arguments.py 109 KB
Newer Older
liangjing's avatar
liangjing committed
1
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
Raul Puri's avatar
Raul Puri committed
2

Mohammad's avatar
Mohammad committed
3
"""Megatron arguments."""
Raul Puri's avatar
Raul Puri committed
4
5

import argparse
liangjing's avatar
v1  
liangjing committed
6
import dataclasses
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
7
import json
liangjing's avatar
liangjing committed
8
import logging
Raul Puri's avatar
Raul Puri committed
9
import os
10
import torch
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
11
12
import types

liangjing's avatar
v1  
liangjing committed
13
import torch.nn.functional as F
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
14

liangjing's avatar
liangjing committed
15
16
17
18
19
from megatron.core.dist_checkpointing.validation import StrictHandling
from megatron.core.models.retro.utils import (
    get_config_path as get_retro_config_path,
    get_gpt_data_dir as get_retro_data_dir,
)
liangjing's avatar
v1  
liangjing committed
20
from megatron.core.transformer import TransformerConfig
liangjing's avatar
liangjing committed
21
22
23
from megatron.training.activations import squared_relu
from megatron.training.utils import update_use_dist_ckpt

Raul Puri's avatar
Raul Puri committed
24

25
def parse_args(extra_args_provider=None, ignore_unknown_args=False):
Mohammad's avatar
Mohammad committed
26
    """Parse all arguments."""
27
28
    parser = argparse.ArgumentParser(description='Megatron-LM Arguments',
                                     allow_abbrev=False)
Mohammad's avatar
Mohammad committed
29

Mohammad's avatar
Mohammad committed
30
31
32
33
34
35
36
37
38
39
40
41
    # Standard arguments.
    parser = _add_network_size_args(parser)
    parser = _add_regularization_args(parser)
    parser = _add_training_args(parser)
    parser = _add_initialization_args(parser)
    parser = _add_learning_rate_args(parser)
    parser = _add_checkpointing_args(parser)
    parser = _add_mixed_precision_args(parser)
    parser = _add_distributed_args(parser)
    parser = _add_validation_args(parser)
    parser = _add_data_args(parser)
    parser = _add_autoresume_args(parser)
Mostofa Patwary's avatar
Mostofa Patwary committed
42
    parser = _add_biencoder_args(parser)
43
    parser = _add_vision_args(parser)
liangjing's avatar
liangjing committed
44
    parser = _add_moe_args(parser)
45
    parser = _add_logging_args(parser)
liangjing's avatar
liangjing committed
46
    parser = _add_straggler_detector_args(parser)
mshoeybi's avatar
mshoeybi committed
47
    parser = _add_inference_args(parser)
48
    parser = _add_transformer_engine_args(parser)
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
49
    parser = _add_retro_args(parser)
liangjing's avatar
liangjing committed
50
51
52
53
54
    parser = _add_experimental_args(parser)
    parser = _add_one_logger_args(parser)
    parser = _add_ft_package_args(parser)
    parser = _add_config_logger_args(parser)
    parser = _add_unsloth_args(parser)
Mohammad's avatar
Mohammad committed
55
56
57
58

    # Custom arguments.
    if extra_args_provider is not None:
        parser = extra_args_provider(parser)
Mohammad's avatar
Mohammad committed
59

Mohammad's avatar
Mohammad committed
60
    # Parse.
61
62
63
64
    if ignore_unknown_args:
        args, _ = parser.parse_known_args()
    else:
        args = parser.parse_args()
Mohammad's avatar
Mohammad committed
65

liangjing's avatar
liangjing committed
66
67
68
69
70
71
72
73
    # Experimental yaml
    if args.yaml_cfg is not None:
        from .yaml_arguments import load_yaml
        assert args.yaml_cfg and not args.use_legacy_models, \
            "Yaml config is not supported with legacy models."
        args = load_yaml(args.yaml_cfg)


74
    # Args from environment
liangjing's avatar
v1  
liangjing committed
75
76
77
    #args.rank = int(os.getenv('RANK', '0'))
    #args.world_size = int(os.getenv("WORLD_SIZE", '1'))

78
79
    return args

liangjing's avatar
liangjing committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

def load_retro_config(retro_project_dir):
    '''Load Retro's config.json.'''

    # Retro config path.
    retro_config_path = get_retro_config_path(retro_project_dir)
    assert os.path.exists(retro_config_path), \
        "Retro project dir missing config.json."

    # Load retro config.
    with open(retro_config_path) as f:
        retro_config = types.SimpleNamespace(**json.load(f))

    return retro_config


def load_retro_args(args):
    """Load predefined args from Retro config (if applicable).

    When using Retro (or GPT for comparison purposes), data arguments are
    overridden by the saved config.json within the Retro project directory. This
    is to ensure that the data used for pretraining is consistent with the data
    that was preprocessed using the Retro preprocessing pipeline (see
    `tools/retro/preprocess_data.py`).
    """

    # Return if no project directory is specified.
    if args.retro_project_dir is None:
        return

    # Load retro config.
    retro_config = load_retro_config(args.retro_project_dir)

    # Retro data path is relative to project dir (via hard or soft links).
    data_dir = get_retro_data_dir(args.retro_project_dir)
    data_path = list(retro_config.retro_gpt_data_path)
    if len(data_path) % 2 == 0:
        for i in range(len(data_path) - 1, -1, -2):
            data_path[i] = os.path.join(data_dir, data_path[i])
    else:
        assert len(data_path) == 1
        data_path[0] = os.path.join(data_dir, data_path[0])

    # Update args.
    args.data_cache_path = retro_config.retro_gpt_data_cache_path
    args.data_path = data_path if args.data_path is None else args.data_path
    args.eval_interval = retro_config.retro_gpt_eval_interval
    args.eval_iters = retro_config.retro_gpt_eval_iters
    args.global_batch_size = retro_config.retro_gpt_global_batch_size
    args.max_position_embeddings = retro_config.retro_gpt_seq_length
    args.merge_file = os.path.join(
        args.retro_project_dir,
        retro_config.retro_gpt_merge_file,
    ) if retro_config.retro_gpt_merge_file is not None else None
    args.seed = retro_config.retro_gpt_seed
    args.seq_length = retro_config.retro_gpt_seq_length
    args.tokenizer_model = os.path.join(
        args.retro_project_dir,
        retro_config.retro_gpt_tokenizer_model,
    ) if retro_config.retro_gpt_tokenizer_model is not None else None
    args.tokenizer_type = retro_config.retro_gpt_tokenizer_type
    args.train_samples = retro_config.retro_gpt_train_samples
    args.vocab_file = os.path.join(
        args.retro_project_dir,
        retro_config.retro_gpt_vocab_file,
    ) if retro_config.retro_gpt_vocab_file is not None else None

    # Retro-specific args.
    args.retro_block_size = retro_config.retro_block_size
    args.retro_chunk_length = retro_config.retro_gpt_chunk_length
    args.retro_neighbor_dirs = retro_config.retro_neighbor_dirs
    args.retro_split_preprocessing = retro_config.retro_gpt_split
    args.retro_bert_tokenizer_type = retro_config.retro_bert_tokenizer_type
    args.retro_bert_vocab_file = retro_config.retro_bert_vocab_file


156
def validate_args(args, defaults={}):
liangjing's avatar
liangjing committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

    # Temporary
    assert args.non_persistent_ckpt_type in ['global', None], \
        'Currently only global checkpoints are supported'

    # Load saved args from Retro (if applicable).
    load_retro_args(args)

    # Set args.use_dist_ckpt from args.ckpt_format.
    update_use_dist_ckpt(args)

    if args.encoder_tensor_model_parallel_size > 0:
        assert args.encoder_pipeline_model_parallel_size > 0, "encoder_pipeline_model_parallel_size must be defined."
        assert args.num_attention_heads % args.encoder_tensor_model_parallel_size == 0
        assert args.encoder_tensor_model_parallel_size <= args.tensor_model_parallel_size, "We do not support encoders with more TP than the decoder."

    if args.encoder_pipeline_model_parallel_size > 0 and args.encoder_tensor_model_parallel_size == 0:
        args.encoder_tensor_model_parallel_size = args.tensor_model_parallel_size

    encoder_model_size = args.encoder_tensor_model_parallel_size * args.encoder_pipeline_model_parallel_size * args.context_parallel_size
    decoder_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size
    total_model_size = encoder_model_size + decoder_model_size

    # Total model size.
    assert args.world_size % total_model_size == 0, (
        f"world size ({args.world_size}) is not divisible by total_model_size ({encoder_model_size=} + {decoder_model_size=})"
    )

mohammad's avatar
mohammad committed
185
    # Pipeline model parallel size.
186
187
    args.transformer_pipeline_model_parallel_size = (
        args.pipeline_model_parallel_size - 1
188
        if args.standalone_embedding_stage else
189
190
        args.pipeline_model_parallel_size
    )
liangjing's avatar
liangjing committed
191
192
193

    args.data_parallel_size = args.world_size // total_model_size

mohammad's avatar
mohammad committed
194
    # Checks.
Mohammad's avatar
Mohammad committed
195
    if args.rank == 0:
liangjing's avatar
liangjing committed
196
197
        print('using world size: {}, data-parallel size: {}, '
              'context-parallel size: {}, '
mohammad's avatar
mohammad committed
198
              'tensor-model-parallel size: {}, '
liangjing's avatar
liangjing committed
199
200
201
              'encoder-tensor-model-parallel size: {}, '
              'pipeline-model-parallel size: {}, '
              'encoder-pipeline-model-parallel size: {}'.format(
mohammad's avatar
mohammad committed
202
                  args.world_size, args.data_parallel_size,
liangjing's avatar
liangjing committed
203
                  args.context_parallel_size,
mohammad's avatar
mohammad committed
204
                  args.tensor_model_parallel_size,
liangjing's avatar
liangjing committed
205
206
207
208
209
210
211
212
213
214
215
216
                  args.encoder_tensor_model_parallel_size,
                  args.pipeline_model_parallel_size,
                  args.encoder_pipeline_model_parallel_size), flush=True)

    # backwards compatibility.
    if args.pipeline_model_parallel_split_rank is not None:
        args.encoder_pipeline_model_parallel_size = args.pipeline_model_parallel_split_rank
        args.pipeline_model_parallel_size -= args.encoder_pipeline_model_parallel_size
        assert args.pipeline_model_parallel_size > 0

    if args.tp_comm_overlap:
        assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled'
mohammad's avatar
mohammad committed
217

218
219
220
221
222
223
224
225
226
227
    # Deprecated arguments
    assert args.batch_size is None, '--batch-size argument is no longer ' \
        'valid, use --micro-batch-size instead'
    del args.batch_size
    assert args.warmup is None, '--warmup argument is no longer valid, use ' \
        '--lr-warmup-fraction instead'
    del args.warmup
    assert args.model_parallel_size is None, '--model-parallel-size is no ' \
        'longer valid, use --tensor-model-parallel-size instead'
    del args.model_parallel_size
Vijay Korthikanti's avatar
Vijay Korthikanti committed
228

229
    if args.checkpoint_activations:
slym's avatar
slym committed
230
        if args.rank == 0:
liangjing's avatar
v1  
liangjing committed
231
232
233
            print('--checkpoint-activations is no longer valid, use --recompute-activations, '
                  'or, for more control, --recompute-granularity and --recompute-method.')
        exit()
234
    del args.checkpoint_activations
235

Vijay Korthikanti's avatar
Vijay Korthikanti committed
236
237
238
239
    if args.recompute_activations:
        args.recompute_granularity = 'selective'
    del args.recompute_activations

Jared Casper's avatar
Jared Casper committed
240
241
242
243
244
    # Set input defaults.
    for key in defaults:
        # For default to be valid, it should not be provided in the
        # arguments that are passed to the program. We check this by
        # ensuring the arg is set to None.
liangjing's avatar
v1  
liangjing committed
245
        if getattr(args, key, None) is not None:
Jared Casper's avatar
Jared Casper committed
246
247
248
249
250
251
252
253
            if args.rank == 0:
                print('WARNING: overriding default arguments for {key}:{v} \
                       with {key}:{v2}'.format(key=key, v=defaults[key],
                                               v2=getattr(args, key)),
                                               flush=True)
        else:
            setattr(args, key, defaults[key])

liangjing's avatar
liangjing committed
254
255
256
257
258
259
260
    if args.data_path is not None and args.split is None:
        legacy_default_split_value = '969, 30, 1'
        if args.rank == 0:
            print('WARNING: Please specify --split when using --data-path. Using legacy default value '
                  f'of "{legacy_default_split_value}"')
        args.split = legacy_default_split_value

mohammad's avatar
mohammad committed
261
262
263
264
265
266
267
268
269
    # Batch size.
    assert args.micro_batch_size is not None
    assert args.micro_batch_size > 0
    if args.global_batch_size is None:
        args.global_batch_size = args.micro_batch_size * args.data_parallel_size
        if args.rank == 0:
            print('setting global batch size to {}'.format(
                args.global_batch_size), flush=True)
    assert args.global_batch_size > 0
270
    if args.num_layers_per_virtual_pipeline_stage is not None:
liangjing's avatar
liangjing committed
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
        if args.overlap_p2p_comm:
            assert args.pipeline_model_parallel_size > 1, \
                'when interleaved schedule is used, pipeline-model-parallel size '\
                'should be greater than 1'
        else:
            assert args.pipeline_model_parallel_size > 2, \
                'when interleaved schedule is used and p2p communication overlap is disabled, '\
                'pipeline-model-parallel size should be greater than 2 to avoid having multiple '\
                'p2p sends and recvs between same 2 ranks per communication batch'
        assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \
            'number of layers should be divisible by the pipeline parallel size'
        num_layers_per_pipeline_stage = args.num_layers // args.transformer_pipeline_model_parallel_size
        assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \
            'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage'
        args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \
