Skip to content

Commit 8aee157

Browse files
lxd-cumtlxd-cumt
committed
add assert
1 parent fe33d23 commit 8aee157

1 file changed

Lines changed: 24 additions & 22 deletions

File tree

‎flagscale/backends/Megatron-LM/megatron/training/arguments.py.patch‎

Lines changed: 24 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
2-
index e36803fc..f82dc6ca 100644
2+
index e36803fc..cbde6421 100644
33
--- a/megatron/training/arguments.py
44
+++ b/megatron/training/arguments.py
55
@@ -59,6 +59,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
@@ -196,7 +196,7 @@ index e36803fc..f82dc6ca 100644
196196
else:
197197
args.ffn_hidden_size = 4 * args.hidden_size
198198

199-
@@ -1085,6 +1110,70 @@ def validate_args(args, defaults={}):
199+
@@ -1085,6 +1110,72 @@ def validate_args(args, defaults={}):
200200
if args.external_cuda_graph:
201201
assert args.te_rng_tracker, "--te-rng-tracker must be enabled when using CUDA Graphs."
202202

@@ -263,11 +263,13 @@ index e36803fc..f82dc6ca 100644
263263
+ assert args.transformer_impl == 'transformer_engine', \
264264
+ 'PEFT is only supported with transformer_engine implementation'
265265
+ assert args.num_experts is None, "PEFT is not tested with MoE currently"
266+
+ assert args.recompute_method is None and args.recompute_granularity is None and args.recompute_num_layers is None, "PEFT will raise comfilcts with recompute currently"
267+
+ assert args.ckpt_format == 'torch', "PEFT is only tested with torch format checkpoint"
266268
+
267269
# Print arguments.
268270
_print_args("arguments", args)
269271

270-
@@ -1435,6 +1524,8 @@ def _add_network_size_args(parser):
272+
@@ -1435,6 +1526,8 @@ def _add_network_size_args(parser):
271273
help='Which normalization technique to use.')
272274
group.add_argument('--norm-epsilon', type=float, default=1e-5,
273275
help='Epsilon for layer norm and RMS norm.')
@@ -276,7 +278,7 @@ index e36803fc..f82dc6ca 100644
276278
group.add_argument('--apply-layernorm-1p', action='store_true',
277279
help='Adjust LayerNorm weights such that they are centered '
278280
'around zero. This improves numerical stability.')
279-
@@ -1450,6 +1541,10 @@ def _add_network_size_args(parser):
281+
@@ -1450,6 +1543,10 @@ def _add_network_size_args(parser):
280282
help='Use squared relu activation instead of default gelu')
281283
group.add_argument('--swiglu', action='store_true',
282284
help='Use gated linear units and SiLU activation instead of default gelu')
@@ -287,7 +289,7 @@ index e36803fc..f82dc6ca 100644
287289
group.add_argument('--onnx-safe', type=bool, required=False,
288290
help='Use workarounds for known problems with '
289291
'Torch ONNX exporter')
290-
@@ -1657,6 +1752,14 @@ def _add_logging_args(parser):
292+
@@ -1657,6 +1754,14 @@ def _add_logging_args(parser):
291293
help='The wandb experiment name.')
292294
group.add_argument('--wandb-save-dir', type=str, default='',
293295
help='Path to save the wandb results locally.')
@@ -302,7 +304,7 @@ index e36803fc..f82dc6ca 100644
302304
group.add_argument('--logging-level', type=int, default=None,
303305
help='Set default logging level')
304306
return parser
305-
@@ -1780,6 +1883,25 @@ def _add_training_args(parser):
307+
@@ -1780,6 +1885,25 @@ def _add_training_args(parser):
306308
'"moe": recompute the MoE layer.'
307309
'"moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, '
308310
'"core_attn", "mlp", and "moe" uses normal checkpointing.')
@@ -328,7 +330,7 @@ index e36803fc..f82dc6ca 100644
328330
group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false',
329331
help='If not set, clone the output of the scatter in embedding layer to GC original tensor.',
330332
dest='clone_scatter_output_in_embedding')
331-
@@ -1866,6 +1988,10 @@ def _add_training_args(parser):
333+
@@ -1866,6 +1990,10 @@ def _add_training_args(parser):
332334
help='Total number of samples to train over all '
333335
'training runs. Note that either train-iters or '
334336
'train-samples should be provided.')
@@ -339,7 +341,7 @@ index e36803fc..f82dc6ca 100644
339341
group.add_argument('--log-interval', type=int, default=100,
340342
help='Report loss and timing interval.')
341343
group.add_argument('--exit-interval', type=int, default=None,
342-
@@ -1981,6 +2107,10 @@ def _add_training_args(parser):
344+
@@ -1981,6 +2109,10 @@ def _add_training_args(parser):
343345
'If None, the default backend will be used.')
344346
group.add_argument('--high-priority-stream-groups', nargs='*', type=str, default=[],
345347
help='The communicator group names to use high priority streams.')
@@ -350,7 +352,7 @@ index e36803fc..f82dc6ca 100644
350352

