in modules/SwissArmyTransformer/sat/arguments.py [0:0]
def add_training_args(parser):
"""Training arguments."""
group = parser.add_argument_group('train', 'training configurations')
# --------------- Core hyper-parameters ---------------
group.add_argument('--experiment-name', type=str, default="MyModel",
help="The experiment name for summary and checkpoint."
"Will load the previous name if mode==pretrain and with --load ")
group.add_argument('--train-iters', type=int, default=None,
help='total number of iterations to train over all training runs')
group.add_argument('--batch-size', type=int, default=4,
help='batch size on a single GPU. batch-size * world_size = total batch_size.')
group.add_argument('--lr', type=float, default=1.0e-4,
help='initial learning rate')
group.add_argument('--mode', type=str,
default='pretrain',
choices=['pretrain', # from_scratch / load ckpt for continue pretraining.
'finetune', # finetuning, auto-warmup 100 iters, new exp name.
'inference' # don't train.
],
help='what type of task to use, will influence auto-warmup, exp name, iteration')
group.add_argument('--seed', type=int, default=1234, help='random seed')
group.add_argument('--zero-stage', type=int, default=0, choices=[0, 1, 2, 3],
help='deepspeed ZeRO stage. 0 means no ZeRO.')
# --------------- Optional hyper-parameters ---------------
# Efficiency.
group.add_argument('--checkpoint-activations', action='store_true',
help='checkpoint activation to allow for training '
'with larger models and sequences. become slow (< 1.5x), save CUDA memory.')
# Inessential
group.add_argument('--checkpoint-num-layers', type=int, default=1,
help='chunk size (number of layers) for checkpointing. ')
group.add_argument('--checkpoint-skip-layers', type=int, default=0,
help='skip the last N layers for checkpointing.')
group.add_argument('--fp16', action='store_true',
help='Run model in fp16 mode')
group.add_argument('--bf16', action='store_true',
help='Run model in bf16 mode')
group.add_argument('--gradient-accumulation-steps', type=int, default=1,
help='run optimizer after every gradient-accumulation-steps backwards.')
group.add_argument('--profiling', type=int, default=-1,
help='profiling, -1 means no profiling, otherwise means warmup args.profiling iters then profiling.')
group.add_argument('--epochs', type=int, default=None,
help='number of train epochs')
group.add_argument('--log-interval', type=int, default=50,
help='report interval')
group.add_argument('--summary-dir', type=str, default="", help="The directory to store the summary")
group.add_argument('--save-args', action='store_true',
help='save args corresponding to the experiment-name')
# Learning rate & weight decay.
group.add_argument('--lr-decay-iters', type=int, default=None,
help='number of iterations to decay LR over,'
' If None defaults to `--train-iters`*`--epochs`')
group.add_argument('--lr-decay-style', type=str, default='linear',
choices=['constant', 'linear', 'cosine', 'exponential'],
help='learning rate decay function')
group.add_argument('--lr-decay-ratio', type=float, default=0.1)
group.add_argument('--warmup', type=float, default=0.01,
help='percentage of data to warmup on (.01 = 1% of all '
'training iters). Default 0.01')
group.add_argument('--weight-decay', type=float, default=0.01,
help='weight decay coefficient for L2 regularization')
# model checkpointing
group.add_argument('--save', type=str, default=None,
help='Output directory to save checkpoints to.')
group.add_argument('--load', type=str, default=None,
help='Path to a directory containing a model checkpoint.')
group.add_argument('--force-train', action='store_true',
help='Force training even with missing keys.')
group.add_argument('--save-interval', type=int, default=5000,
help='number of iterations between saves')
group.add_argument('--no-save-rng', action='store_true',
help='Do not save current rng state.')
group.add_argument('--no-load-rng', action='store_true',
help='Do not load rng state when loading checkpoint.')
group.add_argument('--resume-dataloader', action='store_true',
help='Resume the dataloader when resuming training. ')
# distributed training related, don't use them.
group.add_argument('--distributed-backend', default='nccl',
help='which backend to use for distributed '
'training. One of [gloo, nccl]')
group.add_argument('--local_rank', type=int, default=None,
help='local rank passed from distributed launcher')
# exit, for testing the first period of a long training
group.add_argument('--exit-interval', type=int, default=None,
help='Exit the program after this many new iterations.')
group.add_argument('--wandb', action="store_true", help='whether to use wandb')
group.add_argument('--wandb-project-name', type=str, default="default_project",
help="The project name in wandb.")
return parser