286
287
288
            args.num_layers_per_virtual_pipeline_stage
    else:
        args.virtual_pipeline_model_parallel_size = None
liangjing's avatar
liangjing committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
        # Overlap P2P communication is disabled if not using the interleaved schedule.
        args.overlap_p2p_comm = False
        args.align_param_gather = False
        if args.rank == 0:
            print('WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False '
                  'since non-interleaved schedule does not support overlapping p2p communication '
                  'and aligned param AG')

    if args.overlap_param_gather:
        assert args.use_distributed_optimizer, \
            '--overlap-param-gather only supported with distributed optimizer'
        assert args.overlap_grad_reduce, \
            'Must use --overlap-param-gather with --overlap-grad-reduce'
        assert not args.use_legacy_models, \
            '--overlap-param-gather only supported with MCore models'

    if args.overlap_param_gather_with_optimizer_step:
        assert args.use_distributed_optimizer, \
            '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer'
        assert args.overlap_param_gather, \
            'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather'
        assert args.virtual_pipeline_model_parallel_size is not None, \
            '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism'
        assert not args.use_dist_ckpt, \
            '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet'

    if args.fp8_param_gather:
        assert args.use_distributed_optimizer, \
            '--fp8-param-gather only supported with distributed optimizer'
Mohammad's avatar
Mohammad committed
318

319
320
321
    # Parameters dtype.
    args.params_dtype = torch.float
    if args.fp16:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
322
        assert not args.bf16
323
        args.params_dtype = torch.half
liangjing's avatar
liangjing committed
324
325
326
327
328
329
330
        # Turn off checking for NaNs in loss and grads if using dynamic loss scaling,
        # where NaNs in grads / loss are signal to the loss scaler.
        if not args.loss_scale:
            args.check_for_nan_in_loss_and_grad = False
            if args.rank == 0:
                print('WARNING: Setting args.check_for_nan_in_loss_and_grad to False since '
                      'dynamic loss scaling is being used')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
331
332
333
    if args.bf16:
        assert not args.fp16
        args.params_dtype = torch.bfloat16
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
334
335
336
337
338
339
340
        # bfloat16 requires gradient accumulation and all-reduce to
        # be done in fp32.
        if not args.accumulate_allreduce_grads_in_fp32:
            args.accumulate_allreduce_grads_in_fp32 = True
            if args.rank == 0:
                print('accumulate and all-reduce gradients in fp32 for '
                      'bfloat16 data type.', flush=True)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
341

342
343
344
345
    if args.rank == 0:
        print('using {} for parameters ...'.format(args.params_dtype),
              flush=True)

346
347
348
    if args.dataloader_type is None:
        args.dataloader_type = 'single'

liangjing's avatar
liangjing committed
349
350
351
    # data
    assert args.num_dataset_builder_threads > 0

352
353
    # Consumed tokens.
    args.consumed_train_samples = 0
liangjing's avatar
liangjing committed
354
    args.skipped_train_samples = 0
355
    args.consumed_valid_samples = 0
356

357
358
359
360
361
362
363
    # Support for variable sequence lengths across batches/microbatches.
    # set it if the dataloader supports generation of variable sequence lengths
    # across batches/microbatches. Due to additional communication overhead
    # during pipeline parallelism, it should not be set if sequence length
    # is constant during training.
    args.variable_seq_lengths = False

364
365
366
367
368
369
370
371
372
    # Iteration-based training.
    if args.train_iters:
        # If we use iteration-based training, make sure the
        # sample-based options are off.
        assert args.train_samples is None, \
            'expected iteration-based training'
        assert args.lr_decay_samples is None, \
            'expected iteration-based learning rate decay'
        assert args.lr_warmup_samples == 0, \
373
            'expected iteration-based learning rate warmup'
374
375
        assert args.rampup_batch_size is None, \
            'expected no batch-size rampup for iteration-based training'
376
        if args.lr_warmup_fraction is not None:
377
            assert args.lr_warmup_iters == 0, \
378
                'can only specify one of lr-warmup-fraction and lr-warmup-iters'
379
380
381
382
383
384
385
386
387
388
389

    # Sample-based training.
    if args.train_samples:
        # If we use sample-based training, make sure the
        # iteration-based options are off.
        assert args.train_iters is None, \
            'expected sample-based training'
        assert args.lr_decay_iters is None, \
            'expected sample-based learning rate decay'
        assert args.lr_warmup_iters == 0, \
            'expected sample-based learnig rate warmup'
390
        if args.lr_warmup_fraction is not None:
391
            assert args.lr_warmup_samples == 0, \
392
393
                'can only specify one of lr-warmup-fraction ' \
                'and lr-warmup-samples'
394

395
    if args.num_layers is not None:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
396
397
        assert args.encoder_num_layers is None, \
            'cannot have both num-layers and encoder-num-layers specified'
398
399
        args.encoder_num_layers = args.num_layers
    else:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
400
401
        assert args.encoder_num_layers is not None, \
            'either num-layers or encoder-num-layers should be specified'
402
403
        args.num_layers = args.encoder_num_layers

404
    # Check required arguments.
Mohammad's avatar
Mohammad committed
405
406
    required_args = ['num_layers', 'hidden_size', 'num_attention_heads',
                     'max_position_embeddings']
407
    for req_arg in required_args:
Mohammad's avatar
Mohammad committed
408
        _check_arg_is_not_none(args, req_arg)
409

Mohammad's avatar
Mohammad committed
410
    # Checks.
411
    if args.ffn_hidden_size is None:
liangjing's avatar
liangjing committed
412
413
414
415
416
417
418
419
420
        if args.swiglu:
            # reduce the dimnesion for MLP since projections happens on
            # two linear layers. this keeps the number of paramters in
            # the same ballpark as the counterpart with 4*h size
            # we keep it a multiple of 64, which means the actual tensor size
            # will be a multiple of 64 / tp_size
            args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64
        else:
            args.ffn_hidden_size = 4 * args.hidden_size
421

422
423
424
425
    if args.kv_channels is None:
        assert args.hidden_size % args.num_attention_heads == 0
        args.kv_channels = args.hidden_size // args.num_attention_heads

liangjing's avatar
liangjing committed
426
427
428
429
430
    if args.seq_length is not None and args.context_parallel_size > 1:
        assert args.seq_length % (args.context_parallel_size * 2) == 0, \
            'seq-length should be a multiple of 2 * context-parallel-size ' \
            'if context-parallel-size > 1.'

431
432
433
434
435
436
    if args.seq_length is not None:
        assert args.encoder_seq_length is None
        args.encoder_seq_length = args.seq_length
    else:
        assert args.encoder_seq_length is not None
        args.seq_length = args.encoder_seq_length
437

Mohammad's avatar
Mohammad committed
438
439
    if args.seq_length is not None:
        assert args.max_position_embeddings >= args.seq_length
Jared Casper's avatar
Jared Casper committed
440
441
    if args.decoder_seq_length is not None:
        assert args.max_position_embeddings >= args.decoder_seq_length
Mohammad's avatar
Mohammad committed
442
443
    if args.lr is not None:
        assert args.min_lr <= args.lr
Mohammad's avatar
Mohammad committed
444
445
    if args.save is not None:
        assert args.save_interval is not None
mohammad's avatar
mohammad committed
446
447
448
    # Mixed precision checks.
    if args.fp16_lm_cross_entropy:
        assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.'
449
    if args.fp32_residual_connection:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
450
451
        assert args.fp16 or args.bf16, \
            'residual connection in fp32 only supported when using fp16 or bf16.'
Vijay Korthikanti's avatar
Vijay Korthikanti committed
452

liangjing's avatar
liangjing committed
453
454
455
456
457
    if args.moe_grouped_gemm:
        assert args.bf16, 'Currently GroupedGEMM for MoE only supports bf16 dtype.'
        dc = torch.cuda.get_device_capability()
        assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels."

Vijay Korthikanti's avatar
Vijay Korthikanti committed
458
459
460
461
462
    if args.weight_decay_incr_style == 'constant':
        assert args.start_weight_decay is None
        assert args.end_weight_decay is None
        args.start_weight_decay = args.weight_decay
        args.end_weight_decay = args.weight_decay
Vijay Korthikanti's avatar
Vijay Korthikanti committed
463
    else:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
464
465
        assert args.start_weight_decay is not None
        assert args.end_weight_decay is not None
466

Sangkug Lym's avatar
Sangkug Lym committed
467
468
469
470
471
472
473
474
475
476
    TORCH_MAJOR = int(torch.__version__.split('.')[0])
    TORCH_MINOR = int(torch.__version__.split('.')[1])
    # Persistent fused layer norm.
    if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 11):
        args.no_persist_layer_norm = True
        if args.rank == 0:
            print('Persistent fused layer norm kernel is supported from '
                  'pytorch v1.11 (nvidia pytorch container paired with v1.11). '
                  'Defaulting to no_persist_layer_norm=True')

Vijay Korthikanti's avatar
Vijay Korthikanti committed
477
    # Activation recomputing.
Vijay Korthikanti's avatar
Vijay Korthikanti committed
478
    if args.distribute_saved_activations:
mshoeybi's avatar
mshoeybi committed
479
        assert args.tensor_model_parallel_size > 1, 'can distribute ' \
Vijay Korthikanti's avatar
Vijay Korthikanti committed
480
            'recomputed activations only across tensor model ' \
mshoeybi's avatar
mshoeybi committed
481
            'parallel groups'
Vijay Korthikanti's avatar
Vijay Korthikanti committed
482
483
484
485
486
487
        assert args.recompute_granularity == 'full', \
            'distributed recompute activations is only '\
            'application to full recompute granularity'
        assert args.recompute_method is not None, \
            'for distributed recompute activations to work you '\
            'need to use a recompute method '
liangjing's avatar
v1  
liangjing committed
488
        assert (TORCH_MAJOR, TORCH_MINOR) >= (1, 10), \
Vijay Korthikanti's avatar
Vijay Korthikanti committed
489
            'distributed recompute activations are supported for pytorch ' \
490
491
            'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \
            'pytorch version is v%s.%s.' % (TORCH_MAJOR, TORCH_MINOR)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
492

Vijay Korthikanti's avatar
Vijay Korthikanti committed
493
494
495
496
    if args.recompute_granularity == 'selective':
        assert args.recompute_method is None, \
            'recompute method is not yet supported for ' \
            'selective recomputing granularity'
Vijay Korthikanti's avatar
Vijay Korthikanti committed
497
498
499
500
501
502
503

    # disable sequence parallelism when tp=1
    # to avoid change in numerics when
    # sequence_parallelism is enabled.
    if args.tensor_model_parallel_size == 1:
        args.sequence_parallel = False

Vijay Korthikanti's avatar
Vijay Korthikanti committed
504
    # disable async_tensor_model_parallel_allreduce when
Vijay Korthikanti's avatar
Vijay Korthikanti committed
505
    # model parallel memory optimization is enabled
Vijay Korthikanti's avatar
Vijay Korthikanti committed
506
507
    if args.sequence_parallel:
        args.async_tensor_model_parallel_allreduce = False
Vijay Korthikanti's avatar
Vijay Korthikanti committed
508

509
510
511
512
513
514
515
516
517
518
    if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1":
        if args.sequence_parallel:
            raise RuntimeError(
                "Using sequence parallelism requires setting the environment variable "
                "CUDA_DEVICE_MAX_CONNECTIONS to 1")
        if args.async_tensor_model_parallel_allreduce:
            raise RuntimeError(
                "Using async gradient all reduce requires setting the environment "
                "variable CUDA_DEVICE_MAX_CONNECTIONS to 1")

519
520
521
522
    # Disable bias gelu fusion if we are disabling bias altogether
    if not args.add_bias_linear:
        args.bias_gelu_fusion = False

liangjing's avatar
v1  
liangjing committed
523
524
525
    # Retro checks.
    if args.retro_add_retriever:

liangjing's avatar
liangjing committed
526
527
528
529
        # Train samples should be auto-loaded.
        assert args.train_samples is not None, \
            "args.train_samples should be auto-loaded from the retro config."

liangjing's avatar
v1  
liangjing committed
530
531
532
533
534
535
536
537
        # Sequence parallelism unsupported.
        assert not args.sequence_parallel, \
            "retro currently does not support sequence parallelism."

        # Pipeline parallelism unsupported.
        assert args.pipeline_model_parallel_size == 1, \
            "retro currently does not support pipeline parallelism."

liangjing's avatar
liangjing committed
538
539
540
541
542
    if args.decoupled_lr is not None or args.decoupled_min_lr is not None:
        assert not args.use_legacy_models, \
            '--decoupled-lr and --decoupled-min-lr is not supported in legacy models.'
    # FlashAttention
    args.use_flash_attn = args.use_flash_attn_ck or args.use_flash_attn_triton
liangjing's avatar
v1  
liangjing committed
543
544
545
546

    # Legacy RoPE arguments
    if args.use_rotary_position_embeddings:
        args.position_embedding_type = 'rope'
liangjing's avatar
liangjing committed
547
548
549
550
    if args.rotary_interleaved and args.apply_rope_fusion:
        raise RuntimeError('--rotary-interleaved does not work with rope_fusion.')
    if args.rotary_interleaved and args.use_legacy_models:
        raise RuntimeError('--rotary-interleaved is not supported in legacy models.')
liangjing's avatar
v1  
liangjing committed
551
552
553
554
555

    # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now
    # don't allow it to keep things simple
    if not args.add_position_embedding and args.position_embedding_type != 'rope':
        raise RuntimeError('--no-position-embedding is deprecated, use --position-embedding-type')
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
556