351353
return parser
352354

353-
@@ -2029,11 +2159,26 @@ def _add_learning_rate_args(parser):
355+
@@ -2029,11 +2161,26 @@ def _add_learning_rate_args(parser):
354356
'and initial warmup, the learning rate at each '
355357
'iteration would be different.')
356358
group.add_argument('--lr-decay-style', type=str, default='linear',
@@ -378,7 +380,7 @@ index e36803fc..f82dc6ca 100644
378380
group.add_argument('--lr-decay-iters', type=int, default=None,
379381
help='number of iterations to decay learning rate over,'
380382
' If None defaults to `--train-iters`')
381-
@@ -2089,6 +2234,8 @@ def _add_checkpointing_args(parser):
383+
@@ -2089,6 +2236,8 @@ def _add_checkpointing_args(parser):
382384
help='Output directory to save checkpoints to.')
383385
group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
384386
help='Number of iterations between persistent checkpoint saves.')
@@ -387,7 +389,7 @@ index e36803fc..f82dc6ca 100644
387389
group.add_argument('--no-save-optim', action='store_true', default=None,
388390
help='Do not save current optimizer.')
389391
group.add_argument('--no-save-rng', action='store_true', default=None,
390-
@@ -2136,6 +2283,8 @@ def _add_checkpointing_args(parser):
392+
@@ -2136,6 +2285,8 @@ def _add_checkpointing_args(parser):
391393
group.add_argument('--no-use-tokenizer-model-from-checkpoint-args', action='store_false',
392394
dest='use_tokenizer_model_from_checkpoint_args',
393395
help='If set, do not use tokenizer model path from checkpoint')
@@ -396,7 +398,7 @@ index e36803fc..f82dc6ca 100644
396398
group.add_argument('--exit-on-missing-checkpoint', action='store_true',
397399
help="If '--load' is set, but checkpoint is not found "
398400
"(e.g., path typo), then exit instead of random "
399-
@@ -2293,7 +2442,7 @@ def _add_distributed_args(parser):
401+
@@ -2293,7 +2444,7 @@ def _add_distributed_args(parser):
400402
default=False, help='if set, overlap pipeline parallel communication in warmup and flush',
401403
dest='overlap_p2p_comm_warmup_flush')
402404
group.add_argument('--distributed-backend', default='nccl',
@@ -405,7 +407,7 @@ index e36803fc..f82dc6ca 100644
405407
help='Which backend to use for distributed training.')
406408
group.add_argument('--distributed-timeout-minutes', type=int, default=10,
407409
help='Timeout minutes for torch.distributed.')
408-
@@ -2344,6 +2493,11 @@ def _add_distributed_args(parser):
410+
@@ -2344,6 +2495,11 @@ def _add_distributed_args(parser):
409411
'complete it instead. Also turns on '
410412
'--use-cpu-initialization flag. This is for '
411413
'external DDP manager.' )
@@ -417,7 +419,7 @@ index e36803fc..f82dc6ca 100644
417419
group.add_argument('--account-for-embedding-in-pipeline-split', action='store_true',
418420
default=False, help='If set, *input* embedding layer will be treated as a standard transformer'
419421
'layer in the context of partition and placement for pipeline parallelism.')
420-
@@ -2378,6 +2532,10 @@ def _add_distributed_args(parser):
422+
@@ -2378,6 +2534,10 @@ def _add_distributed_args(parser):
421423
'and performance requirements.')
422424
group.add_argument('--keep-fp8-transpose-cache-when-using-custom-fsdp', action='store_true',
423425
help='If set, keep the fp8 transpose cache when using custom FSDP.')
@@ -428,7 +430,7 @@ index e36803fc..f82dc6ca 100644
428430
group.add_argument('--num-distributed-optimizer-instances', type=int, default=1,
429431
help='Number of Distributed Optimizer copies across Data Parallel domain.')
430432
group.add_argument('--use-torch-fsdp2', action='store_true',
431-
@@ -2430,6 +2588,9 @@ def _add_validation_args(parser):
433+
@@ -2430,6 +2590,9 @@ def _add_validation_args(parser):
432434
group.add_argument('--eval-interval', type=int, default=1000,
433435
help='Interval between running evaluation on '
434436
'validation set.')
@@ -438,7 +440,7 @@ index e36803fc..f82dc6ca 100644
438440
group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.')
439441
group.add_argument('--skip-train', action='store_true',
440442
default=False, help='If set, bypass the training loop, '
441-
@@ -2444,6 +2605,8 @@ def _add_tokenizer_args(parser):
443+
@@ -2444,6 +2607,8 @@ def _add_tokenizer_args(parser):
442444
help='Size of vocab before EOD or padding.')
443445
group.add_argument('--vocab-file', type=str, default=None,
444446
help='Path to the vocab file.')
@@ -447,7 +449,7 @@ index e36803fc..f82dc6ca 100644
447449
group.add_argument('--merge-file', type=str, default=None,
448450
help='Path to the BPE merge file.')
449451
group.add_argument('--vocab-extra-ids', type=int, default=0,
450-
@@ -2462,8 +2625,17 @@ def _add_tokenizer_args(parser):
452+
@@ -2462,8 +2627,17 @@ def _add_tokenizer_args(parser):
451453
'MultimodalTokenizer',
452454
'NullTokenizer',
453455
'NullMultimodalTokenizer',
@@ -466,7 +468,7 @@ index e36803fc..f82dc6ca 100644
466468
group.add_argument('--tokenizer-model', type=str, default=None,
467469
help='Sentencepiece tokenizer model.')
468470
group.add_argument('--tiktoken-pattern', type=str, default=None,
469-
@@ -2497,6 +2669,11 @@ def _add_data_args(parser):
471+
@@ -2497,6 +2671,11 @@ def _add_data_args(parser):
470472
group.add_argument('--valid-data-path', nargs='*', default=None,
471473
help='The weight and prefix list for an independent validation dataset. '
472474
'Follows the same pattern rules as --data-path.')
@@ -478,7 +480,7 @@ index e36803fc..f82dc6ca 100644
478480
group.add_argument('--test-data-path', nargs='*', default=None,
479481
help='The weight and prefix list for an independent test dataset. '
480482
'Follows the same pattern rules as --data-path.')
481-
@@ -2545,11 +2722,18 @@ def _add_data_args(parser):
483+
@@ -2545,11 +2724,18 @@ def _add_data_args(parser):
482484
'end-of-document token.')
483485
group.add_argument('--eod-mask-loss', action='store_true',
484486
help='Mask loss for the end of document tokens.')
@@ -497,7 +499,7 @@ index e36803fc..f82dc6ca 100644
497499
group.add_argument('--object-storage-cache-path', type=str, default=None,
498500
help='Path to cache index files when using s3 or msc dataloader')
499501
group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005,
500-
@@ -2626,6 +2810,19 @@ def _add_biencoder_args(parser):
502+
@@ -2626,6 +2812,19 @@ def _add_biencoder_args(parser):
501503
return parser
502504

503505

@@ -517,7 +519,7 @@ index e36803fc..f82dc6ca 100644
517519
def _add_vision_args(parser):
518520
group = parser.add_argument_group(title="vision")
519521

520-
@@ -2696,6 +2893,8 @@ def _add_vision_args(parser):
522+
@@ -2696,6 +2895,8 @@ def _add_vision_args(parser):
521523
help='Whether to layer normalize the q and k attention embeddings.')
522524
group.add_argument('--qk-l2-norm', action='store_true',
523525
help='Use llama 4 qk l2 norm')
@@ -526,7 +528,7 @@ index e36803fc..f82dc6ca 100644
526528

527529
return parser
528530

529-
@@ -3000,3 +3199,75 @@ def _add_sft_args(parser):
531+
@@ -3000,3 +3201,75 @@ def _add_sft_args(parser):
530532
group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned",
531533
help='SFT prompt format.')
532534
return parser

0 commit comments

Comments
 (0)