Training arguments.
(parser)
| 94 | |
| 95 | |
| 96 | def add_training_args(parser): |
| 97 | """Training arguments.""" |
| 98 | |
| 99 | group = parser.add_argument_group('train', 'training configurations') |
| 100 | |
| 101 | group.add_argument('--batch-size', type=int, default=4, |
| 102 | help='Data Loader batch size') |
| 103 | group.add_argument('--weight-decay', type=float, default=0.01, |
| 104 | help='weight decay coefficient for L2 regularization') |
| 105 | group.add_argument('--checkpoint-activations', action='store_true', |
| 106 | help='checkpoint activation to allow for training ' |
| 107 | 'with larger models and sequences') |
| 108 | group.add_argument('--checkpoint-num-layers', type=int, default=1, |
| 109 | help='chunk size (number of layers) for checkpointing') |
| 110 | group.add_argument('--partition-activations', action='store_true', |
| 111 | help='partition the checkpoint activation across model parallel process') |
| 112 | group.add_argument('--clip-grad', type=float, default=1.0, |
| 113 | help='gradient clipping') |
| 114 | group.add_argument('--train-iters', type=int, default=1000000, |
| 115 | help='total number of iterations to train over all training runs') |
| 116 | group.add_argument('--log-interval', type=int, default=100, |
| 117 | help='report interval') |
| 118 | group.add_argument('--exit-interval', type=int, default=None, |
| 119 | help='Exit the program after this many new iterations.') |
| 120 | |
| 121 | group.add_argument('--seed', type=int, default=1234, |
| 122 | help='random seed') |
| 123 | # Batch prodecuer arguments |
| 124 | group.add_argument('--reset-position-ids', action='store_true', |
| 125 | help='Reset posistion ids after end-of-document token.') |
| 126 | group.add_argument('--reset-attention-mask', action='store_true', |
| 127 | help='Reset self attention maske after ' |
| 128 | 'end-of-document token.') |
| 129 | |
| 130 | # Learning rate. |
| 131 | group.add_argument('--lr-decay-iters', type=int, default=None, |
| 132 | help='number of iterations to decay LR over,' |
| 133 | ' If None defaults to `--train-iters`*`--epochs`') |
| 134 | group.add_argument('--lr-decay-style', type=str, default='linear', |
| 135 | choices=['constant', 'linear', 'cosine', 'exponential'], |
| 136 | help='learning rate decay function') |
| 137 | group.add_argument('--lr', type=float, default=1.0e-4, |
| 138 | help='initial learning rate') |
| 139 | group.add_argument('--warmup', type=float, default=0.01, |
| 140 | help='percentage of data to warmup on (.01 = 1% of all ' |
| 141 | 'training iters). Default 0.01') |
| 142 | # model checkpointing |
| 143 | group.add_argument('--save', type=str, default=None, |
| 144 | help='Output directory to save checkpoints to.') |
| 145 | group.add_argument('--save-interval', type=int, default=5000, |
| 146 | help='number of iterations between saves') |
| 147 | group.add_argument('--no-save-optim', action='store_true', |
| 148 | help='Do not save current optimizer.') |
| 149 | group.add_argument('--no-save-rng', action='store_true', |
| 150 | help='Do not save current rng state.') |
| 151 | group.add_argument('--load', type=str, default=None, |
| 152 | help='Path to a directory containing a model checkpoint.') |
| 153 | group.add_argument('--no-load-optim', action='store_true', |