liangjing's avatar
liangjing committed
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
    # MoE Spec check
    if args.num_experts == 0:
        args.num_experts = None
    if args.num_experts is not None:
        assert args.spec is None, "Model Spec must be None when using MoEs"

    # Context parallel
    if args.context_parallel_size > 1:
        assert not args.use_legacy_models, "Context parallelism is not supported in legacy models."

    # Expert parallelism check
    if args.expert_model_parallel_size  > 1:
        assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism"
        assert args.num_experts % args.expert_model_parallel_size == 0, \
            "Number of experts should be a multiple of expert model parallel_size."
        assert not args.fp16, \
            "Expert parallelism is not supported with fp16 training."

    # Distributed checkpointing checks
    args.use_dist_ckpt = False
    if args.use_dist_ckpt and args.use_legacy_models:
        raise RuntimeError('--use-dist-ckpt is not supported in legacy models.')

    # Data blend checks
    assert args.mock_data + \
           bool(args.data_path) + \
           any([args.train_data_path, args.valid_data_path, args.test_data_path]) \
           <= 1, "A single data source must be provided in training mode, else None"

    if args.use_tp_pp_dp_mapping:
        assert args.context_parallel_size * args.expert_model_parallel_size <= 1, \
            "context_parallel and expert_model_parallel can't be used with tp-pp-dp mapping."

    # Deterministic mode
    if args.deterministic_mode:
        assert not args.use_flash_attn, "Flash attention can not be used in deterministic mode."
        assert not args.cross_entropy_loss_fusion, "Cross Entropy Fusion is currently not deterministic."

        all_reduce_choices = ["Tree", "Ring", "CollnetDirect", "CollnetChain", "^NVLS"]
        assert os.getenv("NCCL_ALGO", -1) != -1 and os.getenv("NCCL_ALGO") in all_reduce_choices, \
            f"NCCL_ALGO must be one of {all_reduce_choices}."

        torch.use_deterministic_algorithms(True)

    # Update the printed args to reflect that `apply_query_key_layer_scaling` also controls `attention_softmax_in_fp32`
    if args.apply_query_key_layer_scaling:
        args.attention_softmax_in_fp32 = True

    # Checkpointing
    if args.ckpt_fully_parallel_save_deprecated and args.rank == 0:
        print('--ckpt-fully-parallel-save flag is deprecated and has no effect.'
              ' Use --no-ckpt-fully-parallel-save to disable parallel save.')
    if (
        args.use_dist_ckpt
        and not args.ckpt_fully_parallel_save
        and args.use_distributed_optimizer
        and args.rank == 0
    ):
        print('Warning: With non-parallel ckpt save and DistributedOptimizer,'
              ' it will be impossible to resume training with different parallelism.'
              ' Consider removing flag --no-ckpt-fully-parallel-save.')
    if args.use_dist_ckpt_deprecated and args.rank == 0:
        print('--use-dist-ckpt is deprecated and has no effect.'
              ' Use --ckpt-format to select the checkpoint format.')
    if args.dist_ckpt_format_deprecated and args.rank == 0:
        print('--dist-ckpt-format is deprecated and has no effect.'
              ' Use --ckpt-format to select the checkpoint format.')

    # MoE upcycling check
    if args.moe_use_upcycling:
        assert args.save is not None, "When using upcycling, the --save option must be specified."
        if not args.no_load_optim:
            args.no_load_optim = True
            print('Warning: disabling --no-load-optim for upcycling.')
        if not args.no_load_rng:
            args.no_load_rng = True
            print('Warning: disabling --no-load-rng for upcycling.')

Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
635
636
    # Print arguments.
    _print_args("arguments", args)
637

Mohammad's avatar
Mohammad committed
638
    return args
Mohammad's avatar
Mohammad committed
639
640


Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
641
def _print_args(title, args):
Mohammad's avatar
Mohammad committed
642
643
    """Print arguments."""
    if args.rank == 0:
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
644
        print(f'------------------------ {title} ------------------------',
mohammad's avatar
mohammad committed
645
              flush=True)
Mohammad's avatar
Mohammad committed
646
647
        str_list = []
        for arg in vars(args):
mohammad's avatar
mohammad committed
648
            dots = '.' * (48 - len(arg))
Mohammad's avatar
Mohammad committed
649
650
651
            str_list.append('  {} {} {}'.format(arg, dots, getattr(args, arg)))
        for arg in sorted(str_list, key=lambda x: x.lower()):
            print(arg, flush=True)
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
652
        print(f'-------------------- end of {title} ---------------------',
mohammad's avatar
mohammad committed
653
              flush=True)
Mohammad's avatar
Mohammad committed
654
655


656
657
658
def _check_arg_is_not_none(args, arg):
    assert getattr(args, arg) is not None, '{} argument is None'.format(arg)

liangjing's avatar
liangjing committed
659
660
661
662
663

def core_transformer_config_from_args(args, config_class=None):

    # Config class.
    config_class = config_class or TransformerConfig
liangjing's avatar
v1  
liangjing committed
664
665
666

    # Translate args to core transformer configuration
    kw_args = {}
liangjing's avatar
liangjing committed
667
    for f in dataclasses.fields(config_class):
liangjing's avatar
v1  
liangjing committed
668
669
670
671
        if hasattr(args, f.name):
            kw_args[f.name] = getattr(args, f.name)
    kw_args['persist_layer_norm'] = not args.no_persist_layer_norm
    kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p
liangjing's avatar
liangjing committed
672
    kw_args['layernorm_epsilon'] = args.norm_epsilon
liangjing's avatar
v1  
liangjing committed
673
674
675
    kw_args['deallocate_pipeline_outputs'] = True
    kw_args['pipeline_dtype'] = args.params_dtype
    kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm
liangjing's avatar
liangjing committed
676
677
678
679
    kw_args['num_moe_experts'] = args.num_experts
    kw_args['rotary_interleaved'] = args.rotary_interleaved
    kw_args['first_pipeline_num_layers']= args.decoder_first_pipeline_num_layers
    kw_args['last_pipeline_num_layers']= args.decoder_last_pipeline_num_layers
liangjing's avatar
v1  
liangjing committed
680
681
682
    if args.swiglu:
        kw_args['activation_func'] = F.silu
        kw_args['gated_linear_unit'] = True
liangjing's avatar
liangjing committed
683
684
685
686
687
688
        kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion
    else:
        kw_args['bias_activation_fusion'] = args.bias_gelu_fusion
    if args.squared_relu:
        assert not args.swiglu
        kw_args['activation_func'] = squared_relu
liangjing's avatar
v1  
liangjing committed
689
690
691
692
693
694
695
    if args.init_method_xavier_uniform:
        kw_args['init_method'] = torch.nn.init.xavier_uniform_
        kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_
    if args.group_query_attention:
        kw_args['num_query_groups'] = args.num_query_groups
    else:
        kw_args['num_query_groups'] = None
liangjing's avatar
liangjing committed
696
697
698
699
    kw_args['config_logger_dir'] = args.config_logger_dir

    # Return config.
    return config_class(**kw_args)
liangjing's avatar
v1  
liangjing committed
700

701

702
703
704
def _add_transformer_engine_args(parser):
    group = parser.add_argument_group(title='Transformer-Engine')

liangjing's avatar
v1  
liangjing committed
705
706
707
708
    group.add_argument('--fp8-format', default=None,
                       choices=['e4m3', 'hybrid'],
                       help='Which fp8 format scheme to use for FP8 tensors in the forward and backward pass',
                       dest='fp8')
709
    group.add_argument('--fp8-margin', type=int, default=0,
liangjing's avatar
v1  
liangjing committed
710
711
                       help='Scaling margin for fp8',
                       dest='fp8_margin')
712
    group.add_argument('--fp8-interval', type=int, default=1,
liangjing's avatar
liangjing committed
713
                       help='DEPRECATED. This flag is ignored. Scaling update interval for fp8',
liangjing's avatar
v1  
liangjing committed
714
                       dest='fp8_interval')
715
    group.add_argument('--fp8-amax-history-len', type=int, default=1,
liangjing's avatar
v1  
liangjing committed
716
717
                       help='Number of steps for which amax history is recorded per tensor',
                       dest='fp8_amax_history_len')
718
719
720
721
    group.add_argument('--fp8-amax-compute-algo', default='most_recent',
                       choices=['most_recent', 'max'],
                       help='Algorithm for computing amax from history',
                       dest='fp8_amax_compute_algo')
liangjing's avatar
v1  
liangjing committed
722
723
724
    group.add_argument('--no-fp8-wgrad', action='store_false',
                       help='Execute wgrad in higher precision even for FP8 runs',
                       dest='fp8_wgrad')
liangjing's avatar
liangjing committed
725
    group.add_argument('--transformer-impl', default='transformer_engine',
liangjing's avatar
v1  
liangjing committed
726
                       choices=['local', 'transformer_engine'],
liangjing's avatar
liangjing committed
727
728
729
730
                       help='Which Transformer implementation to use.')
    group.add_argument('--fp8-param-gather', action='store_true',
                       help='Keep the compute param in fp8 (do not use any other intermediate '
                            'dtype) and perform the param all-gather in fp8.')
731
732
733

    return parser

mshoeybi's avatar
mshoeybi committed
734
735
736
737
738
739
740
741
def _add_inference_args(parser):
    group = parser.add_argument_group(title='inference')

    group.add_argument('--inference-batch-times-seqlen-threshold',
                       type=int, default=512,
                       help='During inference, if batch-size times '
                       'sequence-length is smaller than this threshold '
                       'then we will not use pipelining, otherwise we will.')
742
743
744
745
746
    group.add_argument('--max-tokens-to-oom',
                       type=int, default=12000,
                       help='Maximum number of tokens during inference'
                       'tokens here is # in prompt + # to generate'
                       'Allows us to throw an error before OOM crashes server')
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
747
748
749
750
751
752
753
754
755
    group.add_argument('--output-bert-embeddings', action='store_true',
                       help='Output Bert embeddings (via mean pooling) from '
                       'model, rather than its binary head output or entire '
                       'hidden batch.')
    group.add_argument('--bert-embedder-type', default="megatron",
                       choices=["megatron", "huggingface"],
                       help='Select either Megatron or Huggingface as the '
                       'Bert embedder.')

mshoeybi's avatar
mshoeybi committed
756
757
    return parser

Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
758
759
760
761

def _add_retro_args(parser):
    group = parser.add_argument_group(title='retro')

liangjing's avatar
liangjing committed
762
763
764
    group.add_argument('--retro-project-dir', default=None,
                       help='Retro project directory, which contains the '
                       'preprocessed data for pretraining. This directory '
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
                       'is built during preprocessing (see '
                       'tools/retro/README.md), and contains subdirectories '
                       'for the chunk database and pretraining neighbors.')
    group.add_argument('--retro-add-retriever',
                       action='store_true', default=False,
                       help='Add a retriever to the transformer, for use in '
                       'pretraining a Retro model.')
    group.add_argument('--retro-cyclic-train-iters', type=int, default=None,
                       help='Set number of training iterations for cyclic '
                       'Retro training.')
    group.add_argument('--retro-encoder-layers', type=int, default=2,
                       help='Number of layers to use for the retrieval '
                       'encoder.')
    group.add_argument('--retro-encoder-hidden-dropout',
                       type=float, default=0.1, help='Hidden dropout for '
                       'retrieval encoder.')
    group.add_argument('--retro-encoder-attention-dropout',
                       type=float, default=0.1, help='Attention dropout for '
                       'retrieval encoder.')
    group.add_argument("--retro-num-neighbors", type=int, default=2,
                       help='Number of neighbors to retrieve during '
                       'pretraining.')
    group.add_argument("--retro-num-retrieved-chunks", type=int, default=2,
                       help='Number of chunks to retrieve from the retrieval '
                       'database.')
liangjing's avatar
liangjing committed
790
791
792
793
794
795
    group.add_argument("--retro-attention-gate", type=float, default=1,
                       help="Gated cross attention.")
    group.add_argument("--retro-no-verify-neighbor-count", action="store_false",
                       dest="retro_verify_neighbor_count",
                       help="Skip verifying that len(GPT dataset) == len(saved "
                       "neighbors).")
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
796
797
798
799
800
801
802
803
804
805
806

    # Enforce argument naming convention.
    for action in group._group_actions:
        prefix = action.dest.split("_")[0]
        assert prefix == "retro", \
            "Retro args must be prefixed with '--retro-*', for consistent " \
            "styling. Please fix '%s'." % ", ".join(action.option_strings)

    return parser


Mohammad's avatar
Mohammad committed
807
def _add_network_size_args(parser):
Mohammad's avatar
Mohammad committed
808
    group = parser.add_argument_group(title='network size')
Mohammad's avatar
Mohammad committed
809

810
    group.add_argument('--num-layers', type=int, default=None,
Mohammad's avatar
Mohammad committed
811
                       help='Number of transformer layers.')
812
813
814
815
    group.add_argument('--encoder-num-layers', type=int, default=None,
                       help='Number of encoder transformer layers.')
    group.add_argument('--decoder-num-layers', type=int, default=None,
                       help='Number of decoder transformer layers.')
816
    group.add_argument('--hidden-size', type=int, default=None,
Mohammad's avatar
Mohammad committed
817
                       help='Tansformer hidden size.')
818
    group.add_argument('--ffn-hidden-size', type=int, default=None,
819
820
                       help='Transformer Feed-Forward Network hidden size. '
                       'This is set to 4*hidden-size if not provided')
821
    group.add_argument('--num-attention-heads', type=int, default=None,
Mohammad's avatar
Mohammad committed
822
                       help='Number of transformer attention heads.')
823
    group.add_argument('--kv-channels', type=int, default=None,
824
825
826
827
                       help='Projection weights dimension in multi-head '
                       'attention. This is set to '
                       '   args.hidden_size // args.num_attention_heads '
                       'if not provided.')
liangjing's avatar
v1  
liangjing committed
828
829
830
831
    group.add_argument('--group-query-attention', action='store_true',
                          help='Use group-query attention.')
    group.add_argument('--num-query-groups', type=int, default=1)

