| 587 | |
| 588 | |
| 589 | def _add_training_args(parser): |
| 590 | group = parser.add_argument_group(title="training") |
| 591 | |
| 592 | group.add_argument( |
| 593 | "--micro-batch-size", |
| 594 | type=int, |
| 595 | default=None, |
| 596 | help="Batch size per model instance (local batch size). " |
| 597 | "Global batch size is local batch size times data " |
| 598 | "parallel size times number of micro batches.", |
| 599 | ) |
| 600 | group.add_argument( |
| 601 | "--batch-size", |
| 602 | type=int, |
| 603 | default=None, |
| 604 | help="Old batch size parameter, do not use. " "Use --micro-batch-size instead", |
| 605 | ) |
| 606 | group.add_argument( |
| 607 | "--global-batch-size", |
| 608 | type=int, |
| 609 | default=None, |
| 610 | help="Training batch size. If set, it should be a " |
| 611 | "multiple of micro-batch-size times data-parallel-size. " |
| 612 | "If this value is None, then " |
| 613 | "use micro-batch-size * data-parallel-size as the " |
| 614 | "global batch size. This choice will result in 1 for " |
| 615 | "number of micro-batches.", |
| 616 | ) |
| 617 | group.add_argument( |
| 618 | "--rampup-batch-size", |
| 619 | nargs="*", |
| 620 | default=None, |
| 621 | help="Batch size ramp up with the following values:" |
| 622 | " --rampup-batch-size <start batch size> " |
| 623 | " <batch size incerement> " |
| 624 | " <ramp-up samples> " |
| 625 | "For example:" |
| 626 | " --rampup-batch-size 16 8 300000 \ " |
| 627 | " --global-batch-size 1024" |
| 628 | "will start with global batch size 16 and over " |
| 629 | " (1024 - 16) / 8 = 126 intervals will increase" |
| 630 | "the batch size linearly to 1024. In each interval" |
| 631 | "we will use approximately 300000 / 126 = 2380 samples.", |
| 632 | ) |
| 633 | group.add_argument( |
| 634 | "--checkpoint-activations", |
| 635 | action="store_true", |
| 636 | help="Checkpoint activation to allow for training " |
| 637 | "with larger models, sequences, and batch sizes.", |
| 638 | ) |
| 639 | group.add_argument( |
| 640 | "--distribute-checkpointed-activations", |
| 641 | action="store_true", |
| 642 | help="If set, distribute checkpointed activations " |
| 643 | "across model parallel group.", |
| 644 | ) |
| 645 | group.add_argument( |
| 646 | "--checkpoint-num-layers", |