832
    group.add_argument('--max-position-embeddings', type=int, default=None,
Mohammad's avatar
Mohammad committed
833
834
                       help='Maximum number of position embeddings to use. '
                       'This is the size of position embedding.')
liangjing's avatar
v1  
liangjing committed
835
    group.add_argument('--position-embedding-type', type=str, default='learned_absolute',
liangjing's avatar
liangjing committed
836
                       choices=['learned_absolute', 'rope', 'none'],
liangjing's avatar
v1  
liangjing committed
837
                       help='Position embedding type.')
Mostofa Patwary's avatar
Mostofa Patwary committed
838
    group.add_argument('--use-rotary-position-embeddings', action='store_true',
liangjing's avatar
v1  
liangjing committed
839
840
                       help='Use rotary positional embeddings or not. '
                       'Deprecated: use --position-embedding-type')
liangjing's avatar
liangjing committed
841
842
    group.add_argument('--rotary-base', type=int, default=10000,
                       help='Base to use for rotary positional embeddings, default 10000')
Mostofa Patwary's avatar
Mostofa Patwary committed
843
    group.add_argument('--rotary-percent', type=float, default=1.0,
liangjing's avatar
v1  
liangjing committed
844
                       help='Percent of rotary dimension to use, default 100%%')
liangjing's avatar
liangjing committed
845
846
    group.add_argument('--rotary-interleaved', action='store_true',
                          help='Use interleaved rotary embedding.')
liangjing's avatar
v1  
liangjing committed
847
848
    group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None,
                       help='Sequence length interpolation factor for rotary embeddings.')
Mostofa Patwary's avatar
Mostofa Patwary committed
849
850
    group.add_argument('--no-position-embedding',
                       action='store_false',
liangjing's avatar
v1  
liangjing committed
851
                       help='Disable position embedding. Deprecated: use --position-embedding-type',
Mostofa Patwary's avatar
Mostofa Patwary committed
852
                       dest='add_position_embedding')
Mohammad's avatar
Mohammad committed
853
854
855
    group.add_argument('--make-vocab-size-divisible-by', type=int, default=128,
                       help='Pad the vocab size to be divisible by this value.'
                       'This is added for computational efficieny reasons.')
liangjing's avatar
liangjing committed
856
857
858
859
860
    group.add_argument('--normalization', default='LayerNorm',
                       choices=['LayerNorm', 'RMSNorm'],
                       help='Which normalization technique to use.')
    group.add_argument('--norm-epsilon', type=float, default=1e-5,
                       help='Epsilon for layer norm and RMS norm.')
Mostofa Patwary's avatar
Mostofa Patwary committed
861
    group.add_argument('--apply-layernorm-1p', action='store_true',
862
863
                       help='Adjust LayerNorm weights such that they are centered '
                       'around zero. This improves numerical stability.')
Mohammad's avatar
Mohammad committed
864
865
866
867
    group.add_argument('--apply-residual-connection-post-layernorm',
                       action='store_true',
                       help='If set, use original BERT residula connection '
                       'ordering.')
868
869
870
871
    group.add_argument('--openai-gelu', action='store_true',
                       help='Use OpenAIs GeLU implementation. This option'
                       'should not be used unless for backward compatibility'
                       'reasons.')
872
873
874
875
    group.add_argument('--squared-relu', action='store_true',
                       help='Use squared relu activation instead of default gelu')
    group.add_argument('--swiglu', action='store_true',
                       help='Use gated linear units and SiLU activation instead of default gelu')
876
    group.add_argument('--onnx-safe', type=bool, required=False,
877
878
                       help='Use workarounds for known problems with '
                       'Torch ONNX exporter')
879
880
881
    group.add_argument('--bert-no-binary-head', action='store_false',
                       help='Disable BERT binary head.',
                       dest='bert_binary_head')
882
883
    group.add_argument('--untie-embeddings-and-output-weights', action='store_true',
                       help='Untie embeddings and output weights.'),
liangjing's avatar
liangjing committed
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
    return parser


def _add_straggler_detector_args(parser):
    group = parser.add_argument_group(title='straggler')
    group.add_argument('--log-straggler', action='store_true',
                       help='If set, tracks and logs straggler per GPU.')
    group.add_argument('--disable-straggler-on-startup', action='store_true',
                       help='If set, StragglerDetector is disabled on startup.')
    group.add_argument('--straggler-ctrlr-port', type=int, default=65535,
                       help='Port number to toggle StragglerDetector on/off at runtime')
    group.add_argument('--straggler-minmax-count', type=int, default=1,
                       help='Number of ranks to report with high/low estimated throughput')
    return parser


def _add_one_logger_args(parser):
    group = parser.add_argument_group(title='one logger')
    group.add_argument('--no-one-logger', action='store_false',
                       help='If set, disable using one_logger to track E2E metrics'
                       'Note that one_logger is an internal tool and not '
                       'available externally. For installation, please go to '
                       'https://confluence.nvidia.com/display/MLWFO/Package+Repositories'
                       'for more details',
                       dest='enable_one_logger')
    group.add_argument('--one-logger-project', type=str, default='megatron-lm',
                       help='The one-logger project name. Will ignore if '
                       '--no-one-logger is set')
    group.add_argument('--one-logger-run-name', type=str, default=None,
                       help='The one-logger run name displayed. Will ignore if '
                       '--no-one-logger is set')
    group.add_argument('--one-logger-async', action='store_true',
                       help='If set, forces one_logger to use async mode.')
    group.add_argument('--app-tag-run-name', type=str, default=None,
                       help='Jobs belonging to same training run, suppose to '
                       'have the same name. It will be used to track progress of '
                       'a training done over multiple different jobs')
    group.add_argument('--app-tag-run-version', type=str, default='0.0.0',
                       help='The version of the training of which current job is '
                       'part of. It will be used to track the changes in the '
                       'application side which might change the performance '
                       'baseline')
    return parser


def _add_ft_package_args(parser):
    group = parser.add_argument_group(title='ft_package')
    group.add_argument('--enable-ft-package', action='store_true',
                       help='If set, Fault Tolerance package is enabled. '
                       'Note: This feature is for Nvidia internal use only.')
    return parser


def _add_config_logger_args(parser):
    group = parser.add_argument_group(title='config logger')
    group.add_argument('--config-logger-dir', type=str, default='',
                       help='If set, will dump all configs to --config-logger-dir',
                       dest='config_logger_dir')
Mohammad's avatar
Mohammad committed
942
943
944
    return parser


945
946
947
948
949
def _add_logging_args(parser):
    group = parser.add_argument_group(title='logging')

    group.add_argument('--log-params-norm', action='store_true',
                       help='If set, calculate and log parameters norm.')
950
    group.add_argument('--log-num-zeros-in-grad', action='store_true',
Rewon Child's avatar
Rewon Child committed
951
                       help='If set, calculate and log the number of zeros in gradient.')
liangjing's avatar
liangjing committed
952
953
954
955
956
957
    group.add_argument('--log-throughput', action='store_true',
                       help='If set, calculate and log throughput per GPU.')
    group.add_argument('--log-progress', action='store_true',
                       help='If set, log progress (in terms of number of processed tokens and '
                       'number of floating-point operations) to progress.txt file in checkpoint '
                       'directory.')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
    group.add_argument('--timing-log-level', type=int,
                       default=0, choices=range(0,3),
                       help='Granularity level to measure and report timing. '
                       '   0: report only iteration time and make sure timing '
                       '      does not introduce extra overhead.'
                       '   1: report timing for operations that are executed '
                       '      very limited times (basically once) during '
                       '      each iteration (such as gradient all-reduce) '
                       '   2: report timing for operations that migh be '
                       '      executed numerous times during each iteration. '
                       'Note that setting the level to 1 or 2 might '
                       'cause increase in iteration time.')
    group.add_argument('--no-barrier-with-level-1-timing', action='store_false',
                       help='If not set, use barrier with level 1 time '
                       'measurements. Note that this is up to the user '
                       'to make sure calling barrier with their timers '
                       'will not result in hangs. This can happen if for '
                       'example the user adds a level 1 timer that is not '
                       'called by all ranks.',
                       dest='barrier_with_L1_time')
    group.add_argument('--timing-log-option', type=str, default='minmax',
                       choices=['max', 'minmax', 'all'],
                       help='Options for logging timing:'
                       '  max: report the max timing across all ranks'
                       '  minmax: report min and max timings across all ranks'
                       '  all: report timings of all ranks.')
984
985
    group.add_argument('--tensorboard-log-interval', type=int, default=1,
                       help='Report to tensorboard interval.')
986
987
988
989
    group.add_argument('--tensorboard-queue-size', type=int, default=1000,
                       help='Size of the tensorboard queue for pending events '
                       'and summaries before one of the ‘add’ calls forces a '
                       'flush to disk.')
990
991
992
993
994
995
996
997
998
999
    group.add_argument('--log-timers-to-tensorboard', action='store_true',
                       help='If set, write timers to tensorboard.')
    group.add_argument('--no-log-loss-scale-to-tensorboard',
                       action='store_false',
                       help='Disable loss-scale logging to tensorboard.',
                       dest='log_loss_scale_to_tensorboard')
    group.add_argument('--log-validation-ppl-to-tensorboard',
                       action='store_true',
                       help='If set, write validation perplexity to '
                       'tensorboard.')
1000
1001
    group.add_argument('--log-memory-to-tensorboard',
                       action='store_true',
1002
                       help='Enable memory logging to tensorboard.')
1003
1004
1005
    group.add_argument('--log-world-size-to-tensorboard',
                       action='store_true',
                       help='Enable world size logging to tensorboard.')
liangjing's avatar
liangjing committed
1006
1007
1008
1009
1010
1011
1012
1013
    group.add_argument('--wandb-project', type=str, default='',
                       help='The wandb project name. Ignore wandb by default.')
    group.add_argument('--wandb-exp-name', type=str, default='',
                       help='The wandb experiment name.')
    group.add_argument('--wandb-save-dir', type=str, default='',
                       help='Path to save the wandb results locally.')
    group.add_argument('--logging-level', type=int, default=None,
                       help='Set default logging level')
1014
1015
1016
    return parser


Mohammad's avatar
Mohammad committed
1017
def _add_regularization_args(parser):
Mohammad's avatar
Mohammad committed
1018
1019
1020
    group = parser.add_argument_group(title='regularization')

    group.add_argument('--attention-dropout', type=float, default=0.1,
1021
                       help='Post attention dropout probability.')
Mohammad's avatar
Mohammad committed
1022
1023
1024
1025
    group.add_argument('--hidden-dropout', type=float, default=0.1,
                       help='Dropout probability for hidden state transformer.')
    group.add_argument('--weight-decay', type=float, default=0.01,
                       help='Weight decay coefficient for L2 regularization.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1026
    group.add_argument('--start-weight-decay', type=float,
1027
                       help='Initial weight decay coefficient for L2 regularization.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1028
    group.add_argument('--end-weight-decay', type=float,
1029
                       help='End of run weight decay coefficient for L2 regularization.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1030
    group.add_argument('--weight-decay-incr-style', type=str, default='constant',
1031
1032
                       choices=['constant', 'linear', 'cosine'],
                       help='Weight decay increment function.')
Mohammad's avatar
Mohammad committed
1033
1034
    group.add_argument('--clip-grad', type=float, default=1.0,
                       help='Gradient clipping based on global L2 norm.')
1035
    group.add_argument('--adam-beta1', type=float, default=0.9,
1036
1037
                       help='First coefficient for computing running averages '
                       'of gradient and its square')
1038
    group.add_argument('--adam-beta2', type=float, default=0.999,
1039
1040
                       help='Second coefficient for computing running averages '
                       'of gradient and its square')
1041
    group.add_argument('--adam-eps', type=float, default=1e-08,
1042
                       help='Term added to the denominator to improve'
1043
                       'numerical stability')
1044
1045
    group.add_argument('--sgd-momentum', type=float, default=0.9,
                       help='Momentum factor for sgd')
Mohammad's avatar
Mohammad committed
1046
1047
    return parser

Mohammad's avatar
Mohammad committed
1048
1049

def _add_training_args(parser):
Mohammad's avatar
Mohammad committed
1050
1051
    group = parser.add_argument_group(title='training')

1052
    group.add_argument('--micro-batch-size', type=int, default=None,
Mohammad's avatar
Mohammad committed
1053
1054
                       help='Batch size per model instance (local batch size). '
                       'Global batch size is local batch size times data '
mohammad's avatar
mohammad committed
1055
                       'parallel size times number of micro batches.')
1056
1057
1058
    group.add_argument('--batch-size', type=int, default=None,
                       help='Old batch size parameter, do not use. '
                       'Use --micro-batch-size instead')
mohammad's avatar
mohammad committed
1059
    group.add_argument('--global-batch-size', type=int, default=None,
mohammad's avatar
mohammad committed
1060
1061
1062
                       help='Training batch size. If set, it should be a '
                       'multiple of micro-batch-size times data-parallel-size. '
                       'If this value is None, then '
mohammad's avatar
mohammad committed
1063
                       'use micro-batch-size * data-parallel-size as the '
mohammad's avatar
mohammad committed
1064
1065
                       'global batch size. This choice will result in 1 for '
                       'number of micro-batches.')
mohammad's avatar
mohammad committed
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
    group.add_argument('--rampup-batch-size', nargs='*', default=None,
                       help='Batch size ramp up with the following values:'
                       '  --rampup-batch-size <start batch size> '
                       '                      <batch size incerement> '
                       '                      <ramp-up samples> '
                       'For example:'
                       '   --rampup-batch-size 16 8 300000 \ '
                       '   --global-batch-size 1024'
                       'will start with global batch size 16 and over '
                       ' (1024 - 16) / 8 = 126 intervals will increase'
                       'the batch size linearly to 1024. In each interval'
                       'we will use approximately 300000 / 126 = 2380 samples.')
liangjing's avatar
liangjing committed
1078
1079
1080
1081
1082
1083
1084
    group.add_argument('--decrease-batch-size-if-needed', action='store_true', default=False,
                       help='If set, decrease batch size if microbatch_size * dp_size'
                       'does not divide batch_size. Useful for KSO (Keep Soldiering On)'
                       'to continue making progress if number of healthy GPUs (and'
                       'corresponding dp_size) does not support current batch_size.'
                       'Old batch_size will be restored if training is re-started with'
                       'dp_size that divides batch_size // microbatch_size.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1085
1086
    group.add_argument('--recompute-activations', action='store_true',
                       help='recompute activation to allow for training '
Mohammad's avatar
Mohammad committed
1087
                       'with larger models, sequences, and batch sizes.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1088
    group.add_argument('--recompute-granularity', type=str, default=None,
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1089
                       choices=['full', 'selective'],
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1090
                       help='Checkpoint activations to allow for training '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1091
1092
                       'with larger models, sequences, and batch sizes. '
                       'It is supported at two granularities 1) full: '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1093
                       'whole transformer layer is recomputed, '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1094
                       '2) selective: core attention part of the transformer '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1095
                       'layer is recomputed.')
liangjing's avatar
liangjing committed
1096
1097
1098
    group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false',
                       help='Check for NaNs in loss and grad',
                       dest='check_for_nan_in_loss_and_grad')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1099
    group.add_argument('--distribute-saved-activations',
1100
                       action='store_true',
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1101
                       help='If set, distribute recomputed activations '
1102
                       'across model parallel group.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1103
    group.add_argument('--recompute-method', type=str, default=None,
1104
1105
                       choices=['uniform', 'block'],
                       help='1) uniform: uniformly divide the total number of '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1106
                       'Transformer layers and recompute the input activation of '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1107
                       'each divided chunk at specified granularity, '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1108
                       '2) recompute the input activations of only a set number of '
slym's avatar
slym committed
1109
                       'individual Transformer layers per pipeline stage and do the '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1110
1111
                       'rest without any recomputing at specified granularity'
                       'default) do not apply activations recompute to any layers')
liangjing's avatar
v1  
liangjing committed
1112
    group.add_argument('--recompute-num-layers', type=int, default=None,
1113
                       help='1) uniform: the number of Transformer layers in each '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1114
                       'uniformly divided recompute unit, '
1115
                       '2) block: the number of individual Transformer layers '
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1116
                       'to recompute within each pipeline stage.')
liangjing's avatar
liangjing committed
1117
1118
1119
    group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false',
                       help='If not set, clone the output of the scatter in embedding layer to GC original tensor.',
                       dest='clone_scatter_output_in_embedding')
liangjing's avatar
v1  
liangjing committed
1120
1121
1122
1123
1124
1125
1126
1127
    group.add_argument('--profile', action='store_true',
                       help='Enable nsys profiling. When using this option, nsys '
                       'options should be specified in commandline. An example '
                       'nsys commandline is `nsys profile -s none -t nvtx,cuda '
                       '-o <path/to/output_file> --force-overwrite true '
                       '--capture-range=cudaProfilerApi '
                       '--capture-range-end=stop`.')
    group.add_argument('--profile-step-start', type=int, default=10,
liangjing's avatar
liangjing committed
1128
                       help='Global step to start profiling.')
liangjing's avatar
v1  
liangjing committed
1129
    group.add_argument('--profile-step-end', type=int, default=12,
liangjing's avatar
liangjing committed
1130
1131
1132
1133
1134
                       help='Global step to stop profiling.')
    group.add_argument('--use-pytorch-profiler', action='store_true',
                       help='Use the built-in pytorch profiler. '
                       'Useful if you wish to view profiles in tensorboard.',
                       dest='use_pytorch_profiler')
liangjing's avatar
v1  
liangjing committed
1135
1136
    group.add_argument('--profile-ranks', nargs='+', type=int, default=[0],
                       help='Global ranks to profile.')
liangjing's avatar
liangjing committed
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
    group.add_argument('--tp-comm-overlap', action='store_true', help='Enables the '
                       ' overlap of Tensor parallel communication and GEMM kernels.')
    group.add_argument('--tp-comm-overlap-cfg', type=str, default=None,
                       help='Config file when tp_comm_overlap is enabled.')
    group.add_argument('--disable-tp-comm-overlap-ag', action='store_false',
                       help=('Disables the All-Gather overlap with GEMM by '
                             'pipelining the GEMM and All-Gather.'),
                       dest='tp_comm_overlap_ag')
    group.add_argument('--disable-tp-comm-overlap-rs', action='store_false',
                       help=('Disables the Reduce-Scatter overlap with GEMM by '
                             'pipelining the GEMM and Reduce-Scatter.'),
                       dest='tp_comm_overlap_rs')
    group.add_argument('--tp-comm-overlap-rs-dgrad', action='store_true',
                       help = 'Enables the Reduce-Scatter overlap with dgrad GEMM.',
                       dest='tp_comm_overlap_rs_dgrad')
    group.add_argument('--disable-tp-comm-bulk-dgrad', action='store_false',
                       help='Disables the All-Gather overlap with bprop activation gradient GEMM.',
                       dest='tp_comm_bulk_dgrad')
    group.add_argument('--disable-tp-comm-bulk-wgrad', action='store_false',
                       help='Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.',
                       dest='tp_comm_bulk_wgrad')
    group.add_argument('--use-cpu-initialization', action='store_true',
                       default=None,
                       help='If set, initialize weights on the CPU. This eliminates init differences based on tensor parallelism.')
    group.add_argument('--empty-unused-memory-level', default=0, type=int,
                       choices=[0, 1, 2],
                       help='Call torch.cuda.empty_cache() each iteration '
                       '(training and eval), to reduce fragmentation.'
                       '0=off, 1=moderate, 2=aggressive.')
    group.add_argument('--deterministic-mode', action='store_true',
                       help='Choose code that has deterministic execution. This usually '
                       'means slower execution, but is good for debugging and testing.')
    group.add_argument('--check-weight-hash-across-dp-replicas-interval', type=int, default=None,
                       help='Interval to check weight hashes are same across DP replicas. If not specified, weight hashes not checked.')
    group.add_argument('--calculate-per-token-loss', action='store_true',
                       help=('Scale cross entropy loss by the number of non-padded tokens in the '
                             'global batch, versus the default behavior of assuming all tokens are non-padded.'))
    group.add_argument('--train-sync-interval', type=int, default=None,
                       help='Training CPU-GPU synchronization interval, to ensure that CPU is not running too far ahead of GPU.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1176
1177
1178
1179
1180

    # deprecated
    group.add_argument('--checkpoint-activations', action='store_true',
                       help='Checkpoint activation to allow for training '
                       'with larger models, sequences, and batch sizes.')
Mohammad's avatar
Mohammad committed
1181
    group.add_argument('--train-iters', type=int, default=None,
Mohammad's avatar
Mohammad committed
1182
                       help='Total number of iterations to train over all '
1183
1184
1185
1186
1187
1188
                       'training runs. Note that either train-iters or '
                       'train-samples should be provided.')
    group.add_argument('--train-samples', type=int, default=None,
                       help='Total number of samples to train over all '
                       'training runs. Note that either train-iters or '
                       'train-samples should be provided.')
Mohammad's avatar
Mohammad committed
1189
1190
1191
1192
1193
    group.add_argument('--log-interval', type=int, default=100,
                       help='Report loss and timing interval.')
    group.add_argument('--exit-interval', type=int, default=None,
                       help='Exit the program after the iteration is divisible '
                       'by this value.')
1194
1195
    group.add_argument('--exit-duration-in-mins', type=int, default=None,
                       help='Exit the program after this many minutes.')
1196
1197
1198
    group.add_argument('--exit-signal-handler', action='store_true',
                       help='Dynamically save the checkpoint and shutdown the '
                       'training if SIGTERM is received')
Mohammad's avatar
Mohammad committed
1199
1200
    group.add_argument('--tensorboard-dir', type=str, default=None,
                       help='Write TensorBoard logs to this directory.')
1201
    group.add_argument('--no-masked-softmax-fusion',
1202
1203
1204
                       action='store_false',
                       help='Disable fusion of query_key_value scaling, '
                       'masking, and softmax.',
1205
                       dest='masked_softmax_fusion')
1206
1207
1208
    group.add_argument('--no-bias-gelu-fusion', action='store_false',
                       help='Disable bias and gelu fusion.',
                       dest='bias_gelu_fusion')
liangjing's avatar
liangjing committed
1209
1210
1211
1212
    group.add_argument('--no-bias-swiglu-fusion', action='store_false',
                       help='Disable bias and swiglu fusion, the fusion is '
                       'available only when using megatron-core.',
                       dest='bias_swiglu_fusion')
1213
1214
1215
    group.add_argument('--no-bias-dropout-fusion', action='store_false',
                       help='Disable bias and dropout fusion.',
                       dest='bias_dropout_fusion')
liangjing's avatar
liangjing committed
1216
1217
1218
1219
1220
1221
1222
1223
    group.add_argument('--no-rope-fusion', action='store_false',
                       help='Disable rope fusion, the fusion is available '
                       'only when using megatron-core.',
                       dest='apply_rope_fusion')
    group.add_argument('--cross-entropy-loss-fusion', action='store_true',
                       help='Enabled fusion of cross entropy loss calculation.',
                       dest='cross_entropy_loss_fusion')
    group.add_argument('--use-flash-attn-ck', action='store_true',
1224
1225
                       help='use FlashAttention implementation of attention. '
                       'https://arxiv.org/abs/2205.14135')
liangjing's avatar
liangjing committed
1226
1227
    group.add_argument('--use-flash-attn-triton', action='store_true',
                       help='use FlashAttention implementation of attention using Triton.')
1228
1229
1230
    group.add_argument('--disable-bias-linear', action='store_false',
                       help='Disable bias in the linear layers',
                       dest='add_bias_linear')
liangjing's avatar
liangjing committed
1231
1232
1233
    group.add_argument('--add-qkv-bias', action='store_true',
                       help='Enable bias only in the QKV linear layers',
                       dest='add_qkv_bias')
1234
1235
1236
    group.add_argument('--optimizer', type=str, default='adam',
                       choices=['adam', 'sgd'],
                       help='Optimizer function')
1237
    group.add_argument('--dataloader-type', type=str, default=None,
liangjing's avatar
liangjing committed
1238
                       choices=['single', 'cyclic', 'external'],
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1239
                       help='Single pass vs multiple pass data loader')
slym's avatar
slym committed
1240
    group.add_argument('--no-async-tensor-model-parallel-allreduce',
Sangkug Lym's avatar
Sangkug Lym committed
1241
                       action='store_false',
liangjing's avatar
liangjing committed
1242
                       help='DEPRECATED. This flag is ignored.',
Sangkug Lym's avatar
Sangkug Lym committed
1243
                       dest='async_tensor_model_parallel_allreduce')
Sangkug Lym's avatar
Sangkug Lym committed
1244
1245
1246
1247
1248
    group.add_argument('--no-persist-layer-norm', action='store_true',
                       help='Disable using persistent fused layer norm kernel. '
                       'This kernel supports only a set of hidden sizes. Please '
                       'check persist_ln_hidden_sizes if your hidden '
                       'size is supported.')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1249
    group.add_argument('--sequence-parallel', action='store_true',
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1250
                       help='Enable sequence parallel optimization.')
Sangkug Lym's avatar
Sangkug Lym committed
1251
1252
    group.add_argument('--no-gradient-accumulation-fusion',
                       action='store_false',
1253
                       help='Disable fusing gradient accumulation to weight '
Sangkug Lym's avatar
Sangkug Lym committed
1254
1255
                       'gradient computation of linear layers',
                       dest='gradient_accumulation_fusion')
liangjing's avatar
liangjing committed
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
    group.add_argument('--use-mcore-models', action='store_true',
                       dest='deprecated_use_mcore_models',
                       help='DEPRECATED. Use the implementation from megatron core.'
                       'Now ignored and mcore models are the default, use '
                       '--use-legacy-models to not use core models.')
    group.add_argument('--use-legacy-models', action='store_true',
                       help='Use the legacy Megatron models, not Megatron-Core models.')
    group.add_argument('--manual-gc', action='store_true',
                       help='Disable the threshold-based default garbage '
                       'collector and trigger the garbage collection manually. '
                       'Manual garbage collection helps to align the timing of '
                       'the collection across ranks which mitigates the impact '
                       'of CPU-associated jitters. When the manual gc is enabled, '
                       'garbage collection is performed only at the start and the '
                       'end of the validation routine by default.')
    group.add_argument('--manual-gc-interval', type=int, default=0,
                       help='Training step interval to trigger manual garbage '
                       'collection. When the value is set to 0, garbage '
                       'collection is not triggered between training steps.')
    group.add_argument('--no-manual-gc-eval', action='store_false',
                       help='When using manual garbage collection, disable '
                       'garbage collection at the start and the end of each '
                       'evaluation run.', dest='manual_gc_eval')
    group.add_argument('--disable-tp-comm-split-ag', action='store_false',
                       help='Disables the All-Gather overlap with fprop GEMM.',
                       dest='tp_comm_split_ag')
    group.add_argument('--disable-tp-comm-split-rs', action='store_false',
                       help='Disables the Reduce-Scatter overlap with fprop GEMM.',
                       dest='tp_comm_split_rs')

Mohammad's avatar
Mohammad committed
1286
1287
1288
    return parser


Mohammad's avatar
Mohammad committed
1289
def _add_initialization_args(parser):
Mohammad's avatar
Mohammad committed
1290
1291
1292
1293
1294
    group = parser.add_argument_group(title='initialization')

    group.add_argument('--seed', type=int, default=1234,
                       help='Random seed used for python, numpy, '
                       'pytorch, and cuda.')
1295
1296
1297
    group.add_argument('--data-parallel-random-init', action='store_true',
                       help='Enable random initialization of params '
                       'across data parallel ranks')
Mohammad's avatar
Mohammad committed
1298
1299
1300
    group.add_argument('--init-method-std', type=float, default=0.02,
                       help='Standard deviation of the zero mean normal '
                       'distribution used for weight initialization.')
1301
1302
    group.add_argument('--init-method-xavier-uniform', action='store_true',
                       help='Enable Xavier uniform parameter initialization')
Mohammad's avatar
Mohammad committed
1303

Mohammad's avatar
Mohammad committed
1304
1305
1306
    return parser


Mohammad's avatar
Mohammad committed
1307
def _add_learning_rate_args(parser):
Mohammad's avatar
Mohammad committed
1308
1309
    group = parser.add_argument_group(title='learning rate')

Mohammad's avatar
Mohammad committed
1310
    group.add_argument('--lr', type=float, default=None,
Mohammad's avatar
Mohammad committed
1311
                       help='Initial learning rate. Depending on decay style '
liangjing's avatar
liangjing committed
1312
                       'and initial warmup, the learning rate at each '
Mohammad's avatar
Mohammad committed
1313
1314
                       'iteration would be different.')
    group.add_argument('--lr-decay-style', type=str, default='linear',
liangjing's avatar
liangjing committed
1315
                       choices=['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD'],
Mohammad's avatar
Mohammad committed
1316
                       help='Learning rate decay function.')
liangjing's avatar
liangjing committed
1317
1318
1319
    group.add_argument('--lr-wsd-decay-style', type=str, default='exponential',
                       choices=['exponential', 'linear', 'cosine'],
                       help='Decay style for the annealing phase of WSD'),
Mohammad's avatar
Mohammad committed
1320
1321
1322
    group.add_argument('--lr-decay-iters', type=int, default=None,
                       help='number of iterations to decay learning rate over,'
                       ' If None defaults to `--train-iters`')
1323
1324
1325
    group.add_argument('--lr-decay-samples', type=int, default=None,
                       help='number of samples to decay learning rate over,'
                       ' If None defaults to `--train-samples`')
liangjing's avatar
liangjing committed
1326
1327
1328
1329
    group.add_argument('--lr-wsd-decay-samples', type=int, default=None,
                       help='number of samples for the annealing phase in the wsd schedule')
    group.add_argument('--lr-wsd-decay-iters', type=int, default=None,
                       help='number of iterations for the annealing phase in the wsd schedule')
1330
1331
1332
    group.add_argument('--lr-warmup-fraction', type=float, default=None,
                       help='fraction of lr-warmup-(iters/samples) to use '
                       'for warmup (as a float)')
1333
1334
1335
1336
1337
1338
    group.add_argument('--lr-warmup-iters', type=int, default=0,
                       help='number of iterations to linearly warmup '
                       'learning rate over.')
    group.add_argument('--lr-warmup-samples', type=int, default=0,
                       help='number of samples to linearly warmup '
                       'learning rate over.')
liangjing's avatar
v1  
liangjing committed
1339
1340
1341
    group.add_argument('--lr-warmup-init', type=float, default=0.0,
                       help='Initial value for learning rate warmup. The '
                       'scheduler starts warmup from this value.')
1342
    group.add_argument('--warmup', type=int, default=None,
1343
                       help='Old lr warmup argument, do not use. Use one of the'
1344
                       '--lr-warmup-* arguments above')
Mohammad's avatar
Mohammad committed
1345
    group.add_argument('--min-lr', type=float, default=0.0,
liangjing's avatar
liangjing committed
1346
                       help='Minimum value for learning rate. The scheduler'
Mohammad's avatar
Mohammad committed
1347
                       'clip values below this threshold.')
1348
    group.add_argument('--override-opt_param-scheduler', action='store_true',
Mohammad's avatar
Mohammad committed
1349
1350
1351
1352
1353
                       help='Reset the values of the scheduler (learning rate,'
                       'warmup iterations, minimum learning rate, maximum '
                       'number of iterations, and decay style from input '
                       'arguments and ignore values from checkpoints. Note'
                       'that all the above values will be reset.')
1354
    group.add_argument('--use-checkpoint-opt_param-scheduler', action='store_true',
Mohammad's avatar
Mohammad committed
1355
1356
1357
1358
                       help='Use checkpoint to set the values of the scheduler '
                       '(learning rate, warmup iterations, minimum learning '
                       'rate, maximum number of iterations, and decay style '
                       'from checkpoint and ignore input arguments.')
liangjing's avatar
liangjing committed
1359
1360
1361
1362
1363
    group.add_argument('--decoupled-lr', type=float, default=None,
                       help='Separate learning rate for the input and output layer')
    group.add_argument('--decoupled-min-lr', type=float, default=None,
                       help='Minimum value for learning rate for the input and output layer. The scheduler'
                       'clip values below this threshold')
Mohammad's avatar
Mohammad committed
1364
1365
1366
1367

    return parser


Mohammad's avatar
Mohammad committed
1368
def _add_checkpointing_args(parser):
Mohammad's avatar
Mohammad committed
1369
1370
1371
1372
    group = parser.add_argument_group(title='checkpointing')

    group.add_argument('--save', type=str, default=None,
                       help='Output directory to save checkpoints to.')
liangjing's avatar
liangjing committed
1373
1374
    group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
                       help='Number of iterations between persistent checkpoint saves.')
1375
    group.add_argument('--no-save-optim', action='store_true', default=None,
Mohammad's avatar
Mohammad committed
1376
                       help='Do not save current optimizer.')
1377
    group.add_argument('--no-save-rng', action='store_true', default=None,
Mohammad's avatar
Mohammad committed
1378
1379
1380
                       help='Do not save current rng state.')
    group.add_argument('--load', type=str, default=None,
                       help='Directory containing a model checkpoint.')
Jared Casper's avatar
Jared Casper committed
1381
    group.add_argument('--no-load-optim', action='store_true', default=None,
Mohammad's avatar
Mohammad committed
1382
                       help='Do not load optimizer when loading checkpoint.')
Jared Casper's avatar
Jared Casper committed
1383
    group.add_argument('--no-load-rng', action='store_true', default=None,
Mohammad's avatar
Mohammad committed
1384
                       help='Do not load rng state when loading checkpoint.')
liangjing's avatar
liangjing committed
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
    group.add_argument('--non-persistent-save-interval', type=int, default=None,
                       help='Number of iterations between non-persistent saves.')
    group.add_argument('--non-persistent-ckpt-type', type=str, default=None,
                       choices=['global', 'local', 'in_memory', None],
                       help='Type of non-persistent model checkpoints. '
                           '"global" - Saved as a standard checkpoint (e.g., on Lustre) with old checkpoints being removed. '
                           '"local" - [TBD] Each rank saves a portion of the checkpoint locally (e.g., on SSD/ramdisk). '
                           '"in_memory" - [TBD] A special kind of local checkpoint that avoids serialization. '
                           'None - No non-persistent checkpointing (default option).')
    group.add_argument('--non-persistent-global-ckpt-dir', type=str, default=None,
                       help='Directory containing global non-persistent model checkpoints.')
    group.add_argument('--non-persistent-local-ckpt-dir', type=str, default=None,
                       help='Directory containing local non-persistent model checkpoints.')
    group.add_argument('--non-persistent-local-ckpt-algo', type=str, default='fully_parallel',
                       choices=['fully_parallel', 'atomic'],
                       help='Algorithm for local non-persistent checkpointing.')
Mohammad's avatar
Mohammad committed
1401
1402
1403
1404
    group.add_argument('--finetune', action='store_true',
                       help='Load model for finetuning. Do not load optimizer '
                       'or rng state from checkpoint and set iteration to 0. '
                       'Assumed when loading a release checkpoint.')
liangjing's avatar
liangjing committed
1405
1406
1407
1408
    group.add_argument('--pretrained-checkpoint', type=str, default=None,
                       help='Directory containing a pretrained model checkpoint for finetuning.')
    group.add_argument('--ckpt-step', type=int, default=None,
                       help='Checkpoint step to load model from.')
1409
1410
1411
1412
1413
    group.add_argument('--no-initialization', action='store_false',
                       help='Do not perform initialization when building model, '
                       'can reduce startup time when definitely loading from a '
                       'checkpoint',
                       dest='perform_initialization')
1414
1415
1416
    group.add_argument('--use-checkpoint-args', action='store_true',
                       help='Override any command line arguments with arguments '
                       'from the checkpoint')
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
1417
1418
1419
1420
    group.add_argument('--exit-on-missing-checkpoint', action='store_true',
                       help="If '--load' is set, but checkpoint is not found "
                       "(e.g., path typo), then exit instead of random "
                       "initialization.")
liangjing's avatar
liangjing committed
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
    group.add_argument('--use-dist-ckpt', action='store_true',
                       dest='use_dist_ckpt_deprecated',
                       help='Deprecated: see --ckpt-format.')
    group.add_argument('--auto-detect-ckpt-format', action='store_true',
                       help='Determine if the checkpoint format is in legacy or distributed format.'
                            ' If False, expects distributed checkpoint iff args.ckpt_format != "torch".'
                            ' Might slow down loading a bit (double rank0 ckpt load).')
    group.add_argument('--dist-ckpt-format',
                       dest='dist_ckpt_format_deprecated',
                       help='Deprecated: see --ckpt-format.')
    group.add_argument('--ckpt-format', default='torch_dist',
                       choices=['torch', 'torch_dist', 'zarr'],
                       help='Checkpoint format to use.')
    group.add_argument('--ckpt-convert-format', default=None,
                       choices=['torch', 'torch_dist', 'zarr'],
                       help='Checkpoint format for conversion.')
    group.add_argument('--ckpt-convert-save', default=None,
                       help='Save directory for converted checkpoint.')
    group.add_argument('--ckpt-convert-update-legacy-dist-opt-format', action='store_true',
                       help='When loading a checkpoint, update the legacy format '
                       'for the distributed optimizer, which previously used a '
                       'merged param/grad buffer and a different bucket mapping. '
                       'The legacy format was deprecated on Feb 13, 2024.')
    group.add_argument('--ckpt-fully-parallel-save', action='store_true',
                       dest='ckpt_fully_parallel_save_deprecated',
                       help='Deprecated: see --no-ckpt-fully-parallel-save.')
    group.add_argument('--no-ckpt-fully-parallel-save', action='store_false',
                       dest='ckpt_fully_parallel_save',
                       help='Disable applying full save parallelization across DP for'
                            ' distributed checkpoints. Depending on ckpt format'
                            ' might decrease the number of files in the checkpoint.'
                            ' Makes DistributedOptimizer checkpoint non-reshardable.')
    group.add_argument('--async-save', action='store_true', default=None,
                       help='Apply async checkpointing save. Currently works only with'
                            '`torch_dist` distributed checkpoint format.')
    group.add_argument('--ckpt-fully-parallel-load', action='store_true',
                       help='Apply full load parallelization across DP for'
                            ' distributed checkpoints.')
    group.add_argument('--ckpt-assume-constant-structure', action='store_true',
                       help='If the model and optimizer state dict structure is'
                            'constant throughout a *single training job*, it allows for'
                            'different checkpointing performance optimizations.')
    group.add_argument('--dist-ckpt-strictness', type=str, default='assume_ok_unexpected',
                       choices=[e.value for e in StrictHandling],
                       help='Determine handling of key mismatch during checkpoint load.'
                            ' Check StrictHandling docs for flags meaning.'
                            ' NOTE: This flag controls only distributed checkpoint'
                            ' load from storage, not loading state dict into the model.')
Mohammad's avatar
Mohammad committed
1469
1470
1471
    return parser


Mohammad's avatar
Mohammad committed
1472
def _add_mixed_precision_args(parser):
Mohammad's avatar
Mohammad committed
1473
1474
1475
1476
    group = parser.add_argument_group(title='mixed precision')

    group.add_argument('--fp16', action='store_true',
                       help='Run model in fp16 mode.')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1477
1478
    group.add_argument('--bf16', action='store_true',
                       help='Run model in bfloat16 mode.')
mohammad's avatar
mohammad committed
1479
1480
1481
1482
1483
1484
1485
    group.add_argument('--loss-scale', type=float, default=None,
                       help='Static loss scaling, positive power of 2 '
                       'values can improve fp16 convergence. If None, dynamic'
                       'loss scaling is used.')
    group.add_argument('--initial-loss-scale', type=float, default=2**32,
                       help='Initial loss-scale for dynamic loss scaling.')
    group.add_argument('--min-loss-scale', type=float, default=1.0,
liangjing's avatar
liangjing committed
1486
                       help='Minimum loss scale for dynamic loss scaling.')
mohammad's avatar
mohammad committed
1487
1488
1489
1490
    group.add_argument('--loss-scale-window', type=float, default=1000,
                       help='Window over which to raise/lower dynamic scale.')
    group.add_argument('--hysteresis', type=int, default=2,
                       help='hysteresis for dynamic loss scaling')
1491
1492
    group.add_argument('--fp32-residual-connection', action='store_true',
                       help='Move residual connections to fp32.')
liangjing's avatar
liangjing committed
1493
1494
1495
    group.add_argument('--apply-query-key-layer-scaling', action='store_true',
                       help='Scale Q * K^T by 1 / layer-number. '
                       'Useful for fp16 training. Also sets `attention_softmax_in_fp32` to True.')
Mohammad's avatar
Mohammad committed
1496
    group.add_argument('--attention-softmax-in-fp32', action='store_true',
liangjing's avatar
liangjing committed
1497
                       help='Run attention masking and softmax in fp32.')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1498
1499
1500
    group.add_argument('--accumulate-allreduce-grads-in-fp32',
                       action='store_true',
                       help='Gradient accumulation and all-reduce in fp32.')
1501
1502
1503
1504
    group.add_argument('--fp16-lm-cross-entropy', action='store_true',
                       help='Move the cross entropy unreduced loss calculation'
                       'for lm head to fp16.')

Mohammad's avatar
Mohammad committed
1505
1506
1507
    return parser


Mohammad's avatar
Mohammad committed
1508
def _add_distributed_args(parser):
1509
1510
    group = parser.add_argument_group(title='distributed')

1511
1512
    group.add_argument('--tensor-model-parallel-size', type=int, default=1,
                       help='Degree of tensor model parallelism.')
liangjing's avatar
liangjing committed
1513
1514
    group.add_argument('--encoder-tensor-model-parallel-size', type=int, default=0,
                       help='Degree of tensor model parallelism for the encoder.')
1515
1516
    group.add_argument('--pipeline-model-parallel-size', type=int, default=1,
                       help='Degree of pipeline model parallelism.')
liangjing's avatar
liangjing committed
1517
1518
1519
    group.add_argument('--encoder-pipeline-model-parallel-size', type=int, default=0,
                       help=('Degree of pipeline model parallelism in the encoder. This is '
                             'independent of the amount of pipeline in the decoder.'))
1520
1521
    group.add_argument('--pipeline-model-parallel-split-rank',
                       type=int, default=None,
liangjing's avatar
liangjing committed
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
                       help=('Rank where encoder and decoder should be split. '
                             'Deprecated; use --encoder-pipeline-model-parallel-size instead.'))
    group.add_argument('--decoder-first-pipeline-num-layers',
                       type=int, default=None,
                       help=('The number of transformer layers on the first pipeline stage of the decoder. '
                       'Default None is even split of transformer layers across all pipeline stages'))
    group.add_argument('--decoder-last-pipeline-num-layers',
                       type=int, default=None,
                       help=('The number of transformer layers on the last pipeline stage of the decoder. '
                       'Default None is even split of transformer layers across all pipeline stages'))
1532
1533
1534
    group.add_argument('--model-parallel-size', type=int, default=None,
                       help='Old model parallel argument, do not use. Use '
                       '--tensor-model-parallel-size instead.')
1535
1536
    group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,
                       help='Number of layers per virtual pipeline stage')
liangjing's avatar
liangjing committed
1537
    group.add_argument('--no-overlap-p2p-communication', action='store_false',
liangjing's avatar
v1  
liangjing committed
1538
1539
                       help='overlap pipeline parallel communication with forward and backward chunks',
                       dest='overlap_p2p_comm')
Mohammad's avatar
Mohammad committed
1540
1541
1542
    group.add_argument('--distributed-backend', default='nccl',
                       choices=['nccl', 'gloo'],
                       help='Which backend to use for distributed training.')
Lawrence McAfee's avatar
Retro  
Lawrence McAfee committed
1543
1544
    group.add_argument('--distributed-timeout-minutes', type=int, default=10,
                       help='Timeout minutes for torch.distributed.')
liangjing's avatar
liangjing committed
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
    group.add_argument('--overlap-grad-reduce', action='store_true',
                       default=False, help='If set, overlap DDP grad reduce.')
    group.add_argument('--defer-embedding-wgrad-compute', action='store_true',
                       default=False, help='If set, defers the vocabulary projection linear layer weight'
                       'gradient compute to pipeline flush.', dest='defer_embedding_wgrad_compute')
    group.add_argument('--wgrad-deferral-limit', type=int, default=0, help='Number of micro-batches for which'
                       'weight gradient computation of vocabulary projection is deferred, defaults to 0 which'
                       'means all the micro-batches are deferred. Invalid if `defer-embedding-wgrad-compute`'
                       'is not set')
    group.add_argument('--no-align-grad-reduce', action='store_false',
                       help='If not set, all PP stages will launch gradient reduces simultaneously. '
                       'Otherwise, each PP stage will independently launch as needed.',
                       dest='align_grad_reduce')
    group.add_argument('--ddp-bucket-size', type=int, default=None,
                       help='Bucket size for data-parallel communication')
    group.add_argument('--ddp-average-in-collective', action='store_true',
                       default=False, help='If set, average directly in data-parallel communication collective.')
    group.add_argument('--overlap-param-gather', action='store_true',
                       default=False, help='If set, overlap param all-gather in distributed optimizer.')
    group.add_argument('--overlap-param-gather-with-optimizer-step', action='store_true',
                       default=False, help='If set, overlap param all-gather of first bucket with optimizer step.')
    group.add_argument('--no-align-param-gather', action='store_false',
                       help='If not set, all PP stages will launch param all-gathers simultaneously. '
                       'Otherwise, each PP stage will independently launch as needed.',
                       dest='align_param_gather')
1570
    group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false',
liangjing's avatar
liangjing committed
1571
                       help='If not set, use scatter/gather to optimize communication of tensors in pipeline.',
1572
                       dest='scatter_gather_tensors_in_pipeline')
1573
1574
1575
1576
    group.add_argument('--use-ring-exchange-p2p', action='store_true',
                       default=False, help='If set, use custom-built ring exchange '
                       'for p2p communications. Note that this option will require '
                       'a custom built image that support ring-exchange p2p.')
liangjing's avatar
liangjing committed
1577
1578
#    group.add_argument('--local-rank', type=int, default=int(os.getenv('LOCAL_RANK', '0')),
#                       help='local rank passed from distributed launcher.')
Mohammad's avatar
Mohammad committed
1579
1580
    group.add_argument('--local_rank', type=int, default=None,
                       help='local rank passed from distributed launcher.')
1581
    group.add_argument('--lazy-mpu-init', type=bool, required=False,
1582
1583
1584
1585
1586
                       help='If set to True, initialize_megatron() '
                       'skips DDP initialization and returns function to '
                       'complete it instead.Also turns on '
                       '--use-cpu-initialization flag. This is for '
                       'external DDP manager.' )
1587
    group.add_argument('--standalone-embedding-stage', action='store_true',
Lawrence McAfee's avatar
Lawrence McAfee committed
1588
1589
                       default=False, help='If set, *input* embedding layer '
                       'is placed on its own pipeline stage, without any '
Lawrence McAfee's avatar
Lawrence McAfee committed
1590
1591
                       'transformer layers. (For T5, this flag currently only '
                       'affects the encoder embedding.)')
1592
1593
    group.add_argument('--use-distributed-optimizer', action='store_true',
                       help='Use distributed optimizer.')
liangjing's avatar
liangjing committed
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
    group.add_argument('--context-parallel-size', type=int, default=1,
                       help='Degree of context parallelism.')
    group.add_argument('--nccl-communicator-config-path', type=str, default=None,
                       help='Path to the yaml file with NCCL communicator '
                       'configurations. The number of min/max thread groups and thread '
                       'group cluster size of each communicator can be configured by '
                       'setting `min_ctas`, `max_ctas`, and `cga_cluster_size`.')
    group.add_argument('--use-tp-pp-dp-mapping', action='store_true', default=False,
                        help='If set, distributed ranks initialize order is changed '
                        'from tp-dp-pp to tp-pp-dp. Make sure EP and CP aren\'t used '
                        'with this option enabled')
liangjing's avatar
v1  
liangjing committed
1605
1606
    group.add_argument('--rank', default=-1, type=int,
                       help='node rank for distributed training')
liangjing's avatar
liangjing committed
1607
    group.add_argument('--world_size', type=int, default=8,
liangjing's avatar
v1  
liangjing committed
1608
1609
1610
                       help='number of nodes for distributed training')
    group.add_argument('--dist_url',
                       help='Which master node url for distributed training.')
Mohammad's avatar
Mohammad committed
1611
1612
1613
    return parser


Mohammad's avatar
Mohammad committed
1614
def _add_validation_args(parser):
Mohammad's avatar
Mohammad committed
1615
1616
1617
1618
1619
1620
1621
1622
    group = parser.add_argument_group(title='validation')

    group.add_argument('--eval-iters', type=int, default=100,
                       help='Number of iterations to run for evaluation'
                       'validation/test for.')
    group.add_argument('--eval-interval', type=int, default=1000,
                       help='Interval between running evaluation on '
                       'validation set.')
liangjing's avatar
liangjing committed
1623
    group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.')
liangjing's avatar
v1  
liangjing committed
1624
1625
1626
    group.add_argument('--skip-train', action='store_true',
                       default=False, help='If set, bypass the training loop, '
                       'optionally do evaluation for validation/test, and exit.')
Mohammad's avatar
Mohammad committed
1627

Mohammad's avatar
Mohammad committed
1628
1629
1630
    return parser


Mohammad's avatar
Mohammad committed
1631
def _add_data_args(parser):
Mohammad's avatar
Mohammad committed
1632
1633
    group = parser.add_argument_group(title='data and dataloader')

mohammad's avatar
mohammad committed
1634
    group.add_argument('--data-path', nargs='*', default=None,
liangjing's avatar
liangjing committed
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
                       help='The weight and prefix list for a set of train, validation, and test'
                       'datasets which split according to --split. The accepted formats are: '
                       '(1) a single prefix, '
                       '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, '
                       '(3) a list of prefixes e.g. prefix1 prefix2. '
                       'For (3), weights are inferred from the lengths of the contributing datasets. '
                       'This argument is exclusive to the other independent --*-data-path arguments.')
    group.add_argument('--renormalize-blend-weights', action='store_true',
                       help='Renormalize the blend weights to account for the mid-level dataset '
                       'oversampling done to ensure fulfillment of the requested number of '
                       'samples. Use this option if prompted. Defaults to False for backward '
                       'comparability in the data sample order.')
    group.add_argument('--split', type=str, default=None,
Mohammad's avatar
Mohammad committed
1648
1649
                       help='Comma-separated list of proportions for training,'
                       ' validation, and test split. For example the split '
1650
1651
                       '`90,5,5` will use 90%% of data for training, 5%% for '
                       'validation and 5%% for test.')
1652
    group.add_argument('--train-data-path', nargs='*', default=None,
liangjing's avatar
liangjing committed
1653
1654
                       help='The weight and prefix list for an independent train dataset. '
                       'Follows the same pattern rules as --data-path.')
1655
    group.add_argument('--valid-data-path', nargs='*', default=None,
liangjing's avatar
liangjing committed
1656
1657
                       help='The weight and prefix list for an independent validation dataset. '
                       'Follows the same pattern rules as --data-path.')
1658
    group.add_argument('--test-data-path', nargs='*', default=None,
liangjing's avatar
liangjing committed
1659
1660
                       help='The weight and prefix list for an independent test dataset. '
                       'Follows the same pattern rules as --data-path.')
liangjing's avatar
v1  
liangjing committed
1661
1662
    group.add_argument('--data-cache-path', default=None,
                       help='Path to a directory to hold cached index files.')
liangjing's avatar
liangjing committed
1663
1664
1665
1666
1667
1668
    group.add_argument('--no-mmap-bin-files', action='store_false',
                       help='Disable mmap-ing of .bin files.',
                       dest='mmap_bin_files')
    group.add_argument('--mock-data', action='store_true',
                       help='Skip data loading and validation and opt for artificial '
                       'generation of mock data when an implementation is available.')
liangjing's avatar
v1  
liangjing committed
1669
1670
    group.add_argument('--vocab-size', type=int, default=None,
                       help='Size of vocab before EOD or padding.')
Mohammad's avatar
Mohammad committed
1671
    group.add_argument('--vocab-file', type=str, default=None,
Mohammad's avatar
Mohammad committed
1672
                       help='Path to the vocab file.')
Mohammad's avatar
Mohammad committed
1673
1674
    group.add_argument('--merge-file', type=str, default=None,
                       help='Path to the BPE merge file.')
1675
1676
1677
    group.add_argument('--vocab-extra-ids', type=int, default=0,
                       help='Number of additional vocabulary tokens. '
                            'They are used for span masking in the T5 model')
Mohammad's avatar
Mohammad committed
1678
    group.add_argument('--seq-length', type=int, default=None,
1679
                       help='Maximum sequence length to process.')
1680
    group.add_argument('--encoder-seq-length', type=int, default=None,
1681
1682
                       help='Maximum encoder sequence length to process.'
                       'This should be exclusive of --seq-length')
1683
1684
    group.add_argument('--decoder-seq-length', type=int, default=None,
                       help="Maximum decoder sequence length to process.")
Mostofa Patwary's avatar
Mostofa Patwary committed
1685
1686
    group.add_argument('--retriever-seq-length', type=int, default=256,
                       help='Maximum sequence length for the biencoder model '
1687
                       'for retriever')
1688
1689
1690
    group.add_argument('--sample-rate', type=float, default=1.0,
                       help='sample rate for training data. Supposed to be 0 '
                            ' < sample_rate < 1')
Mohammad's avatar
Mohammad committed
1691
1692
1693
1694
1695
1696
    group.add_argument('--mask-prob', type=float, default=0.15,
                       help='Probability of replacing a token with mask.')
    group.add_argument('--short-seq-prob', type=float, default=0.1,
                       help='Probability of producing a short sequence.')
    group.add_argument('--num-workers', type=int, default=2,
                       help="Dataloader number of workers.")
Mohammad's avatar
Mohammad committed
1697
1698
1699
    group.add_argument('--tokenizer-type', type=str,
                       default=None,
                       choices=['BertWordPieceLowerCase',
Raul Puri's avatar
Raul Puri committed
1700
                                'BertWordPieceCase',
1701
                                'GPT2BPETokenizer',
1702
                                'SentencePieceTokenizer',
liangjing's avatar
v1  
liangjing committed
1703
                                'GPTSentencePieceTokenizer',
liangjing's avatar
liangjing committed
1704
1705
                                'HuggingFaceTokenizer',
                                'Llama2Tokenizer',
liangjing's avatar
liangjing committed
1706
                                'QwenTokenizer',
liangjing's avatar
liangjing committed
1707
                                'TikTokenizer',
liangjing's avatar
v1  
liangjing committed
1708
                                'NullTokenizer'],
Mohammad's avatar
Mohammad committed
1709
                       help='What type of tokenizer to use.')
1710
    group.add_argument('--tokenizer-model', type=str, default=None,
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1711
                       help='Sentencepiece tokenizer model.')
liangjing's avatar
liangjing committed
1712
1713
1714
1715
1716
1717
    group.add_argument('--tiktoken-pattern', type=str, default=None,
                       help='Which tiktoken pattern to use. Options: [v1, v2]')
    group.add_argument('--tiktoken-num-special-tokens', type=int, default=1000,
                       help='Number of special tokens in tiktoken tokenizer')
    group.add_argument('--tiktoken-special-tokens', type=str, nargs='+', default=None,
                       help='List of tiktoken special tokens, needs to have ["<unk>", "<s>", "</s>"]')
1718
1719
1720
1721
1722
1723
1724
    group.add_argument('--reset-position-ids', action='store_true',
                       help='Reset posistion ids after end-of-document token.')
    group.add_argument('--reset-attention-mask', action='store_true',
                       help='Reset self attention maske after '
                       'end-of-document token.')
    group.add_argument('--eod-mask-loss', action='store_true',
                       help='Mask loss for the end of document tokens.')
liangjing's avatar
liangjing committed
1725
1726
1727
1728
1729
1730
1731
    group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false',
                       help='If set, do not create attention_masks in dataloader.',
                       dest='create_attention_mask_in_dataloader')
    group.add_argument('--num-dataset-builder-threads', type=int, default=1,
                       help='Number of parallel threads per rank for dataset builder')
    group.add_argument('--s3-cache-path', type=str, default=None,
                       help='Path to cache index files when using s3 dataloader')
Mohammad's avatar
Mohammad committed
1732
1733
    return parser

Raul Puri's avatar
Raul Puri committed
1734

Mohammad's avatar
Mohammad committed
1735
1736
def _add_autoresume_args(parser):
    group = parser.add_argument_group(title='autoresume')
Raul Puri's avatar
Raul Puri committed
1737

Mohammad's avatar
Mohammad committed
1738
1739
1740
1741
1742
    group.add_argument('--adlr-autoresume', action='store_true',
                       help='Enable autoresume on adlr cluster.')
    group.add_argument('--adlr-autoresume-interval', type=int, default=1000,
                       help='Intervals over which check for autoresume'
                       'termination signal')
Raul Puri's avatar
Raul Puri committed
1743

Mohammad's avatar
Mohammad committed
1744
    return parser
Neel Kant's avatar
Neel Kant committed
1745
1746


Mostofa Patwary's avatar
Mostofa Patwary committed
1747
1748
def _add_biencoder_args(parser):
    group = parser.add_argument_group(title='biencoder')
Neel Kant's avatar
Neel Kant committed
1749
1750
1751

    # network size
    group.add_argument('--ict-head-size', type=int, default=None,
1752
                       help='Size of block embeddings to be used in ICT and '
Mostofa Patwary's avatar
Mostofa Patwary committed
1753
                        'REALM (paper default: 128)')
1754
    group.add_argument('--biencoder-projection-dim', type=int, default=0,
Mostofa Patwary's avatar
Mostofa Patwary committed
1755
1756
                       help='Size of projection head used in biencoder (paper'
                        ' default: 128)')
1757
    group.add_argument('--biencoder-shared-query-context-model', action='store_true',
Mostofa Patwary's avatar
Mostofa Patwary committed
1758
1759
                        help='Whether to share the parameters of the query '
                        'and context models or not')
Neel Kant's avatar
Neel Kant committed
1760
1761
1762
1763
1764

    # checkpointing
    group.add_argument('--ict-load', type=str, default=None,
                       help='Directory containing an ICTBertModel checkpoint')
    group.add_argument('--bert-load', type=str, default=None,
1765
1766
                       help='Directory containing an BertModel checkpoint '
                       '(needed to start ICT and REALM)')
Neel Kant's avatar
Neel Kant committed
1767
1768
1769
1770
1771

    # data
    group.add_argument('--titles-data-path', type=str, default=None,
                       help='Path to titles dataset used for ICT')
    group.add_argument('--query-in-block-prob', type=float, default=0.1,
1772
1773
                       help='Probability of keeping query in block for '
                       'ICT dataset')
Neel Kant's avatar
Neel Kant committed
1774
    group.add_argument('--use-one-sent-docs', action='store_true',
Neel Kant's avatar
Neel Kant committed
1775
                       help='Whether to use one sentence documents in ICT')
1776
1777
    group.add_argument('--evidence-data-path', type=str, default=None,
                       help='Path to Wikipedia Evidence frm DPR paper')
Neel Kant's avatar
Neel Kant committed
1778

1779
    # training
1780
    group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int,
Mostofa Patwary's avatar
Mostofa Patwary committed
1781
1782
                        default=[], help="Which top-k accuracies to report "
                        "(e.g. '1 5 20')")
Mostofa Patwary's avatar
Mostofa Patwary committed
1783
    group.add_argument('--retriever-score-scaling', action='store_true',
Mostofa Patwary's avatar
Mostofa Patwary committed
1784
1785
                       help='Whether to scale retriever scores by inverse '
                        'square root of hidden size')
1786

Neel Kant's avatar
Neel Kant committed
1787
    # faiss index
Neel Kant's avatar
Neel Kant committed
1788
    group.add_argument('--block-data-path', type=str, default=None,
Neel Kant's avatar
Neel Kant committed
1789
                       help='Where to save/load BlockData to/from')
Mostofa Patwary's avatar
Mostofa Patwary committed
1790
1791
1792
    group.add_argument('--embedding-path', type=str, default=None,
                       help='Where to save/load Open-Retrieval Embedding'
                        ' data to/from')
Neel Kant's avatar
Neel Kant committed
1793
1794
1795

    # indexer
    group.add_argument('--indexer-batch-size', type=int, default=128,
1796
1797
                       help='How large of batches to use when doing indexing '
                       'jobs')
Neel Kant's avatar
Neel Kant committed
1798
    group.add_argument('--indexer-log-interval', type=int, default=1000,
1799
1800
                       help='After how many batches should the indexer '
                       'report progress')
Neel Kant's avatar
Neel Kant committed
1801
    return parser
1802
1803


1804
1805
def _add_vision_args(parser):
    group = parser.add_argument_group(title="vision")
1806

1807
    # general vision arguements
1808
1809
    group.add_argument('--num-classes', type=int, default=1000,
                       help='num of classes in vision classificaiton task')
1810
1811
1812
1813
    group.add_argument('--img-h', type=int, default=224,
                       help='Image height for vision classification task')
    group.add_argument('--img-w', type=int, default=224,
                       help='Image height for vision classification task')
1814
1815
1816
    group.add_argument('--num-channels', type=int, default=3,
                       help='Number of channels in input image data')
    group.add_argument('--patch-dim', type=int, default=16,
1817
                       help='patch dimension')
1818
1819
1820
1821
1822
1823
1824
    group.add_argument('--classes-fraction', type=float, default=1.0,
                       help='training with fraction of classes.')
    group.add_argument('--data-per-class-fraction', type=float, default=1.0,
                       help='training with fraction of data per class.')
    group.add_argument('--no-data-sharding', action='store_false',
                       help='Disable data sharding.',
                       dest='data_sharding')
1825
1826
1827
1828
    group.add_argument('--head-lr-mult', type=float, default=1.0,
                       help='learning rate multiplier for head during finetuning')

    # pretraining type and backbone selection`
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1829
1830
    group.add_argument('--vision-pretraining', action='store_true',
                       help='flag to indicate vision pretraining')
1831
    group.add_argument('--vision-pretraining-type', type=str, default='classify',
Vijay Korthikanti's avatar
Vijay Korthikanti committed
1832
                       choices=['classify', 'inpaint', 'dino'],
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
                       help='pretraining objectives')
    group.add_argument('--vision-backbone-type', type=str, default='vit',
                       choices=['vit', 'mit', 'swin'],
                       help='backbone types types')
    group.add_argument('--swin-backbone-type', type=str, default='tiny',
                       choices=['tiny', 'base', 'h3'],
                       help='pretraining objectives')
    # inpainting arguments
    group.add_argument('--mask-type', type=str, default='random',
                       choices=['random', 'row'],
                       help='mask types')
    group.add_argument('--mask-factor', type=float, default=1.0,
                       help='mask size scaling parameter')
liangjing's avatar
v1  
liangjing committed
1846

1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
    # dino arguments
    group.add_argument('--iter-per-epoch', type=int, default=1250,
                       help='iterations per epoch')
    group.add_argument('--dino-local-img-size', type=int, default=96,
                       help='Image size for vision classification task')
    group.add_argument('--dino-local-crops-number', type=int, default=10,
                       help='Number of local crops')
    group.add_argument('--dino-head-hidden-size', type=int, default=2048,
                       help='Hidden dimension size in dino head')
    group.add_argument('--dino-bottleneck-size', type=int, default=256,
                       help='Bottle neck dimension in dino head ')
    group.add_argument('--dino-freeze-last-layer', type=float, default=1,
                       help='Freezing last layer weights')
    group.add_argument('--dino-norm-last-layer', action='store_true',
                       help='Disable Norm in last layer.')
    group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04,
                       help='warump teacher temperature')
    group.add_argument('--dino-teacher-temp', type=float, default=0.07,
                       help='teacher temperature')
    group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30,
                       help='warmup teacher temperaure epochs')
1868

liangjing's avatar
liangjing committed
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
    # regularization arguments
    group.add_argument('--qk-layernorm', action='store_true',
                       help='Whether to layer normalize the q and k attention embeddings.')

    return parser

def _add_moe_args(parser):
    group = parser.add_argument_group(title="moe")
    group.add_argument('--expert-model-parallel-size', type=int, default=1,
                       help='Degree of expert model parallelism.')
    group.add_argument('--num-experts', type=int, default=None,
                       help='Number of Experts in MoE (None means no MoE)')
    group.add_argument('--moe-router-load-balancing-type', type=str,
                       choices=['aux_loss', 'sinkhorn', 'none'],
                       default='aux_loss',
                       help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer, "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".')
    group.add_argument('--moe-router-topk', type=int, default=2,
                       help='Number of experts to route to for each token. The default is 2.')
    group.add_argument('--moe-router-pre-softmax', action='store_true',
                       help='Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.')
    group.add_argument('--moe-grouped-gemm', action='store_true',
                       help='When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine.')
    group.add_argument('--moe-aux-loss-coeff', type=float, default=0.0,
                       help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
    group.add_argument('--moe-z-loss-coeff', type=float, default=None,
                       help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.')
    group.add_argument('--moe-input-jitter-eps', type=float, default=None,
                       help='Add noise to the input tensor by applying jitter with a specified epsilon value.')
    group.add_argument('--moe-token-dispatcher-type', type=str,
                       choices=['allgather', 'alltoall', 'alltoall_seq'],
                       default='allgather',
                       help="The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather', 'alltoall' and 'alltoall_seq'. We recommend using 'alltoall' when applying expert parallelism. For more information, please refer to the documentation in core/moe/README.")
    group.add_argument('--moe-per-layer-logging', action='store_true',
                       help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.')
    # Token dropping arguments
    group.add_argument('--moe-expert-capacity-factor', type=float, default=None,
                       help='The capacity factor for each expert, None means no token will be dropped.')
    group.add_argument('--moe-pad-expert-input-to-capacity', action='store_true',
                       help='Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set.')
    group.add_argument('--moe-token-drop-policy', type=str, default='probs', choices=['probs', 'position'],
                       help='The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped.')
    group.add_argument('--moe-layer-recompute', action='store_true',
                       help='Enable checkpointing for moe_layer, should be used when memory is not sufficient.')
    group.add_argument('--moe-extended-tp', action='store_true',
                       help='Alternative to expert parallelism, all experts are sharded across TPXEP domain.')
    group.add_argument('--moe-use-upcycling', action='store_true',
                       help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. '
                       'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.')

    return parser

def _add_experimental_args(parser):
    group = parser.add_argument_group(title='experimental')

    group.add_argument('--spec', type=str, default=None, nargs='*',
                       help='Specify the <module_location function_name> pair '
                       'that returns a spec to customize a model, transformer '
                       'block, or transformer layer, depending on the use case.'
                       'To use local spec specify local as the argument.'
                       'For more details, see the model class, '
                       '`transformer_block.py`, or `transformer_layer.py`')
    group.add_argument('--hybrid-attention-ratio', type=float, default=0.0,
                       help='Ratio of attention layers to total layers, in the '
                       'range [0.0, 1.0].')
    group.add_argument('--hybrid-mlp-ratio', type=float, default=0.0,
                       help='Ratio of mlp layers to total layers, in the '
                       'range [0.0, 1.0].')
    group.add_argument('--hybrid-override-pattern', type=str, default=None,
                       help='Force a specific hybrid layer pattern. If a value'
                       'greater than 0.0 is supplied to any of the hybrid ratio'
                       'arguments, then the number of each type of layer in the'
                       'override pattern must match number in the overidden'
                       'pattern')
    group.add_argument('--yaml-cfg', type=str, default=None,
                       help = 'Config file to add additional arguments')
    return parser

def _add_unsloth_args(parser):
    group = parser.add_argument_group(title='unsloth')

    group.add_argument('--use-fast-cross-entropy-loss', action='store_true',
                       help='Use fast_cross_entropy_loss of unsloth more faster in calculating loss')
    group.add_argument('--use-fast-rms-layernorm', action='store_true',
                       help='Use fast_rms_layernorm of unsloth more faster in Layer Normalization')

1954
    return parser