diff --git a/.gitignore b/.gitignore index a40ad60e..5463348e 100644 --- a/.gitignore +++ b/.gitignore @@ -82,38 +82,6 @@ target/ profile_default/ ipython_config.py -# pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -# Pipfile.lock - -# UV -# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# uv.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -# poetry.lock -# poetry.toml - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. -# https://pdm-project.org/en/latest/usage/project/#working-with-version-control -# pdm.lock -# pdm.toml .pdm-python .pdm-build/ @@ -182,25 +150,14 @@ dmypy.json # Cython debug symbols cython_debug/ -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -# .idea/ +# Idea +.idea/ # Abstra -# Abstra is an AI-powered process automation framework. -# Ignore directories containing user credentials, local state, and settings. -# Learn more at https://abstra.io/docs .abstra/ -# Visual Studio Code -# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore -# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore -# and can be added to the global gitignore or merged into this file. However, if you prefer, -# you could uncomment the following to ignore the entire vscode folder -# .vscode/ +# VSCode +.vscode/ # Ruff stuff: .ruff_cache/ @@ -237,8 +194,6 @@ dump/ # SLURM artefacts left in the working tree by ad-hoc sbatch runs slurm-*.out -*_[0-9]*.out -*_[0-9]*.err # Weights & Biases local run cache wandb/ diff --git a/config/backend/megadlm_torchrun.yaml b/config/backend/megadlm_torchrun.yaml new file mode 100644 index 00000000..c9b4d898 --- /dev/null +++ b/config/backend/megadlm_torchrun.yaml @@ -0,0 +1,39 @@ +defaults: + - megatron: base + - _self_ + +class_name: MegatronBackend + +python_cmd: "python" # not used anyways +full_schema_validation: false +launcher_script: "./submodules/MegaDLMs/pretrain_difflm.py" +dist_cmd: "torchrun ${oc.join:'',${oc.mapcondtmpl:'\\-\\-.*','\\\n %',' %',\ + ${oc.cliargs:${.torchrun_args}}}}\\\n ${.launcher_script}" +torchrun_args: + nproc_per_node: ${slurm.sbatch.gpus_per_node} + nnodes: ${slurm.sbatch.nodes} + max_restarts: 0 + tee: 3 + rdzv_endpoint: "$MASTER_ADDR:$MASTER_PORT" + rdzv_backend: "static" + node_rank: "$SLURM_NODEID" + master_addr: "$MASTER_ADDR" + master_port: "$MASTER_PORT" +env: + PYTHONPATH: ".:submodules/MegaDLMs" + RUN_DIR: ${.PROJECT_DIR}/submodules/MegaDLMs + PROJECT_DIR: ${oc.env:PROJECT_DIR,.} + HF_HOME: ${oc.env:HF_HOME,"/scratch-shared/${oc.env:USER}/hf"} + +# MegaDLMs-only CLI flags, appended verbatim to full_cmd. They cannot go through +# the megatron arg mapping: build_cmdline_args drops any key that is not in the +# stock Megatron argparser, and these are added by MegaDLMs' +# custom_args/difflm.py extra_args_provider. `aux` is the declared free-form +# slot on MegatronBackendConfig for exactly this kind of untyped variable -- +# putting the key at the top level only works by grace of NonStrictDataclass. +# Override per-experiment via backend: {aux: {difflm_args: "..."}} +aux: + difflm_args: "" + +# map megatron command to +full_cmd: "${.dist_cmd} ${oc.join:'',${oc.mapcondtmpl:'\\-\\-.*','\\\n %',' %',${oc.argsmegatron:${backend.megatron}}}} ${.aux.difflm_args}" diff --git a/config/backend/megatron/baby/models/qwen3_dense_9b.yaml b/config/backend/megatron/baby/models/qwen3_dense_9b.yaml index 7db9cc71..7a964ef8 100755 --- a/config/backend/megatron/baby/models/qwen3_dense_9b.yaml +++ b/config/backend/megatron/baby/models/qwen3_dense_9b.yaml @@ -19,7 +19,6 @@ apply_rope_fusion: true gradient_accumulation_fusion: true bias_dropout_fusion: true masked_softmax_fusion: true -async_tensor_model_parallel_allreduce: true accumulate_allreduce_grads_in_fp32: true # Attention @@ -38,7 +37,7 @@ untie_embeddings_and_output_weights: true # Normalization normalization: "RMSNorm" -norm_epsilon: 1.e-5 # rmsnorm epsilon +layernorm_epsilon: 1.e-5 # rmsnorm epsilon qk_layernorm: True # enable qk-norm bf16: True diff --git a/config/backend/megatron/base_defaults.yaml b/config/backend/megatron/base_defaults.yaml index 2bc8645a..5f7a7612 100644 --- a/config/backend/megatron/base_defaults.yaml +++ b/config/backend/megatron/base_defaults.yaml @@ -1,13 +1,13 @@ -account_for_embedding_in_pipeline_split: false # If set, *input* embedding layer will be treated as a standard transformerlayer in the context of partition and placement for pipeline parallelism. -account_for_loss_in_pipeline_split: false # If set, loss layer will be treated as a standard transformerlayer in the context of partition and placement for pipeline parallelism. +account_for_embedding_in_pipeline_split: false # If set, the embedding layer will be treated as a standard transformer layer in the context of partition and placement for pipeline parallelism. +account_for_loss_in_pipeline_split: false # If set, the loss layer will be treated as a standard transformer layer in the context of partition and placement for pipeline parallelism. accumulate_allreduce_grads_in_fp32: false # Gradient accumulation and all-reduce in fp32. activation_func_clamp_value: null # Clamp the output of the linear_fc1 in the activation function. Only used when activation_func is quick_gelu. adam_beta1: 0.9 # First coefficient for computing running averages of gradient and its square adam_beta2: 0.999 # Second coefficient for computing running averages of gradient and its square adam_eps: 1.0e-08 # Term added to the denominator to improvenumerical stability -add_bias_linear: true # Disable bias in the linear layers +add_bias_linear: true # Include/exclude a bias term in all linear layers (QKV projections, after core attention, and two in MLP layer). add_position_embedding: true # Disable position embedding. Deprecated: use --position-embedding-type -add_qkv_bias: false # Enable bias only in the QKV linear layers +add_qkv_bias: false # Add a bias term only for QKV projections. adlr_autoresume: false # Enable autoresume on adlr cluster. adlr_autoresume_interval: 1000 # Intervals over which check for autoresumetermination signal align_grad_reduce: true # If not set, all PP stages will launch gradient reduces simultaneously. Otherwise, each PP stage will independently launch as needed. @@ -15,22 +15,27 @@ align_param_gather: true # If not set, all PP stages will launch param all-gath allow_ambiguous_pad_tokens: false # Whether to prevent pad tokens already present in the dataset from being masked out when the pad token incorrectly shares the same id with other special tokens in the tokenizer. Note that this argument has no effect when the tokenizer correctly provides a unique id for the pad. Masking out such ambiguous pad tokens results in training instability. Such a scenario is best resolved by fixing the tokenizer; leaving this option as False provides a workaround. When left to the default of False, any token ids that collide with the pad token id - as provided by the tokenizer - will not be masked out of the loss calculation: it cannot be determined whether they are truly pad. If instead this argument is set, the training flow will treat all tokens that share the same id as the pad token as true pad tokens, potentially causing severe training instability. app_tag_run_name: null # 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 app_tag_run_version: 0.0.0 # 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 -apply_layernorm_1p: false # Adjust LayerNorm weights such that they are centered around zero. This improves numerical stability. -apply_query_key_layer_scaling: false # Scale Q * K^T by 1 / layer-number. Useful for fp16 training. Also sets `attention_softmax_in_fp32` to True. -apply_residual_connection_post_layernorm: false # If set, use original BERT residula connection ordering. -apply_rope_fusion: true # Disable rope fusion, the fusion is available only when using megatron-core. -async_save: null # Apply async checkpointing save. Currently works only with`torch_dist` distributed checkpoint format. -async_tensor_model_parallel_allreduce: true # DEPRECATED. This flag is ignored. -attention_backend: auto # choices: [1, 2, 3, 4, 5] | Attention backend to use (flash,fused,unfused,local,auto). Defaults to auto +apply_query_key_layer_scaling: false # If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with fp16. Also sets `attention_softmax_in_fp32` to True. +apply_residual_connection_post_layernorm: false # If True, uses the original BERT residule connection ordering. +apply_rope_fusion: true # Requires TE >= 1.4 or Apex. Override to false in configs that don't have TE (e.g. DLM). +apply_wd_to_qk_layernorm: false # Apply weight decay to qk layernorm as a special case. +async_ckpt_cpu_priority: 10 # CPU nice value target (0-19, higher = lower priority) for the async checkpoint writer process. If it exceeds 19, it will be set to 19. If the current nice value is greater than the target, it will be left unchanged. Only applies when using persistent ckpt worker. +async_ckpt_io_priority: 3 # I/O scheduling class (0-3, 3=idle) for the async checkpoint writer process. +async_ckpt_use_cpu_shm: false # Copy GPU tensors to CPU shared-memory in the training process before handing off to the async checkpoint worker. Avoids CUDA IPC / NVLink fabric handles in the worker subprocess. Useful on MNNVL systems where fabric resources are exhausted. Only applies with the nvrx async strategy. +async_save: false # Apply async checkpointing save. Currently works only with `torch_dist` distributed checkpoint format. +async_strategy: nvrx # choices: ['nvrx', 'mcore'] | Which async save strategy to use. Available strategies: nvrx, mcore. +attention_backend: auto # choices: [1, 2, 3, 4, 5] | Attention backend to run. By default we let transformer engine decide the best backend to run (except in the case of local). If attention backend is local we use the local pytorch implementation in mcore. Users can specify exact backend by changing this config. attention_dropout: 0.1 # Post attention dropout probability. +attention_output_gate: false # Whether to apply output gate to the attention layers. attention_softmax_in_fp32: false # Run attention masking and softmax in fp32. auto_detect_ckpt_format: false # 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). -barrier_with_L1_time: true # 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. +barrier_with_L1_time: true # If not disabled, 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. +batch_invariant_mode: false # If true, uses batch-invariant kernels that provide deterministic forward execution regardless of batch size. This ensures bitwise identical results when the same inputs are processed in different batch configurations. This will significantly affect speed of training and inference as the kernels are not full optimized. Defaults to False. batch_size: null # Old batch size parameter, do not use. Use --micro-batch-size instead bert_binary_head: true # Disable BERT binary head. bert_embedder_type: megatron # choices: ['megatron', 'huggingface'] | Select either Megatron or Huggingface as the Bert embedder. bert_load: null # Directory containing an BertModel checkpoint (needed to start ICT and REALM) -bf16: false # Run model in bfloat16 mode. +bf16: false # If true, train with bf16 mixed precision training. bias_dropout_fusion: true # Disable bias and dropout fusion. bias_gelu_fusion: true # Disable bias and gelu fusion. bias_swiglu_fusion: true # Disable bias and swiglu fusion, the fusion is available only when using megatron-core. @@ -39,47 +44,58 @@ biencoder_shared_query_context_model: false # Whether to share the parameters o block_data_path: null # Where to save/load BlockData to/from cache_mla_latents: false # If set caches the mla down projected latents with mla flash decode. calc_ft_timeouts: false # If set, FT package will try to automatically compute the timeouts. Note: This feature is for Nvidia internal use only. -calculate_per_token_loss: false # 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. +calculate_per_token_loss: false # Whether cross entropy loss is calculated over the actual number of non-padded tokens in the global batch, versus the default behavior of assuming all tokens are non-padded. +chat_template: null # Custom chat template in jinja format for conversation formatting. check_for_large_grads: false # Check for unexpectedly large grads check_for_nan_in_loss_and_grad: true # Check for NaNs in loss and grad -check_for_spiky_loss: false # Check for spiky loss +check_for_spiky_loss: false # Check for spiky loss. check_weight_hash_across_dp_replicas_interval: null # Interval to check weight hashes are same across DP replicas. If not specified, weight hashes not checked. checkpoint_activations: false # Checkpoint activation to allow for training with larger models, sequences, and batch sizes. -ckpt_assume_constant_structure: false # If the model and optimizer state dict structure isconstant throughout a *single training job*, it allows fordifferent checkpointing performance optimizations. -ckpt_convert_format: null # choices: ['torch', 'torch_dist', 'zarr'] | Checkpoint format for conversion. +ckpt_assume_constant_structure: false # Assume the checkpoint structure is constant across saves to enable optimizations. +ckpt_convert_format: null # choices: ['torch', 'torch_dist'] | Checkpoint format for conversion. ckpt_convert_save: null # Save directory for converted checkpoint. ckpt_convert_update_legacy_dist_opt_format: false # 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. -ckpt_format: torch_dist # choices: ['torch', 'torch_dist', 'zarr', 'torch_dcp', 'fsdp_dtensor'] | Checkpoint format to use. torch is the format used by torch.save/load. torch_dist is a megatron built-in distributed checkpointing format. torch_dcp is the torch.distributed.checkpoint format. fsdp_dtensor is a torch DCP native, Megatron FSDP training-specific checkpoint format. +ckpt_format: torch_dist # choices: ['torch', 'torch_dist', 'torch_dcp', 'fsdp_dtensor'] | Checkpoint format to use. torch is the format used by torch.save/load. torch_dist is a megatron built-in distributed checkpointing format. torch_dcp is the torch.distributed.checkpoint format. fsdp_dtensor is a torch DCP native, Megatron FSDP training-specific checkpoint format. ckpt_fully_parallel_load: false # Apply full load parallelization across DP for distributed checkpoints. +ckpt_fully_parallel_load_exchange_algo: broadcast # choices: ['broadcast', 'gather_rounds', 'gather_object'] | Algorithm for fully parallel load of distributed checkpoints. "broadcast"(default): Broadcast the checkpoint from rank 0 to all other ranks. "gather_rounds": Gather the checkpoint from all ranks in rounds. "gather_object": Gather the checkpoint from all ranks in a single operation. +ckpt_fully_parallel_load_process_group: dp # choices: ['dp', 'ep_dp'] | Process group for fully parallel load of distributed checkpoints. "dp"(default): Data parallel process group. "ep_dp": Expert data parallel process group. ckpt_fully_parallel_save: true # 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. ckpt_fully_parallel_save_deprecated: false # Deprecated: see --no-ckpt-fully-parallel-save. +ckpt_fully_parallel_save_process_group: dp # choices: ['dp', 'ep_dp'] | Process group for fully parallel save of distributed checkpoints. "dp"(default): Data parallel process group. "ep_dp": Expert data parallel process group. +ckpt_load_validate_sharding_integrity: true # Whether to validate sharding access integrity when loading a distributed checkpoint. When True (default), each tensor shard is checked to be accessed exactly once as main replica by some rank. Disabling skips this validation ckpt_step: null # Checkpoint step to load model from. classes_fraction: 1.0 # training with fraction of classes. clip_grad: 1.0 # Gradient clipping based on global L2 norm. -clone_scatter_output_in_embedding: true # If not set, clone the output of the scatter in embedding layer to GC original tensor. -config_logger_dir: '' # If set, will dump all configs to --config-logger-dir -context_parallel_size: 1 # Degree of context parallelism. +clone_scatter_output_in_embedding: true # When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer to facilitate garbage collection of input. +config_logger_dir: '' # When non-empty, dumps entry-point configs to config_logger_dir +context_parallel_size: 1 # Splits network input along sequence dimension across GPU ranks. cp_comm_type: # Inter-gpu communication type for context parallelism: p2p, a2a, allgather or a2a+p2p. If a single string is provided, all layers will share the same communication type. Users can also specify separated types for each layer like --cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p - p2p -cpu_offloading_num_layers: 0 # The number of Transformer layers to offload to CPU. +cpu_offloading_num_layers: 0 # Tells the number of transformer layers for which activations has to be offloaded. +cpu_offloading_retain_pinned_cpu_buffers: false # If True, the pinned CPU buffers are retained after offloading and reused for the next iteration. It is useful for cuda graphs capture. +create_all_gather_group: false # Create a separate process group for all-gather operations to overlap reduce-scatter and all-gather operations. create_attention_mask_in_dataloader: true # If set, do not create attention_masks in dataloader. -cross_entropy_fusion_impl: native # choices: ['native', 'te'] | Implementation of cross entropy loss calculation. -cross_entropy_loss_fusion: false # Enabled fusion of cross entropy loss calculation. -cuda_graph_impl: none # choices: ['none', 'local', 'transformer_engine'] | Determines the CUDA graph capture implementation. "none": no CUDA graph. "local": capture the CUDA graph using MCore local implementation. --cuda-graph-scope="full_iteration" enables whole iteration CUDA graph. "transformer_engine": capture the CUDA graph using TE make_graphed_callables(). -cuda_graph_scope: full # choices: ['full', 'attn', 'full_iteration'] | Determines the CUDA graphs capturing scope. Valid values are "full", "attn" and "full_iteration". "Full" scope captures a whole Transformer layer. "Attn" scope only captures operations in TransformerLayer._forward_attention(). "ful_iteration" scope captures a whole iteration. full_iteration scope is only supported with --cuda-graph-impl=local, attn scope is only supported with --cuda-graph-impl=transformer_engine. -cuda_graph_warmup_steps: 3 # Number of CUDA graph warmup steps +cross_entropy_fusion_impl: native # choices: ['native', 'te'] | If 'native', MCore based CE loss fusion is used, if 'te', Parallel CE loss from Transformer Engine library is used. Defaults to 'native'. +cross_entropy_loss_fusion: false # If this is enabled, the fused cross entropy implementation would be used. Defaults to False. +cuda_graph_impl: none # choices: ['none', 'local', 'transformer_engine', 'full_iteration'] | Determines the CUDA graph capture implementation. "none": no CUDA graph. "local": MCore CUDA graph implementation. During training, graphable modules own per-layer CUDA graphs controlled by cuda_graph_modules. During inference, graph ownership is controlled separately by inference_cuda_graph_scope. "transformer_engine": Transformer Engine CUDA graph implementation. During training, TE make_graphed_callables() creates per-layer CUDA graphs controlled by cuda_graph_modules. Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none". "full_iteration": full-iteration CUDA graph implementation for the training iteration (1 CUDA graph for the whole forward-backward path excluding the optimizer step). Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none". cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when cuda_graph_impl="full_iteration". +cuda_graph_modules: [] # Selects training capture coverage within per-layer CUDA graphs (local and transformer_engine implementations). Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba": "attn": captures operations in TransformerLayer._forward_attention(). "mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. "moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. "moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), including the shared experts if they are not overlapped with EP comm. "moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". "mamba": captures the mamba layer. An empty list means capturing the whole Transformer layer. This field is meaningless when --cuda-graph-impl=full_iteration and must be empty. Backward compatibility: "full" is deprecated but kept for backward compatibility; it is transformed to an empty list in validate_args. The deprecated values "full_iteration" and "full_iteration_inference" are also accepted and migrated to the new API in validate_args. +cuda_graph_scope_deprecated: null # ==SUPPRESS== +cuda_graph_warmup_steps: 3 # Number of warmup steps for CUDA graphs data_args_path: null # Path to data-args. Instead of feeding `--data-path` with weighted dataset, we pass in a file path from which we read that argument. This is useful when the list of data is too big. data_cache_path: null # Path to a directory to hold cached index files. data_parallel_random_init: false # Enable random initialization of params across data parallel ranks -data_parallel_sharding_strategy: no_shard # choices: ['no_shard', 'optim', 'optim_grads', 'optim_grads_params'] | Sharding strategy of data parallelism. +data_parallel_sharding_strategy: optim_grads_params # choices: ['no_shard', 'optim', 'optim_grads', 'optim_grads_params'] | Sharding strategy of data parallelism. data_path: null # The weight and prefix list for a set of train, validation, and testdatasets 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. data_per_class_fraction: 1.0 # training with fraction of data per class. data_sharding: true # Disable data sharding. +dataloader_defer_npy_index_mmap: false # Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path. +dataloader_fast_cache_load: false # Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path. dataloader_type: null # choices: ['single', 'cyclic', 'external'] | Single pass vs multiple pass data loader ddp_average_in_collective: false # If set, average directly in data-parallel communication collective. ddp_bucket_size: null # Bucket size for data-parallel communication ddp_num_buckets: null # Number of buckets for data-parallel communication ddp_pad_buckets_for_high_nccl_busbw: false # If set, make sure the bucket size is divisible by a large power of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, since NCCL message size (which for ring algorithms is bucket_size / dp_size) apparently needs to be divisible by a power of 2 for high busbw. +ddp_param_name_patterns_for_fp32_local_accumulation: [] # List of param_name patterns (in Python's fnmatch format) to match against to do local gradient accumulation in FP32. The special pattern 'all' matches every parameter. ddp_reduce_scatter_with_fp32_accumulation: false # If set, use a reduce-scatter implementation which sends lower-precision values over the wire (using an all-to-all to keep total communication overhead in line with the standard ring implementation) but performs accumulation locally in FP32. decode_only_cuda_graphs: false # Only use cuda graphs for decode-only steps, not prefill and mixed steps. decoder_first_pipeline_num_layers: null # 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 @@ -88,11 +104,11 @@ decoder_num_layers: null # Number of decoder transformer layers. decoder_seq_length: null # Maximum decoder sequence length to process. decoupled_lr: null # Separate learning rate for the input and output layer decoupled_min_lr: null # Minimum value for learning rate for the input and output layer. The schedulerclip values below this threshold -decrease_batch_size_if_needed: false # If set, decrease batch size if microbatch_size * dp_sizedoes not divide batch_size. Useful for KSO (Keep Soldiering On)to continue making progress if number of healthy GPUs (andcorresponding dp_size) does not support current batch_size.Old batch_size will be restored if training is re-started withdp_size that divides batch_size // microbatch_size. -defer_embedding_wgrad_compute: false # If set, defers the vocabulary projection linear layer weightgradient compute to pipeline flush. -delay_wgrad_compute: false # Delay the wgrad compute for batch-level overlapping -deprecated_use_mcore_models: false # DEPRECATED. Use the implementation from megatron core.Now ignored and mcore models are the default, use --use-legacy-models to not use core models. -deterministic_mode: false # Choose code that has deterministic execution. This usually means slower execution, but is good for debugging and testing. +decrease_batch_size_if_needed: false # If set, decrease batch size if microbatch_size * dp_size does not divide batch_size. Old batch_size will be restored if training is re-started with dp_size that divides batch_size // microbatch_size. Not supported with step-batch-size-schedule. +defer_embedding_wgrad_compute: false # If true, defers the embedding WGRAD GEMMs while pipeline flush is taking place enabling us to hide pipeline flush latency. Defaults to False. +delay_wgrad_compute: false # Delay the weight gradient computation to improve batch-level communication overlapping +deprecated_use_mcore_models: false # DEPRECATED. Use the implementation from megatron core.Now ignored and mcore models are the default. +deterministic_mode: false # If true, code that has deterministic execution will be chosen. This usually means slower execution, but is good for debugging and testing. Defaults to False. dino_bottleneck_size: 256 # Bottle neck dimension in dino head dino_freeze_last_layer: 1 # Freezing last layer weights dino_head_hidden_size: 2048 # Hidden dimension size in dino head @@ -103,51 +119,72 @@ dino_teacher_temp: 0.07 # teacher temperature dino_warmup_teacher_temp: 0.04 # warump teacher temperature dino_warmup_teacher_temp_epochs: 30 # warmup teacher temperaure epochs disable_bf16_reduced_precision_matmul: false # If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to prevent matmul from using reduced precision accumulation when using BF16. -disable_chunked_prefill: false # Disable chunked prefill (chunked prefill is enabled by default). -disable_mamba_mem_eff_path: false # Disable Mamba efficient path. +disable_jit_fuser: false # Disable the JIT fuser. +disable_msc_deprecated: false # [DEPRECATED] MSC is disabled by default; this flag is a no-op and will be removed in a future release. disable_straggler_on_startup: false # If set, StragglerDetector is disabled on startup. disable_symmetric_registration: false # Disable symmetric (window) registration for NCCL userbuffer registration.This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set. dist_ckpt_format_deprecated: null # Deprecated: see --ckpt-format. dist_ckpt_optim_fully_reshardable: false # Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as opposed to plain DP reshardability. dist_ckpt_save_pre_mcore_014: false # Revert checkpointing simplifications introduced in Megatron-Core v0.14. This option affects only checkpoint saving format and will be removed soon (checkpoint load format is determined based on checkpoint metadata). dist_ckpt_strictness: assume_ok_unexpected # choices: ['assume_ok_unexpected', 'log_unexpected', 'log_all', 'raise_unexpected', 'raise_all', 'return_unexpected', 'return_all', 'ignore_all'] | 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. -distrib_optim_fully_reshardable_mem_efficient: false # During distributed optimizer checkpoint save and load tries to use as little memory as possible by using Gloo (instead of NCCL) and only one rank for saving. Turn on only if experiencing host or device memory issues. Has affect only with `--dist-ckpt-optim-fully-reshardable` flag. -distribute_saved_activations: false # If set, distribute recomputed activations across model parallel group. +dist_ckpt_workers: 1 # Number of workers for distributed checkpointing. Only used for async save. If set to 1, the checkpointing is performed in a single process. +distrib_optim_fully_reshardable_mem_efficient: false # During distributed optimizer checkpoint save and load tries to use as little memory as possible by using Gloo (instead of NCCL) and only one rank for saving. Turn on only if experiencing host or device memory issues. Has affect only with `dist_ckpt_optim_fully_reshardable` flag. +distribute_saved_activations: false # If True, distribute recomputed activations across the model parallel group. distributed_backend: nccl # choices: ['nccl', 'gloo'] | Which backend to use for distributed training. -distributed_timeout_minutes: 10 # Default timeout minutes for torch.distributed. -distributed_timeout_seconds_after_init: null # Timeout seconds for process groups after initialization.This timeout is applied to all process groups after initialization. +distributed_timeout_minutes: 10 # Timeout minutes for torch.distributed. +distributed_timeout_seconds_after_init: null # Timeout in seconds for process groups after initialization. This timeout is applied to all process groups after initialization and the first iteration completes. +dsa_indexer_head_dim: null # Dimension per DSA indexer head. +dsa_indexer_loss_coeff: null # Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss. +dsa_indexer_n_heads: null # Number of DSA indexer heads. +dsa_indexer_topk: null # Number of top-k tokens to select in DSA indexer. +dsa_indexer_use_sparse_loss: false # Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices. dump_param_to_param_group_map: null # Path to a file containing parameter-to-parameter-group mapping. Provide a JSON file that specifies which parameters belong to which parameter group for global coordination. -embedding_init_method_std: null # Standard deviation of the zero mean normal distribution used for embedding weight initialization. If unset, embeddings will be initialized the same way as other weights. Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying weight decay on embedding weights to avoid shrinkage towards zero. See https://arxiv.org/abs/2312.16903 for more details. +embedding_init_method_std: null # Standard deviation of the zero mean normal for the default initialization method for the embedding layer. If None, will be set to init_method_std. Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying weight decay on embedding weights to avoid shrinkage towards zero. See https://arxiv.org/abs/2312.16903 for more details. embedding_path: null # Where to save/load Open-Retrieval Embedding data to/from -empty_unused_memory_level: 0 # choices: [0, 1, 2] | Call torch.cuda.empty_cache() each iteration (training and eval), to reduce fragmentation.0=off, 1=moderate, 2=aggressive. -enable_cuda_graph: false # Deprecated. Use --cuda-graph-impl=local instead. Use local implementation of CUDA graph capture and replay. --cuda-graph-scope="full_iteration" enables whole iteration CUDA graph. +empty_unused_memory_level: 0 # choices: [0, 1, 2] | Call torch.cuda.empty_cache() each iteration (training and eval), to reduce fragmentation. 0=off, 1=moderate, 2=aggressive. +enable_chunked_prefill: false # Enable chunked prefill (disabled by default) +enable_cuda_graph: false # DEPRECATED and replaced by cuda_graph_impl. When set to true, either partial CUDA graph (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole iteration excluding optimizer) is enabled. cuda_graph_modules determines the scope of graph capture. enable_experimental: false # Enable experimental features. enable_ft_package: false # If set, Fault Tolerance package is enabled. Note: This feature is for Nvidia internal use only. enable_full_sharding_in_hsdp: false # If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode. -enable_gloo_process_groups: true # Disables creation and usage of Gloo process groups. -enable_msc: true # Disable the usage of Multi-Storage Client (MSC) in Megatron Core. +enable_msc: false # Enable the usage of Multi-Storage Client (MSC) in Megatron Core. Disabled by default; pass this flag to opt in. enable_one_logger: true # If set, disable using one_logger to track E2E metricsNote that one_logger is an internal tool and not available externally. For installation, please go to https://confluence.nvidia.com/display/MLWFO/Package+Repositoriesfor more details encoder_num_layers: null # Number of encoder transformer layers. encoder_seq_length: null # Maximum encoder sequence length to process.This should be exclusive of --seq-length end_weight_decay: null # End of run weight decay coefficient for L2 regularization. eod_mask_loss: false # Mask loss for the end of document tokens. +ep_overlap_early_attn_memory_release: false # Enable early memory release of attention activations during EP overlap. EP overlap can increase peak memory usage when the overlapped forward module allocates more memory than what is freed by the backward module. This flag addresses this by reordering the attention backward pass to occur earlier in the schedule. Specifically: - Without this flag: attn_bwd executes after moe_combine_fwd - With this flag: attn_bwd executes before mlp_fwd The earlier execution releases attention activations sooner, reducing peak memory. Note: This may impact performance as moe_combine_fwd and moe_dispatch_bwd become exposed (not overlapped with other computation). error_injection_rate: 0 # Rate at which to inject unexpected results, e.g. 1000 means once every 1000 result validations error_injection_type: transient_error # choices: ['correct_result', 'transient_error', 'persistent_error'] | Type of error to inject. -eval_interval: 1000 # Interval between running evaluation on validation set. -eval_iters: 100 # Number of iterations to run for evaluationvalidation/test for. +eval_global_batch_size: null # Global batch size to use during evaluation. If not set, defaults to global_batch_size. Must be divisible by (eval_micro_batch_size * data_parallel_size). +eval_interval: null # Interval between running evaluation on validation set. If not set, evaluation will not run during training. +eval_iters: 100 # Number of iterations to run for evaluation. Used for both validation and test. If not set, evaluation will not run. +eval_micro_batch_size: null # Micro batch size to use during evaluation. If not set, defaults to micro_batch_size. Changing this affects per-device memory usage during eval and the number of microbatches per eval step. evidence_data_path: null # Path to Wikipedia Evidence frm DPR paper exit_duration_in_mins: null # Exit the program after this many minutes. exit_interval: null # Exit the program after the iteration is divisible by this value. -exit_on_missing_checkpoint: false # If '--load' is set, but checkpoint is not found (e.g., path typo), then exit instead of random initialization. -exit_signal: SIGTERM # choices: ['SIGTERM', 'SIGINT', 'SIGUSR1', 'SIGUSR2'] | Signal to use for exit signal handler. If not specified, defaults to SIGTERM. -exit_signal_handler: false # Dynamically save the checkpoint and shutdown the training if signal is received +exit_on_missing_checkpoint: false # If 'load' is set, but checkpoint is not found (e.g., path typo), then exit instead of random initialization. +exit_signal: '15' # choices: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 34, 64] | Signal for the signal handler to detect. +exit_signal_handler: false # Dynamically save the checkpoint and shutdown the training if SIGTERM is received +exit_signal_handler_for_dataloader: false # Use signal handler for dataloader workers +exit_signal_handler_for_training: false # Shutdown the training when SIGINT or SIGTERM received to avoid unclear traceback exp_avg_dtype: fp32 # choices: ['fp32', 'fp16', 'bf16', 'fp8'] | Dtype of exp_avg (1st moment in adam optimizer) when enabling precision-aware-optimizer. This dtype is used for storing the optimizer state in memory during training but does not affect the precision in the kernel computation. exp_avg_sq_dtype: fp32 # choices: ['fp32', 'fp16', 'bf16', 'fp8'] | Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling precision-aware-optimizer. This dtype is used for storing the optimizer state in memory during training but does not affect the precision in the kernel computation. -expert_model_parallel_size: 1 # Degree of expert model parallelism. -expert_tensor_parallel_size: null # Degree of expert model parallelism. Default is None, which will be set to the value of --tensor-model-paralle-size. -external_cuda_graph: false # Deprecated. Use --cuda-graph-impl=transformer_engine instead. Use TE make_graphed_callables() to capture the CUDA graph. +experimental_attention_variant: null # choices: ['gated_delta_net', 'dsa'] | Type of attention variant to use. Currently support gated_delta_net and dsa. +expert_model_parallel_size: 1 # Distributes Moe Experts across sub data parallel dimension. +expert_tensor_parallel_size: null # Intra-layer tensor model parallelism for expert layer. Splits tensors across GPU ranks. Default is None, which will be set to the value of tensor_model_parallel_size. +external_cuda_graph: false # DEPRECATED and replaced by cuda_graph_impl. When set to true, TransformerLayer layers are swapped with user provided CUDA graphs. fake_process_group: false # If set, initialize with fake distributed process group and all distributed communication operations will be skipped. This is quite useful for profiling memory usage of distributed training with just one GPU. Setting WORLD_SIZE and RANK to the specific values for target distribtued scale. -ffn_hidden_size: null # Transformer Feed-Forward Network hidden size. This is set to 4*hidden-size if not provided +fault_injector_delay_start_iteration: null # Start the fault delay timer after iteration N completes. If unset, fault delay timing starts from the beginning of training. +fault_injector_fault_delay: null # Force a specific fault delay in seconds from training start or delay_start_iteration. +fault_injector_fault_probabilities: null # Comma-separated list of fault probabilities (normalized at runtime). +fault_injector_fault_types: null # Comma-separated list of fault types to inject (e.g. 'hang,crash'). +fault_injector_mtti_seconds: null # Mean time to inject (MTTI) in seconds; used when fault_delay is None. +fault_injector_num_ranks: null # Number of ranks to inject faults on (random selection). +fault_injector_offset_seconds: null # Offset seconds added to the sampled fault delay. +fault_injector_ranks: null # Comma-separated list of ranks to inject faults on. +fault_injector_seed: null # RNG seed for the fault injector. +ffn_hidden_size: null # Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size if not provided. fim_data: false # Whether to use the FIM dataset. fim_eod_token: <|endoftext|> # FIM EOD token fim_fragment_rate: null # Rate of FIM on each fragment when --fim-split-sample is not None. @@ -159,37 +196,48 @@ fim_rate: 0.5 # Probability to convert a training sample into a FIM format. fim_split_sample: null # String around which to split the sample for FIM. fim_spm_rate: 0.5 # Probability that the a FIM sample uses the SPM format over the PSM format. fim_suffix_token: # FIM suffix token +fine_grained_activation_offloading: false # If True, offload the input of the specified modules to the CPU. Fine-grained activation offloading is a module-level offloading method instead of a layer-level offloading method like cpu_offloading. +fine_grained_offloading_max_inflight_offloads: null # Per fine-grained offloading group name, max number of inflight offloads for that name not yet joined on the main stream (wait_event on D2H). The same cap applies to every name (e.g., ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait after every offload for that name. 1 = at most one not-yet-waited offload per name, etc. None = do not insert these joins. This feature is particularly useful when using with full-iteration CUDA graphs finetune: false # Load model for finetuning. Do not load optimizer or rng state from checkpoint and set iteration to 0. Assumed when loading a release checkpoint. -first_last_layers_bf16: false # Construct first and last layers in bf16 when doing FP8 training. -flash_decode: false # Whether to use the flash decoding kernel. -fp16: false # Run model in fp16 mode. +first_last_layers_bf16: false # If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8. +flash_decode: false # Use the optimized flash decoding kernel during inference. +flight_recorder_dump_on_timeout: true # Dump flight recorder traces on NCCL timeout (TORCH_NCCL_DUMP_ON_TIMEOUT). +flight_recorder_dump_path: null # Path for NCCL flight recorder trace dumps. Sets TORCH_FR_DUMP_TEMP_FILE and TORCH_NCCL_DEBUG_INFO_TEMP_FILE env variables before distributed init. +flight_recorder_extra_dump_on_exec: true # Enable extra flight recorder dump on execution (TORCH_NCCL_EXTRA_DUMP_ON_EXEC). +flight_recorder_include_only_active: true # Include only active operations in flight recorder dumps (TORCH_INCLUDE_ONLY_ACTIVE). +flight_recorder_include_stack_trace: false # Include stack traces in flight recorder dumps (TORCH_INCLUDE_STACK_TRACE). +flight_recorder_trace_buffer_size: 2000 # Size of the NCCL flight recorder trace buffer (TORCH_NCCL_TRACE_BUFFER_SIZE). +fp16: false # If true, train with fp16 mixed precision training. fp16_lm_cross_entropy: false # Move the cross entropy unreduced loss calculationfor lm head to fp16. -fp32_residual_connection: false # Move residual connections to fp32. -fp4: null # choices: ['e2m1'] | Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass -fp4_param: false # Keep the compute param in fp4 (do not use any other intermediate dtype) and perform the param all-gather in fp4. -fp4_quantizer_factory: null # Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. -fp4_recipe: nvfp4 # choices: ['nvfp4', 'custom'] | Which fp4 recipe to use for FP4 tensors in the forward and backward pass -fp8: null # choices: ['e4m3', 'hybrid'] | Which fp8 format scheme to use for FP8 tensors in the forward and backward pass -fp8_amax_compute_algo: most_recent # choices: ['most_recent', 'max'] | Algorithm for computing amax from history -fp8_amax_history_len: 1 # Number of steps for which amax history is recorded per tensor -fp8_interval: 1 # DEPRECATED. This flag is ignored. Scaling update interval for fp8 -fp8_margin: 0 # Scaling margin for fp8 +fp32_residual_connection: false # If true, move residual connections to fp32. +fp4: null # choices: ['e2m1'] | If set, enables the use of FP4 precision through Transformer Engine. Currently only supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0). +fp4_param_gather: false # Keep the compute param in fp4 (do not use any other intermediate dtype) and perform the param all-gather in fp4. +fp4_quantizer_factory: null # Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. Required when fp4_recipe is custom. +fp4_recipe: nvfp4 # choices: ['nvfp4', 'custom'] | If set, enables the use of FP4 precision through Transformer Engine. Currently only 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture. +fp8: null # choices: ['e4m3', 'hybrid'] | If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8 activation and weight tensors and e5m2 for all FP8 output activation gradient tensors. +fp8_amax_compute_algo: most_recent # choices: ['most_recent', 'max'] | Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2 predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent` always chooses the most recently seen value. +fp8_amax_history_len: 1 # The length of the amax history window used for scaling factor computation. +fp8_interval: 1 # DEPRECATED from TransformerEngine v1.8.0. This flag is ignored. Controls how often the scaling factor is recomputed. +fp8_margin: 0 # Margin for the scaling factor computation. +fp8_output_proj: false # If True, run the LM-head output projection with a TE ColumnParallelLinear under the MXFP8 autocast context. Only active when fp8=True and fp8_recipe='mxfp8'. fp8_param_gather: false # Keep the compute param in fp8 (do not use any other intermediate dtype) and perform the param all-gather in fp8. -fp8_quantizer_factory: null # Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. -fp8_recipe: delayed # choices: ['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom'] | Which fp8 recipe to use for FP8 tensors in the forward and backward pass -fp8_wgrad: true # Execute wgrad in higher precision even for FP8 runs +fp8_quantizer_factory: null # Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. Required when fp8_recipe is custom. +fp8_recipe: delayed # choices: ['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom'] | If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed' uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, 4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe. +fp8_wgrad: true # When set to False, override FP8 config options and do the wgrad computation in higher precision. fsdp_double_buffer: false # Enable double buffering for temporary memory needed for Megatron FSDP communications. Double-buffering the communication memory improves memory management efficiency by reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. This is required for user buffer registration and is enabled by default when using NCCL user buffers. +fsdp_manual_registration: false # Manually register the FSDP communication buffers to NCCL user buffer.This option is only effective when use-megatron-fsdp and use-nccl-ub is set. +ft_num_warmup_iters: 5 # Number of warmup iterations before monitoring step section and out-of-section timeouts. The first N iterations are excluded from timeout monitoring as they can be significantly slower than steady-state. Default: 5. Note: This feature is for Nvidia internal use only. full_validation: false # If set, each time validation occurs it uses the full validation dataset(s). This currently only works for GPT datasets! +fused_residual_rmsnorm: false # If True, fuses residual connection and RMSNorm backward pass when TE is used. global_batch_size: null # Training batch size. If set, it should be a multiple of micro-batch-size times data-parallel-size. If this value is None, then use micro-batch-size * data-parallel-size as the global batch size. This choice will result in 1 for number of micro-batches. glu_linear_offset: 0.0 # Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only used when gated_linear_unit is True +gpu_sniff_test_interval: null # Interval (in iterations) to run GPU performance sniff tests (GEMMs, all-to-all, send/recv). Flags any rank whose throughput differs from the mean by more than one standard deviation. If not specified, sniff tests are not run. grad_reduce_in_bf16: false # Reduce gradients in bfloat16. gradient_accumulation_fusion: true # Disable fusing gradient accumulation to weight gradient computation of linear layers gradient_reduce_div_fusion: true # If not set, fuse the division in gradient reduce. group_query_attention: false # Use group-query attention. grpo_clamp_eps_lower: 0.01 # Lower GRPO clipping bound. grpo_clamp_eps_upper: 0.01 # Upper GRPO clipping bound. In vanilla implementation, equals to the lower one. -grpo_default_temperature: 1.0 # Default temperature for model inference. -grpo_default_top_p: 0 # Default top-p for model inference. grpo_entropy_term_weight: 0.0 # Entropy term weight in GRPO loss. grpo_filter_groups_with_same_reward: false # Filter groups with same reward. grpo_group_size: 2 # Number of samples per a GRPO group. @@ -199,13 +247,14 @@ grpo_prompts_per_step: 32 # Number of GRPO groups (G in the paper). head_lr_mult: 1.0 # learning rate multiplier for head during finetuning heterogeneous_layers_config_encoded_json: null # This is encoded json string of the heterogeneous model configuration. Used to keep the content of the heterogeneous model specification in args when the model is loaded from a checkpoint. Use the format of the HuggingFace config files in llama nemotron models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json. heterogeneous_layers_config_path: null # Path to json file containing heterogeneous model configuration. Use the format of the HuggingFace config files in llama nemotron models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json. -hidden_dropout: 0.1 # Dropout probability for hidden state transformer. +hidden_dropout: 0.1 # Dropout probability for transformer hidden state. hidden_size: null # Transformer hidden size. -hierarchical_context_parallel_sizes: null # Degrees of the hierarchical context parallelism. Users should provide a list to specify the sizes for different levels. --hierarchical-context-parallel-sizes 2 4 indicates every two adjacent gpus forms the first level of cp groups and the cp ranks with the same odevity forms the second level of cp groups. -high_priority_stream_groups: [] # The communicator group names to use high priority streams. -hybrid_attention_ratio: 0.0 # Ratio of attention layers to total layers, in the range [0.0, 1.0]. -hybrid_mlp_ratio: 0.0 # Ratio of mlp layers to total layers, in the range [0.0, 1.0]. -hybrid_override_pattern: null # Force a specific hybrid layer pattern. The valueshould be a string of characters chosen fromcore.ssm.mamba_hybrid_layer_allocation.Symbols.If a value greater than 0.0 is supplied to any of the hybrid ratio arguments, then the number of each typeof layer in the override pattern must match number inthe overidden pattern +hierarchical_context_parallel_sizes: null # Degrees of the hierarchical context parallelism. Users should provide a list to specify the sizes for different levels. Taking the a2a+p2p cp comm type as example, it contains groups of two levels, so the first value of the list indicates the group size of the a2a communication type, and the second value indicates the group size of the p2p communication type. +high_priority_a2a_comm_stream: false # If True, the communication stream created by set_streams for combined 1f1b a2a overlap is created with CUDA high priority. +high_priority_stream_groups: [] # Specify which communicator groups should use high priority streams during creation. Assigning high priority to communication streams ensures that communication kernels are scheduled with higher priority, minimizing the exposed communication when it is overlapped with other computation kernels. +hybrid_context_parallel: false # If true, enables hybrid context parallel. This is used to balance the workload of each CP rank when we use packed samples with variable sequence lengths. Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. +hybrid_layer_pattern: null # Specify a hybrid layer pattern using M (mamba), G (gdn), * (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline stage boundaries for flexible virtual pipeline parallel (fVPP). Use / to separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". When this flag is used, it is the sole indicator that a hybrid model is being run. +hybrid_override_pattern: null # Deprecated. Use --hybrid-layer-pattern instead. If specified, its value will be forwarded to --hybrid-layer-pattern. hysteresis: 2 # hysteresis for dynamic loss scaling ict_head_size: null # Size of block embeddings to be used in ICT and REALM (paper default: 128) ict_load: null # Directory containing an ICTBertModel checkpoint @@ -214,23 +263,43 @@ img_w: 224 # Image height for vision classification task indexer_batch_size: 128 # How large of batches to use when doing indexing jobs indexer_log_interval: 1000 # After how many batches should the indexer report progress inference_batch_times_seqlen_threshold: -1 # If (batch-size * sequence-length) is smaller than this thresholdthen batches will not be split up for pipelining.Requires setting --pipeline-model-parallel-size > 1.Setting this to -1 indicates that batch pipelining is not used. -inference_coordinator_port: 12346 # This port will be used to setup the inference coordinator on node-0 +inference_coordinator_port: null # This port will be used to setup the inference coordinator on node-0 +inference_cuda_graph_all_prefills: false # Extend prefill/mixed CUDA graph capture up to `max_tokens`. By default, all graphs are limited by the decode limit of `max_requests * (num_speculative_tokens + 1)`. +inference_cuda_graph_scope: null # choices: ['none', 'layer', 'block'] | Controls the CUDA graph scope during inference. When unset, the effective default is derived from cuda_graph_impl: "local" -> "layer", all other impls -> "none". "none": inference runs in eager mode (no CUDA graphs). "layer": inference graphs are owned at the module/layer boundary, e.g. TransformerLayer or MambaLayer. "block": inference graphs are owned by the enclosing block, e.g. TransformerBlock or HybridBlock. Currently supported combinations are: cuda_graph_impl="local" -> "layer" or "block"; all other cuda_graph_impl values -> "none". +inference_disable_ep_consensus: false # Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. Pause/unpause take effect as soon as the signal is delivered to a rank. Only safe when EP coordination is not required (e.g. ep_world_size == 1). +inference_disable_triton_nvls_kernels: false # If true, disables the use of Triton NVLS kernels during inference. inference_dynamic_batching: false # Enable dynamic batching mode. inference_dynamic_batching_block_size: 256 # KV cache block size. It should be a multiple of 256 -inference_dynamic_batching_buffer_size_gb: 40.0 # Amount of on-GPU memory allocated for the KV cache. The total amount of memory allocated for the KV cache (CPU + GPU memory) depends on the value set for the unified virtual memory (UVM) level (via `--inference-dynamic-batching-unified-memory-level`).If the UVM level is 0, then only GPU memory is used and the total memory equals `buffer_size_gb`. If the UVM level is 1, then additional memory is utilized on the CPU and the total memory equals `2 * buffer_size_gb`. -inference_dynamic_batching_cuda_graph_max_tokens: 16384 # Maximum number of tokens to capture in a cuda graph. +inference_dynamic_batching_buffer_size_gb: 40.0 # Amount of on-GPU memory allocated for the KV cache. The total amount of memory allocated for the KV cache (CPU + GPU memory) depends on the value set for the unified virtual memory (UVM) level (via `--inference-dynamic-batching-unified-memory-level`).If the UVM level is 0, then only GPU memory is used and the total memory equals `buffer_size_gb`. If the UVM level is 1, then additional memory is utilized on the CPU and the total memory equals `buffer_size_gb + paused_buffer_size_gb`. inference_dynamic_batching_cuda_graph_mixed_prefill_count: 16 # Number of mixed prefill requests to capture in a cuda graph. +inference_dynamic_batching_enable_prefix_caching: false # Enable/disable prefix caching for dynamic batching inference. When disabled, KV cache blocks cannot be shared between requests with identical prompt prefixes. +inference_dynamic_batching_mamba_memory_ratio: null # Percentage of memory buffer to allocate for Mamba states. If not specified, allocates Mamba state tensors for each KV cache block. Only used for hybrid models. +inference_dynamic_batching_max_requests: null # Override the inference context's `max_requests`. By default, `max_requests` is set to the number of blocks in the context's memory buffer. inference_dynamic_batching_max_tokens: null # Override the inference context's default `max_tokens`. -inference_dynamic_batching_num_cuda_graphs: 16 # Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to `max_requests`. (See `dynamic_context.py` for details on how `max_requests` is computed). Due to rounding, the actual number of cuda graphs may not equal this argument. +inference_dynamic_batching_num_cuda_graphs: 16 # Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to `max_requests`. (See `dynamic_context.py` for details on how `max_requests` is computed). Due to rounding, the actual number of cuda graphs may not equal this argument.The user can also pass -1, in which case we automatically determine the number of graphs to capture based on the `max_requests`. +inference_dynamic_batching_paused_buffer_size_gb: null # Amount of memory reserved for paused requests in the dynamic inference context. Active requests are paused when there are not enough active blocks available to continue generating a request. +inference_dynamic_batching_prefix_caching_coordinator_policy: first_prefix_block # choices: ['longest_prefix', 'first_prefix_block', 'round_robin'] | Coordinator routing policy for prefix caching. "first_prefix_block" (default) routes based on the first block hash only. "longest_prefix" routes to the rank with the longest matching prefix. "round_robin" ignores prefix affinity and cycles through ranks. +inference_dynamic_batching_prefix_caching_eviction_policy: ref_zero # choices: ['ref_zero', 'lru'] | Eviction policy for prefix caching blocks. "ref_zero" (default) immediately returns blocks to the free pool when ref_count hits 0. "lru" keeps blocks cached and evicts via LRU only when space is needed. +inference_dynamic_batching_prefix_caching_mamba_gb: null # GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid models. When set, Mamba states at block boundaries are cached for reuse. +inference_dynamic_batching_prefix_caching_routing_alpha: 0.5 # Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * normalized_load. Higher alpha favors prefix cache hits; lower alpha favors load balance. Default: 0.5. +inference_dynamic_batching_sampling_backend: torch # choices: ['torch', 'flashinfer'] | Which sampling kernels to use during inference. Falls back to "torch" with a warning if "flashinfer" is requested but the package is not installed. +inference_dynamic_batching_track_generated_token_events: false # Track per-token events with timestamps for each generated token. When enabled, each generated token creates a GENERATED_TOKEN event with a timestamp, useful for per-token latency analysis. inference_dynamic_batching_track_paused_request_events: false # Track paused request ids by adding 'paused' events to each request's event history. This has a very minor impact on latency. -inference_dynamic_batching_unified_memory_level: 1 # choices: [0, 1] | Set unified memory usage within the dynamic inference context. The levels are: 0) no unified memory, 1) allocate `memory_buffer` in unified memory. Eventually, additional levels will be included to control other tensors within the context. -inference_max_batch_size: 8 # Maximum batch size for inference. +inference_dynamic_batching_unified_memory_level: 0 # choices: [0, 1] | Set unified memory usage within the dynamic inference context. The levels are: 0) no unified memory, 1) allocate `memory_buffer` in unified memory. Eventually, additional levels will be included to control other tensors within the context. +inference_fuse_tp_communication: false # If true, uses a fused reduce-scatter-residual-norm-allgather kernel during inference. +inference_grouped_gemm_backend: vllm # choices: ['flashinfer', 'torch', 'vllm'] | Specifies the backend to use for grouped GEMM operations during inference. Options: - 'flashinfer': Uses FlashInfer cutlass_fused_moe. Not compatible with MXFP8. - 'torch': Uses torch.nn.functional.grouped_mm (mcore_fused_moe with Triton kernels). Supports both BF16 and MXFP8. - 'vllm': Uses vLLM's Triton fused MoE kernel (BF16). Avoids physical token permutation via indirect addressing. +inference_logging_step_interval: 0 # Step interval for logging inference metrics. Default to 0 to disable inference logging. +inference_max_requests: 8 # Maximum number of requests for inference. inference_max_seq_length: 2560 # Maximum sequence length expected for inference (prefill + decode). +inference_moe_disable_fused_quant_kernels: false # When False (default), use fused kernels that combine permute/activation with MXFP8 quantization + swizzle into a single kernel launch. Only applies when fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel launches (useful for debugging). +inference_moe_token_dispatcher_type: nvls # choices: ['nccl', 'nvls'] | Token dispatcher to use for MoE expert parallelism during inference. - 'nccl': AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires decode-only CUDA graphs (forced automatically). - 'nvls': Variable-count AllGather-V/ReduceScatter-V via NVLS multimem kernels. Requires Hopper+ GPUs with NVLink and symmetric memory. Default. Only applies when transformer_impl='inference_optimized' and EP > 1. inference_rng_tracker: false # Use a random number generator configured for inference. -inference_wandb_logging_step_interval: 0 # Step interval for logging inference metrics to wandb. Default to 0 to disable inference wandb logging. -init_method_std: 0.02 # Standard deviation of the zero mean normal distribution used for weight initialization. +inference_text_gen_server_logging: false # Enable per-request logging in the inference text generation server. +inference_use_synchronous_zmq_collectives: false # Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs. +inference_wandb_logging: false # Enable inference wandb logging. +init_method_std: 0.02 # Standard deviation of the zero mean normal for the default initialization method, not used if init_method and output_layer_init_method are provided. init_method_xavier_uniform: false # Enable Xavier uniform parameter initialization -init_model_with_meta_device: false +init_model_with_meta_device: false # If True, initializes the model with the meta device. This is helpful for training of very large models. This feature is only works when megatron fsdp is turned on. initial_loss_scale: 4294967296 # Initial loss-scale for dynamic loss scaling. inprocess_active_world_size: 1 # The number of ranks initially executing the workload. The remaining ranks from the allocation are set aside as warm reserve. inprocess_barrier_timeout: 120 # Timeout (in seconds) for internal distributed barrier @@ -248,26 +317,34 @@ inprocess_progress_watchdog_interval: 1.0 # Interval (in seconds) for automatic inprocess_restart: false # Enables in-process restart. inprocess_soft_timeout: 60 # Soft progress timeout (in seconds). inprocess_termination_grace_time: 1 # Interval (in seconds) between SIGTERM and SIGKILL issued on hard timeout -is_hybrid_model: false # Indicates whether the model is a hybrid model. +is_hybrid_model: false # Indicates whether this is a hybrid model. iter_per_epoch: 1250 # iterations per epoch -iterations_to_skip: [] # List of iterations to skip, empty by default. +iterations_to_skip: [] # List of 1-indexed iterations to skip during training, empty by default. keep_fp8_transpose_cache: false # If set, keep the fp8 transpose cache when using Megatron FSDP. -kitchen_config_file: null # Use the config .yaml file at the specified location to configure kitchen quantization. -kitchen_recipe_number: null # Use a default kitchen recipe for all layers as defined by QAT_PARAMS index -kv_channels: null # Projection weights dimension in multi-head attention. This is set to args.hidden_size // args.num_attention_heads if not provided. +kitchen_attention_backend: sdpa # choices: ['sdpa', 'fa'] | Which kitchen attention backend to use when use_kitchen_attention=True. "sdpa" for KitchenDotProductAttention, "fa" for KitchenFlashAttention. +kv_channels: null # Projection weights dimension in multi-head attention. This is set to hidden_size // num_attention_heads if not provided. kv_lora_rank: 32 # Rank of Key and Value tensors' low rank representation. langrl_env_config: null # Path to YAML config file for RL environment configuration. -langrl_external_server: false -langrl_inference_server_conversation_template: null # Conversation template, if using a chat server. -langrl_inference_server_type: inplace_megatron # choices: ['inplace_megatron', 'inplace_megatron_chat'] | Type of inference server to use. -lazy_mpu_init: null # 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. -legacy_tokenizer: false # To use legacy tokenizer system. +layernorm_epsilon: 1.0e-05 # Epsilon value for any LayerNorm/RMSNorm operations. +layernorm_zero_centered_gamma: false # If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves numerical stability. +lazy_mpu_init: false # 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. +linear_attention_freq: null # Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. Accepts either: - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer - A string containing a Python list expression that defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an LA layer and 0 indicates a SDPA layer. Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, "([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice. +linear_conv_kernel_dim: 4 # Conv kernel dimension for the gated delta net. +linear_key_head_dim: 128 # Query and key head dimension for the gated delta net. +linear_num_key_heads: 16 # Number of query and key heads for the gated delta net. +linear_num_value_heads: 32 # Number of value and gate heads for the gated delta net. +linear_value_head_dim: 128 # Value and gate head dimension for the gated delta net. +lion_beta1: 0.95 # First beta coefficient for Lion optimizer (used in sign update). Default: 0.95. +lion_beta2: 0.98 # Second beta coefficient for Lion optimizer (used in momentum EMA update). Default: 0.98. load: null # Directory containing a model checkpoint. -load_main_params_from_ckpt: null # Load main parameters from checkpoint directly. +load_main_params_from_ckpt: false # Load main parameters from checkpoint. When loading a model from a checkpoint without loading the optimizer, the model parameters are updated but for fp16 optimizer with main parameters, the main parameters need to also be updated. local_rank: 0 # local rank passed from distributed launcher. -log_energy: false # If set, log energy consumption (in Joules) +log_device_memory_used: false # Log device memory used (as reported by nvidia-smi). +log_energy: false # If set, log energy consumption (in Joules). log_interval: 100 # Report loss and timing interval. log_loss_scale_to_tensorboard: true # Disable loss-scale logging to tensorboard. +log_max_attention_logit: false # Enable max attention logit logging to tensorboard. +log_memory_interval: null # Report memory interval. log_memory_to_tensorboard: false # Enable memory logging to tensorboard. log_num_zeros_in_grad: false # If set, calculate and log the number of zeros in gradient. log_params_norm: false # If set, calculate and log parameters norm. @@ -281,8 +358,8 @@ logging_level: null # Set default logging level loss_scale: null # Static loss scaling, positive power of 2 values can improve fp16 convergence. If None, dynamicloss scaling is used. loss_scale_window: 1000 # Window over which to raise/lower dynamic scale. lr: null # Initial learning rate. Depending on decay style and initial warmup, the learning rate at each iteration would be different. -lr_decay_iters: null # number of iterations to decay learning rate over, If None defaults to `--train-iters` -lr_decay_samples: null # number of samples to decay learning rate over, If None defaults to `--train-samples` +lr_decay_iters: null # number of iterations to decay learning rate over, If None defaults to train iters +lr_decay_samples: null # number of samples to decay learning rate over, If None defaults to train samples lr_decay_style: linear # choices: ['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD'] | Learning rate decay function. lr_warmup_fraction: null # fraction of lr-warmup-(iters/samples) to use for warmup (as a float) lr_warmup_init: 0.0 # Initial value for learning rate warmup. The scheduler starts warmup from this value. @@ -291,80 +368,124 @@ lr_warmup_samples: 0 # number of samples to linearly warmup learning rate over. lr_wsd_decay_iters: null # number of iterations for the annealing phase in the wsd schedule lr_wsd_decay_samples: null # number of samples for the annealing phase in the wsd schedule lr_wsd_decay_style: exponential # choices: ['exponential', 'linear', 'cosine', 'minus_sqrt'] | Decay style for the annealing phase of WSD -main_grads_dtype: fp32 # choices: ['fp32', 'bf16'] | Dtype of main grads when enabling precision-aware-optimizer -main_params_dtype: fp32 # choices: ['fp32', 'fp16'] | Dtype of main params when enabling precision-aware-optimizer +main_grads_dtype: fp32 # choices: ['fp32', 'bf16'] | Dtype of main grads when enabling precision-aware-optimizer. +main_params_dtype: fp32 # choices: ['fp32', 'fp16'] | Dtype of main params when enabling precision-aware-optimizer. make_vocab_size_divisible_by: 128 # Pad the vocab size to be divisible by this value.This is added for computational efficieny reasons. -mamba_head_dim: 64 # Head dimension for Mamba layers. -mamba_num_groups: 8 # Number of groups for Mamba layers. -mamba_num_heads: null # Number of heads for Mamba layers.If not set, then the number of heads will be --hidden-size * expand // --mamba-head-dim -mamba_state_dim: 128 # State dimension for Mamba layers. +mamba_head_dim: 64 # The dimensionality of the heads in the Mamba layers. +mamba_inference_conv_states_dtype: bf16 # choices: ['bf16', 'fp16', 'fp32'] | Dtype for the Mamba inference conv states tensor +mamba_inference_ssm_states_dtype: bf16 # choices: ['bf16', 'fp16', 'fp32'] | Dtype for the Mamba inference SSM states tensor +mamba_num_groups: 8 # The number of groups used in Mamba layers. +mamba_num_heads: null # The number of heads used in Mamba layers. If None, the number of heads will be hidden_size * expand // mamba_head_dim. +mamba_state_dim: 128 # The dimensionality of the state representation in Mamba layers. manual_gc: false # 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. -manual_gc_eval: true # When using manual garbage collection, disable garbage collection at the start and the end of each evaluation run. -manual_gc_interval: 0 # Training step interval to trigger manual garbage collection. When the value is set to 0, garbage collection is not triggered between training steps. +manual_gc_eval: true # When using manual garbage collection, this controls garbage collection at the start and the end of each evaluation run. +manual_gc_interval: 0 # Training step interval to trigger manual garbage collection. Values > 0 will trigger garbage collections between training steps. mask_factor: 1.0 # mask size scaling parameter mask_prob: 0.15 # Probability of replacing a token with mask. mask_type: random # choices: ['random', 'row'] | mask types masked_softmax_fusion: true # Disable fusion of query_key_value scaling, masking, and softmax. max_position_embeddings: null # Maximum number of position embeddings to use. This is the size of position embedding. +max_seqlen_per_dp_cp_rank: null # Maximum sequence length per DPxCP rank. This is the maximum sequence length each rank can handle without overflowing the memory. Typically, a good starting point is to set this to maximum sequence length / context parallel size. This is used to calculate the number and length of sub-samples assigned to each rank when using hybrid_context_parallel. max_tokens_to_oom: 12000 # Maximum number of tokens during inferencetokens here is # in prompt + # to generateAllows us to throw an error before OOM crashes server +megatron_fsdp_enable_fine_grained_param_gather: false # If set, enables fine-grained parameter gathering for Megatron-FSDP. This allows greater overlap between parameter all-gather operations and forward computation, at the cost of additional communication calls. For MXFP8, this helps save memory during fine-grained activation recomputation, because MXFP8 forward and backward passes use different parameter representations (rowwise data for forward, colwise data for backward). Only the rowwise parameters of modules involved in recomputation will be unsharded. +megatron_fsdp_grad_comm_dtype: auto # choices: ['fp32', 'fp16', 'bf16', 'auto'] | When using Megatron-FSDP, this controls the data-type used when communicating model gradients during FSDP. If 'auto', then the main gradient data-type will be used for the gradient communication / reduction data-type. When using NCCL v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels. +megatron_fsdp_main_grads_dtype: auto # choices: ['fp32', 'bf16', 'fp16', 'auto'] | Data type for the main gradient buffer utilized for distributed optimization with Megatron-FSDP. If 'auto', then the native model gradient data-type will be used for the main gradient / accumulation data-type. Replaces --main-grads-dtype. +megatron_fsdp_main_params_dtype: fp32 # choices: ['fp32', 'bf16', 'fp16', 'auto'] | Data type for the main weight buffer utilized for distributed optimization and quantization with Megatron-FSDP. If 'auto', then the native model parameter data-type will be used for the main weight data-type. Replaces --main-params-dtype. memory_snapshot_path: snapshot.pickle # Specifies where to dump the memory history pickle. merge_file: null # Path to the BPE merge file. +metadata_path: null # Path to the tokenizer metadata file in json format. micro_batch_size: null # Batch size per model instance (local batch size). Global batch size is local batch size times data parallel size times number of micro batches. -microbatch_group_size_per_vp_stage: null # Number of contiguous microbatches per virtual pipeline stage +microbatch_group_size_per_vp_stage: null # This value specifies the number of micro-batches that are executed at a time for a given virtual stage (both forward and backward). Default (in __post_init__() method below) to pipeline_parallel_size which specifies a depth-first schedule. Example: for PP=2 VP=2, when microbatch_group_size_per_vp_stage=2, num_microbatches = 4, we have rank 0 | 0 1 0 1 2 3 2 3 rank 1 | 0 1 0 1 2 3 2 3 When microbatch_group_size_per_vp_stage=3, num_microbatches = 5, we have rank 0 | 0 1 2 0 1 2 3 4 3 4 rank 1 | 0 1 2 0 1 2 3 4 3 4 mid_level_dataset_surplus: 0.005 # The sample surplus to build for the mid-level datasets(s) min_loss_scale: 1.0 # Minimum loss scale for dynamic loss scaling. min_lr: 0.0 # Minimum value for learning rate. The schedulerclip values below this threshold. -mlp_chunks_for_prefill: 1 # Number of chunks along sequence dimension for MLP computation during prefill +min_offloaded_tensor_size: 1048576 # The minimum size of the tensor to be offloaded. +mla_down_proj_fusion: false # Enable fused q/kv down-projection and fused input layernorm when backend supports. Otherwise fall back to the unfused MLA. +mlp_chunks_for_prefill: 1 # The number of chunks along the sequence dimension to use for MLP computation during prefill. +mlp_chunks_for_training: 1 # The number of chunks along the sequence dimension to use for MLP computation during training. mmap_bin_files: true # Disable mmap-ing of .bin files. mock_data: false # Skip data loading and validation and opt for artificial generation of mock data when an implementation is available. model_parallel_size: null # Old model parallel argument, do not use. Use --tensor-model-parallel-size instead. -moe_apply_probs_on_input: false # Apply probs before mlp activation for moe routing. +moe_apply_probs_on_input: false # Apply probs on input of experts instead of applying after activation and glu. moe_aux_loss_coeff: 0.0 # Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended. moe_deepep_num_sms: 20 # Number of SMs to use for DeepEP. -moe_enable_deepep: false # DEPRECATED: Please use --moe-flex-dispatcher-backend=deepep instead. -moe_expert_capacity_factor: null # The capacity factor for each expert, None means no token will be dropped. -moe_extended_tp: false # Deprecated. Use --expert-tensor-parallel-size instead. -moe_ffn_hidden_size: null # The hidden size of each expert's feed-forward network (ffn). If not specified, defaults to the ffn_hidden_size. -moe_flex_dispatcher_backend: deepep # choices: ['deepep', 'hybridep'] | The backend to use for flex token dispatcher. The default is "deepep". Options are "deepep" and "hybridep". -moe_grouped_gemm: false # 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. -moe_hybridep_num_sms: 16 # Number of SMs to use for HybridEP. +moe_enable_deepep: false # [Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models. +moe_enable_routing_replay: false # If True, enable the routing replay feature for MoE layers. +moe_expert_capacity_factor: null # moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token will be dropped. The default is None. +moe_expert_rank_capacity_factor: null # moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens exceeding this budget will be dropped. None means no token will be dropped. The default is None. +moe_ffn_hidden_size: null # MoE Feed-Forward Network hidden size. If not specified, defaults to the ffn_hidden_size. +moe_flex_dispatcher_backend: deepep # choices: ['deepep', 'hybridep'] | [Experimental] The backend to use for flex token dispatcher. The default is "deepep". Options are "deepep" and "hybridep". Currently only "hybridep" backend supports the MNNVL case. +moe_grouped_gemm: false # When there are multiple experts per rank, compress multiple local (potentially small) gemms in a single kernel launch to improve the utilization and performance by leveraging the Grouped GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). +moe_hybridep_num_blocks_permute: null # Number of CUDA thread blocks for the permute part in HybridEP. When permute_fusion_into_hybridep is True, this sets the number of SMs for the permute part (only 1 block per SM). +moe_hybridep_num_blocks_unpermute: null # Number of CUDA thread blocks for the unpermute part in HybridEP. When permute_fusion_into_hybridep is True, this sets the number of SMs for the unpermute part (only 1 block per SM). +moe_hybridep_num_sms: null # Number of SMs to use for HybridEP. None uses the default from DeepEP. In pure NVL scenarios, 16 SMs can generally achieve good bandwidth. +moe_hybridep_num_sms_preprocessing: 108 # Number of SMs to use for HybridEP preprocessing (metadata scan kernel). moe_input_jitter_eps: null # Add noise to the input tensor by applying jitter with a specified epsilon value. -moe_layer_freq: 1 # Frequency between MoE layers and Dense layers. Accepts either: - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers - A string containing a Python list expression that defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an expert layer and 0 indicates a dense layer. Examples: "([0]+[1]*23)": 1 dense layer followed by 23 experts layers, "([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice. -moe_layer_recompute: false # Enable checkpointing for moe_layer, should be used when memory is not sufficient. Deprecated. Use "--recompute-granularity selective --recompute-modules moe" instead. -moe_pad_expert_input_to_capacity: false # Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set. -moe_pad_experts_for_cuda_graph_inference: false # some MoE routers have a D2H sync that will break cuda graphs. If this flag is set the router will switch to dropping and padding during decode time which does not have a D2H sync. The capacity factor is set to the max that an expert could see during inference so no tokens are actually dropped. +moe_latent_size: null # Latent projection dimension for MoE. If None, MoE latent projections are not used. +moe_layer_freq: 1 # Frequency between MoE layers and Dense layers. Accepts either: - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers - A string containing a Python list expression that defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an expert layer and 0 indicates a dense layer. Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, "([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice. +moe_layer_recompute: false # Memory optimization: checkpointing moe_layer to save actiavtion memory. +moe_mlp_glu_interleave_size: null # When set, GLU activations in the MoE grouped MLP layer will use a block interleaved format. Instead of interpreting the input tensor as a concatenation of gates and linear units, it will be interpreted as alternating blocks of gates and linear units. This data format is experimental and primarily intended to enable advanced fused kernels. +moe_pad_expert_input_to_capacity: false # moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match the expert capacity length, effective only after the moe_expert_capacity_factor is set. The default setting is False. +moe_pad_experts_for_cuda_graph_inference: false # moe_pad_experts_for_cuda_graph_inference (bool): If True, the router will switch to dropping and padding during decode time which does not have a D2H sync. The capacity factor is set to the max that an expert could see during inference so no tokens are actually dropped. The default setting is False. +moe_paged_stash: false # If True, enable paged stash for all routed-expert activations needed for backward +moe_paged_stash_buffer_size_factor_cpu: 0.0 # Scale factor for paged stash host buffer. 0 disables host buffer. Same sign convention as moe_paged_stash_buffer_size_factor_cuda: positive = avg-based, negative = actual-max; scale = abs(factor). +moe_paged_stash_buffer_size_factor_cuda: 1.1 # Scale factor for paged stash CUDA buffer allocation. Sign selects sizing: positive = avg-based, negative = actual-max. Magnitude is headroom (e.g. 1.10 = 10%). +moe_paged_stash_page_size: 64 # Number of tokens per page for paged stash memory management. moe_per_layer_logging: false # Enable per-layer logging for MoE, currently supports auxiliary loss and z loss. moe_permute_fusion: false # Fuse token rearrangement ops during token dispatching. -moe_router_bias_update_rate: 0.001 # Expert bias update rate in the aux-loss-free load balancing strategy. The expert bias is updated based on the number of assigned tokens to each expert in a global batch, where the bias is increased for the experts with less assigned tokens and decreased for the experts with more assigned tokens. The default value 1e-3 is same as that used in DeepSeekV3. -moe_router_dtype: null # choices: ['fp32', 'fp64'] | Data type for routing computation and expert output weighted averaging. Fp32/fp64 enhances numerical stability, especially with numerous experts. The perf impact should be negligible when used with permute fusion. None means no changes for dtype. -moe_router_enable_expert_bias: false # TopK routing with dynamic expert bias in the aux-loss-free load balancing strategy. The routing decision is based on the sum of the routing scores and the expert bias. See https://arxiv.org/abs/2408.15664 for details. -moe_router_force_load_balancing: false # [Experimental] Force override routing to balance token distribution using random logits for MoE routers, supporting naive top-k and group-limited top-k. This experimental feature is for benchmarking purposes only! +moe_permute_fusion_into_hybridep: false # Fuse token rearrangement ops during token dispatching for HybridEP. +moe_router_bias_update_rate: 0.001 # The expert bias is updated based on the number of assigned tokens to each expert in a global batch, where the bias is increased for the experts with less assigned tokens and decreased for the experts with more assigned tokens. The default value 1e-3 is same as that used in DeepSeekV3. +moe_router_dtype: null # choices: ['fp32', 'fp64'] | Data type for routing and expert output weighted averaging. Using fp32 or fp64 can improve stability especially when the number of experts is large (e.g. finegrained-moe). None means no changes for dtype. +moe_router_enable_expert_bias: false # TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. The routing decision is based on the sum of the routing scores and the expert bias. See https://arxiv.org/abs/2408.15664 for details. +moe_router_force_biased: null # Apply random expert bias in normal distribution with specified std to router logits. Shared seed across all ranks ensures identical bias. If positive, generates new random bias each forward pass. If negative, generates bias once per layer and reuses it (abs value is std). This is an experimental feature for benchmarking purposes. +moe_router_force_load_balancing: false # [Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark. moe_router_fusion: false # Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above. moe_router_group_topk: null # Number of selected groups for group-limited routing. moe_router_load_balancing_type: aux_loss # choices: ['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'] | Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss". -moe_router_num_groups: null # Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1) Experts are divided into equal-sized groups, 2) For each token, a subset of groups are selected based on routing scores (sum of top-2 expert scores within each group), 3) From these selected groups, moe_router_topk experts are chosen.Two common use cases: 1) Device-limited routing: Set equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437) -moe_router_padding_for_fp8: false # [Compatibility alias for --moe-router-padding-for-quantization] Enabling this will also enable --moe-router-padding-for-quantization. -moe_router_padding_for_quantization: false # Pad the routing_map to make sure the number of tokens each expert received is a multiple of 16/32 for FP8/FP4 precision. It is suggested to enable this for dropless training with FP8/FP4 precision when num_local_experts > 1. This is a more efficient way to pad for FP8/FP4 which eliminates the explicit padding in the GroupedMLP layer. -moe_router_pre_softmax: false # Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k. -moe_router_score_function: softmax # choices: ['softmax', 'sigmoid'] | Score function for MoE TopK routing. Can be "softmax" or "sigmoid". -moe_router_topk: 2 # Number of experts to route to for each token. The default is 2. -moe_router_topk_scaling_factor: null # Scaling factor for routing score in top-k selection, only works when --moe-router-pre-softmax enabled. Defaults to None, which means no scaling. -moe_shared_expert_intermediate_size: null # Shared expert total ffn hidden size. It should be equal to "num_shared_experts * ffn_size_of_each_shared_expert" if there are multiple shared experts. None means no shared expert. By default, the shared experts execute before the router. However, when --moe-shared-expert-overlap or --overlap-moe-expert-parallel-comm is set, the shared experts execute after the router, before the routed experts. This makes the gradients from the router and the shared experts added in different orders to the hidden_states, causing minor numerical differences in the hidden_states gradient. +moe_router_num_groups: null # Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1. Experts are divided into 'moe_router_num_groups' equal-sized groups 2. For each token, 'moe_router_group_topk' groups are selected based on sum of top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group 3. From these selected groups, 'moe_router_topk' individual experts are chosen Two common use cases: - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) - Node-limited routing: Set 'moe_router_num_groups' equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437) +moe_router_padding_for_fp8: false # [Compatibility alias for moe_router_padding_for_quantization] Enabling this will also enable moe_router_padding_for_quantization. +moe_router_padding_for_quantization: false # Whether to pad the routing_map to make sure the number of tokens each expert receives is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit padding in the GroupedMLP layer. +moe_router_pre_softmax: false # Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k. +moe_router_score_function: softmax # choices: ['softmax', 'sigmoid', 'sqrtsoftplus'] | Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus". +moe_router_topk: 2 # Number of experts to route to for each token. +moe_router_topk_scaling_factor: null # Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax enabled. Defaults to None, which means no scaling. +moe_shared_expert_gate: false # Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set. +moe_shared_expert_intermediate_size: null # Shared expert total ffn hidden size. It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if there are multiple shared experts. None means no shared expert. By default, the shared experts execute before the router. However, when moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set, the shared experts execute after the router, before the routed experts. This makes the gradients from the router and the shared experts added in different orders to the hidden_states, causing minor numerical differences in the hidden_states gradient. moe_shared_expert_overlap: false # Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute before the router. Only effective when moe-shared-expert-intermediate-size is set. -moe_token_dispatcher_type: allgather # choices: ['allgather', 'alltoall', 'flex'] | The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather', 'alltoall'. We recommend using 'alltoall' when applying expert parallelism. For more information, please refer to the documentation in core/moe/README. +moe_single_grouped_bias: false # When using TE GroupedLinear for MoE experts, store expert biases as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and ``add_bias_linear=True``. +moe_single_grouped_weight: false # When using TE GroupedLinear for MoE experts, store expert weights as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``. +moe_token_dispatcher_type: allgather # choices: ['allgather', 'alltoall', 'flex'] | The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather','alltoall' and 'flex'. moe_token_drop_policy: probs # choices: ['probs', 'position'] | 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. moe_upcycling_granularity: 1 # This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy. -moe_use_legacy_grouped_gemm: false # Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will be deprecated soon. moe_use_upcycling: false # 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. -moe_z_loss_coeff: null # Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended. -mrope_section: null # Multimodal rope section is for channel dimension, empty by default. +moe_z_loss_coeff: null # Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended. +mrope_section: null # Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. mscale: 1.0 # Mscale for YaRN RoPE in multi-latent attention. mscale_all_dim: 0.0 # Mscale all dimensions for YaRN RoPE in multi-latent attention. -mtp_loss_scaling_factor: 0.1 # Scaling factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP losses across all depths, and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. -mtp_num_layers: null # Number of Multi-Token Prediction (MTP) Layers.MTP extends the prediction scope to multiple future tokens at each position.This MTP implementation sequentially predict additional tokens by using D sequential modules to predict D additional tokens. -multi_latent_attention: false # Use multi-latent attention for model. -multiple_validation_sets: false # If set, multiple datasets listed in the validation split are evaluated independently with a separate loss for each dataset in the list. This argument requires that no weights are included in the list -nccl_all_reduce_for_prefill: false # When using symmeric all reduce kernels this will use regular nccl kernels for prefill. This can be more effecient when prefill is large as the nccl kernels can be more bandwith optimized +mtp_hybrid_override_pattern: null # DEPRECATED: Use unified hybrid_layer_pattern instead. Legacy argument for loading old checkpoints. Force a specific hybrid layer pattern for MTP layers. +mtp_loss_scaling_factor: 0.1 # Weighting factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP losses across all depths, and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. +mtp_num_layers: null # Number of Multi-Token Prediction (MTP) Layers. MTP extends the prediction scope to multiple future tokens at each position. This MTP implementation sequentially predict additional tokens by using D sequential modules to predict D additional tokens. +mtp_standalone: false # This will be set automatically according to the pipeline layout, and will be set to True if MTP is in a separate vpp stage. +mtp_use_repeated_layer: false # Use a single MTP layer repeatedly instead of multiple separate layers. +multi_latent_attention: false # Whether to use multi-latent attention. +multiple_validation_sets: false # If set, multiple datasets listed in the validation split are evaluated independently with a separate loss for each dataset in the list. This argument requires that no weights are included in the list. +muon_coefficient_type: quintic # Newton-Schulz coefficient type for the Muon optimizer. Valid types are discovered from the installed emerging_optimizers package (e.g. simple, quintic, polar_express, aol). Validated at optimizer creation time. +muon_extra_scale_factor: 1.0 # Additional scale factor for the muon update +muon_fp32_matmul_prec: medium # choices: ['low', 'medium', 'high'] | FP32 matmul precision for Newton-Schulz iteration +muon_momentum: 0.9 # Momentum factor for Muon optimizer +muon_nesterov: false # Whether to use Nesterov-style momentum in the internal SGD +muon_num_ns_steps: 5 # Number of Newton-Schulz steps for Muon optimizer +muon_scalar_optimizer: adam # choices: ['adam', 'lion'] | Optimizer for scalar parameters (embeddings, biases, norms) when using muon. Defaults to adam. +muon_scale_mode: spectral # choices: ['spectral', 'unit_rms_norm', 'shape_scaling'] | Scale mode for Muon optimizer. With MuP, set --muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, or set --muon-scale-mode spectral to keep spectral scaling. +muon_split_qkv: true # Whether to split QKV parameters for Muon optimizer +muon_tp_mode: blockwise # choices: ['blockwise', 'duplicated', 'distributed'] | How to perform NS calculation for tensor model parallel weights +mup_attn_scale_power: 1.0 # Power for attention scaling: softmax_scale = 1 / (kv_channels ** mup_attn_scale_power). 0.5 = standard attention (1/sqrt(d_head)), 1.0 = MuP attention (1/d_head). Default: 1.0 (MuP scaling when use_mup is True). Set to 0.5 for standard scaling. +mup_base_head_dim: null # Base head dimension for MuP attention scaling. When set, softmax_scale = sqrt(mup_base_head_dim) / (kv_channels ** mup_attn_scale_power). Set to base model's d_head (e.g., 64) to match standard 1/sqrt(d_head) scaling at the base width, ensuring non-MuP compatibility for that specific value. +mup_base_hidden_size: null # Base hidden size for MuP width scaling. This is the reference width from which scaling factors are computed. Defaults to hidden_size if not specified (base model case where width_mult=1.0). Set this to your base/proxy model's hidden size when scaling up. +mup_embedding_mult: 1.0 # Multiplier for embedding layer output. Applied after the embedding lookup. Default: 1.0 (no scaling). +mup_output_mult: 1.0 # Multiplier for output logits before softmax. When MuP is enabled and this is left at 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable across widths. Override to customize output scaling. Default: 1.0. +mup_width_mult: 1.0 # Width multiplier for MuP scaling, computed as hidden_size / mup_base_hidden_size. This value is automatically computed in __post_init__ when use_mup is enabled. +nccl_all_reduce_for_prefill: false # If True, use NCCL all-reduce kernels when symmetric all-reduce is enabled. nccl_communicator_config_path: null # 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`. nccl_ub: false # Use the userbuffer registration for DP/FSDP communication buffers.This option will reduce GPU SM usage for the DP/FSDP communication,which is improving the performance of the overlapped computation. no_load_optim: null # Do not load optimizer when loading checkpoint. @@ -373,75 +494,93 @@ no_persist_layer_norm: false # Disable using persistent fused layer norm kernel no_rope_freq: null # Controls which layers to skip performing Rotary Position Embedding. Accepts either: - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string containing a Python list expression that defines a custom pattern, e.g.: "([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.By default this is disabled and set to None, indicating RoPE will be performedon every layer. no_save_optim: null # Do not save current optimizer. no_save_rng: null # Do not save current rng state. -non_persistent_ckpt_type: null # choices: ['global', 'local', 'in_memory', None] | Type of non-persistent model checkpoints. "global" - Saved as a standard checkpoint (e.g., on Lustre) with old checkpoints being removed. "local" - Each rank saves a portion of the checkpoint locally (e.g., on SSD/ramdisk). None - No non-persistent checkpointing (default option). +no_weight_decay_cond_type: null # choices: ['apply_wd_to_qk_layernorm'] | Type of no weight decay condition. Choices: None (default): apply weight decay to 1D weights and biases."apply_wd_to_qk_layernorm": additionally apply weight decay to qk layernorm as a special case.DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. +non_persistent_ckpt_type: null # choices: ['global', 'local', 'in_memory'] | 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). non_persistent_global_ckpt_dir: null # Directory containing global non-persistent model checkpoints. non_persistent_local_ckpt_algo: fully_parallel # choices: ['fully_parallel', 'atomic'] | Algorithm for local non-persistent checkpointing. non_persistent_local_ckpt_dir: null # Directory containing local non-persistent model checkpoints. non_persistent_save_interval: null # Number of iterations between non-persistent saves. -norm_epsilon: 1.0e-05 # Epsilon for layer norm and RMS norm. -normalization: LayerNorm # choices: ['LayerNorm', 'RMSNorm'] | Which normalization technique to use. +normalization: LayerNorm # choices: ['LayerNorm', 'RMSNorm'] | Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`. +null_tokenizer_eod_id: null # EOD token id for NullTokenizer. Defaults to `vocab_size - 1`. +null_tokenizer_pad_id: -1 # Pad token id for NullTokenizer. Defaults to -1 (no pad token). Set to a value outside the dataset to avoid masking real tokens. num_attention_heads: null # Number of transformer attention heads. num_channels: 3 # Number of channels in input image data num_classes: 1000 # num of classes in vision classificaiton task num_dataset_builder_threads: 1 # Number of parallel threads per rank for dataset builder num_distributed_optimizer_instances: 1 # Number of Distributed Optimizer copies across Data Parallel domain. num_experts: null # Number of Experts in MoE (None means no MoE) -num_layers: null # Number of transformer layers. -num_layers_at_end_in_bf16: 1 # Number of layers at end to construct in bf16 when --first-last-layers-bf16 is enabled. -num_layers_at_start_in_bf16: 1 # Number of layers at start to construct in bf16 when --first-last-layers-bf16 is enabled. +num_layers: null # Number of transformer layers in a transformer block. +num_layers_at_end_in_bf16: 1 # Number of layers at the end of the model to keep in BF16 precision when first_last_layers_bf16 is True. +num_layers_at_start_in_bf16: 1 # Number of layers at the start of the model to keep in BF16 precision when first_last_layers_bf16 is True. num_layers_per_virtual_pipeline_stage: null # Number of layers per virtual pipeline stage -num_query_groups: 1 +num_query_groups: 1 # Number of query groups for group query attention. If None, normal attention is used. +num_speculative_tokens: 0 # Number of speculative tokens generated during decode num_virtual_stages_per_pipeline_rank: null # Number of virtual pipeline stages per pipeline parallelism rank num_workers: 2 # Dataloader number of workers. +nvtx_ranges: false # Enable NVTX range annotations for profiling. When enabled, inserts NVTX markers to categorize execution in profiler output. object_storage_cache_path: null # Path to cache index files when using s3 or msc dataloader +offload_modules: [] # The submodules to offload its input. choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act". "attn_norm": offload the input of the normalization in the attention part. "qkv_linear": offload the input of the qkv linear part. "core_attn": offload the input of the core attention part. "attn_proj": offload the input of the attn linear projection part. "mlp_norm": offload the input of the normalization in the mlp part. "expert_fc1": offload the input of the expert fc1 part. "moe_act": offload the input of the moe act part. one_logger_async: false # If set, forces one_logger to use async mode. one_logger_project: megatron-lm # The one-logger project name. Will ignore if --no-one-logger is set one_logger_run_name: null # The one-logger run name displayed. Will ignore if --no-one-logger is set onnx_safe: null # Use workarounds for known problems with Torch ONNX exporter openai_gelu: false # Use OpenAIs GeLU implementation. This optionshould not be used unless for backward compatibilityreasons. -optimizer: adam # choices: ['adam', 'sgd'] | Optimizer function +optimizer: adam # choices: ['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'] | Optimizer function. Note: dist_muon is deprecated; use --optimizer muon with --use-distributed-optimizer instead. optimizer_cpu_offload: false # Offload optimizer state to CPU +optimizer_cuda_graph: false # Enable CUDA graph for optimizer step optimizer_offload_fraction: 1.0 # Ratio of optimizer state to offload to CPU +outer_dp_sharding_strategy: no_shard # choices: ['no_shard', 'optim'] | Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). Default: "no_shard". output_bert_embeddings: false # Output Bert embeddings (via mean pooling) from model, rather than its binary head output or entire hidden batch. overlap_cpu_optimizer_d2h_h2d: false # Overlap CPU optimizer step, gradients D2H and updated parameters H2D. +overlap_dispatch_backward_with_experts_wgrad: false # Delay the weight gradient computation for TE Grouped GEMM MoE experts. When enabled with FSDP, the expert weight gradients are computed on a separate CUDA stream after the data gradients finish, allowing overlap of wgrad compute with EP A2A communication. The FSDP gradient reduce-scatter for expert parameters is deferred until the delayed wgrad computation completes. This requires transformer_engine with GroupedLinear support (TE >= 2.3.0). overlap_grad_reduce: false # If set, overlap DDP grad reduce. -overlap_moe_expert_parallel_comm: false # Overlap the EP A2A communication by batch-level overlapping in 1f1b stage. +overlap_moe_expert_parallel_comm: false # Overlap EP A2A communications with independent computations of different micro-batches in 1f1b phase of pipelining or non-pipelining schedule. overlap_p2p_comm: true # overlap pipeline parallel communication with forward and backward chunks in 1F1B -overlap_p2p_comm_warmup_flush: false # if set, overlap pipeline parallel communication in warmup and flush +overlap_p2p_comm_warmup_flush: false # If true, overlap communication and computation in warm up and flush phase. Only valid when overlap_p2p_comm is True and batch_p2p_comm is False. Defaults to False. overlap_param_gather: false # If set, overlap param all-gather in distributed optimizer. overlap_param_gather_with_optimizer_step: false # If set, overlap param all-gather of first bucket with optimizer step. -override_opt_param_scheduler: false # 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. Notethat all the above values will be reset. +override_opt_param_scheduler: false # 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. +pad_vocab_size: true # Whether to pad vocab size of the model automatically if padded_vocab_size is not provided. padded_vocab_size: null # Vocabulary size of the model (padded to be divisible by tensor model parallel size). If not provided, it will be automatically calculated from vocab-size. patch_dim: 16 # patch dimension +per_dataset_sequences_path: null # Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file. per_split_data_args_path: null # Path to per-split-data-args. Instead of feeding `--(train|valid|test)-data-path` with weighted dataset, we pass in a file path from which we read those arguments. This is useful when the list of data is too big. Format is a json file with `train`, `valid, `test` keys -perform_initialization: true # Do not perform initialization when building model, can reduce startup time when definitely loading from a checkpoint +perform_initialization: true # Controls weights initialization. This option can be useful when you know you are going to load values from a checkpoint. perform_rl_step: false # Use the RL training step. +phase_transition_iterations: null # Comma-separated list of iterations where phase transitions occur. Requires fixed global batch size across phases. pin_cpu_grads: true # Disable pinning of CPU memory for gradients. pin_cpu_params: true # Disable pinning of CPU memory for parameters. -pipeline_model_parallel_comm_backend: null # choices: ['nccl', 'ucc'] | Select a communicator backend for pipeline parallel communication. If None, the default backend will be used. +pipeline_model_parallel_comm_backend: null # choices: ['nccl', 'ucc'] | Configuring backend option of pipeline parallel communication (e.g., nccl, ucc) If None, the default backend will be used. pipeline_model_parallel_layout: null # A string that describes a custom pipeline model parallel layout. e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer decoder layer, and mtp layer, respectively. Stages are split by "|". Replicated stages or layers can be described with multiplication. Commas can be used cosmetically. Default None is not using this argument to set the layout. -pipeline_model_parallel_size: 1 # Degree of pipeline model parallelism. -position_embedding_type: learned_absolute # choices: ['learned_absolute', 'rope', 'mrope', 'relative', 'none'] | Position embedding type. +pipeline_model_parallel_size: 1 # Inter-layer model parallelism. Splits transformer layers across GPU ranks. +position_embedding_type: learned_absolute # choices: ['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'] | Position embedding type. pretrained_checkpoint: null # Directory containing a pretrained model checkpoint for finetuning. profile: false # 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 --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop`. -profile_ranks: # Global ranks to profile. -- 0 +profile_ranks: [] # Global ranks to profile. profile_step_end: 12 # Global step to stop profiling. profile_step_start: 10 # Global step to start profiling. +pytorch_profiler_collect_callstack: false # Collect callstack in pytorch profiler. +pytorch_profiler_collect_chakra: false # Collect chakra trace in pytorch profiler. +pytorch_profiler_collect_shapes: false # Collect tensor shape in pytorch profiler. q_lora_rank: null # Rank of Query tensor's low rank representation. +qk_clip: false # Whether to clip the query and key weights. Needed for Muon MLA Model training. +qk_clip_alpha: 0.5 # The balancing alpha for qk-clip. Q = Q * (eta ** alpha) +qk_clip_threshold: 100 # The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0) qk_head_dim: 128 # Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim -qk_l2_norm: false # Use llama 4 qk l2 norm -qk_layernorm: false # Whether to layer normalize the q and k attention embeddings. +qk_l2_norm: false # Whether to apply llama 4-style qk L2 norm. +qk_layernorm: false # Whether to apply `normalization` type of normalization to the query and key embeddings. qk_pos_emb_head_dim: 64 # Dimension of the position embedding in the QK projection. query_in_block_prob: 0.1 # Probability of keeping query in block for ICT dataset quick_geglu: false # Use quick geglu activation instead of default gelu -rampup_batch_size: null # Batch size ramp up with the following values: --rampup-batch-size For example: --rampup-batch-size 16 8 300000 \ --global-batch-size 1024will start with global batch size 16 and over (1024 - 16) / 8 = 126 intervals will increasethe batch size linearly to 1024. In each intervalwe will use approximately 300000 / 126 = 2380 samples. +rampup_batch_size: null # Batch size ramp up with the following values: , , 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. Deprecated. Use step_batch_size_schedule instead. recompute_activations: false # recompute activation to allow for training with larger models, sequences, and batch sizes. -recompute_granularity: null # choices: ['full', 'selective'] | Checkpoint activations to allow for training with larger models, sequences, and batch sizes. It is supported at two granularities 1) full: whole transformer layer is recomputed, 2) selective: submodules set in --recompute-modules are recomputed, default is core_attn. -recompute_method: null # choices: ['uniform', 'block'] | 1) uniform: uniformly divide the total number of Transformer layers and recompute the input activation of each divided chunk at specified granularity, 2) recompute the input activations of only a set number of individual Transformer layers per pipeline stage and do the rest without any recomputing at specified granularitydefault) do not apply activations recompute to any layers -recompute_modules: null # The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts". default: ["core_attn"]."core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. "layernorm": recompute the input_layernorm and pre_mlp_layernorm. "mla_up_proj": recompute the MLA up projection and RoPE applying parts."mlp": recompute the dense MLP layer."moe": recompute the MoE layer."shared_experts": recompute the shared experts in the MoE layer."moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. -recompute_num_layers: null # 1) uniform: the number of Transformer layers in each uniformly divided recompute unit, 2) block: the number of individual Transformer layers to recompute within each pipeline stage. +recompute_granularity: null # choices: ['full', 'selective'] | Determines which type of activation recompute to use. Megatron-core supports 'selective' activation checkpointing where the submodules set in --recompute-modules is checkpointed. The default is "core_attn" which is the memory intensive part of attention. These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+). See Reducing Activation Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198) for more details. 'full' will checkpoint the entire transformer layer. If None, no recompute is performed and all activations are saved. If set, must be 'selective' or 'full'. 'selective' always uses all layers. +recompute_method: null # choices: ['uniform', 'block'] | Determines which transformer layers will be recomputed. uniform will uniformly divide the total number of transformer layers in a transformer block and recompute the input activation of each divided chunk at the specified granularity. block will recompute the input activations for only a set number of transformer layers per pipeline stage. The rest of the layers in the pipeline stage will not have any activations recomputed. If None, and recompute is enabled, all layers will do recomputation. If set, must be 'uniform' or 'block'. +recompute_modules: null # The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. "layernorm": recompute the input_layernorm and pre_mlp_layernorm. "mla_up_proj": recompute the MLA up projection and RoPE applying parts. "mlp": recompute the dense MLP submodule. "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. +recompute_num_layers: null # When recompute_method is uniform, recompute_num_layers is the number of transformer layers in each uniformly divided recompute unit. When recompute_method is block, recompute_num_layers is the number of transformer layers to recompute within each pipeline stage. Must be None for 'selective' activation checkpointing. record_memory_history: false # Record memory history in last rank. +record_shapes: false # Record shapes of tensors in `torch.autograd.profiler.emit_nvtx` for the Nsys profiler. +refit_method: gloo # choices: ['nccl', 'gloo', 'nvshmem'] | Method to refit the model weights between training and inference models during RL. nccl: use NCCLCopyService to refit using NCCL; gloo: use GlooCopyService over CPU; nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM. relative_attention_max_distance: 128 # Maximum distance for relative position embeddings calculation. relative_attention_num_buckets: 32 # Number of buckets for relative position embeddings. replication: false # If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks. @@ -453,139 +592,161 @@ reset_position_ids: false # Reset posistion ids after end-of-document token. result_rejected_tracker_filename: null # Optional name of file tracking `result_rejected` events. retriever_report_topk_accuracies: [] # Which top-k accuracies to report (e.g. '1 5 20') retriever_score_scaling: false # Whether to scale retriever scores by inverse square root of hidden size -retriever_seq_length: 256 # Maximum sequence length for the biencoder model for retriever -retro_add_retriever: false # Add a retriever to the transformer, for use in pretraining a Retro model. -retro_attention_gate: 1 # Gated cross attention. -retro_cyclic_train_iters: null # Set number of training iterations for cyclic Retro training. -retro_encoder_attention_dropout: 0.1 # Attention dropout for retrieval encoder. -retro_encoder_hidden_dropout: 0.1 # Hidden dropout for retrieval encoder. -retro_encoder_layers: 2 # Number of layers to use for the retrieval encoder. -retro_num_neighbors: 2 # Number of neighbors to retrieve during pretraining. -retro_num_retrieved_chunks: 2 # Number of chunks to retrieve from the retrieval database. -retro_project_dir: null # Retro project directory, which contains the preprocessed data for pretraining. This directory is built during preprocessing (see tools/retro/README.md), and contains subdirectories for the chunk database and pretraining neighbors. -retro_verify_neighbor_count: true # Skip verifying that len(GPT dataset) == len(saved neighbors). reuse_grad_buf_for_mxfp8_param_ag: false # If True, reuse the grad buffer for MXFP8 parameter all-gather. -rl_calculate_intra_group_similarity: false # If set, calculate the intra-group similarity of rollouts. +rl_default_temperature: 1.0 # Default temperature for model inference. +rl_default_top_k: -1 # Default top-k for model inference. +rl_default_top_p: 0 # Default top-p for model inference. +rl_generation_batch_size: null # Override the number of groups per generation batch. Defaults to grpo_prompts_per_step when --rl-num-parallel-generation-batches is set. rl_importance_sampling_truncation_coef: null # If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss. +rl_inference_expert_model_parallel_size: null # Degree of expert model parallelism for inference for RL. +rl_inference_expert_tensor_model_parallel_size: null # Degree of expert tensor model parallelism for inference for RL. For MoE models, this controls the TP size for expert layers specifically. Defaults to training expert_tensor_parallel_size if not specified. rl_inference_logprobs_is_correction: false # If set, use inference logprobs in importance sampling correction of the loss. -rl_offload_kv_cache_during_training: false # Offload KV cache to CPU during training to save GPU memory +rl_inference_model_unified_memory_level: 0 # choices: [0, 1] | Allocate the separate RL inference model parameters from a unified virtual memory (UVM) CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe device pointers. +rl_inference_parsers: [] # List of response parsers to enable for RL inference (e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool). +rl_inference_pipeline_model_parallel_size: null # Degree of pipeline model parallelism for inference for RL. +rl_inference_tensor_model_parallel_size: null # Degree of tensor model parallelism for inference for RL. +rl_kv_cache_management_mode: persist # choices: ['persist', 'offload', 'recompute'] | KV cache management mode during RL training: persist: leave KV cache in GPU memory (default), offload: offload KV cache to CPU during training, recompute: deallocate KV cache and recompute from scratch each cycle +rl_num_parallel_generation_batches: null # Number of generation batches in flight. Set to L+1 to allow for L steps of staleness between the inference and training policies. Each batch contains grpo_prompts_per_step groups by default. Requires --rl-partial-rollouts. Mutually exclusive with --rl-num-parallel-generations. +rl_num_parallel_generations: null # Number of rollouts being generated by the inference engine simultaneously. Internally divided by grpo_group_size. Requires --rl-partial-rollouts. Mutually exclusive with --rl-num-parallel-generation-batches. +rl_offload_inference_model_weights_when_idle: false # When using a separate RL inference model, offload its weights to CPU when not doing rollout inference, and restore to GPU right before inference. Works with two backends: 1) UVM (when --rl-inference-model-unified-memory-level=1), or 2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed). rl_offload_optimizer_during_inference: false # Offload optimizer state to CPU during inference/rollout to save GPU memory -rl_partial_rollouts: false # If set, use partial rollouts. +rl_parallel_generation_tasks: null # Deprecated: use --rl-num-parallel-generations instead. +rl_partial_rollouts: false # Allow inference to continue generating rollouts while training updates the policy weights. This enables off-policy training where rollouts may be generated with a stale version of the policy. Use --rl-num-parallel-generations or --rl-num-parallel-generation-batches to control the degree of staleness. +rl_persist_cuda_graphs: false # Persist CUDA graphs when the inference engine is suspended. If False, CUDA graphs are deleted on suspend and re-captured on resume. rl_prompts_per_eval: 32 # Number of prompts to evaluate for for each RL task.This evaluation can be very expensive when using environmentsthat evaluate pass@k so we default to a lower number. -rl_remove_kv_cache_during_training: false # Remove KV cache during training to save GPU memory -rl_reset_cuda_graphs: false # Reset CUDA graphs between inference/training to save GPU memory rl_sequence_packing_algo: fifo # choices: ['fifo', 'round-robin'] | Algorithm for distributing packed bins across ranks. fifo: first-in-first-out sequential distribution, round-robin: distribute bins cyclically across ranks for better load balancing -rl_sequence_packing_bin_size: 8192 # Override bin size for sequence packing. +rl_sequence_packing_max_sequences_per_bin: 50 # Maximum number of sequences that can be packed into a single bin. +rl_skip_bos_token: false # Skip BOS token at the beginning of the sequences. Default is False. +rl_training_cuda_graphs: false # If set, do not toggle CUDA graphs on/off between inference and training phases. rl_use_sequence_packing: false # Enable sequence packing +rl_verify_model_weights_swap: false # If set, verify that the model weights were correctly transferred by comparing forward pass outputs onthe first swap of model weights. rope_scaling_factor: 8.0 # Rope scaling factor in llama3.x models rope_type: null # choices: ['rope', 'yarn'] | Type of rope to use. Note that MLA takes yarn by default, and common attention takes rope by default. rotary_base: 10000 # Base to use for rotary positional embeddings, default 10000 -rotary_interleaved: false # Use interleaved rotary embedding. +rotary_interleaved: false # True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of first half and second half (LLaMa style). Default to False. rotary_percent: 1.0 # Percent of rotary dimension to use, default 100%% rotary_scaling_factor: 1.0 # Rotary scaling factor for the rotary embeddings. rotary_seq_len_interpolation_factor: null # Sequence length interpolation factor for rotary embeddings. run_workload_inspector_server: false # If set, enables workload inspector server for on-demand profiling. sample_rate: 1.0 # sample rate for training data. Supposed to be 0 < sample_rate < 1 save: null # Output directory to save checkpoints to. +save_activations_interval: null # Number of iterations between act.name->act.data mapping saves. +save_dgrads_interval: null # Number of iterations between dgrad saves. save_extra_steps: [] # Specific training iterations at which to save checkpoints save_interval: null # Number of iterations between persistent checkpoint saves. -save_retain_interval: null # Number of iterations between retained checkpoints (othercheckpoints _except the last checkpoint_ are automatically deleted). -scatter_gather_tensors_in_pipeline: true # If not set, use scatter/gather to optimize communication of tensors in pipeline. +save_params_interval: null # Number of iterations between param.name->param.data mapping saves. +save_retain_interval: null # Number of iterations between retained checkpoints (other checkpoints except the last checkpoint are automatically deleted). +save_tokens_per_expert_interval: null # Number of iterations between tokens-per-expert routing metadata saves. +save_wgrads_interval: null # Number of iterations between wgrad (main_grad) saves. seed: 1234 # Random seed used for python, numpy, pytorch, and cuda. seq_length: null # Maximum sequence length to process. -sequence_parallel: false # Enable sequence parallel optimization. +sequence_parallel: false # Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially. See Reducing Activation Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198) for more details. sft: false # Megatron SFT training sft_tokenizer_prompt_format: nemotron-h-aligned # SFT prompt format. sgd_momentum: 0.9 # Momentum factor for sgd -sharp_enabled_group: null # choices: ['dp', 'dp_replica'] | IB SHARP can be enabled from only one communication group. By default, it is enabled from dp group. Available options: [dp, dp_replica] +sharp_enabled_group: null # choices: ['dp', 'dp_replica'] | IB SHARP can be enabled from only one communication group. By default, it is enabled from dp group if not specified and use_sharp=True. Available options: [dp, dp_replica] short_seq_prob: 0.1 # Probability of producing a short sequence. -skip_train: false # If set, bypass the training loop, optionally do evaluation for validation/test, and exit. -softmax_type: vanilla # choices: ['learnable', 'vanilla', 'off-by-one'] | Type of softmax to use for the attention. Supports both a fixed offset and learnable offset. +skip_train: false # If set, bypass the training loop, perform evaluation for validation/test, and exit. +softmax_type: vanilla # choices: ['vanilla', 'off-by-one', 'learnable'] | Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and and learnable offset. spec: null # Specify the 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` +special_tokens: null # List of special tokens. For TikTokenizer needs to have ["", "", "", "", "", "", ""] split: null # Comma-separated list of proportions for training, validation, and test split. For example the split `90,5,5` will use 90%% of data for training, 5%% for validation and 5%% for test. squared_relu: false # Use squared relu activation instead of default gelu +start_eval_at_iter: null # If set, evaluation will only start after this iteration number. Useful for skipping evaluation during early training iterations when the model is not yet meaningful. If not set, evaluation starts from the first eval_interval. start_weight_decay: null # Initial weight decay coefficient for L2 regularization. +step_batch_size_schedule: null # Step-wise batch size schedule in format "THRESHOLD:BS THRESHOLD:BS ...". Thresholds support suffixes: K (1e3), M (1e6), B (1e9), T (1e12). If sequence length is provided, thresholds are interpreted as tokens; otherwise as samples. Example: step_batch_size_schedule = "0:768 250B:1536 500B:3072 750B:6144" Cannot be used together with decrease_batch_size_if_needed. straggler_ctrlr_port: 65535 # Port number to toggle StragglerDetector on/off at runtime straggler_minmax_count: 1 # Number of ranks to report with high/low estimated throughput -strict_fsdp_dtensor_load: true # Do not strict loading for fsdp_dtensor checkpoint format. +strict_fsdp_dtensor_load: true # Whether to enforce strict loading for FSDP DTensor checkpoints. When False, allows partial loading. suggested_communication_unit_size: null # Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system's memory and performance requirements. swiglu: false # Use gated linear units and SiLU activation instead of default gelu swin_backbone_type: tiny # choices: ['tiny', 'base', 'h3'] | pretraining objectives -symmetric_ar_type: null # choices: ['two_shot', 'one_shot', 'multimem_all_reduce', None] | What type of symmetric all reduce to use. The default is none which is no use of symetric memory +symmetric_ar_type: null # choices: ['two_shot', 'one_shot', 'multimem_all_reduce'] | What type of symmetric all reduce to use. The default is None which is no use of symmetric memory. +te_precision_config_file: null # Configuration file to select per-module precision overrides. See TransformerEngineMixedPrecision.md te_rng_tracker: false # Use the Transformer Engine version of the random number generator. Required for CUDA graphs support. -tensor_model_parallel_size: 1 # Degree of tensor model parallelism. +tensor_model_parallel_size: 1 # Intra-layer model parallelism. Splits tensors across GPU ranks. tensorboard_dir: null # Write TensorBoard logs to this directory. tensorboard_log_interval: 1 # Report to tensorboard interval. -tensorboard_queue_size: 1000 # Size of the tensorboard queue for pending events and summaries before one of the "add" calls forces a flush to disk. +tensorboard_queue_size: 1000 # Size of the tensorboard queue for pending events and summaries before one of the 'add' calls forces a flush to disk. test_data_path: null # The weight and prefix list for an independent test dataset. Follows the same pattern rules as --data-path. test_mode: false # Run all real-time test alongside the experiment. -tiktoken_num_special_tokens: 1000 # Number of special tokens in tiktoken tokenizer -tiktoken_pattern: null # Which tiktoken pattern to use. Options: [v1, v2] -tiktoken_special_tokens: null # List of tiktoken special tokens, needs to have ["", "", "", "", "", "", ""] -timing_log_level: 0 # choices: [0, 1, 2] | 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. -timing_log_option: minmax # choices: ['max', 'minmax', 'all'] | 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. +tiktoken_num_special_tokens: 1000 # Number of special tokens in tiktoken tokenizer. +tiktoken_pattern: null # choices: ['v1', 'v2'] | Which tiktoken pattern to use. Options: [v1, v2] +timing_log_level: 0 # choices: [0, 1, 2] | 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. +timing_log_option: minmax # choices: ['max', 'minmax', 'all'] | 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. titles_data_path: null # Path to titles dataset used for ICT -tokenizer_metadata: null # Path to tokenizer metadata in json format. -tokenizer_model: null # Sentencepiece tokenizer model. +tokenizer_hf_no_include_special_tokens: false # Converting text to ids will not include special for HuggingFace tokenizer. +tokenizer_hf_no_use_fast: false # Whether to use fast HuggingFace tokenizer. +tokenizer_model: null # Path to the tokenizer model. +tokenizer_sentencepiece_ignore_extra_whitespaces: true # Whether to ignore extra whitespaces in the input text while encoding. +tokenizer_sentencepiece_legacy: false # SentencePiece tokenizer wrapper legacy behavior. Allows special tokens usage. tokenizer_type: null # choices: ['BertWordPieceLowerCase', 'BertWordPieceCase', 'GPT2BPETokenizer', 'SentencePieceTokenizer', 'GPTSentencePieceTokenizer', 'HuggingFaceTokenizer', 'Llama2Tokenizer', 'TikTokenizer', 'MultimodalTokenizer', 'NullTokenizer', 'NullMultimodalTokenizer', 'SFTTokenizer'] | What type of tokenizer to use. torch_fsdp2_reshard_after_forward: true # Whether to reshard weights after forward pass when using PyTorch FSDP2. Set to enable FSDP ZeRO-2. tp_comm_bootstrap_backend: nccl # choices: ['nccl', 'mpi', 'gloo'] | Set the bootstrapping backend of Tensor parallel communications. -tp_comm_bulk_dgrad: true # Disables the All-Gather overlap with bprop activation gradient GEMM. -tp_comm_bulk_wgrad: true # Disables the Reduce-Scatter overlap with bprop weight gradient GEMM. -tp_comm_overlap: false # Enables the overlap of Tensor parallel communication and GEMM kernels. -tp_comm_overlap_ag: true # Disables the All-Gather overlap with GEMM by pipelining the GEMM and All-Gather. +tp_comm_bulk_dgrad: true # Controls Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if tp_comm_overlap is False. +tp_comm_bulk_wgrad: true # Controls All-Gather overlap with Bprop activation gradient GEMM. Don't care if tp_comm_overlap is False. +tp_comm_overlap: false # If true, allows overlapping of Linear layer execution with tensor parallel communication collectives like AllGather/ReduceScatter. Overlapping is done for the linear layers wherever possible during the forward and the backward pass. +tp_comm_overlap_ag: true # Controls All-Gather overlap with GEMM by pipelining the GEMM and All-Gather. Don't care if tp_comm_overlap is False. tp_comm_overlap_cfg: null # Config file when tp_comm_overlap is enabled. -tp_comm_overlap_rs: true # Disables the Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter. -tp_comm_overlap_rs_dgrad: false # Enables the Reduce-Scatter overlap with dgrad GEMM. -tp_comm_split_ag: true # Disables the All-Gather overlap with fprop GEMM. -tp_comm_split_rs: true # Disables the Reduce-Scatter overlap with fprop GEMM. +tp_comm_overlap_rs: true # Controls Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter. Don't care if tp_comm_overlap is False. +tp_comm_overlap_rs_dgrad: false # If true, allows Reduce-Scatter overlap with DGRAD GEMM by pipelining the GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is False. +tp_comm_split_ag: true # Deprecated from TransformerEngine v1.6.0. Controls All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather splits. Don't care if tp_comm_overlap is False. +tp_comm_split_rs: true # Deprecated from TransformerEngine v1.6.0. Controls Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is False. train_data_path: null # The weight and prefix list for an independent train dataset. Follows the same pattern rules as --data-path. -train_iters: null # Total number of iterations to train over all training runs. Note that either train-iters or train-samples should be provided. -train_samples: null # Total number of samples to train over all training runs. Note that either train-iters or train-samples should be provided. +train_iters: null # Total number of iterations to train over all training runs. Note that either train_iters or train_samples should be provided. +train_samples: null # Total number of samples to train over all training runs. Note that either train_iters or train_samples should be provided. train_sync_interval: null # Training CPU-GPU synchronization interval, to ensure that CPU is not running too far ahead of GPU. -transformer_impl: transformer_engine # choices: ['local', 'transformer_engine', 'inference_optimized'] | Which Transformer implementation to use. -trust_remote_code: false # Whether or not to allow PreTrainedTokenizer to execute remote code +transformer_impl: transformer_engine # choices: ['local', 'transformer_engine', 'inference_optimized'] | Transformer implementation to use. Options are 'transformer_engine' for Transformer Engine and 'local' for MCore. +trust_remote_code: false # Whether or not to allow PreTrainedTokenizer to execute remote code. untie_embeddings_and_output_weights: false # Untie embeddings and output weights. use_checkpoint_args: false # Override model-related command-line arguments with arguments from checkpoint -use_checkpoint_opt_param_scheduler: false # 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. -use_cpu_initialization: null # If set, initialize weights on the CPU. This eliminates init differences based on tensor parallelism. +use_checkpoint_opt_param_scheduler: false # 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. +use_cpu_initialization: null # When set to False, we initialize the weights directly on the GPU. CPU initialization is the same regardless of tensor model parallelism, but GPU initialization is not. Transferring weights from CPU to GPU can take a significant amount of time for large models. use_dist_ckpt_deprecated: false # Deprecated: see --ckpt-format. use_distributed_optimizer: false # Use distributed optimizer. use_flash_attn: false # use FlashAttention implementation of attention. https://arxiv.org/abs/2205.14135 -use_fused_weighted_squared_relu: false # Use fused weighted squared relu when using MoE. -use_legacy_models: false # Use the legacy Megatron models, not Megatron-Core models. +use_fused_weighted_squared_relu: false # If True, uses fused weighted squared relu kernel when using MoE. +use_gloo_process_groups: true # If enabled, create Gloo process groups for communications. +use_grouped_gemm_for_dense_mlp: false # Use GroupedLinear(num_groups=1) for dense MLP to trigger the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. Requires ``use_te_op_fuser=True`` and SwiGLU activation. +use_kitchen_attention: false # Use the kitchen extension for attention (instead of TE's attention). +use_layer_wise_param_layout: true # Opt out of the precomputed LayerWise param layout. When set, falls back to the legacy LayerWise ping-pong path: all params (including non-Muon embeddings, biases, layernorm) live in a single LayerWise buffer and the optimizer uses the allgather_params() codepath. The default (precomputed layout) routes non-Muon params through a separate DistributedOptimizer with byte-level sharding, which is faster and uses less padding but produces different bf16 reduction ordering and so will not match legacy-path loss curves bit-for-bit. use_legacy_static_engine: false # Use legacy static engine. (Current static engine uses dynamic engine under the hood) -use_megatron_fsdp: false # Use the Megatron FSDP code path in DDP. +use_mamba_mem_eff_path: true # Controls usage of the memory efficient path for Mamba layers. +use_megatron_fsdp: false # Use Megatron's Fully Sharded Data Parallel. Cannot be used together with use_torch_fsdp2. use_mp_args_from_checkpoint_args: false # Copy model parallelism command-line arguments from checkpoint +use_mup: false # Enable Maximal Update Parameterization (MuP) for hyperparameter transfer across model widths. When enabled, learning rates and initialization are scaled according to the width multiplier to ensure consistent training dynamics. use_one_sent_docs: false # Whether to use one sentence documents in ICT -use_persistent_ckpt_worker: false # Enables a persitent checkpoint worker for async save +use_persistent_ckpt_worker: false # Use a persistent background worker for async checkpoint saves. When enabled, creates a dedicated worker thread/process for handling async saves. When disabled, uses temporal workers that are created and destroyed for each save operation. use_precision_aware_optimizer: false # Use the precision-aware optimizer in TransformerEngine, which allows setting the main params and optimizer states to lower precision, such as fp16, bf16 and fp8. use_pytorch_profiler: false # Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard. -use_ring_exchange_p2p: false # 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. +use_ring_exchange_p2p: false # Use custom ring_exchange kernel instead of torch.distributed.batch_isend_irecv(). Requires custom built torch with torch.distributed.ring_exchange. use_rope_scaling: false # Apply rope scaling as used in llama3.x use_rotary_position_embeddings: false # Use rotary positional embeddings or not. Deprecated: use --position-embedding-type -use_sharp: false # Required to enable SHARP communication. -use_te_activation_func: false # Use activation function kernel from Transformer Engine in MLP module. +use_sharp: false # Set the use of SHARP for the collective communications of data-parallel process groups. When `True`, run barrier within each data-parallel process group, which specifies the SHARP application target groups. +use_te_activation_func: false # Whether to use ffn activation functions implemented by TransformerEngine use_tokenizer_model_from_checkpoint_args: true # If set, do not use tokenizer model path from checkpoint -use_torch_fsdp2: false # Use the torch FSDP2 implementation. FSDP2 has not been tested with pipeline parallelism, and may contain bugs. +use_torch_fsdp2: false # Use the torch FSDP2 implementation. FSDP2 is not currently working with Pipeline Parallel. It is still not in a stable release stage, and may therefore contain bugs or other potential issues. use_torch_optimizer_for_cpu_offload: false # Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode. use_tp_pp_dp_mapping: false # If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep-pp-dp. +use_transformer_engine_op_fuser: false # If True, submodules may use Transformer Engine's operation fuser API to enable advanced fusions. v_head_dim: 128 # Dimension of the head in the V projection. valid_data_path: null # The weight and prefix list for an independent validation dataset. Follows the same pattern rules as --data-path. +validation_set_names: null # Optional list of names for multiple validation sets. When provided with --multiple-validation-sets, these names are used instead of numeric indices (e.g. 'validation-wiki' instead of 'validation-0'). The number of names must match the number of validation datasets. +verify_integrity: false # Whether to hash checkpointing files during save and validate their integrity during load. vision_backbone_type: vit # choices: ['vit', 'mit', 'swin'] | backbone types types vision_pretraining: false # flag to indicate vision pretraining vision_pretraining_type: classify # choices: ['classify', 'inpaint', 'dino'] | pretraining objectives -vocab_extra_ids: 0 # Number of additional vocabulary tokens. They are used for span masking in the T5 model +vocab_extra_ids: 0 # Number of additional vocabulary tokens. They are used for span masking in the T5 model. vocab_file: null # Path to the vocab file. vocab_size: null # Size of vocab before EOD or padding. -wandb_entity: '' # The wandb entity name. It is useful when there are multiple sub-projects in a project. https://community.wandb.ai/t/how-do-i-decide-which-account-private-or-team-to-upload-the-run-to/5704 Ignore wandb by default. -wandb_exp_name: '' # The wandb experiment name. -wandb_project: '' # The wandb project name. Ignore wandb by default. -wandb_save_dir: '' # Path to save the wandb results locally. +wandb_entity: null # The wandb entity name. It is useful when there are multiple sub-projects in a project. +wandb_exp_name: null # The wandb experiment name. +wandb_project: null # The wandb project name. Ignore wandb by default. +wandb_save_dir: null # Path to save the wandb results locally. warmup: null # Old lr warmup argument, do not use. Use one of the--lr-warmup-* arguments above weight_decay: 0.01 # Weight decay coefficient for L2 regularization. weight_decay_incr_style: constant # choices: ['constant', 'linear', 'cosine'] | Weight decay increment function. -wgrad_deferral_limit: 0 # Number of micro-batches for whichweight gradient computation of vocabulary projection is deferred, defaults to 0 whichmeans all the micro-batches are deferred. Invalid if `defer-embedding-wgrad-compute`is not set +wgrad_deferral_limit: 0 # This value tunes the number of micro-batches for which the embedding weight gradient compute needs to be deferred to pipeline flush, this argument is invalid if `defer_embedding_wgrad_compute` is False. Defaults to 0, which means all micro-batches are deferred. window_attn_skip_freq: null # Frequency of layers to skip window attention. Accepts either: - An integer N: Represents a (N-1):1 ratio, meaning one full attention layer after (N-1) SWA layers. - A string containing a Python list expression that defines a custom pattern, e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates SWA and 0 indicates full attention. window_size: null # Window size for window attention. If not provided, window attention will be disabled. yaml_cfg: null # Config file to add additional arguments diff --git a/config/backend/megatron/base_llama.yaml b/config/backend/megatron/base_llama.yaml index dd7334d7..dcfb6b68 100644 --- a/config/backend/megatron/base_llama.yaml +++ b/config/backend/megatron/base_llama.yaml @@ -18,7 +18,7 @@ num_attention_heads: null tensor_model_parallel_size: 1 sequence_parallel: true normalization: RMSNorm -norm_epsilon: 1e-6 +layernorm_epsilon: 1e-6 swiglu: true bf16: true use_flash_attn: true diff --git a/config/backend/megatron/base_llama_spaj.yaml b/config/backend/megatron/base_llama_spaj.yaml index e9033a31..8384c460 100644 --- a/config/backend/megatron/base_llama_spaj.yaml +++ b/config/backend/megatron/base_llama_spaj.yaml @@ -18,7 +18,7 @@ num_attention_heads: 16 tensor_model_parallel_size: 1 sequence_parallel: true normalization: RMSNorm -norm_epsilon: 1e-6 +layernorm_epsilon: 1e-6 swiglu: true bf16: true use_flash_attn: true diff --git a/config/backend/megatron/base_reference.yaml b/config/backend/megatron/base_reference.yaml index 7f06685a..c8770d37 100644 --- a/config/backend/megatron/base_reference.yaml +++ b/config/backend/megatron/base_reference.yaml @@ -23,7 +23,6 @@ gradient_accumulation_fusion: false hidden_dropout: 0.0 hidden_size: 384 kv_channels: 64 -legacy_tokenizer: true log_interval: 100 log_progress: true log_throughput: true @@ -36,7 +35,7 @@ max_position_embeddings: 4096 micro_batch_size: 16 min_lr: 1e-05 # moe_ffn_hidden_size: 1536 -norm_epsilon: 1e-05 +layernorm_epsilon: 1e-05 normalization: RMSNorm num_attention_heads: 6 num_dataset_builder_threads: 6 diff --git a/config/backend/megatron/base_reference_base.yaml b/config/backend/megatron/base_reference_base.yaml index 082b1cc7..49c98a90 100644 --- a/config/backend/megatron/base_reference_base.yaml +++ b/config/backend/megatron/base_reference_base.yaml @@ -17,12 +17,11 @@ ckpt_format: torch_dist dataloader_type: cyclic gradient_accumulation_fusion: false hidden_dropout: 0.0 -legacy_tokenizer: true log_interval: 100 log_progress: true log_throughput: true lr_decay_style: WSD -norm_epsilon: 1e-05 +layernorm_epsilon: 1e-05 normalization: RMSNorm num_dataset_builder_threads: 6 num_workers: 7 diff --git a/config/backend/megatron/megadlm-8b.yaml b/config/backend/megatron/megadlm-8b.yaml index fbb8e453..572e4076 100644 --- a/config/backend/megatron/megadlm-8b.yaml +++ b/config/backend/megatron/megadlm-8b.yaml @@ -14,7 +14,7 @@ ffn_hidden_size: 14336 swiglu: true vocab_size: 50289 normalization: RMSNorm -norm_epsilon: 1.0e-6 +layernorm_epsilon: 1.0e-6 rotary_base: 500000 add_bias_linear: false position_embedding_type: rope diff --git a/config/backend/megatron/qwen3_125M.yaml b/config/backend/megatron/qwen3_125M.yaml index 6f33a8a4..96359ffe 100644 --- a/config/backend/megatron/qwen3_125M.yaml +++ b/config/backend/megatron/qwen3_125M.yaml @@ -8,3 +8,26 @@ num_attention_heads: 4 num_query_groups: 4 kv_channels: 128 untie_embeddings_and_output_weights: False + +seq_length: 4096 +max_position_embeddings: 4096 + +train_iters: 9536 +micro_batch_size: 8 +lr: 6.0e-4 + +tokenizer_type: SentencePieceTokenizer +tokenizer_model: /projects/0/prjs2013/oellm/tokenizer/oellm-tokenizer-128k-v2/tokenizer.model + +data_path: + - /projects/0/prjs2013/preprocessed/oellm-tokenizer-128k-v2/HPLT/HPLT2.0_cleaned/eng_Latn_67/eng_Latn_67_text_document + +save: ${job.base_output_dir}/checkpoints +wandb_save_dir: ${job.base_output_dir}/wandb +tensorboard_dir: ${job.base_output_dir}/tensorboard +save_interval: 1000 +eval_interval: 1001 +eval_iters: 101 +log_interval: 10 +log_throughput: true +use_flash_attn: true diff --git a/config/backend/megatron/qwen3_125M_dlm.yaml b/config/backend/megatron/qwen3_125M_dlm.yaml new file mode 100644 index 00000000..8c2b2d3a --- /dev/null +++ b/config/backend/megatron/qwen3_125M_dlm.yaml @@ -0,0 +1,44 @@ +defaults: + - qwen3_base + +num_layers: 16 +hidden_size: 512 +ffn_hidden_size: 2048 +num_attention_heads: 4 +num_query_groups: 4 +kv_channels: 128 +untie_embeddings_and_output_weights: False + +seq_length: 4097 +max_position_embeddings: 4097 + +train_iters: 9536 +micro_batch_size: 8 +lr: 6.0e-4 + +# Disable fusions that conflict with diffusion or require TE >= 1.4 +bias_gelu_fusion: false +bias_dropout_fusion: false +bias_swiglu_fusion: false +apply_rope_fusion: false + +aux: + model_name: "megadlm-125m" + mask_token_id: 147 + +tokenizer_type: SentencePieceTokenizer +tokenizer_model: /projects/0/prjs2013/oellm/tokenizer/oellm-tokenizer-128k-v2/tokenizer.model + +data_path: + - /projects/0/prjs2013/preprocessed/oellm-tokenizer-128k-v2/HPLT/HPLT2.0_cleaned/eng_Latn_67/eng_Latn_67_text_document + +save: ${job.base_output_dir}/checkpoints +wandb_save_dir: ${job.base_output_dir}/wandb +tensorboard_dir: ${job.base_output_dir}/tensorboard +save_interval: 1000 +eval_interval: 1001 +eval_iters: 101 +log_interval: 10 +log_throughput: true +log_timers_to_tensorboard: true +use_flash_attn: true diff --git a/config/backend/megatron/qwen3_base.yaml b/config/backend/megatron/qwen3_base.yaml index 319106e5..50d15d3f 100644 --- a/config/backend/megatron/qwen3_base.yaml +++ b/config/backend/megatron/qwen3_base.yaml @@ -44,7 +44,7 @@ add_bias_linear: false rotary_percent: 1.0 rotary_base: 1000000 normalization: RMSNorm -norm_epsilon: 1e-6 +layernorm_epsilon: 1e-6 position_embedding_type: rope use_rotary_position_embeddings: true attention_dropout: 0.0 diff --git a/config/backend/megatron/qwen3_moe.yaml b/config/backend/megatron/qwen3_moe.yaml index ecffb5ec..dc7b86cd 100644 --- a/config/backend/megatron/qwen3_moe.yaml +++ b/config/backend/megatron/qwen3_moe.yaml @@ -21,7 +21,6 @@ ckpt_format: torch_dist cross_entropy_loss_fusion: true manual_gc: true manual_gc_interval: 100 -cuda_graph_scope: "full" # value irrelevant; cuda_graph_impl=none disables graphs entirely cuda_graph_impl: "none" cuda_graph_warmup_steps: 3 cross_entropy_fusion_impl: "te" @@ -45,7 +44,7 @@ add_bias_linear: false rotary_percent: 1.0 rotary_base: 1000000 normalization: RMSNorm -norm_epsilon: 1e-6 +layernorm_epsilon: 1e-6 position_embedding_type: rope use_rotary_position_embeddings: true attention_dropout: 0.0 diff --git a/config/container/snellius.yaml b/config/container/snellius.yaml index 5b81fa29..28c87d50 100644 --- a/config/container/snellius.yaml +++ b/config/container/snellius.yaml @@ -14,7 +14,10 @@ class_name: Container # tokenize correctly under 5.14.1. Repo code that touches transformers # internals (postprocess/) runs on the host venv, not in this container. # -image: ${oc.env:CONTAINER_CACHE_DIR, /projects/0/prjs2013/containers/megatron-torch-2.13-nvcr.26-06-transformers-5.sif} +# The oc.env default covers only the *directory*; the filename must stay outside +# the interpolation, otherwise a set CONTAINER_CACHE_DIR resolves `image` to a +# directory instead of a .sif. Matches the pattern used by the other machines. +image: ${oc.env:CONTAINER_CACHE_DIR,/projects/0/prjs2013/containers}/megatron-torch-2.13-nvcr.26-06-transformers-5.sif runtime: singularity bind: [ diff --git a/config/experiments/korbi/chain_opensci_130M_train_eval_jupiter.yaml b/config/experiments/korbi/chain_opensci_130M_train_eval_jupiter.yaml index b1055515..4fd4d066 100644 --- a/config/experiments/korbi/chain_opensci_130M_train_eval_jupiter.yaml +++ b/config/experiments/korbi/chain_opensci_130M_train_eval_jupiter.yaml @@ -37,7 +37,6 @@ job: # Default backend.megatron tweaks shared by the train stage. backend: megatron: - legacy_tokenizer: false train_iters: 1000 lr_warmup_iters: 50 lr_decay_iters: 1000 diff --git a/config/experiments/korbi/chain_qwen3_bridge_train_eval_jupiter.yaml b/config/experiments/korbi/chain_qwen3_bridge_train_eval_jupiter.yaml index 699a9984..bd7b46c2 100644 --- a/config/experiments/korbi/chain_qwen3_bridge_train_eval_jupiter.yaml +++ b/config/experiments/korbi/chain_qwen3_bridge_train_eval_jupiter.yaml @@ -43,7 +43,6 @@ job: backend: megatron: - legacy_tokenizer: false train_iters: 1000 lr_warmup_iters: 50 lr_decay_iters: 1000 diff --git a/config/experiments/korbi/chain_qwen3_bridge_train_eval_juwels.yaml b/config/experiments/korbi/chain_qwen3_bridge_train_eval_juwels.yaml index 98dff6b4..febcc8ae 100644 --- a/config/experiments/korbi/chain_qwen3_bridge_train_eval_juwels.yaml +++ b/config/experiments/korbi/chain_qwen3_bridge_train_eval_juwels.yaml @@ -28,7 +28,6 @@ job: backend: megatron: - legacy_tokenizer: false trust_remote_code: true train_iters: 50 lr_warmup_iters: 5 diff --git a/config/experiments/korbi/chain_qwen3_bridge_train_eval_leonardo.yaml b/config/experiments/korbi/chain_qwen3_bridge_train_eval_leonardo.yaml index 875a1b45..878201d4 100644 --- a/config/experiments/korbi/chain_qwen3_bridge_train_eval_leonardo.yaml +++ b/config/experiments/korbi/chain_qwen3_bridge_train_eval_leonardo.yaml @@ -25,7 +25,6 @@ job: backend: megatron: - legacy_tokenizer: false trust_remote_code: true train_iters: 50 lr_warmup_iters: 5 diff --git a/config/experiments/korbi/chain_qwen3_bridge_train_eval_lumi.yaml b/config/experiments/korbi/chain_qwen3_bridge_train_eval_lumi.yaml index d468e1c0..ad53ff67 100644 --- a/config/experiments/korbi/chain_qwen3_bridge_train_eval_lumi.yaml +++ b/config/experiments/korbi/chain_qwen3_bridge_train_eval_lumi.yaml @@ -25,7 +25,6 @@ job: backend: megatron: - legacy_tokenizer: false trust_remote_code: true train_iters: 50 lr_warmup_iters: 5 diff --git a/config/experiments/korbi/megatron_1.8B_multinode_speed_sweep_jupiter.yaml b/config/experiments/korbi/megatron_1.8B_multinode_speed_sweep_jupiter.yaml index 8b211012..495459d3 100644 --- a/config/experiments/korbi/megatron_1.8B_multinode_speed_sweep_jupiter.yaml +++ b/config/experiments/korbi/megatron_1.8B_multinode_speed_sweep_jupiter.yaml @@ -46,7 +46,6 @@ backend: sequence_parallel: false # default; overridden to true for TP=2 configs # CUDA graph settings (require garbage_collection_threshold, not expandable_segments) cuda_graph_impl: local - cuda_graph_scope: full check_for_nan_in_loss_and_grad: false use_distributed_optimizer: true micro_batch_size: 6 diff --git a/config/experiments/korbi/megatron_niccolo_130M_scaling_jupiter.yaml b/config/experiments/korbi/megatron_niccolo_130M_scaling_jupiter.yaml index 2177fe84..cfd555ec 100644 --- a/config/experiments/korbi/megatron_niccolo_130M_scaling_jupiter.yaml +++ b/config/experiments/korbi/megatron_niccolo_130M_scaling_jupiter.yaml @@ -48,8 +48,6 @@ backend: head_lr_mult: 1.0 hidden_dropout: 0 hidden_size: 576 - hybrid_attention_ratio: 0.0 - hybrid_mlp_ratio: 0.0 kv_channels: 64 log_interval: 10 log_progress: true diff --git a/config/experiments/korbi/megatron_niccolo_130M_smoke_jupiter.yaml b/config/experiments/korbi/megatron_niccolo_130M_smoke_jupiter.yaml index 5198ea0a..c3506fdd 100644 --- a/config/experiments/korbi/megatron_niccolo_130M_smoke_jupiter.yaml +++ b/config/experiments/korbi/megatron_niccolo_130M_smoke_jupiter.yaml @@ -46,8 +46,6 @@ backend: head_lr_mult: 1.0 hidden_dropout: 0 hidden_size: 576 - hybrid_attention_ratio: 0.0 - hybrid_mlp_ratio: 0.0 kv_channels: 64 log_interval: 10 log_progress: true diff --git a/config/experiments/korbi/megatron_niccolo_130M_smoke_juwels.yaml b/config/experiments/korbi/megatron_niccolo_130M_smoke_juwels.yaml index 840e14d3..3f1d96f8 100644 --- a/config/experiments/korbi/megatron_niccolo_130M_smoke_juwels.yaml +++ b/config/experiments/korbi/megatron_niccolo_130M_smoke_juwels.yaml @@ -46,8 +46,6 @@ backend: head_lr_mult: 1.0 hidden_dropout: 0 hidden_size: 576 - hybrid_attention_ratio: 0.0 - hybrid_mlp_ratio: 0.0 kv_channels: 64 log_interval: 10 log_progress: true diff --git a/config/experiments/korbi/speed_tests/llama1_8b_jupiter_tp1_mbs6_nosp_cg.yaml b/config/experiments/korbi/speed_tests/llama1_8b_jupiter_tp1_mbs6_nosp_cg.yaml index d9eec256..285bde4b 100644 --- a/config/experiments/korbi/speed_tests/llama1_8b_jupiter_tp1_mbs6_nosp_cg.yaml +++ b/config/experiments/korbi/speed_tests/llama1_8b_jupiter_tp1_mbs6_nosp_cg.yaml @@ -30,8 +30,7 @@ backend: cross_entropy_loss_fusion: true sequence_parallel: false # TP=1: sequence_parallel=true adds useless all-gathers cuda_graph_impl: local # CUDA graphs: reduces kernel launch overhead - cuda_graph_scope: full # layer-level scope - check_for_nan_in_loss_and_grad: false # required for cuda_graph_scope=full + check_for_nan_in_loss_and_grad: false # required when CUDA graphs are captured env: TOKENIZERS_PARALLELISM: "false" PYTORCH_CUDA_ALLOC_CONF: "garbage_collection_threshold:0.8" # expandable_segments incompatible with CUDA graphs diff --git a/config/experiments/korbi/test_diana.yaml b/config/experiments/korbi/test_diana.yaml index 4c88c19b..e988585d 100644 --- a/config/experiments/korbi/test_diana.yaml +++ b/config/experiments/korbi/test_diana.yaml @@ -53,7 +53,7 @@ backend: attention_dropout: 0.0 hidden_dropout: 0.0 normalization: RMSNorm - norm_epsilon: 1.e-5 + layernorm_epsilon: 1.e-5 qk_layernorm: True bf16: True swiglu: True diff --git a/config/experiments/korbi/test_sweep.yaml b/config/experiments/korbi/test_sweep.yaml index e0a5dcf6..df4d0c6e 100644 --- a/config/experiments/korbi/test_sweep.yaml +++ b/config/experiments/korbi/test_sweep.yaml @@ -55,7 +55,7 @@ backend: attention_dropout: 0.0 hidden_dropout: 0.0 normalization: RMSNorm - norm_epsilon: 1.e-5 + layernorm_epsilon: 1.e-5 qk_layernorm: True bf16: True swiglu: True diff --git a/config/experiments/megatron_sanity_check.yaml b/config/experiments/megatron_sanity_check.yaml index d26a6215..7e26cfb4 100644 --- a/config/experiments/megatron_sanity_check.yaml +++ b/config/experiments/megatron_sanity_check.yaml @@ -21,7 +21,6 @@ backend: context_parallel_size: 1 sequence_parallel: true use_distributed_optimizer: true - async_tensor_model_parallel_allreduce: false masked_softmax_fusion: false gradient_accumulation_fusion: false bias_dropout_fusion: false diff --git a/config/experiments/opensci_perf_test_1.7b_jupiter.yaml b/config/experiments/opensci_perf_test_1.7b_jupiter.yaml index 35620093..33ab3be0 100644 --- a/config/experiments/opensci_perf_test_1.7b_jupiter.yaml +++ b/config/experiments/opensci_perf_test_1.7b_jupiter.yaml @@ -52,7 +52,6 @@ backend: init_method_std: 0.02 clip_grad: 1.0 add_qkv_bias: false - legacy_tokenizer: false override_opt_param_scheduler: true # (2) Enable all fusions @@ -70,7 +69,7 @@ backend: use_flash_attn: true qk_layernorm: true normalization: RMSNorm - norm_epsilon: 1.0e-05 + layernorm_epsilon: 1.0e-05 swiglu: true position_embedding_type: rope rotary_base: 100000 diff --git a/config/experiments/opensci_ref_0.13b_mixvitae_jupiter.yaml b/config/experiments/opensci_ref_0.13b_mixvitae_jupiter.yaml index c05817c9..0869040e 100644 --- a/config/experiments/opensci_ref_0.13b_mixvitae_jupiter.yaml +++ b/config/experiments/opensci_ref_0.13b_mixvitae_jupiter.yaml @@ -45,7 +45,6 @@ backend: clip_grad: 1.0 add_qkv_bias: false - legacy_tokenizer: false override_opt_param_scheduler: true # Fusions @@ -60,7 +59,7 @@ backend: use_flash_attn: true qk_layernorm: true normalization: RMSNorm - norm_epsilon: 1.0e-05 + layernorm_epsilon: 1.0e-05 swiglu: true position_embedding_type: rope rotary_base: 100000 diff --git a/config/experiments/opensci_ref_0.13b_test_e2e.yaml b/config/experiments/opensci_ref_0.13b_test_e2e.yaml index 880b8cf3..b356d66d 100644 --- a/config/experiments/opensci_ref_0.13b_test_e2e.yaml +++ b/config/experiments/opensci_ref_0.13b_test_e2e.yaml @@ -22,7 +22,6 @@ job: backend: megatron: - legacy_tokenizer: false train_iters: 5 lr_warmup_iters: 2 lr_decay_iters: 5 diff --git a/config/experiments/opensci_ref_0.4b_mixvitae_jupiter.yaml b/config/experiments/opensci_ref_0.4b_mixvitae_jupiter.yaml index 2f581762..a022410b 100644 --- a/config/experiments/opensci_ref_0.4b_mixvitae_jupiter.yaml +++ b/config/experiments/opensci_ref_0.4b_mixvitae_jupiter.yaml @@ -43,7 +43,6 @@ backend: clip_grad: 1.0 add_qkv_bias: false - legacy_tokenizer: false override_opt_param_scheduler: true # Fusions @@ -58,7 +57,7 @@ backend: use_flash_attn: true qk_layernorm: true normalization: RMSNorm - norm_epsilon: 1.0e-05 + layernorm_epsilon: 1.0e-05 swiglu: true position_embedding_type: rope rotary_base: 100000 diff --git a/config/experiments/opensci_ref_1.3b_mixvitae_jupiter.yaml b/config/experiments/opensci_ref_1.3b_mixvitae_jupiter.yaml index 9ccab8b2..ec8003aa 100644 --- a/config/experiments/opensci_ref_1.3b_mixvitae_jupiter.yaml +++ b/config/experiments/opensci_ref_1.3b_mixvitae_jupiter.yaml @@ -44,7 +44,6 @@ backend: # Override base_reference to match JUWELS baseline add_qkv_bias: false - legacy_tokenizer: false override_opt_param_scheduler: true # Fusions @@ -59,7 +58,7 @@ backend: use_flash_attn: true qk_layernorm: true normalization: RMSNorm - norm_epsilon: 1.0e-05 + layernorm_epsilon: 1.0e-05 swiglu: true position_embedding_type: rope rotary_base: 100000 diff --git a/config/experiments/opensci_ref_1.7b_mixvitae_jupiter.yaml b/config/experiments/opensci_ref_1.7b_mixvitae_jupiter.yaml index 696d6f5b..4a25353f 100644 --- a/config/experiments/opensci_ref_1.7b_mixvitae_jupiter.yaml +++ b/config/experiments/opensci_ref_1.7b_mixvitae_jupiter.yaml @@ -43,7 +43,6 @@ backend: clip_grad: 1.0 add_qkv_bias: false - legacy_tokenizer: false override_opt_param_scheduler: true # Fusions @@ -58,7 +57,7 @@ backend: use_flash_attn: true qk_layernorm: true normalization: RMSNorm - norm_epsilon: 1.0e-05 + layernorm_epsilon: 1.0e-05 swiglu: true position_embedding_type: rope rotary_base: 100000 diff --git a/config/experiments/qwen3_1.7b_mixvitae_jupiter_fullrun.yaml b/config/experiments/qwen3_1.7b_mixvitae_jupiter_fullrun.yaml index a2c77be5..85aa788f 100644 --- a/config/experiments/qwen3_1.7b_mixvitae_jupiter_fullrun.yaml +++ b/config/experiments/qwen3_1.7b_mixvitae_jupiter_fullrun.yaml @@ -53,7 +53,6 @@ backend: tokenizer_type: HuggingFaceTokenizer tokenizer_model: /e/data1/datasets/playground/mmlaion/shared/oellm_shared_evals/models/Qwen/Qwen3-1.7B-Base vocab_size: 151936 - legacy_tokenizer: false split: "989,10,1" num_workers: 4 diff --git a/config/experiments/qwen3_1.7b_mixvitae_jupiter_repro.yaml b/config/experiments/qwen3_1.7b_mixvitae_jupiter_repro.yaml index e381bad4..4a1d2742 100644 --- a/config/experiments/qwen3_1.7b_mixvitae_jupiter_repro.yaml +++ b/config/experiments/qwen3_1.7b_mixvitae_jupiter_repro.yaml @@ -44,7 +44,6 @@ backend: tokenizer_type: HuggingFaceTokenizer tokenizer_model: /e/data1/datasets/playground/mmlaion/shared/oellm_shared_evals/models/Qwen/Qwen3-1.7B-Base vocab_size: 151936 - legacy_tokenizer: false split: "989,10,1" num_workers: 4 diff --git a/config/experiments/thomas/qwen3_125M_dlm.yaml b/config/experiments/thomas/qwen3_125M_dlm.yaml new file mode 100644 index 00000000..bda8086d --- /dev/null +++ b/config/experiments/thomas/qwen3_125M_dlm.yaml @@ -0,0 +1,43 @@ +# @package _global_ + +defaults: + - /backend: megadlm_torchrun + - /container: snellius + - /slurm: snellius + - /job: default + - /sweep: none + - override /backend/megatron: qwen3_125M_dlm + - _self_ + +job: + name: "qwen3_125M_dlm_baseline" + base_output_dir: "${oc.env:OUTPUT_DIR,'./output'}/${job.name}" + log_path: ${.base_output_dir}/logs/slurm-%j.log + log_path_current: ${.base_output_dir}/logs/current.log + config_path: ${.base_output_dir}/logs/config-%j.yaml + config_path_current: ${.base_output_dir}/logs/current.yaml + +backend: + # MegaDLMs-only flags (submodules/MegaDLMs/custom_args/difflm.py), appended + # verbatim after the mapped Megatron args. See backend/megadlm_torchrun.yaml. + # + # --mask-token is interpolated from the model config so 147 has a single + # source of truth; 147 is in oellm-tokenizer-128k-v2. + # --mask-prob 0.15 is NOT a diffusion parameter. It is stock Megatron's + # BERT-style masking arg (same default, 0.15) and MegaDLMs never reads it, + # so it is inert. The diffusion masking rate is sampled per-sequence by + # pretrain_difflm.py and is not exposed as a CLI flag. Kept only to make the + # no-op explicit -- drop it if it reads as a knob that does something. + aux: + difflm_args: "--attention-mask-type no_mask --mask-token ${backend.megatron.aux.mask_token_id} --base-model vanilla --model-running-mode difflm-noshift --difflm-varilen-prob 0.01 --mask-prob 0.15 --wandb-exp-name ${job.name}" + megatron: + global_batch_size: 512 + wandb_project: "hplt-megadlm-baseline" + data_cache_path: "${job.base_output_dir}/data_cache" + +slurm: + sbatch: + _non_strict: true + nodes: 1 + partition: gpu_h100 + time: "15:00:00" diff --git a/config/experiments/throughput_sweep_qwen3_1.7b_jupiter_repro.yaml b/config/experiments/throughput_sweep_qwen3_1.7b_jupiter_repro.yaml index b4961c53..d02dee4e 100644 --- a/config/experiments/throughput_sweep_qwen3_1.7b_jupiter_repro.yaml +++ b/config/experiments/throughput_sweep_qwen3_1.7b_jupiter_repro.yaml @@ -40,7 +40,6 @@ backend: tokenizer_type: HuggingFaceTokenizer tokenizer_model: /e/data1/datasets/playground/mmlaion/shared/oellm_shared_evals/models/Qwen/Qwen3-1.7B-Base vocab_size: 151936 - legacy_tokenizer: false split: "989,10,1" num_workers: 1 diff --git a/oellm_autoexp/backends/megatron/cli_metadata.py b/oellm_autoexp/backends/megatron/cli_metadata.py index 0fcac74d..b71ca267 100644 --- a/oellm_autoexp/backends/megatron/cli_metadata.py +++ b/oellm_autoexp/backends/megatron/cli_metadata.py @@ -1,5 +1,6 @@ """Megatron CLI metadata (auto-generated).""" +from typing import Any from collections.abc import Mapping from oellm_autoexp.backends.megatron_args import MegatronArgMetadata, MegatronActionSpec @@ -8,11 +9,7 @@ "account_for_embedding_in_pipeline_split": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, *input* embedding layer will be treated as a standard" - " transformerlayer in the context of partition and placement for pipeline" - " parallelism." - ), + help="If set, the embedding layer will be treated as a standard transformer layer in the context of partition and placement for pipeline parallelism.", choices=None, nargs=0, element_type=None, @@ -20,10 +17,7 @@ "account_for_loss_in_pipeline_split": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, loss layer will be treated as a standard transformerlayer in the" - " context of partition and placement for pipeline parallelism." - ), + help="If set, the loss layer will be treated as a standard transformer layer in the context of partition and placement for pipeline parallelism.", choices=None, nargs=0, element_type=None, @@ -39,10 +33,7 @@ "activation_func_clamp_value": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Clamp the output of the linear_fc1 in the activation function. Only used" - " when activation_func is quick_gelu." - ), + help="Clamp the output of the linear_fc1 in the activation function. Only used when activation_func is quick_gelu.", choices=None, nargs=None, element_type=None, @@ -50,7 +41,7 @@ "adam_beta1": MegatronArgMetadata( arg_type=float, default=0.9, - help=("First coefficient for computing running averages of gradient and its square"), + help="First coefficient for computing running averages of gradient and its square", choices=None, nargs=None, element_type=None, @@ -58,7 +49,7 @@ "adam_beta2": MegatronArgMetadata( arg_type=float, default=0.999, - help=("Second coefficient for computing running averages of gradient and its square"), + help="Second coefficient for computing running averages of gradient and its square", choices=None, nargs=None, element_type=None, @@ -74,7 +65,7 @@ "add_bias_linear": MegatronArgMetadata( arg_type=bool, default=True, - help="Disable bias in the linear layers", + help="Include/exclude a bias term in all linear layers (QKV projections, after core attention, and two in MLP layer).", choices=None, nargs=0, element_type=None, @@ -90,7 +81,7 @@ "add_qkv_bias": MegatronArgMetadata( arg_type=bool, default=False, - help="Enable bias only in the QKV linear layers", + help="Add a bias term only for QKV projections.", choices=None, nargs=0, element_type=None, @@ -114,10 +105,7 @@ "align_grad_reduce": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "If not set, all PP stages will launch gradient reduces simultaneously." - " Otherwise, each PP stage will independently launch as needed." - ), + help="If not set, all PP stages will launch gradient reduces simultaneously. Otherwise, each PP stage will independently launch as needed.", choices=None, nargs=0, element_type=None, @@ -125,10 +113,7 @@ "align_param_gather": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "If not set, all PP stages will launch param all-gathers simultaneously." - " Otherwise, each PP stage will independently launch as needed." - ), + help="If not set, all PP stages will launch param all-gathers simultaneously. Otherwise, each PP stage will independently launch as needed.", choices=None, nargs=0, element_type=None, @@ -136,20 +121,7 @@ "allow_ambiguous_pad_tokens": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Whether to prevent pad tokens already present in the dataset from being" - " masked out when the pad token incorrectly shares the same id with other" - " special tokens in the tokenizer. Note that this argument has no effect" - " when the tokenizer correctly provides a unique id for the pad. Masking" - " out such ambiguous pad tokens results in training instability. Such a" - " scenario is best resolved by fixing the tokenizer; leaving this option as" - " False provides a workaround. When left to the default of False, any token" - " ids that collide with the pad token id - as provided by the tokenizer -" - " will not be masked out of the loss calculation: it cannot be determined" - " whether they are truly pad. If instead this argument is set, the training" - " flow will treat all tokens that share the same id as the pad token as" - " true pad tokens, potentially causing severe training instability." - ), + help="Whether to prevent pad tokens already present in the dataset from being masked out when the pad token incorrectly shares the same id with other special tokens in the tokenizer. Note that this argument has no effect when the tokenizer correctly provides a unique id for the pad. Masking out such ambiguous pad tokens results in training instability. Such a scenario is best resolved by fixing the tokenizer; leaving this option as False provides a workaround. When left to the default of False, any token ids that collide with the pad token id - as provided by the tokenizer - will not be masked out of the loss calculation: it cannot be determined whether they are truly pad. If instead this argument is set, the training flow will treat all tokens that share the same id as the pad token as true pad tokens, potentially causing severe training instability.", choices=None, nargs=0, element_type=None, @@ -157,11 +129,7 @@ "app_tag_run_name": MegatronArgMetadata( arg_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" - ), + 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", choices=None, nargs=None, element_type=None, @@ -169,76 +137,87 @@ "app_tag_run_version": MegatronArgMetadata( arg_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" - ), + 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", choices=None, nargs=None, element_type=None, ), - "apply_layernorm_1p": MegatronArgMetadata( + "apply_query_key_layer_scaling": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Adjust LayerNorm weights such that they are centered around zero. This" - " improves numerical stability." - ), + help="If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with fp16. Also sets `attention_softmax_in_fp32` to True.", choices=None, nargs=0, element_type=None, ), - "apply_query_key_layer_scaling": MegatronArgMetadata( + "apply_residual_connection_post_layernorm": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Scale Q * K^T by 1 / layer-number. Useful for fp16 training. Also sets" - " `attention_softmax_in_fp32` to True." - ), + help="If True, uses the original BERT residule connection ordering.", choices=None, nargs=0, element_type=None, ), - "apply_residual_connection_post_layernorm": MegatronArgMetadata( + "apply_rope_fusion": MegatronArgMetadata( arg_type=bool, - default=False, - help="If set, use original BERT residula connection ordering.", + default=True, + help="Disable rope fusion, the fusion is available only when using megatron-core.", choices=None, nargs=0, element_type=None, ), - "apply_rope_fusion": MegatronArgMetadata( + "apply_wd_to_qk_layernorm": MegatronArgMetadata( arg_type=bool, - default=True, - help=("Disable rope fusion, the fusion is available only when using megatron-core."), + default=False, + help="Apply weight decay to qk layernorm as a special case.", choices=None, nargs=0, element_type=None, ), - "async_save": MegatronArgMetadata( + "async_ckpt_cpu_priority": MegatronArgMetadata( + arg_type=int, + default=10, + help="CPU nice value target (0-19, higher = lower priority) for the async checkpoint writer process. If it exceeds 19, it will be set to 19. If the current nice value is greater than the target, it will be left unchanged. Only applies when using persistent ckpt worker.", + choices=None, + nargs=None, + element_type=None, + ), + "async_ckpt_io_priority": MegatronArgMetadata( + arg_type=int, + default=3, + help="I/O scheduling class (0-3, 3=idle) for the async checkpoint writer process.", + choices=None, + nargs=None, + element_type=None, + ), + "async_ckpt_use_cpu_shm": MegatronArgMetadata( arg_type=bool, - default=None, - help=( - "Apply async checkpointing save. Currently works only with`torch_dist`" - " distributed checkpoint format." - ), + default=False, + help="Copy GPU tensors to CPU shared-memory in the training process before handing off to the async checkpoint worker. Avoids CUDA IPC / NVLink fabric handles in the worker subprocess. Useful on MNNVL systems where fabric resources are exhausted. Only applies with the nvrx async strategy.", choices=None, nargs=0, element_type=None, ), - "async_tensor_model_parallel_allreduce": MegatronArgMetadata( + "async_save": MegatronArgMetadata( arg_type=bool, - default=True, - help="DEPRECATED. This flag is ignored.", + default=False, + help="Apply async checkpointing save. Currently works only with `torch_dist` distributed checkpoint format.", choices=None, nargs=0, element_type=None, ), + "async_strategy": MegatronArgMetadata( + arg_type=str, + default="nvrx", + help="Which async save strategy to use. Available strategies: nvrx, mcore.", + choices=("nvrx", "mcore"), + nargs=None, + element_type=None, + ), "attention_backend": MegatronArgMetadata( arg_type=None, default="auto", - help=("Attention backend to use (flash,fused,unfused,local,auto). Defaults to auto"), + help="Attention backend to run. By default we let transformer engine decide the best backend to run (except in the case of local). If attention backend is local we use the local pytorch implementation in mcore. Users can specify exact backend by changing this config.", choices=(1, 2, 3, 4, 5), nargs=None, element_type=None, @@ -251,6 +230,14 @@ nargs=None, element_type=None, ), + "attention_output_gate": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to apply output gate to the attention layers.", + choices=None, + nargs=0, + element_type=None, + ), "attention_softmax_in_fp32": MegatronArgMetadata( arg_type=bool, default=False, @@ -262,11 +249,7 @@ "auto_detect_ckpt_format": MegatronArgMetadata( arg_type=bool, default=False, - 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)." - ), + 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).', choices=None, nargs=0, element_type=None, @@ -274,12 +257,15 @@ "barrier_with_L1_time": MegatronArgMetadata( arg_type=bool, default=True, - 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." - ), + help="If not disabled, 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.", + choices=None, + nargs=0, + element_type=None, + ), + "batch_invariant_mode": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If true, uses batch-invariant kernels that provide deterministic forward execution regardless of batch size. This ensures bitwise identical results when the same inputs are processed in different batch configurations. This will significantly affect speed of training and inference as the kernels are not full optimized. Defaults to False.", choices=None, nargs=0, element_type=None, @@ -311,7 +297,7 @@ "bert_load": MegatronArgMetadata( arg_type=str, default=None, - help=("Directory containing an BertModel checkpoint (needed to start ICT and REALM)"), + help="Directory containing an BertModel checkpoint (needed to start ICT and REALM)", choices=None, nargs=None, element_type=None, @@ -319,7 +305,7 @@ "bf16": MegatronArgMetadata( arg_type=bool, default=False, - help="Run model in bfloat16 mode.", + help="If true, train with bf16 mixed precision training.", choices=None, nargs=0, element_type=None, @@ -343,9 +329,7 @@ "bias_swiglu_fusion": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "Disable bias and swiglu fusion, the fusion is available only when using megatron-core." - ), + help="Disable bias and swiglu fusion, the fusion is available only when using megatron-core.", choices=None, nargs=0, element_type=None, @@ -385,10 +369,7 @@ "calc_ft_timeouts": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, FT package will try to automatically compute the timeouts. Note:" - " This feature is for Nvidia internal use only." - ), + help="If set, FT package will try to automatically compute the timeouts. Note: This feature is for Nvidia internal use only.", choices=None, nargs=0, element_type=None, @@ -396,14 +377,19 @@ "calculate_per_token_loss": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + help="Whether cross entropy loss is calculated over the actual number of non-padded tokens in the global batch, versus the default behavior of assuming all tokens are non-padded.", choices=None, nargs=0, element_type=None, ), + "chat_template": MegatronArgMetadata( + arg_type=str, + default=None, + help="Custom chat template in jinja format for conversation formatting.", + choices=None, + nargs=None, + element_type=None, + ), "check_for_large_grads": MegatronArgMetadata( arg_type=bool, default=False, @@ -423,7 +409,7 @@ "check_for_spiky_loss": MegatronArgMetadata( arg_type=bool, default=False, - help="Check for spiky loss", + help="Check for spiky loss.", choices=None, nargs=0, element_type=None, @@ -431,10 +417,7 @@ "check_weight_hash_across_dp_replicas_interval": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Interval to check weight hashes are same across DP replicas. If not" - " specified, weight hashes not checked." - ), + help="Interval to check weight hashes are same across DP replicas. If not specified, weight hashes not checked.", choices=None, nargs=None, element_type=None, @@ -442,10 +425,7 @@ "checkpoint_activations": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Checkpoint activation to allow for training with larger models, sequences," - " and batch sizes." - ), + help="Checkpoint activation to allow for training with larger models, sequences, and batch sizes.", choices=None, nargs=0, element_type=None, @@ -453,25 +433,21 @@ "ckpt_assume_constant_structure": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If the model and optimizer state dict structure isconstant throughout a" - " *single training job*, it allows fordifferent checkpointing performance" - " optimizations." - ), + help="Assume the checkpoint structure is constant across saves to enable optimizations.", choices=None, nargs=0, element_type=None, ), "ckpt_convert_format": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, help="Checkpoint format for conversion.", - choices=("torch", "torch_dist", "zarr"), + choices=("torch", "torch_dist"), nargs=None, element_type=None, ), "ckpt_convert_save": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, help="Save directory for converted checkpoint.", choices=None, @@ -481,26 +457,16 @@ "ckpt_convert_update_legacy_dist_opt_format": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, ), "ckpt_format": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="torch_dist", - help=( - "Checkpoint format to use. torch is the format used by torch.save/load." - " torch_dist is a megatron built-in distributed checkpointing format." - " torch_dcp is the torch.distributed.checkpoint format. fsdp_dtensor is a" - " torch DCP native, Megatron FSDP training-specific checkpoint format." - ), - choices=("torch", "torch_dist", "zarr", "torch_dcp", "fsdp_dtensor"), + help="Checkpoint format to use. torch is the format used by torch.save/load. torch_dist is a megatron built-in distributed checkpointing format. torch_dcp is the torch.distributed.checkpoint format. fsdp_dtensor is a torch DCP native, Megatron FSDP training-specific checkpoint format.", + choices=("torch", "torch_dist", "torch_dcp", "fsdp_dtensor"), nargs=None, element_type=None, ), @@ -512,14 +478,26 @@ nargs=0, element_type=None, ), + "ckpt_fully_parallel_load_exchange_algo": MegatronArgMetadata( + arg_type=str, + default="broadcast", + help='Algorithm for fully parallel load of distributed checkpoints. "broadcast"(default): Broadcast the checkpoint from rank 0 to all other ranks. "gather_rounds": Gather the checkpoint from all ranks in rounds. "gather_object": Gather the checkpoint from all ranks in a single operation.', + choices=("broadcast", "gather_rounds", "gather_object"), + nargs=None, + element_type=None, + ), + "ckpt_fully_parallel_load_process_group": MegatronArgMetadata( + arg_type=str, + default="dp", + help='Process group for fully parallel load of distributed checkpoints. "dp"(default): Data parallel process group. "ep_dp": Expert data parallel process group.', + choices=("dp", "ep_dp"), + nargs=None, + element_type=None, + ), "ckpt_fully_parallel_save": MegatronArgMetadata( arg_type=bool, default=True, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -532,6 +510,22 @@ nargs=0, element_type=None, ), + "ckpt_fully_parallel_save_process_group": MegatronArgMetadata( + arg_type=str, + default="dp", + help='Process group for fully parallel save of distributed checkpoints. "dp"(default): Data parallel process group. "ep_dp": Expert data parallel process group.', + choices=("dp", "ep_dp"), + nargs=None, + element_type=None, + ), + "ckpt_load_validate_sharding_integrity": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Whether to validate sharding access integrity when loading a distributed checkpoint. When True (default), each tensor shard is checked to be accessed exactly once as main replica by some rank. Disabling skips this validation", + choices=None, + nargs=0, + element_type=None, + ), "ckpt_step": MegatronArgMetadata( arg_type=int, default=None, @@ -559,9 +553,7 @@ "clone_scatter_output_in_embedding": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "If not set, clone the output of the scatter in embedding layer to GC original tensor." - ), + help="When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer to facilitate garbage collection of input.", choices=None, nargs=0, element_type=None, @@ -569,7 +561,7 @@ "config_logger_dir": MegatronArgMetadata( arg_type=str, default="", - help="If set, will dump all configs to --config-logger-dir", + help="When non-empty, dumps entry-point configs to config_logger_dir", choices=None, nargs=None, element_type=None, @@ -577,7 +569,7 @@ "context_parallel_size": MegatronArgMetadata( arg_type=int, default=1, - help="Degree of context parallelism.", + help="Splits network input along sequence dimension across GPU ranks.", choices=None, nargs=None, element_type=None, @@ -585,12 +577,7 @@ "cp_comm_type": MegatronArgMetadata( arg_type=list, default=["p2p"], - help=( - "Inter-gpu communication type for context parallelism: p2p, a2a, allgather" - " or a2a+p2p. If a single string is provided, all layers will share the" - " same communication type. Users can also specify separated types for each" - " layer like --cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p" - ), + help="Inter-gpu communication type for context parallelism: p2p, a2a, allgather or a2a+p2p. If a single string is provided, all layers will share the same communication type. Users can also specify separated types for each layer like --cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p", choices=None, nargs="+", element_type=str, @@ -598,11 +585,27 @@ "cpu_offloading_num_layers": MegatronArgMetadata( arg_type=int, default=0, - help="The number of Transformer layers to offload to CPU.", + help="Tells the number of transformer layers for which activations has to be offloaded.", choices=None, nargs=None, element_type=None, ), + "cpu_offloading_retain_pinned_cpu_buffers": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, the pinned CPU buffers are retained after offloading and reused for the next iteration. It is useful for cuda graphs capture.", + choices=None, + nargs=0, + element_type=None, + ), + "create_all_gather_group": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Create a separate process group for all-gather operations to overlap reduce-scatter and all-gather operations.", + choices=None, + nargs=0, + element_type=None, + ), "create_attention_mask_in_dataloader": MegatronArgMetadata( arg_type=bool, default=True, @@ -614,7 +617,7 @@ "cross_entropy_fusion_impl": MegatronArgMetadata( arg_type=str, default="native", - help="Implementation of cross entropy loss calculation.", + help="If 'native', MCore based CE loss fusion is used, if 'te', Parallel CE loss from Transformer Engine library is used. Defaults to 'native'.", choices=("native", "te"), nargs=None, element_type=None, @@ -622,7 +625,7 @@ "cross_entropy_loss_fusion": MegatronArgMetadata( arg_type=bool, default=False, - help="Enabled fusion of cross entropy loss calculation.", + help="If this is enabled, the fused cross entropy implementation would be used. Defaults to False.", choices=None, nargs=0, element_type=None, @@ -630,37 +633,26 @@ "cuda_graph_impl": MegatronArgMetadata( arg_type=str, default="none", - help=( - 'Determines the CUDA graph capture implementation. "none": no CUDA graph.' - ' "local": capture the CUDA graph using MCore local implementation.' - ' --cuda-graph-scope="full_iteration" enables whole iteration CUDA graph.' - ' "transformer_engine": capture the CUDA graph using TE' - " make_graphed_callables()." - ), - choices=("none", "local", "transformer_engine"), + help='Determines the CUDA graph capture implementation. "none": no CUDA graph. "local": MCore CUDA graph implementation. During training, graphable modules own per-layer CUDA graphs controlled by cuda_graph_modules. During inference, graph ownership is controlled separately by inference_cuda_graph_scope. "transformer_engine": Transformer Engine CUDA graph implementation. During training, TE make_graphed_callables() creates per-layer CUDA graphs controlled by cuda_graph_modules. Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none". "full_iteration": full-iteration CUDA graph implementation for the training iteration (1 CUDA graph for the whole forward-backward path excluding the optimizer step). Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none". cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when cuda_graph_impl="full_iteration".', + choices=("none", "local", "transformer_engine", "full_iteration"), nargs=None, element_type=None, ), - "cuda_graph_scope": MegatronArgMetadata( - arg_type=str, - default="full", - help=( - 'Determines the CUDA graphs capturing scope. Valid values are "full",' - ' "attn" and "full_iteration". "Full" scope captures a whole Transformer' - ' layer. "Attn" scope only captures operations in' - ' TransformerLayer._forward_attention(). "ful_iteration" scope captures a' - " whole iteration. full_iteration scope is only supported with" - " --cuda-graph-impl=local, attn scope is only supported with" - " --cuda-graph-impl=transformer_engine." - ), - choices=("full", "attn", "full_iteration"), - nargs=None, - element_type=None, + "cuda_graph_modules": MegatronArgMetadata( + arg_type=list, + default=[], + help='Selects training capture coverage within per-layer CUDA graphs (local and transformer_engine implementations). Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba": "attn": captures operations in TransformerLayer._forward_attention(). "mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. "moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. "moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), including the shared experts if they are not overlapped with EP comm. "moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". "mamba": captures the mamba layer. An empty list means capturing the whole Transformer layer. This field is meaningless when --cuda-graph-impl=full_iteration and must be empty. Backward compatibility: "full" is deprecated but kept for backward compatibility; it is transformed to an empty list in validate_args. The deprecated values "full_iteration" and "full_iteration_inference" are also accepted and migrated to the new API in validate_args.', + choices=None, + nargs="+", + element_type=Any, + ), + "cuda_graph_scope_deprecated": MegatronArgMetadata( + arg_type=list, default=None, help="==SUPPRESS==", choices=None, nargs="+", element_type=Any ), "cuda_graph_warmup_steps": MegatronArgMetadata( arg_type=int, default=3, - help="Number of CUDA graph warmup steps", + help="Number of warmup steps for CUDA graphs", choices=None, nargs=None, element_type=None, @@ -668,11 +660,7 @@ "data_args_path": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Path to data-args. Instead of feeding `--data-path` with weighted dataset," - " we pass in a file path from which we read that argument. This is useful" - " when the list of data is too big." - ), + help="Path to data-args. Instead of feeding `--data-path` with weighted dataset, we pass in a file path from which we read that argument. This is useful when the list of data is too big.", choices=None, nargs=None, element_type=None, @@ -695,7 +683,7 @@ ), "data_parallel_sharding_strategy": MegatronArgMetadata( arg_type=str, - default="no_shard", + default="optim_grads_params", help="Sharding strategy of data parallelism.", choices=("no_shard", "optim", "optim_grads", "optim_grads_params"), nargs=None, @@ -704,15 +692,7 @@ "data_path": MegatronArgMetadata( arg_type=list, default=None, - help=( - "The weight and prefix list for a set of train, validation, and" - " testdatasets 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." - ), + help="The weight and prefix list for a set of train, validation, and testdatasets 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.", choices=None, nargs="*", element_type=str, @@ -733,6 +713,22 @@ nargs=0, element_type=None, ), + "dataloader_defer_npy_index_mmap": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.", + choices=None, + nargs=0, + element_type=None, + ), + "dataloader_fast_cache_load": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.", + choices=None, + nargs=0, + element_type=None, + ), "dataloader_type": MegatronArgMetadata( arg_type=str, default=None, @@ -768,26 +764,23 @@ "ddp_pad_buckets_for_high_nccl_busbw": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, make sure the bucket size is divisible by a large power of 2" - " (2^16) to ensure NCCL collectives have high bus bandwidth at large DP" - " counts, since NCCL message size (which for ring algorithms is bucket_size" - " / dp_size) apparently needs to be divisible by a power of 2 for high" - " busbw." - ), + help="If set, make sure the bucket size is divisible by a large power of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, since NCCL message size (which for ring algorithms is bucket_size / dp_size) apparently needs to be divisible by a power of 2 for high busbw.", choices=None, nargs=0, element_type=None, ), + "ddp_param_name_patterns_for_fp32_local_accumulation": MegatronArgMetadata( + arg_type=list, + default=[], + help="List of param_name patterns (in Python's fnmatch format) to match against to do local gradient accumulation in FP32. The special pattern 'all' matches every parameter.", + choices=None, + nargs="+", + element_type=str, + ), "ddp_reduce_scatter_with_fp32_accumulation": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, use a reduce-scatter implementation which sends lower-precision" - " values over the wire (using an all-to-all to keep total communication" - " overhead in line with the standard ring implementation) but performs" - " accumulation locally in FP32." - ), + help="If set, use a reduce-scatter implementation which sends lower-precision values over the wire (using an all-to-all to keep total communication overhead in line with the standard ring implementation) but performs accumulation locally in FP32.", choices=None, nargs=0, element_type=None, @@ -803,11 +796,7 @@ "decoder_first_pipeline_num_layers": MegatronArgMetadata( arg_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" - ), + 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", choices=None, nargs=None, element_type=None, @@ -815,11 +804,7 @@ "decoder_last_pipeline_num_layers": MegatronArgMetadata( arg_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" - ), + 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", choices=None, nargs=None, element_type=None, @@ -851,10 +836,7 @@ "decoupled_min_lr": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Minimum value for learning rate for the input and output layer. The" - " schedulerclip values below this threshold" - ), + help="Minimum value for learning rate for the input and output layer. The schedulerclip values below this threshold", choices=None, nargs=None, element_type=None, @@ -862,13 +844,7 @@ "decrease_batch_size_if_needed": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, decrease batch size if microbatch_size * dp_sizedoes not divide" - " batch_size. Useful for KSO (Keep Soldiering On)to continue making" - " progress if number of healthy GPUs (andcorresponding dp_size) does not" - " support current batch_size.Old batch_size will be restored if training is" - " re-started withdp_size that divides batch_size // microbatch_size." - ), + help="If set, decrease batch size if microbatch_size * dp_size does not divide batch_size. Old batch_size will be restored if training is re-started with dp_size that divides batch_size // microbatch_size. Not supported with step-batch-size-schedule.", choices=None, nargs=0, element_type=None, @@ -876,10 +852,7 @@ "defer_embedding_wgrad_compute": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, defers the vocabulary projection linear layer weightgradient" - " compute to pipeline flush." - ), + help="If true, defers the embedding WGRAD GEMMs while pipeline flush is taking place enabling us to hide pipeline flush latency. Defaults to False.", choices=None, nargs=0, element_type=None, @@ -887,7 +860,7 @@ "delay_wgrad_compute": MegatronArgMetadata( arg_type=bool, default=False, - help="Delay the wgrad compute for batch-level overlapping", + help="Delay the weight gradient computation to improve batch-level communication overlapping", choices=None, nargs=0, element_type=None, @@ -895,11 +868,7 @@ "deprecated_use_mcore_models": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + help="DEPRECATED. Use the implementation from megatron core.Now ignored and mcore models are the default.", choices=None, nargs=0, element_type=None, @@ -907,10 +876,7 @@ "deterministic_mode": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Choose code that has deterministic execution. This usually means slower" - " execution, but is good for debugging and testing." - ), + help="If true, code that has deterministic execution will be chosen. This usually means slower execution, but is good for debugging and testing. Defaults to False.", choices=None, nargs=0, element_type=None, @@ -990,28 +956,23 @@ "disable_bf16_reduced_precision_matmul": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If True, sets" - " torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False" - " to prevent matmul from using reduced precision accumulation when using" - " BF16." - ), + help="If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to prevent matmul from using reduced precision accumulation when using BF16.", choices=None, nargs=0, element_type=None, ), - "disable_chunked_prefill": MegatronArgMetadata( + "disable_jit_fuser": MegatronArgMetadata( arg_type=bool, default=False, - help="Disable chunked prefill (chunked prefill is enabled by default).", + help="Disable the JIT fuser.", choices=None, nargs=0, element_type=None, ), - "disable_mamba_mem_eff_path": MegatronArgMetadata( + "disable_msc_deprecated": MegatronArgMetadata( arg_type=bool, default=False, - help="Disable Mamba efficient path.", + help="[DEPRECATED] MSC is disabled by default; this flag is a no-op and will be removed in a future release.", choices=None, nargs=0, element_type=None, @@ -1027,11 +988,7 @@ "disable_symmetric_registration": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Disable symmetric (window) registration for NCCL userbuffer" - " registration.This option will force to use conventional (local)" - " userbuffer registration when use-nccl-ub is set." - ), + help="Disable symmetric (window) registration for NCCL userbuffer registration.This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.", choices=None, nargs=0, element_type=None, @@ -1047,10 +1004,7 @@ "dist_ckpt_optim_fully_reshardable": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as" - " opposed to plain DP reshardability." - ), + help="Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as opposed to plain DP reshardability.", choices=None, nargs=0, element_type=None, @@ -1058,11 +1012,7 @@ "dist_ckpt_save_pre_mcore_014": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Revert checkpointing simplifications introduced in Megatron-Core v0.14." - " This option affects only checkpoint saving format and will be removed" - " soon (checkpoint load format is determined based on checkpoint metadata)." - ), + help="Revert checkpointing simplifications introduced in Megatron-Core v0.14. This option affects only checkpoint saving format and will be removed soon (checkpoint load format is determined based on checkpoint metadata).", choices=None, nargs=0, element_type=None, @@ -1070,12 +1020,7 @@ "dist_ckpt_strictness": MegatronArgMetadata( arg_type=str, default="assume_ok_unexpected", - 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." - ), + 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.", choices=( "assume_ok_unexpected", "log_unexpected", @@ -1089,15 +1034,18 @@ nargs=None, element_type=None, ), + "dist_ckpt_workers": MegatronArgMetadata( + arg_type=int, + default=1, + help="Number of workers for distributed checkpointing. Only used for async save. If set to 1, the checkpointing is performed in a single process.", + choices=None, + nargs=None, + element_type=None, + ), "distrib_optim_fully_reshardable_mem_efficient": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "During distributed optimizer checkpoint save and load tries to use as" - " little memory as possible by using Gloo (instead of NCCL) and only one" - " rank for saving. Turn on only if experiencing host or device memory" - " issues. Has affect only with `--dist-ckpt-optim-fully-reshardable` flag." - ), + help="During distributed optimizer checkpoint save and load tries to use as little memory as possible by using Gloo (instead of NCCL) and only one rank for saving. Turn on only if experiencing host or device memory issues. Has affect only with `dist_ckpt_optim_fully_reshardable` flag.", choices=None, nargs=0, element_type=None, @@ -1105,13 +1053,13 @@ "distribute_saved_activations": MegatronArgMetadata( arg_type=bool, default=False, - help="If set, distribute recomputed activations across model parallel group.", + help="If True, distribute recomputed activations across the model parallel group.", choices=None, nargs=0, element_type=None, ), "distributed_backend": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="nccl", help="Which backend to use for distributed training.", choices=("nccl", "gloo"), @@ -1121,7 +1069,7 @@ "distributed_timeout_minutes": MegatronArgMetadata( arg_type=int, default=10, - help="Default timeout minutes for torch.distributed.", + help="Timeout minutes for torch.distributed.", choices=None, nargs=None, element_type=None, @@ -1129,22 +1077,55 @@ "distributed_timeout_seconds_after_init": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Timeout seconds for process groups after initialization.This timeout is" - " applied to all process groups after initialization." - ), + help="Timeout in seconds for process groups after initialization. This timeout is applied to all process groups after initialization and the first iteration completes.", + choices=None, + nargs=None, + element_type=None, + ), + "dsa_indexer_head_dim": MegatronArgMetadata( + arg_type=int, + default=None, + help="Dimension per DSA indexer head.", + choices=None, + nargs=None, + element_type=None, + ), + "dsa_indexer_loss_coeff": MegatronArgMetadata( + arg_type=float, + default=None, + help="Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss.", + choices=None, + nargs=None, + element_type=None, + ), + "dsa_indexer_n_heads": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of DSA indexer heads.", + choices=None, + nargs=None, + element_type=None, + ), + "dsa_indexer_topk": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of top-k tokens to select in DSA indexer.", choices=None, nargs=None, element_type=None, ), + "dsa_indexer_use_sparse_loss": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices.", + choices=None, + nargs=0, + element_type=None, + ), "dump_param_to_param_group_map": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Path to a file containing parameter-to-parameter-group mapping. Provide a" - " JSON file that specifies which parameters belong to which parameter group" - " for global coordination." - ), + help="Path to a file containing parameter-to-parameter-group mapping. Provide a JSON file that specifies which parameters belong to which parameter group for global coordination.", choices=None, nargs=None, element_type=None, @@ -1152,14 +1133,7 @@ "embedding_init_method_std": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Standard deviation of the zero mean normal distribution used for embedding" - " weight initialization. If unset, embeddings will be initialized the same" - " way as other weights. Setting this to a value around 1.0 may avoid loss" - " spikes in training. Setting this to any value will also skip applying" - " weight decay on embedding weights to avoid shrinkage towards zero. See" - " https://arxiv.org/abs/2312.16903 for more details." - ), + help="Standard deviation of the zero mean normal for the default initialization method for the embedding layer. If None, will be set to init_method_std. Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying weight decay on embedding weights to avoid shrinkage towards zero. See https://arxiv.org/abs/2312.16903 for more details.", choices=None, nargs=None, element_type=None, @@ -1175,68 +1149,55 @@ "empty_unused_memory_level": MegatronArgMetadata( arg_type=int, default=0, - help=( - "Call torch.cuda.empty_cache() each iteration (training and eval), to" - " reduce fragmentation.0=off, 1=moderate, 2=aggressive." - ), + help="Call torch.cuda.empty_cache() each iteration (training and eval), to reduce fragmentation. 0=off, 1=moderate, 2=aggressive.", choices=(0, 1, 2), nargs=None, element_type=None, ), - "enable_cuda_graph": MegatronArgMetadata( + "enable_chunked_prefill": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Deprecated. Use --cuda-graph-impl=local instead. Use local implementation" - ' of CUDA graph capture and replay. --cuda-graph-scope="full_iteration"' - " enables whole iteration CUDA graph. " - ), + help="Enable chunked prefill (disabled by default)", choices=None, nargs=0, element_type=None, ), - "enable_experimental": MegatronArgMetadata( + "enable_cuda_graph": MegatronArgMetadata( arg_type=bool, default=False, - help="Enable experimental features.", + help="DEPRECATED and replaced by cuda_graph_impl. When set to true, either partial CUDA graph (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole iteration excluding optimizer) is enabled. cuda_graph_modules determines the scope of graph capture.", choices=None, nargs=0, element_type=None, ), - "enable_ft_package": MegatronArgMetadata( + "enable_experimental": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, Fault Tolerance package is enabled. Note: This feature is for" - " Nvidia internal use only." - ), + help="Enable experimental features.", choices=None, nargs=0, element_type=None, ), - "enable_full_sharding_in_hsdp": MegatronArgMetadata( + "enable_ft_package": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel" - " (HSDP) mode." - ), + help="If set, Fault Tolerance package is enabled. Note: This feature is for Nvidia internal use only.", choices=None, nargs=0, element_type=None, ), - "enable_gloo_process_groups": MegatronArgMetadata( + "enable_full_sharding_in_hsdp": MegatronArgMetadata( arg_type=bool, - default=True, - help="Disables creation and usage of Gloo process groups.", + default=False, + help="If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.", choices=None, nargs=0, element_type=None, ), "enable_msc": MegatronArgMetadata( arg_type=bool, - default=True, - help="Disable the usage of Multi-Storage Client (MSC) in Megatron Core.", + default=False, + help="Enable the usage of Multi-Storage Client (MSC) in Megatron Core. Disabled by default; pass this flag to opt in.", choices=None, nargs=0, element_type=None, @@ -1244,13 +1205,7 @@ "enable_one_logger": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "If set, disable using one_logger to track E2E metricsNote that one_logger" - " is an internal tool and not available externally. For installation," - " please go to" - " https://confluence.nvidia.com/display/MLWFO/Package+Repositoriesfor more" - " details" - ), + help="If set, disable using one_logger to track E2E metricsNote that one_logger is an internal tool and not available externally. For installation, please go to https://confluence.nvidia.com/display/MLWFO/Package+Repositoriesfor more details", choices=None, nargs=0, element_type=None, @@ -1266,9 +1221,7 @@ "encoder_seq_length": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Maximum encoder sequence length to process.This should be exclusive of --seq-length" - ), + help="Maximum encoder sequence length to process.This should be exclusive of --seq-length", choices=None, nargs=None, element_type=None, @@ -1289,13 +1242,18 @@ nargs=0, element_type=None, ), + "ep_overlap_early_attn_memory_release": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable early memory release of attention activations during EP overlap. EP overlap can increase peak memory usage when the overlapped forward module allocates more memory than what is freed by the backward module. This flag addresses this by reordering the attention backward pass to occur earlier in the schedule. Specifically: - Without this flag: attn_bwd executes after moe_combine_fwd - With this flag: attn_bwd executes before mlp_fwd The earlier execution releases attention activations sooner, reducing peak memory. Note: This may impact performance as moe_combine_fwd and moe_dispatch_bwd become exposed (not overlapped with other computation).", + choices=None, + nargs=0, + element_type=None, + ), "error_injection_rate": MegatronArgMetadata( arg_type=int, default=0, - help=( - "Rate at which to inject unexpected results, e.g. 1000 means once every" - " 1000 result validations" - ), + help="Rate at which to inject unexpected results, e.g. 1000 means once every 1000 result validations", choices=None, nargs=None, element_type=None, @@ -1303,15 +1261,23 @@ "error_injection_type": MegatronArgMetadata( arg_type=str, default="transient_error", - help="Type of error to inject. ", + help="Type of error to inject.", choices=("correct_result", "transient_error", "persistent_error"), nargs=None, element_type=None, ), + "eval_global_batch_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Global batch size to use during evaluation. If not set, defaults to global_batch_size. Must be divisible by (eval_micro_batch_size * data_parallel_size).", + choices=None, + nargs=None, + element_type=None, + ), "eval_interval": MegatronArgMetadata( arg_type=int, - default=1000, - help="Interval between running evaluation on validation set.", + default=None, + help="Interval between running evaluation on validation set. If not set, evaluation will not run during training.", choices=None, nargs=None, element_type=None, @@ -1319,7 +1285,15 @@ "eval_iters": MegatronArgMetadata( arg_type=int, default=100, - help="Number of iterations to run for evaluationvalidation/test for.", + help="Number of iterations to run for evaluation. Used for both validation and test. If not set, evaluation will not run.", + choices=None, + nargs=None, + element_type=None, + ), + "eval_micro_batch_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Micro batch size to use during evaluation. If not set, defaults to micro_batch_size. Changing this affects per-device memory usage during eval and the number of microbatches per eval step.", choices=None, nargs=None, element_type=None, @@ -1351,26 +1325,73 @@ "exit_on_missing_checkpoint": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If '--load' is set, but checkpoint is not found (e.g., path typo), then" - " exit instead of random initialization." - ), + help="If 'load' is set, but checkpoint is not found (e.g., path typo), then exit instead of random initialization.", choices=None, nargs=0, element_type=None, ), "exit_signal": MegatronArgMetadata( - arg_type=str, - default="SIGTERM", - help=("Signal to use for exit signal handler. If not specified, defaults to SIGTERM."), - choices=("SIGTERM", "SIGINT", "SIGUSR1", "SIGUSR2"), + arg_type=None, + default="15", + help="Signal for the signal handler to detect.", + choices=( + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 34, + 64, + ), nargs=None, element_type=None, ), "exit_signal_handler": MegatronArgMetadata( arg_type=bool, default=False, - help=("Dynamically save the checkpoint and shutdown the training if signal is received"), + help="Dynamically save the checkpoint and shutdown the training if SIGTERM is received", + choices=None, + nargs=0, + element_type=None, + ), + "exit_signal_handler_for_dataloader": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Use signal handler for dataloader workers", + choices=None, + nargs=0, + element_type=None, + ), + "exit_signal_handler_for_training": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Shutdown the training when SIGINT or SIGTERM received to avoid unclear traceback", choices=None, nargs=0, element_type=None, @@ -1378,12 +1399,7 @@ "exp_avg_dtype": MegatronArgMetadata( arg_type=None, default="fp32", - help=( - "Dtype of exp_avg (1st moment in adam optimizer) when enabling" - " precision-aware-optimizer. This dtype is used for storing the optimizer" - " state in memory during training but does not affect the precision in the" - " kernel computation." - ), + help="Dtype of exp_avg (1st moment in adam optimizer) when enabling precision-aware-optimizer. This dtype is used for storing the optimizer state in memory during training but does not affect the precision in the kernel computation.", choices=("fp32", "fp16", "bf16", "fp8"), nargs=None, element_type=None, @@ -1391,20 +1407,23 @@ "exp_avg_sq_dtype": MegatronArgMetadata( arg_type=None, default="fp32", - help=( - "Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling" - " precision-aware-optimizer. This dtype is used for storing the optimizer" - " state in memory during training but does not affect the precision in the" - " kernel computation." - ), + help="Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling precision-aware-optimizer. This dtype is used for storing the optimizer state in memory during training but does not affect the precision in the kernel computation.", choices=("fp32", "fp16", "bf16", "fp8"), nargs=None, element_type=None, ), + "experimental_attention_variant": MegatronArgMetadata( + arg_type=str, + default=None, + help="Type of attention variant to use. Currently support gated_delta_net and dsa.", + choices=("gated_delta_net", "dsa"), + nargs=None, + element_type=None, + ), "expert_model_parallel_size": MegatronArgMetadata( arg_type=int, default=1, - help="Degree of expert model parallelism.", + help="Distributes Moe Experts across sub data parallel dimension.", choices=None, nargs=None, element_type=None, @@ -1412,10 +1431,7 @@ "expert_tensor_parallel_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Degree of expert model parallelism. Default is None, which will be set to" - " the value of --tensor-model-paralle-size." - ), + help="Intra-layer tensor model parallelism for expert layer. Splits tensors across GPU ranks. Default is None, which will be set to the value of tensor_model_parallel_size.", choices=None, nargs=None, element_type=None, @@ -1423,10 +1439,7 @@ "external_cuda_graph": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Deprecated. Use --cuda-graph-impl=transformer_engine instead. Use TE" - " make_graphed_callables() to capture the CUDA graph." - ), + help="DEPRECATED and replaced by cuda_graph_impl. When set to true, TransformerLayer layers are swapped with user provided CUDA graphs.", choices=None, nargs=0, element_type=None, @@ -1434,24 +1447,87 @@ "fake_process_group": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, initialize with fake distributed process group and all distributed" - " communication operations will be skipped. This is" - " quite useful for profiling memory usage of distributed training with just" - " one GPU. Setting WORLD_SIZE and RANK to the" - " specific values for target distribtued scale." - ), + help="If set, initialize with fake distributed process group and all distributed communication operations will be skipped. This is quite useful for profiling memory usage of distributed training with just one GPU. Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.", choices=None, nargs=0, element_type=None, ), + "fault_injector_delay_start_iteration": MegatronArgMetadata( + arg_type=int, + default=None, + help="Start the fault delay timer after iteration N completes. If unset, fault delay timing starts from the beginning of training.", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_fault_delay": MegatronArgMetadata( + arg_type=float, + default=None, + help="Force a specific fault delay in seconds from training start or delay_start_iteration.", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_fault_probabilities": MegatronArgMetadata( + arg_type=str, + default=None, + help="Comma-separated list of fault probabilities (normalized at runtime).", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_fault_types": MegatronArgMetadata( + arg_type=str, + default=None, + help="Comma-separated list of fault types to inject (e.g. 'hang,crash').", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_mtti_seconds": MegatronArgMetadata( + arg_type=float, + default=None, + help="Mean time to inject (MTTI) in seconds; used when fault_delay is None.", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_num_ranks": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of ranks to inject faults on (random selection).", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_offset_seconds": MegatronArgMetadata( + arg_type=float, + default=None, + help="Offset seconds added to the sampled fault delay.", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_ranks": MegatronArgMetadata( + arg_type=str, + default=None, + help="Comma-separated list of ranks to inject faults on.", + choices=None, + nargs=None, + element_type=None, + ), + "fault_injector_seed": MegatronArgMetadata( + arg_type=int, + default=None, + help="RNG seed for the fault injector.", + choices=None, + nargs=None, + element_type=None, + ), "ffn_hidden_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Transformer Feed-Forward Network hidden size. This is set to 4*hidden-size" - " if not provided" - ), + help="Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size if not provided.", choices=None, nargs=None, element_type=None, @@ -1531,7 +1607,7 @@ "fim_spm_rate": MegatronArgMetadata( arg_type=float, default=0.5, - help=("Probability that the a FIM sample uses the SPM format over the PSM format."), + help="Probability that the a FIM sample uses the SPM format over the PSM format.", choices=None, nargs=None, element_type=None, @@ -1544,14 +1620,26 @@ nargs=None, element_type=None, ), + "fine_grained_activation_offloading": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, offload the input of the specified modules to the CPU. Fine-grained activation offloading is a module-level offloading method instead of a layer-level offloading method like cpu_offloading.", + choices=None, + nargs=0, + element_type=None, + ), + "fine_grained_offloading_max_inflight_offloads": MegatronArgMetadata( + arg_type=int, + default=None, + help="Per fine-grained offloading group name, max number of inflight offloads for that name not yet joined on the main stream (wait_event on D2H). The same cap applies to every name (e.g., ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait after every offload for that name. 1 = at most one not-yet-waited offload per name, etc. None = do not insert these joins. This feature is particularly useful when using with full-iteration CUDA graphs", + choices=None, + nargs=None, + element_type=None, + ), "finetune": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -1559,7 +1647,7 @@ "first_last_layers_bf16": MegatronArgMetadata( arg_type=bool, default=False, - help="Construct first and last layers in bf16 when doing FP8 training.", + help="If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8.", choices=None, nargs=0, element_type=None, @@ -1567,15 +1655,63 @@ "flash_decode": MegatronArgMetadata( arg_type=bool, default=False, - help="Whether to use the flash decoding kernel.", + help="Use the optimized flash decoding kernel during inference.", + choices=None, + nargs=0, + element_type=None, + ), + "flight_recorder_dump_on_timeout": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Dump flight recorder traces on NCCL timeout (TORCH_NCCL_DUMP_ON_TIMEOUT).", + choices=None, + nargs=0, + element_type=None, + ), + "flight_recorder_dump_path": MegatronArgMetadata( + arg_type=str, + default=None, + help="Path for NCCL flight recorder trace dumps. Sets TORCH_FR_DUMP_TEMP_FILE and TORCH_NCCL_DEBUG_INFO_TEMP_FILE env variables before distributed init.", + choices=None, + nargs=None, + element_type=None, + ), + "flight_recorder_extra_dump_on_exec": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Enable extra flight recorder dump on execution (TORCH_NCCL_EXTRA_DUMP_ON_EXEC).", + choices=None, + nargs=0, + element_type=None, + ), + "flight_recorder_include_only_active": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Include only active operations in flight recorder dumps (TORCH_INCLUDE_ONLY_ACTIVE).", + choices=None, + nargs=0, + element_type=None, + ), + "flight_recorder_include_stack_trace": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Include stack traces in flight recorder dumps (TORCH_INCLUDE_STACK_TRACE).", choices=None, nargs=0, element_type=None, ), + "flight_recorder_trace_buffer_size": MegatronArgMetadata( + arg_type=int, + default=2000, + help="Size of the NCCL flight recorder trace buffer (TORCH_NCCL_TRACE_BUFFER_SIZE).", + choices=None, + nargs=None, + element_type=None, + ), "fp16": MegatronArgMetadata( arg_type=bool, default=False, - help="Run model in fp16 mode.", + help="If true, train with fp16 mixed precision training.", choices=None, nargs=0, element_type=None, @@ -1591,61 +1727,55 @@ "fp32_residual_connection": MegatronArgMetadata( arg_type=bool, default=False, - help="Move residual connections to fp32.", + help="If true, move residual connections to fp32.", choices=None, nargs=0, element_type=None, ), "fp4": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, - help=("Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass"), + help="If set, enables the use of FP4 precision through Transformer Engine. Currently only supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0).", choices=("e2m1",), nargs=None, element_type=None, ), - "fp4_param": MegatronArgMetadata( + "fp4_param_gather": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Keep the compute param in fp4 (do not use any other intermediate dtype)" - " and perform the param all-gather in fp4." - ), + help="Keep the compute param in fp4 (do not use any other intermediate dtype) and perform the param all-gather in fp4.", choices=None, nargs=0, element_type=None, ), "fp4_quantizer_factory": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, - help=( - "Python import path to a callable quantizer factory, e.g.," - " package.module.quantizer_factory." - ), + help="Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. Required when fp4_recipe is custom.", choices=None, nargs=None, element_type=None, ), "fp4_recipe": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="nvfp4", - help="Which fp4 recipe to use for FP4 tensors in the forward and backward pass", + help="If set, enables the use of FP4 precision through Transformer Engine. Currently only 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture.", choices=("nvfp4", "custom"), nargs=None, element_type=None, ), "fp8": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, - help=("Which fp8 format scheme to use for FP8 tensors in the forward and backward pass"), + help="If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8 activation and weight tensors and e5m2 for all FP8 output activation gradient tensors.", choices=("e4m3", "hybrid"), nargs=None, element_type=None, ), "fp8_amax_compute_algo": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="most_recent", - help="Algorithm for computing amax from history", + help="Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2 predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent` always chooses the most recently seen value.", choices=("most_recent", "max"), nargs=None, element_type=None, @@ -1653,7 +1783,7 @@ "fp8_amax_history_len": MegatronArgMetadata( arg_type=int, default=1, - help="Number of steps for which amax history is recorded per tensor", + help="The length of the amax history window used for scaling factor computation.", choices=None, nargs=None, element_type=None, @@ -1661,7 +1791,7 @@ "fp8_interval": MegatronArgMetadata( arg_type=int, default=1, - help="DEPRECATED. This flag is ignored. Scaling update interval for fp8", + help="DEPRECATED from TransformerEngine v1.8.0. This flag is ignored. Controls how often the scaling factor is recomputed.", choices=None, nargs=None, element_type=None, @@ -1669,37 +1799,39 @@ "fp8_margin": MegatronArgMetadata( arg_type=int, default=0, - help="Scaling margin for fp8", + help="Margin for the scaling factor computation.", choices=None, nargs=None, element_type=None, ), + "fp8_output_proj": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, run the LM-head output projection with a TE ColumnParallelLinear under the MXFP8 autocast context. Only active when fp8=True and fp8_recipe='mxfp8'.", + choices=None, + nargs=0, + element_type=None, + ), "fp8_param_gather": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Keep the compute param in fp8 (do not use any other intermediate dtype)" - " and perform the param all-gather in fp8." - ), + help="Keep the compute param in fp8 (do not use any other intermediate dtype) and perform the param all-gather in fp8.", choices=None, nargs=0, element_type=None, ), "fp8_quantizer_factory": MegatronArgMetadata( - arg_type=None, + arg_type=str, default=None, - help=( - "Python import path to a callable quantizer factory, e.g.," - " package.module.quantizer_factory." - ), + help="Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. Required when fp8_recipe is custom.", choices=None, nargs=None, element_type=None, ), "fp8_recipe": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="delayed", - help="Which fp8 recipe to use for FP8 tensors in the forward and backward pass", + help="If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed' uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, 4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe.", choices=("tensorwise", "delayed", "mxfp8", "blockwise", "custom"), nargs=None, element_type=None, @@ -1707,7 +1839,7 @@ "fp8_wgrad": MegatronArgMetadata( arg_type=bool, default=True, - help="Execute wgrad in higher precision even for FP8 runs", + help="When set to False, override FP8 config options and do the wgrad computation in higher precision.", choices=None, nargs=0, element_type=None, @@ -1715,25 +1847,39 @@ "fsdp_double_buffer": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Enable double buffering for temporary memory needed for Megatron FSDP" - " communications. Double-buffering the communication memory improves memory" - " management efficiency by reusing previously allocated buffers, rather" - " than creating new buffers for each FSDP communication. This is required" - " for user buffer registration and is enabled by default when using NCCL" - " user buffers." - ), + help="Enable double buffering for temporary memory needed for Megatron FSDP communications. Double-buffering the communication memory improves memory management efficiency by reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. This is required for user buffer registration and is enabled by default when using NCCL user buffers.", + choices=None, + nargs=0, + element_type=None, + ), + "fsdp_manual_registration": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Manually register the FSDP communication buffers to NCCL user buffer.This option is only effective when use-megatron-fsdp and use-nccl-ub is set.", choices=None, nargs=0, element_type=None, ), + "ft_num_warmup_iters": MegatronArgMetadata( + arg_type=int, + default=5, + help="Number of warmup iterations before monitoring step section and out-of-section timeouts. The first N iterations are excluded from timeout monitoring as they can be significantly slower than steady-state. Default: 5. Note: This feature is for Nvidia internal use only.", + choices=None, + nargs=None, + element_type=None, + ), "full_validation": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, each time validation occurs it uses the full validation" - " dataset(s). This currently only works for GPT datasets!" - ), + help="If set, each time validation occurs it uses the full validation dataset(s). This currently only works for GPT datasets!", + choices=None, + nargs=0, + element_type=None, + ), + "fused_residual_rmsnorm": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, fuses residual connection and RMSNorm backward pass when TE is used.", choices=None, nargs=0, element_type=None, @@ -1741,12 +1887,7 @@ "global_batch_size": MegatronArgMetadata( arg_type=int, default=None, - 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 use" - " micro-batch-size * data-parallel-size as the global batch size. This" - " choice will result in 1 for number of micro-batches." - ), + 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 use micro-batch-size * data-parallel-size as the global batch size. This choice will result in 1 for number of micro-batches.", choices=None, nargs=None, element_type=None, @@ -1754,10 +1895,15 @@ "glu_linear_offset": MegatronArgMetadata( arg_type=float, default=0.0, - help=( - "Offset term in the GLU activation function: activation_func(x[0]) * (x[1]" - " + offset). Only used when gated_linear_unit is True" - ), + help="Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only used when gated_linear_unit is True", + choices=None, + nargs=None, + element_type=None, + ), + "gpu_sniff_test_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Interval (in iterations) to run GPU performance sniff tests (GEMMs, all-to-all, send/recv). Flags any rank whose throughput differs from the mean by more than one standard deviation. If not specified, sniff tests are not run.", choices=None, nargs=None, element_type=None, @@ -1773,9 +1919,7 @@ "gradient_accumulation_fusion": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "Disable fusing gradient accumulation to weight gradient computation of linear layers" - ), + help="Disable fusing gradient accumulation to weight gradient computation of linear layers", choices=None, nargs=0, element_type=None, @@ -1807,23 +1951,7 @@ "grpo_clamp_eps_upper": MegatronArgMetadata( arg_type=float, default=0.01, - help=("Upper GRPO clipping bound. In vanilla implementation, equals to the lower one."), - choices=None, - nargs=None, - element_type=None, - ), - "grpo_default_temperature": MegatronArgMetadata( - arg_type=float, - default=1.0, - help="Default temperature for model inference.", - choices=None, - nargs=None, - element_type=None, - ), - "grpo_default_top_p": MegatronArgMetadata( - arg_type=float, - default=0, - help="Default top-p for model inference.", + help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.", choices=None, nargs=None, element_type=None, @@ -1887,13 +2015,7 @@ "heterogeneous_layers_config_encoded_json": MegatronArgMetadata( arg_type=str, default=None, - help=( - "This is encoded json string of the heterogeneous model configuration. Used" - " to keep the content of the heterogeneous model specification in args when" - " the model is loaded from a checkpoint. Use the format of the HuggingFace" - " config files in llama nemotron models, e.g." - " https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json." - ), + help="This is encoded json string of the heterogeneous model configuration. Used to keep the content of the heterogeneous model specification in args when the model is loaded from a checkpoint. Use the format of the HuggingFace config files in llama nemotron models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.", choices=None, nargs=None, element_type=None, @@ -1901,11 +2023,7 @@ "heterogeneous_layers_config_path": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Path to json file containing heterogeneous model configuration. Use the" - " format of the HuggingFace config files in llama nemotron models, e.g." - " https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json." - ), + help="Path to json file containing heterogeneous model configuration. Use the format of the HuggingFace config files in llama nemotron models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.", choices=None, nargs=None, element_type=None, @@ -1913,7 +2031,7 @@ "hidden_dropout": MegatronArgMetadata( arg_type=float, default=0.1, - help="Dropout probability for hidden state transformer.", + help="Dropout probability for transformer hidden state.", choices=None, nargs=None, element_type=None, @@ -1929,37 +2047,39 @@ "hierarchical_context_parallel_sizes": MegatronArgMetadata( arg_type=list, default=None, - help=( - "Degrees of the hierarchical context parallelism. Users should provide a" - " list to specify the sizes for different levels." - " --hierarchical-context-parallel-sizes 2 4 indicates every two adjacent" - " gpus forms the first level of cp groups and the cp ranks with the same" - " odevity forms the second level of cp groups." - ), + help="Degrees of the hierarchical context parallelism. Users should provide a list to specify the sizes for different levels. Taking the a2a+p2p cp comm type as example, it contains groups of two levels, so the first value of the list indicates the group size of the a2a communication type, and the second value indicates the group size of the p2p communication type.", choices=None, nargs="+", element_type=int, ), + "high_priority_a2a_comm_stream": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, the communication stream created by set_streams for combined 1f1b a2a overlap is created with CUDA high priority.", + choices=None, + nargs=0, + element_type=None, + ), "high_priority_stream_groups": MegatronArgMetadata( arg_type=list, default=[], - help="The communicator group names to use high priority streams.", + help="Specify which communicator groups should use high priority streams during creation. Assigning high priority to communication streams ensures that communication kernels are scheduled with higher priority, minimizing the exposed communication when it is overlapped with other computation kernels.", choices=None, - nargs="*", + nargs="+", element_type=str, ), - "hybrid_attention_ratio": MegatronArgMetadata( - arg_type=float, - default=0.0, - help="Ratio of attention layers to total layers, in the range [0.0, 1.0].", + "hybrid_context_parallel": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If true, enables hybrid context parallel. This is used to balance the workload of each CP rank when we use packed samples with variable sequence lengths. Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel.", choices=None, - nargs=None, + nargs=0, element_type=None, ), - "hybrid_mlp_ratio": MegatronArgMetadata( - arg_type=float, - default=0.0, - help="Ratio of mlp layers to total layers, in the range [0.0, 1.0].", + "hybrid_layer_pattern": MegatronArgMetadata( + arg_type=str, + default=None, + help='Specify a hybrid layer pattern using M (mamba), G (gdn), * (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline stage boundaries for flexible virtual pipeline parallel (fVPP). Use / to separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". When this flag is used, it is the sole indicator that a hybrid model is being run.', choices=None, nargs=None, element_type=None, @@ -1967,13 +2087,7 @@ "hybrid_override_pattern": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Force a specific hybrid layer pattern. The valueshould be a string of" - " characters chosen fromcore.ssm.mamba_hybrid_layer_allocation.Symbols.If a" - " value greater than 0.0 is supplied to any of the hybrid ratio arguments," - " then the number of each typeof layer in the override pattern must match" - " number inthe overidden pattern" - ), + help="Deprecated. Use --hybrid-layer-pattern instead. If specified, its value will be forwarded to --hybrid-layer-pattern.", choices=None, nargs=None, element_type=None, @@ -1989,7 +2103,7 @@ "ict_head_size": MegatronArgMetadata( arg_type=int, default=None, - help=("Size of block embeddings to be used in ICT and REALM (paper default: 128)"), + help="Size of block embeddings to be used in ICT and REALM (paper default: 128)", choices=None, nargs=None, element_type=None, @@ -2037,24 +2151,51 @@ "inference_batch_times_seqlen_threshold": MegatronArgMetadata( arg_type=int, default=-1, - help=( - "If (batch-size * sequence-length) is smaller than this thresholdthen" - " batches will not be split up for pipelining.Requires setting" - " --pipeline-model-parallel-size > 1.Setting this to -1 indicates that" - " batch pipelining is not used." - ), + help="If (batch-size * sequence-length) is smaller than this thresholdthen batches will not be split up for pipelining.Requires setting --pipeline-model-parallel-size > 1.Setting this to -1 indicates that batch pipelining is not used.", choices=None, nargs=None, element_type=None, ), "inference_coordinator_port": MegatronArgMetadata( arg_type=int, - default=12346, + default=None, help="This port will be used to setup the inference coordinator on node-0", choices=None, nargs=None, element_type=None, ), + "inference_cuda_graph_all_prefills": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Extend prefill/mixed CUDA graph capture up to `max_tokens`. By default, all graphs are limited by the decode limit of `max_requests * (num_speculative_tokens + 1)`.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_cuda_graph_scope": MegatronArgMetadata( + arg_type=str, + default=None, + help='Controls the CUDA graph scope during inference. When unset, the effective default is derived from cuda_graph_impl: "local" -> "layer", all other impls -> "none". "none": inference runs in eager mode (no CUDA graphs). "layer": inference graphs are owned at the module/layer boundary, e.g. TransformerLayer or MambaLayer. "block": inference graphs are owned by the enclosing block, e.g. TransformerBlock or HybridBlock. Currently supported combinations are: cuda_graph_impl="local" -> "layer" or "block"; all other cuda_graph_impl values -> "none".', + choices=("none", "layer", "block"), + nargs=None, + element_type=None, + ), + "inference_disable_ep_consensus": MegatronArgMetadata( + arg_type=None, + default=False, + help="Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. Pause/unpause take effect as soon as the signal is delivered to a rank. Only safe when EP coordination is not required (e.g. ep_world_size == 1).", + choices=None, + nargs=0, + element_type=None, + ), + "inference_disable_triton_nvls_kernels": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If true, disables the use of Triton NVLS kernels during inference.", + choices=None, + nargs=0, + element_type=None, + ), "inference_dynamic_batching": MegatronArgMetadata( arg_type=bool, default=False, @@ -2074,31 +2215,39 @@ "inference_dynamic_batching_buffer_size_gb": MegatronArgMetadata( arg_type=float, default=40.0, - help=( - "Amount of on-GPU memory allocated for the KV cache. The total amount of" - " memory allocated for the KV cache (CPU + GPU memory) depends on the value" - " set for the unified virtual memory (UVM) level (via" - " `--inference-dynamic-batching-unified-memory-level`).If the UVM level is" - " 0, then only GPU memory is used and the total memory equals" - " `buffer_size_gb`. If the UVM level is 1, then additional memory is" - " utilized on the CPU and the total memory equals `2 * buffer_size_gb`." - ), + help="Amount of on-GPU memory allocated for the KV cache. The total amount of memory allocated for the KV cache (CPU + GPU memory) depends on the value set for the unified virtual memory (UVM) level (via `--inference-dynamic-batching-unified-memory-level`).If the UVM level is 0, then only GPU memory is used and the total memory equals `buffer_size_gb`. If the UVM level is 1, then additional memory is utilized on the CPU and the total memory equals `buffer_size_gb + paused_buffer_size_gb`.", choices=None, nargs=None, element_type=None, ), - "inference_dynamic_batching_cuda_graph_max_tokens": MegatronArgMetadata( + "inference_dynamic_batching_cuda_graph_mixed_prefill_count": MegatronArgMetadata( arg_type=int, - default=16384, - help="Maximum number of tokens to capture in a cuda graph.", + default=16, + help="Number of mixed prefill requests to capture in a cuda graph.", choices=None, nargs=None, element_type=None, ), - "inference_dynamic_batching_cuda_graph_mixed_prefill_count": MegatronArgMetadata( + "inference_dynamic_batching_enable_prefix_caching": MegatronArgMetadata( + arg_type=None, + default=False, + help="Enable/disable prefix caching for dynamic batching inference. When disabled, KV cache blocks cannot be shared between requests with identical prompt prefixes.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_dynamic_batching_mamba_memory_ratio": MegatronArgMetadata( + arg_type=float, + default=None, + help="Percentage of memory buffer to allocate for Mamba states. If not specified, allocates Mamba state tensors for each KV cache block. Only used for hybrid models.", + choices=None, + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_max_requests": MegatronArgMetadata( arg_type=int, - default=16, - help="Number of mixed prefill requests to capture in a cuda graph.", + default=None, + help="Override the inference context's `max_requests`. By default, `max_requests` is set to the number of blocks in the context's memory buffer.", choices=None, nargs=None, element_type=None, @@ -2114,44 +2263,111 @@ "inference_dynamic_batching_num_cuda_graphs": MegatronArgMetadata( arg_type=int, default=16, - help=( - "Maximum number of cuda graphs to capture, where the cuda graph batch sizes" - " range from 1 to `max_requests`. (See `dynamic_context.py` for details on" - " how `max_requests` is computed). Due to rounding, the actual number of" - " cuda graphs may not equal this argument." - ), + help="Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to `max_requests`. (See `dynamic_context.py` for details on how `max_requests` is computed). Due to rounding, the actual number of cuda graphs may not equal this argument.The user can also pass -1, in which case we automatically determine the number of graphs to capture based on the `max_requests`.", choices=None, nargs=None, element_type=None, ), - "inference_dynamic_batching_track_paused_request_events": MegatronArgMetadata( - arg_type=bool, - default=False, - help=( - "Track paused request ids by adding 'paused' events to each request's event" - " history. This has a very minor impact on latency." - ), + "inference_dynamic_batching_paused_buffer_size_gb": MegatronArgMetadata( + arg_type=float, + default=None, + help="Amount of memory reserved for paused requests in the dynamic inference context. Active requests are paused when there are not enough active blocks available to continue generating a request.", choices=None, - nargs=0, + nargs=None, element_type=None, ), - "inference_dynamic_batching_unified_memory_level": MegatronArgMetadata( - arg_type=int, - default=1, - help=( - "Set unified memory usage within the dynamic inference context. The levels" - " are: 0) no unified memory, 1) allocate `memory_buffer` in unified memory." - " Eventually, additional levels will be included to control other tensors" - " within the context." - ), + "inference_dynamic_batching_prefix_caching_coordinator_policy": MegatronArgMetadata( + arg_type=str, + default="first_prefix_block", + help='Coordinator routing policy for prefix caching. "first_prefix_block" (default) routes based on the first block hash only. "longest_prefix" routes to the rank with the longest matching prefix. "round_robin" ignores prefix affinity and cycles through ranks.', + choices=("longest_prefix", "first_prefix_block", "round_robin"), + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_prefix_caching_eviction_policy": MegatronArgMetadata( + arg_type=str, + default="ref_zero", + help='Eviction policy for prefix caching blocks. "ref_zero" (default) immediately returns blocks to the free pool when ref_count hits 0. "lru" keeps blocks cached and evicts via LRU only when space is needed.', + choices=("ref_zero", "lru"), + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_prefix_caching_mamba_gb": MegatronArgMetadata( + arg_type=float, + default=None, + help="GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid models. When set, Mamba states at block boundaries are cached for reuse.", + choices=None, + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_prefix_caching_routing_alpha": MegatronArgMetadata( + arg_type=float, + default=0.5, + help="Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * normalized_load. Higher alpha favors prefix cache hits; lower alpha favors load balance. Default: 0.5.", + choices=None, + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_sampling_backend": MegatronArgMetadata( + arg_type=str, + default="torch", + help='Which sampling kernels to use during inference. Falls back to "torch" with a warning if "flashinfer" is requested but the package is not installed.', + choices=("torch", "flashinfer"), + nargs=None, + element_type=None, + ), + "inference_dynamic_batching_track_generated_token_events": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Track per-token events with timestamps for each generated token. When enabled, each generated token creates a GENERATED_TOKEN event with a timestamp, useful for per-token latency analysis.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_dynamic_batching_track_paused_request_events": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Track paused request ids by adding 'paused' events to each request's event history. This has a very minor impact on latency.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_dynamic_batching_unified_memory_level": MegatronArgMetadata( + arg_type=int, + default=0, + help="Set unified memory usage within the dynamic inference context. The levels are: 0) no unified memory, 1) allocate `memory_buffer` in unified memory. Eventually, additional levels will be included to control other tensors within the context.", choices=(0, 1), nargs=None, element_type=None, ), - "inference_max_batch_size": MegatronArgMetadata( + "inference_fuse_tp_communication": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If true, uses a fused reduce-scatter-residual-norm-allgather kernel during inference.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_grouped_gemm_backend": MegatronArgMetadata( + arg_type=str, + default="vllm", + help="Specifies the backend to use for grouped GEMM operations during inference. Options: - 'flashinfer': Uses FlashInfer cutlass_fused_moe. Not compatible with MXFP8. - 'torch': Uses torch.nn.functional.grouped_mm (mcore_fused_moe with Triton kernels). Supports both BF16 and MXFP8. - 'vllm': Uses vLLM's Triton fused MoE kernel (BF16). Avoids physical token permutation via indirect addressing.", + choices=("flashinfer", "torch", "vllm"), + nargs=None, + element_type=None, + ), + "inference_logging_step_interval": MegatronArgMetadata( + arg_type=int, + default=0, + help="Step interval for logging inference metrics. Default to 0 to disable inference logging.", + choices=None, + nargs=None, + element_type=None, + ), + "inference_max_requests": MegatronArgMetadata( arg_type=int, default=8, - help="Maximum batch size for inference.", + help="Maximum number of requests for inference.", choices=None, nargs=None, element_type=None, @@ -2164,6 +2380,22 @@ nargs=None, element_type=None, ), + "inference_moe_disable_fused_quant_kernels": MegatronArgMetadata( + arg_type=bool, + default=False, + help="When False (default), use fused kernels that combine permute/activation with MXFP8 quantization + swizzle into a single kernel launch. Only applies when fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel launches (useful for debugging).", + choices=None, + nargs=0, + element_type=None, + ), + "inference_moe_token_dispatcher_type": MegatronArgMetadata( + arg_type=str, + default="nvls", + help="Token dispatcher to use for MoE expert parallelism during inference. - 'nccl': AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires decode-only CUDA graphs (forced automatically). - 'nvls': Variable-count AllGather-V/ReduceScatter-V via NVLS multimem kernels. Requires Hopper+ GPUs with NVLink and symmetric memory. Default. Only applies when transformer_impl='inference_optimized' and EP > 1.", + choices=("nccl", "nvls"), + nargs=None, + element_type=None, + ), "inference_rng_tracker": MegatronArgMetadata( arg_type=bool, default=False, @@ -2172,24 +2404,34 @@ nargs=0, element_type=None, ), - "inference_wandb_logging_step_interval": MegatronArgMetadata( - arg_type=int, - default=0, - help=( - "Step interval for logging inference metrics to wandb. Default to 0 to" - " disable inference wandb logging." - ), + "inference_text_gen_server_logging": MegatronArgMetadata( + arg_type=None, + default=False, + help="Enable per-request logging in the inference text generation server.", choices=None, - nargs=None, + nargs=0, + element_type=None, + ), + "inference_use_synchronous_zmq_collectives": MegatronArgMetadata( + arg_type=None, + default=False, + help="Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.", + choices=None, + nargs=0, + element_type=None, + ), + "inference_wandb_logging": MegatronArgMetadata( + arg_type=None, + default=False, + help="Enable inference wandb logging.", + choices=None, + nargs=0, element_type=None, ), "init_method_std": MegatronArgMetadata( arg_type=float, default=0.02, - help=( - "Standard deviation of the zero mean normal distribution used for weight" - " initialization." - ), + help="Standard deviation of the zero mean normal for the default initialization method, not used if init_method and output_layer_init_method are provided.", choices=None, nargs=None, element_type=None, @@ -2205,7 +2447,7 @@ "init_model_with_meta_device": MegatronArgMetadata( arg_type=bool, default=False, - help=None, + help="If True, initializes the model with the meta device. This is helpful for training of very large models. This feature is only works when megatron fsdp is turned on.", choices=None, nargs=0, element_type=None, @@ -2221,10 +2463,7 @@ "inprocess_active_world_size": MegatronArgMetadata( arg_type=int, default=1, - help=( - "The number of ranks initially executing the workload. The remaining ranks" - " from the allocation are set aside as warm reserve." - ), + help="The number of ranks initially executing the workload. The remaining ranks from the allocation are set aside as warm reserve.", choices=None, nargs=None, element_type=None, @@ -2288,7 +2527,7 @@ "inprocess_last_call_wait": MegatronArgMetadata( arg_type=float, default=1, - help=("Time interval (in seconds) for other ranks to report concurrent terminal failures."), + help="Time interval (in seconds) for other ranks to report concurrent terminal failures.", choices=None, nargs=None, element_type=None, @@ -2352,7 +2591,7 @@ "is_hybrid_model": MegatronArgMetadata( arg_type=bool, default=False, - help="Indicates whether the model is a hybrid model.", + help="Indicates whether this is a hybrid model.", choices=None, nargs=0, element_type=None, @@ -2368,7 +2607,7 @@ "iterations_to_skip": MegatronArgMetadata( arg_type=list, default=[], - help="List of iterations to skip, empty by default.", + help="List of 1-indexed iterations to skip during training, empty by default.", choices=None, nargs="+", element_type=int, @@ -2381,31 +2620,18 @@ nargs=0, element_type=None, ), - "kitchen_config_file": MegatronArgMetadata( + "kitchen_attention_backend": MegatronArgMetadata( arg_type=str, - default=None, - help=( - "Use the config .yaml file at the specified location to configure kitchen quantization." - ), - choices=None, - nargs=None, - element_type=None, - ), - "kitchen_recipe_number": MegatronArgMetadata( - arg_type=int, - default=None, - help=("Use a default kitchen recipe for all layers as defined by QAT_PARAMS index"), - choices=None, + default="sdpa", + help='Which kitchen attention backend to use when use_kitchen_attention=True. "sdpa" for KitchenDotProductAttention, "fa" for KitchenFlashAttention.', + choices=("sdpa", "fa"), nargs=None, element_type=None, ), "kv_channels": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Projection weights dimension in multi-head attention. This is set to " - " args.hidden_size // args.num_attention_heads if not provided." - ), + help="Projection weights dimension in multi-head attention. This is set to hidden_size // num_attention_heads if not provided.", choices=None, nargs=None, element_type=None, @@ -2426,48 +2652,92 @@ nargs=None, element_type=None, ), - "langrl_external_server": MegatronArgMetadata( - arg_type=None, + "layernorm_epsilon": MegatronArgMetadata( + arg_type=float, + default=1e-05, + help="Epsilon value for any LayerNorm/RMSNorm operations.", + choices=None, + nargs=None, + element_type=None, + ), + "layernorm_zero_centered_gamma": MegatronArgMetadata( + arg_type=bool, default=False, - help=None, + help="If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves numerical stability.", choices=None, nargs=0, element_type=None, ), - "langrl_inference_server_conversation_template": MegatronArgMetadata( - arg_type=str, + "lazy_mpu_init": MegatronArgMetadata( + arg_type=bool, + default=False, + 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.", + choices=None, + nargs=0, + element_type=None, + ), + "linear_attention_freq": MegatronArgMetadata( + arg_type=None, default=None, - help="Conversation template, if using a chat server.", + help='Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. Accepts either: - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer - A string containing a Python list expression that defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an LA layer and 0 indicates a SDPA layer. Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, "([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.', choices=None, nargs=None, element_type=None, ), - "langrl_inference_server_type": MegatronArgMetadata( - arg_type=str, - default="inplace_megatron", - help="Type of inference server to use.", - choices=("inplace_megatron", "inplace_megatron_chat"), + "linear_conv_kernel_dim": MegatronArgMetadata( + arg_type=int, + default=4, + help="Conv kernel dimension for the gated delta net.", + choices=None, nargs=None, element_type=None, ), - "lazy_mpu_init": MegatronArgMetadata( - arg_type=bool, - default=None, - 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." - ), + "linear_key_head_dim": MegatronArgMetadata( + arg_type=int, + default=128, + help="Query and key head dimension for the gated delta net.", choices=None, nargs=None, element_type=None, ), - "legacy_tokenizer": MegatronArgMetadata( - arg_type=bool, - default=False, - help="To use legacy tokenizer system.", + "linear_num_key_heads": MegatronArgMetadata( + arg_type=int, + default=16, + help="Number of query and key heads for the gated delta net.", choices=None, - nargs=0, + nargs=None, + element_type=None, + ), + "linear_num_value_heads": MegatronArgMetadata( + arg_type=int, + default=32, + help="Number of value and gate heads for the gated delta net.", + choices=None, + nargs=None, + element_type=None, + ), + "linear_value_head_dim": MegatronArgMetadata( + arg_type=int, + default=128, + help="Value and gate head dimension for the gated delta net.", + choices=None, + nargs=None, + element_type=None, + ), + "lion_beta1": MegatronArgMetadata( + arg_type=float, + default=0.95, + help="First beta coefficient for Lion optimizer (used in sign update). Default: 0.95.", + choices=None, + nargs=None, + element_type=None, + ), + "lion_beta2": MegatronArgMetadata( + arg_type=float, + default=0.98, + help="Second beta coefficient for Lion optimizer (used in momentum EMA update). Default: 0.98.", + choices=None, + nargs=None, element_type=None, ), "load": MegatronArgMetadata( @@ -2480,8 +2750,8 @@ ), "load_main_params_from_ckpt": MegatronArgMetadata( arg_type=bool, - default=None, - help="Load main parameters from checkpoint directly.", + default=False, + help="Load main parameters from checkpoint. When loading a model from a checkpoint without loading the optimizer, the model parameters are updated but for fp16 optimizer with main parameters, the main parameters need to also be updated.", choices=None, nargs=0, element_type=None, @@ -2494,10 +2764,18 @@ nargs=None, element_type=None, ), + "log_device_memory_used": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Log device memory used (as reported by nvidia-smi).", + choices=None, + nargs=0, + element_type=None, + ), "log_energy": MegatronArgMetadata( arg_type=bool, default=False, - help="If set, log energy consumption (in Joules)", + help="If set, log energy consumption (in Joules).", choices=None, nargs=0, element_type=None, @@ -2518,6 +2796,22 @@ nargs=0, element_type=None, ), + "log_max_attention_logit": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable max attention logit logging to tensorboard.", + choices=None, + nargs=0, + element_type=None, + ), + "log_memory_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Report memory interval.", + choices=None, + nargs=None, + element_type=None, + ), "log_memory_to_tensorboard": MegatronArgMetadata( arg_type=bool, default=False, @@ -2545,10 +2839,7 @@ "log_progress": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -2604,10 +2895,7 @@ "loss_scale": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Static loss scaling, positive power of 2 values can improve fp16" - " convergence. If None, dynamicloss scaling is used." - ), + help="Static loss scaling, positive power of 2 values can improve fp16 convergence. If None, dynamicloss scaling is used.", choices=None, nargs=None, element_type=None, @@ -2623,10 +2911,7 @@ "lr": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Initial learning rate. Depending on decay style and initial warmup, the" - " learning rate at each iteration would be different." - ), + help="Initial learning rate. Depending on decay style and initial warmup, the learning rate at each iteration would be different.", choices=None, nargs=None, element_type=None, @@ -2634,9 +2919,7 @@ "lr_decay_iters": MegatronArgMetadata( arg_type=int, default=None, - help=( - "number of iterations to decay learning rate over, If None defaults to `--train-iters`" - ), + help="number of iterations to decay learning rate over, If None defaults to train iters", choices=None, nargs=None, element_type=None, @@ -2644,9 +2927,7 @@ "lr_decay_samples": MegatronArgMetadata( arg_type=int, default=None, - help=( - "number of samples to decay learning rate over, If None defaults to `--train-samples`" - ), + help="number of samples to decay learning rate over, If None defaults to train samples", choices=None, nargs=None, element_type=None, @@ -2670,9 +2951,7 @@ "lr_warmup_init": MegatronArgMetadata( arg_type=float, default=0.0, - help=( - "Initial value for learning rate warmup. The scheduler starts warmup from this value." - ), + help="Initial value for learning rate warmup. The scheduler starts warmup from this value.", choices=None, nargs=None, element_type=None, @@ -2720,7 +2999,7 @@ "main_grads_dtype": MegatronArgMetadata( arg_type=None, default="fp32", - help="Dtype of main grads when enabling precision-aware-optimizer", + help="Dtype of main grads when enabling precision-aware-optimizer.", choices=("fp32", "bf16"), nargs=None, element_type=None, @@ -2728,7 +3007,7 @@ "main_params_dtype": MegatronArgMetadata( arg_type=None, default="fp32", - help="Dtype of main params when enabling precision-aware-optimizer", + help="Dtype of main params when enabling precision-aware-optimizer.", choices=("fp32", "fp16"), nargs=None, element_type=None, @@ -2736,10 +3015,7 @@ "make_vocab_size_divisible_by": MegatronArgMetadata( arg_type=int, default=128, - help=( - "Pad the vocab size to be divisible by this value.This is added for" - " computational efficieny reasons." - ), + help="Pad the vocab size to be divisible by this value.This is added for computational efficieny reasons.", choices=None, nargs=None, element_type=None, @@ -2747,15 +3023,31 @@ "mamba_head_dim": MegatronArgMetadata( arg_type=int, default=64, - help="Head dimension for Mamba layers.", + help="The dimensionality of the heads in the Mamba layers.", choices=None, nargs=None, element_type=None, ), + "mamba_inference_conv_states_dtype": MegatronArgMetadata( + arg_type=str, + default="bf16", + help="Dtype for the Mamba inference conv states tensor", + choices=("bf16", "fp16", "fp32"), + nargs=None, + element_type=None, + ), + "mamba_inference_ssm_states_dtype": MegatronArgMetadata( + arg_type=str, + default="bf16", + help="Dtype for the Mamba inference SSM states tensor", + choices=("bf16", "fp16", "fp32"), + nargs=None, + element_type=None, + ), "mamba_num_groups": MegatronArgMetadata( arg_type=int, default=8, - help="Number of groups for Mamba layers.", + help="The number of groups used in Mamba layers.", choices=None, nargs=None, element_type=None, @@ -2763,10 +3055,7 @@ "mamba_num_heads": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Number of heads for Mamba layers.If not set, then the number of heads will" - " be --hidden-size * expand // --mamba-head-dim" - ), + help="The number of heads used in Mamba layers. If None, the number of heads will be hidden_size * expand // mamba_head_dim.", choices=None, nargs=None, element_type=None, @@ -2774,7 +3063,7 @@ "mamba_state_dim": MegatronArgMetadata( arg_type=int, default=128, - help="State dimension for Mamba layers.", + help="The dimensionality of the state representation in Mamba layers.", choices=None, nargs=None, element_type=None, @@ -2782,14 +3071,7 @@ "manual_gc": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -2797,10 +3079,7 @@ "manual_gc_eval": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "When using manual garbage collection, disable garbage collection at the" - " start and the end of each evaluation run." - ), + help="When using manual garbage collection, this controls garbage collection at the start and the end of each evaluation run.", choices=None, nargs=0, element_type=None, @@ -2808,11 +3087,7 @@ "manual_gc_interval": MegatronArgMetadata( arg_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." - ), + help="Training step interval to trigger manual garbage collection. Values > 0 will trigger garbage collections between training steps.", choices=None, nargs=None, element_type=None, @@ -2852,9 +3127,15 @@ "max_position_embeddings": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Maximum number of position embeddings to use. This is the size of position embedding." - ), + help="Maximum number of position embeddings to use. This is the size of position embedding.", + choices=None, + nargs=None, + element_type=None, + ), + "max_seqlen_per_dp_cp_rank": MegatronArgMetadata( + arg_type=int, + default=None, + help="Maximum sequence length per DPxCP rank. This is the maximum sequence length each rank can handle without overflowing the memory. Typically, a good starting point is to set this to maximum sequence length / context parallel size. This is used to calculate the number and length of sub-samples assigned to each rank when using hybrid_context_parallel.", choices=None, nargs=None, element_type=None, @@ -2862,14 +3143,43 @@ "max_tokens_to_oom": MegatronArgMetadata( arg_type=int, default=12000, - help=( - "Maximum number of tokens during inferencetokens here is # in prompt + # to" - " generateAllows us to throw an error before OOM crashes server" - ), + help="Maximum number of tokens during inferencetokens here is # in prompt + # to generateAllows us to throw an error before OOM crashes server", choices=None, nargs=None, element_type=None, ), + "megatron_fsdp_enable_fine_grained_param_gather": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If set, enables fine-grained parameter gathering for Megatron-FSDP. This allows greater overlap between parameter all-gather operations and forward computation, at the cost of additional communication calls. For MXFP8, this helps save memory during fine-grained activation recomputation, because MXFP8 forward and backward passes use different parameter representations (rowwise data for forward, colwise data for backward). Only the rowwise parameters of modules involved in recomputation will be unsharded.", + choices=None, + nargs=0, + element_type=None, + ), + "megatron_fsdp_grad_comm_dtype": MegatronArgMetadata( + arg_type=None, + default="auto", + help="When using Megatron-FSDP, this controls the data-type used when communicating model gradients during FSDP. If 'auto', then the main gradient data-type will be used for the gradient communication / reduction data-type. When using NCCL v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.", + choices=("fp32", "fp16", "bf16", "auto"), + nargs=None, + element_type=None, + ), + "megatron_fsdp_main_grads_dtype": MegatronArgMetadata( + arg_type=None, + default="auto", + help="Data type for the main gradient buffer utilized for distributed optimization with Megatron-FSDP. If 'auto', then the native model gradient data-type will be used for the main gradient / accumulation data-type. Replaces --main-grads-dtype.", + choices=("fp32", "bf16", "fp16", "auto"), + nargs=None, + element_type=None, + ), + "megatron_fsdp_main_params_dtype": MegatronArgMetadata( + arg_type=None, + default="fp32", + help="Data type for the main weight buffer utilized for distributed optimization and quantization with Megatron-FSDP. If 'auto', then the native model parameter data-type will be used for the main weight data-type. Replaces --main-params-dtype.", + choices=("fp32", "bf16", "fp16", "auto"), + nargs=None, + element_type=None, + ), "memory_snapshot_path": MegatronArgMetadata( arg_type=str, default="snapshot.pickle", @@ -2886,13 +3196,18 @@ nargs=None, element_type=None, ), + "metadata_path": MegatronArgMetadata( + arg_type=str, + default=None, + help="Path to the tokenizer metadata file in json format.", + choices=None, + nargs=None, + element_type=None, + ), "micro_batch_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Batch size per model instance (local batch size). Global batch size is" - " local batch size times data parallel size times number of micro batches." - ), + help="Batch size per model instance (local batch size). Global batch size is local batch size times data parallel size times number of micro batches.", choices=None, nargs=None, element_type=None, @@ -2900,7 +3215,7 @@ "microbatch_group_size_per_vp_stage": MegatronArgMetadata( arg_type=int, default=None, - help="Number of contiguous microbatches per virtual pipeline stage", + help="This value specifies the number of micro-batches that are executed at a time for a given virtual stage (both forward and backward). Default (in __post_init__() method below) to pipeline_parallel_size which specifies a depth-first schedule. Example: for PP=2 VP=2, when microbatch_group_size_per_vp_stage=2, num_microbatches = 4, we have rank 0 | 0 1 0 1 2 3 2 3 rank 1 | 0 1 0 1 2 3 2 3 When microbatch_group_size_per_vp_stage=3, num_microbatches = 5, we have rank 0 | 0 1 2 0 1 2 3 4 3 4 rank 1 | 0 1 2 0 1 2 3 4 3 4", choices=None, nargs=None, element_type=None, @@ -2924,15 +3239,39 @@ "min_lr": MegatronArgMetadata( arg_type=float, default=0.0, - help=("Minimum value for learning rate. The schedulerclip values below this threshold."), + help="Minimum value for learning rate. The schedulerclip values below this threshold.", + choices=None, + nargs=None, + element_type=None, + ), + "min_offloaded_tensor_size": MegatronArgMetadata( + arg_type=int, + default=1048576, + help="The minimum size of the tensor to be offloaded.", choices=None, nargs=None, element_type=None, ), + "mla_down_proj_fusion": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable fused q/kv down-projection and fused input layernorm when backend supports. Otherwise fall back to the unfused MLA.", + choices=None, + nargs=0, + element_type=None, + ), "mlp_chunks_for_prefill": MegatronArgMetadata( arg_type=int, default=1, - help=("Number of chunks along sequence dimension for MLP computation during prefill"), + help="The number of chunks along the sequence dimension to use for MLP computation during prefill.", + choices=None, + nargs=None, + element_type=None, + ), + "mlp_chunks_for_training": MegatronArgMetadata( + arg_type=int, + default=1, + help="The number of chunks along the sequence dimension to use for MLP computation during training.", choices=None, nargs=None, element_type=None, @@ -2948,10 +3287,7 @@ "mock_data": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Skip data loading and validation and opt for artificial generation of mock" - " data when an implementation is available." - ), + help="Skip data loading and validation and opt for artificial generation of mock data when an implementation is available.", choices=None, nargs=0, element_type=None, @@ -2959,7 +3295,7 @@ "model_parallel_size": MegatronArgMetadata( arg_type=int, default=None, - help=("Old model parallel argument, do not use. Use --tensor-model-parallel-size instead."), + help="Old model parallel argument, do not use. Use --tensor-model-parallel-size instead.", choices=None, nargs=None, element_type=None, @@ -2967,7 +3303,7 @@ "moe_apply_probs_on_input": MegatronArgMetadata( arg_type=bool, default=False, - help="Apply probs before mlp activation for moe routing.", + help="Apply probs on input of experts instead of applying after activation and glu.", choices=None, nargs=0, element_type=None, @@ -2975,7 +3311,7 @@ "moe_aux_loss_coeff": MegatronArgMetadata( arg_type=list, default=0.0, - help=("Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended."), + help="Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.", choices=None, nargs="+", element_type=float, @@ -2991,7 +3327,15 @@ "moe_enable_deepep": MegatronArgMetadata( arg_type=bool, default=False, - help="DEPRECATED: Please use --moe-flex-dispatcher-backend=deepep instead.", + help="[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.", + choices=None, + nargs=0, + element_type=None, + ), + "moe_enable_routing_replay": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, enable the routing replay feature for MoE layers.", choices=None, nargs=0, element_type=None, @@ -2999,26 +3343,23 @@ "moe_expert_capacity_factor": MegatronArgMetadata( arg_type=float, default=None, - help=("The capacity factor for each expert, None means no token will be dropped."), + help="moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token will be dropped. The default is None.", choices=None, nargs=None, element_type=None, ), - "moe_extended_tp": MegatronArgMetadata( - arg_type=bool, - default=False, - help="Deprecated. Use --expert-tensor-parallel-size instead.", + "moe_expert_rank_capacity_factor": MegatronArgMetadata( + arg_type=float, + default=None, + help="moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens exceeding this budget will be dropped. None means no token will be dropped. The default is None.", choices=None, - nargs=0, + nargs=None, element_type=None, ), "moe_ffn_hidden_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "The hidden size of each expert's feed-forward network (ffn). If not" - " specified, defaults to the ffn_hidden_size." - ), + help="MoE Feed-Forward Network hidden size. If not specified, defaults to the ffn_hidden_size.", choices=None, nargs=None, element_type=None, @@ -3026,10 +3367,7 @@ "moe_flex_dispatcher_backend": MegatronArgMetadata( arg_type=str, default="deepep", - help=( - 'The backend to use for flex token dispatcher. The default is "deepep".' - ' Options are "deepep" and "hybridep".' - ), + help='[Experimental] The backend to use for flex token dispatcher. The default is "deepep". Options are "deepep" and "hybridep". Currently only "hybridep" backend supports the MNNVL case.', choices=("deepep", "hybridep"), nargs=None, element_type=None, @@ -3037,19 +3375,39 @@ "moe_grouped_gemm": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + help="When there are multiple experts per rank, compress multiple local (potentially small) gemms in a single kernel launch to improve the utilization and performance by leveraging the Grouped GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).", choices=None, nargs=0, element_type=None, ), + "moe_hybridep_num_blocks_permute": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of CUDA thread blocks for the permute part in HybridEP. When permute_fusion_into_hybridep is True, this sets the number of SMs for the permute part (only 1 block per SM).", + choices=None, + nargs=None, + element_type=None, + ), + "moe_hybridep_num_blocks_unpermute": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of CUDA thread blocks for the unpermute part in HybridEP. When permute_fusion_into_hybridep is True, this sets the number of SMs for the unpermute part (only 1 block per SM).", + choices=None, + nargs=None, + element_type=None, + ), "moe_hybridep_num_sms": MegatronArgMetadata( arg_type=int, - default=16, - help="Number of SMs to use for HybridEP.", + default=None, + help="Number of SMs to use for HybridEP. None uses the default from DeepEP. In pure NVL scenarios, 16 SMs can generally achieve good bandwidth.", + choices=None, + nargs=None, + element_type=None, + ), + "moe_hybridep_num_sms_preprocessing": MegatronArgMetadata( + arg_type=int, + default=108, + help="Number of SMs to use for HybridEP preprocessing (metadata scan kernel).", choices=None, nargs=None, element_type=None, @@ -3057,7 +3415,15 @@ "moe_input_jitter_eps": MegatronArgMetadata( arg_type=float, default=None, - help=("Add noise to the input tensor by applying jitter with a specified epsilon value."), + help="Add noise to the input tensor by applying jitter with a specified epsilon value.", + choices=None, + nargs=None, + element_type=None, + ), + "moe_latent_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Latent projection dimension for MoE. If None, MoE latent projections are not used.", choices=None, nargs=None, element_type=None, @@ -3065,16 +3431,7 @@ "moe_layer_freq": MegatronArgMetadata( arg_type=None, default=1, - help=( - "Frequency between MoE layers and Dense layers. Accepts either: - An" - " integer N: Represents a 1:N ratio, meaning one expert layer for every N-1" - " dense layers - A string containing a Python list expression that defines" - ' a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to' - " [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an expert layer and 0" - ' indicates a dense layer. Examples: "([0]+[1]*23)": 1 dense layer followed' - ' by 23 experts layers, "([1]*3+[0]*2)*2": Three expert layers followed by' - " two dense layers, repeated twice." - ), + help='Frequency between MoE layers and Dense layers. Accepts either: - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers - A string containing a Python list expression that defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an expert layer and 0 indicates a dense layer. Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, "([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.', choices=None, nargs=None, element_type=None, @@ -3082,22 +3439,23 @@ "moe_layer_recompute": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Enable checkpointing for moe_layer, should be used when memory is not" - ' sufficient. Deprecated. Use "--recompute-granularity selective' - ' --recompute-modules moe" instead.' - ), + help="Memory optimization: checkpointing moe_layer to save actiavtion memory.", choices=None, nargs=0, element_type=None, ), + "moe_mlp_glu_interleave_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="When set, GLU activations in the MoE grouped MLP layer will use a block interleaved format. Instead of interpreting the input tensor as a concatenation of gates and linear units, it will be interpreted as alternating blocks of gates and linear units. This data format is experimental and primarily intended to enable advanced fused kernels.", + choices=None, + nargs=None, + element_type=None, + ), "moe_pad_expert_input_to_capacity": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Pads the input for each expert to match the expert capacity length," - " effective only after the --moe-expert-capacity-factor is set." - ), + help="moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match the expert capacity length, effective only after the moe_expert_capacity_factor is set. The default setting is False.", choices=None, nargs=0, element_type=None, @@ -3105,21 +3463,47 @@ "moe_pad_experts_for_cuda_graph_inference": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "some MoE routers have a D2H sync that will break cuda graphs. If this" - " flag is set the router will switch to dropping and padding during decode" - " time which does not have a D2H sync. The capacity factor is set to the" - " max that an expert could see during inference so no tokens are actually" - " dropped." - ), + help="moe_pad_experts_for_cuda_graph_inference (bool): If True, the router will switch to dropping and padding during decode time which does not have a D2H sync. The capacity factor is set to the max that an expert could see during inference so no tokens are actually dropped. The default setting is False.", choices=None, nargs=0, element_type=None, ), + "moe_paged_stash": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, enable paged stash for all routed-expert activations needed for backward", + choices=None, + nargs=0, + element_type=None, + ), + "moe_paged_stash_buffer_size_factor_cpu": MegatronArgMetadata( + arg_type=float, + default=0.0, + help="Scale factor for paged stash host buffer. 0 disables host buffer. Same sign convention as moe_paged_stash_buffer_size_factor_cuda: positive = avg-based, negative = actual-max; scale = abs(factor).", + choices=None, + nargs=None, + element_type=None, + ), + "moe_paged_stash_buffer_size_factor_cuda": MegatronArgMetadata( + arg_type=float, + default=1.1, + help="Scale factor for paged stash CUDA buffer allocation. Sign selects sizing: positive = avg-based, negative = actual-max. Magnitude is headroom (e.g. 1.10 = 10%).", + choices=None, + nargs=None, + element_type=None, + ), + "moe_paged_stash_page_size": MegatronArgMetadata( + arg_type=int, + default=64, + help="Number of tokens per page for paged stash memory management.", + choices=None, + nargs=None, + element_type=None, + ), "moe_per_layer_logging": MegatronArgMetadata( arg_type=bool, default=False, - help=("Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."), + help="Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.", choices=None, nargs=0, element_type=None, @@ -3132,17 +3516,18 @@ nargs=0, element_type=None, ), + "moe_permute_fusion_into_hybridep": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Fuse token rearrangement ops during token dispatching for HybridEP.", + choices=None, + nargs=0, + element_type=None, + ), "moe_router_bias_update_rate": MegatronArgMetadata( arg_type=float, default=0.001, - help=( - "Expert bias update rate in the aux-loss-free load balancing strategy. The" - " expert bias is updated based on the number of assigned tokens to each" - " expert in a global batch, where the bias is increased for the experts" - " with less assigned tokens and decreased for the experts with more" - " assigned tokens. The default value 1e-3 is same as that used in" - " DeepSeekV3." - ), + help="The expert bias is updated based on the number of assigned tokens to each expert in a global batch, where the bias is increased for the experts with less assigned tokens and decreased for the experts with more assigned tokens. The default value 1e-3 is same as that used in DeepSeekV3.", choices=None, nargs=None, element_type=None, @@ -3150,12 +3535,7 @@ "moe_router_dtype": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Data type for routing computation and expert output weighted averaging." - " Fp32/fp64 enhances numerical stability, especially with numerous experts." - " The perf impact should be negligible when used with permute fusion. None" - " means no changes for dtype." - ), + help="Data type for routing and expert output weighted averaging. Using fp32 or fp64 can improve stability especially when the number of experts is large (e.g. finegrained-moe). None means no changes for dtype.", choices=("fp32", "fp64"), nargs=None, element_type=None, @@ -3163,23 +3543,23 @@ "moe_router_enable_expert_bias": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "TopK routing with dynamic expert bias in the aux-loss-free load balancing" - " strategy. The routing decision is based on the sum of the routing scores" - " and the expert bias. See https://arxiv.org/abs/2408.15664 for details." - ), + help="TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. The routing decision is based on the sum of the routing scores and the expert bias. See https://arxiv.org/abs/2408.15664 for details.", choices=None, nargs=0, element_type=None, ), + "moe_router_force_biased": MegatronArgMetadata( + arg_type=float, + default=None, + help="Apply random expert bias in normal distribution with specified std to router logits. Shared seed across all ranks ensures identical bias. If positive, generates new random bias each forward pass. If negative, generates bias once per layer and reuses it (abs value is std). This is an experimental feature for benchmarking purposes.", + choices=None, + nargs=None, + element_type=None, + ), "moe_router_force_load_balancing": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "[Experimental] Force override routing to balance token distribution using" - " random logits for MoE routers, supporting naive top-k and group-limited" - " top-k. This experimental feature is for benchmarking purposes only!" - ), + help="[Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark.", choices=None, nargs=0, element_type=None, @@ -3187,10 +3567,7 @@ "moe_router_fusion": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Enable fusion for MoE TopK routing and aux-loss computation. This is only" - " supported in TransformerEngine 2.7.0 and above." - ), + help="Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above.", choices=None, nargs=0, element_type=None, @@ -3206,14 +3583,7 @@ "moe_router_load_balancing_type": MegatronArgMetadata( arg_type=list, 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; "seq_aux_loss" corresponds to the load balancing loss' - " used in DeepSeekV2, which computes the loss for each individual sample;" - ' "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and' - ' "none" implies no load balancing. The default is "aux_loss".' - ), + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".', choices=("aux_loss", "seq_aux_loss", "global_aux_loss", "sinkhorn", "none"), nargs="+", element_type=str, @@ -3221,18 +3591,7 @@ "moe_router_num_groups": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Number of groups to divide experts into for group-limited routing. When" - " using group-limited routing: 1) Experts are divided into equal-sized" - " groups, 2) For each token, a subset of groups are selected based on" - " routing scores (sum of top-2 expert scores within each group), 3) From" - " these selected groups, moe_router_topk experts are chosen.Two common use" - " cases: 1) Device-limited routing: Set equal to expert parallel size (EP)" - " to limit each token to experts on a subset of devices (See DeepSeek-V2:" - " https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to" - " number of nodes in EP group to limit each token to experts on a subset of" - " nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)" - ), + help="Number of groups to divide experts into for group-limited routing. When using group-limited routing: 1. Experts are divided into 'moe_router_num_groups' equal-sized groups 2. For each token, 'moe_router_group_topk' groups are selected based on sum of top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group 3. From these selected groups, 'moe_router_topk' individual experts are chosen Two common use cases: - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) - Node-limited routing: Set 'moe_router_num_groups' equal to number of nodes in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)", choices=None, nargs=None, element_type=None, @@ -3240,10 +3599,7 @@ "moe_router_padding_for_fp8": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "[Compatibility alias for --moe-router-padding-for-quantization] Enabling" - " this will also enable --moe-router-padding-for-quantization." - ), + help="[Compatibility alias for moe_router_padding_for_quantization] Enabling this will also enable moe_router_padding_for_quantization.", choices=None, nargs=0, element_type=None, @@ -3251,13 +3607,7 @@ "moe_router_padding_for_quantization": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Pad the routing_map to make sure the number of tokens each expert received" - " is a multiple of 16/32 for FP8/FP4 precision. It is suggested to enable" - " this for dropless training with FP8/FP4 precision when num_local_experts" - " > 1. This is a more efficient way to pad for FP8/FP4 which eliminates the" - " explicit padding in the GroupedMLP layer." - ), + help="Whether to pad the routing_map to make sure the number of tokens each expert receives is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit padding in the GroupedMLP layer.", choices=None, nargs=0, element_type=None, @@ -3265,10 +3615,7 @@ "moe_router_pre_softmax": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Enable pre-softmax routing for MoE, which means softmax is before the" - " top-k selection. By default, softmax is done after top-k." - ), + help="Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the top-k selection. By default, softmax is done after top-k.", choices=None, nargs=0, element_type=None, @@ -3276,15 +3623,15 @@ "moe_router_score_function": MegatronArgMetadata( arg_type=str, default="softmax", - help='Score function for MoE TopK routing. Can be "softmax" or "sigmoid".', - choices=("softmax", "sigmoid"), + help='Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus".', + choices=("softmax", "sigmoid", "sqrtsoftplus"), nargs=None, element_type=None, ), "moe_router_topk": MegatronArgMetadata( arg_type=int, default=2, - help="Number of experts to route to for each token. The default is 2.", + help="Number of experts to route to for each token.", choices=None, nargs=None, element_type=None, @@ -3292,29 +3639,23 @@ "moe_router_topk_scaling_factor": MegatronArgMetadata( arg_type=float, default=None, - help=( - "Scaling factor for routing score in top-k selection, only works when" - " --moe-router-pre-softmax enabled. Defaults to None, which means no" - " scaling." - ), + help="Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax enabled. Defaults to None, which means no scaling.", choices=None, nargs=None, element_type=None, ), + "moe_shared_expert_gate": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.", + choices=None, + nargs=0, + element_type=None, + ), "moe_shared_expert_intermediate_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Shared expert total ffn hidden size. It should be equal to" - ' "num_shared_experts * ffn_size_of_each_shared_expert" if there are' - " multiple shared experts. None means no shared expert. By default, the" - " shared experts execute before the router. However, when" - " --moe-shared-expert-overlap or --overlap-moe-expert-parallel-comm is set," - " the shared experts execute after the router, before the routed experts." - " This makes the gradients from the router and the shared experts added in" - " different orders to the hidden_states, causing minor numerical" - " differences in the hidden_states gradient." - ), + help="Shared expert total ffn hidden size. It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if there are multiple shared experts. None means no shared expert. By default, the shared experts execute before the router. However, when moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set, the shared experts execute after the router, before the routed experts. This makes the gradients from the router and the shared experts added in different orders to the hidden_states, causing minor numerical differences in the hidden_states gradient.", choices=None, nargs=None, element_type=None, @@ -3322,11 +3663,23 @@ "moe_shared_expert_overlap": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Enable overlapping between shared expert computations and dispatcher" - " communications. Without this, the shared experts execute before the" - " router. Only effective when moe-shared-expert-intermediate-size is set." - ), + help="Enable overlapping between shared expert computations and dispatcher communications. Without this, the shared experts execute before the router. Only effective when moe-shared-expert-intermediate-size is set.", + choices=None, + nargs=0, + element_type=None, + ), + "moe_single_grouped_bias": MegatronArgMetadata( + arg_type=bool, + default=False, + help="When using TE GroupedLinear for MoE experts, store expert biases as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and ``add_bias_linear=True``.", + choices=None, + nargs=0, + element_type=None, + ), + "moe_single_grouped_weight": MegatronArgMetadata( + arg_type=bool, + default=False, + help="When using TE GroupedLinear for MoE experts, store expert weights as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``.", choices=None, nargs=0, element_type=None, @@ -3334,12 +3687,7 @@ "moe_token_dispatcher_type": MegatronArgMetadata( arg_type=str, default="allgather", - help=( - "The type of token dispatcher to use. The default is 'allgather'. Options" - " are 'allgather', 'alltoall'. We recommend using 'alltoall' when applying" - " expert parallelism. For more information, please refer to the" - " documentation in core/moe/README." - ), + help="The type of token dispatcher to use. The default is 'allgather'. Options are 'allgather','alltoall' and 'flex'.", choices=("allgather", "alltoall", "flex"), nargs=None, element_type=None, @@ -3347,11 +3695,7 @@ "moe_token_drop_policy": MegatronArgMetadata( arg_type=str, default="probs", - 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.' - ), + 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.', choices=("probs", "position"), nargs=None, element_type=None, @@ -3359,36 +3703,15 @@ "moe_upcycling_granularity": MegatronArgMetadata( arg_type=int, default=1, - help=( - "This param sepecifics how many times smaller is the expert hidden size" - " compared with the original dense FFN hidden size. For using granular" - " upcycling strategy, please set this param as a positive integer. If this" - " param is set to 1, it means using the default upcycling strategy." - ), + help="This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.", choices=None, nargs=None, element_type=None, ), - "moe_use_legacy_grouped_gemm": MegatronArgMetadata( - arg_type=bool, - default=False, - help=( - "Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will" - " be deprecated soon." - ), - choices=None, - nargs=0, - element_type=None, - ), "moe_use_upcycling": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -3396,7 +3719,7 @@ "moe_z_loss_coeff": MegatronArgMetadata( arg_type=float, default=None, - help=("Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended."), + help="Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended.", choices=None, nargs=None, element_type=None, @@ -3404,7 +3727,7 @@ "mrope_section": MegatronArgMetadata( arg_type=list, default=None, - help="Multimodal rope section is for channel dimension, empty by default.", + help="Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.", choices=None, nargs="+", element_type=int, @@ -3425,15 +3748,18 @@ nargs=None, element_type=None, ), + "mtp_hybrid_override_pattern": MegatronArgMetadata( + arg_type=str, + default=None, + help="DEPRECATED: Use unified hybrid_layer_pattern instead. Legacy argument for loading old checkpoints. Force a specific hybrid layer pattern for MTP layers.", + choices=None, + nargs=None, + element_type=None, + ), "mtp_loss_scaling_factor": MegatronArgMetadata( arg_type=float, default=0.1, - help=( - "Scaling factor of Multi-Token Prediction (MTP) loss. We compute the" - " average of the MTP losses across all depths, and multiply it the scaling" - " factor to obtain the overall MTP loss, which serves as an additional" - " training objective." - ), + help="Weighting factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP losses across all depths, and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective.", choices=None, nargs=None, element_type=None, @@ -3441,20 +3767,31 @@ "mtp_num_layers": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Number of Multi-Token Prediction (MTP) Layers.MTP extends the prediction" - " scope to multiple future tokens at each position.This MTP implementation" - " sequentially predict additional tokens by using D sequential modules to" - " predict D additional tokens." - ), + help="Number of Multi-Token Prediction (MTP) Layers. MTP extends the prediction scope to multiple future tokens at each position. This MTP implementation sequentially predict additional tokens by using D sequential modules to predict D additional tokens.", choices=None, nargs=None, element_type=None, ), + "mtp_standalone": MegatronArgMetadata( + arg_type=bool, + default=False, + help="This will be set automatically according to the pipeline layout, and will be set to True if MTP is in a separate vpp stage.", + choices=None, + nargs=0, + element_type=None, + ), + "mtp_use_repeated_layer": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Use a single MTP layer repeatedly instead of multiple separate layers.", + choices=None, + nargs=0, + element_type=None, + ), "multi_latent_attention": MegatronArgMetadata( arg_type=bool, default=False, - help="Use multi-latent attention for model.", + help="Whether to use multi-latent attention.", choices=None, nargs=0, element_type=None, @@ -3462,23 +3799,143 @@ "multiple_validation_sets": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, multiple datasets listed in the validation split are evaluated" - " independently with a separate loss for each dataset in the list. This" - " argument requires that no weights are included in the list" - ), + help="If set, multiple datasets listed in the validation split are evaluated independently with a separate loss for each dataset in the list. This argument requires that no weights are included in the list.", + choices=None, + nargs=0, + element_type=None, + ), + "muon_coefficient_type": MegatronArgMetadata( + arg_type=str, + default="quintic", + help="Newton-Schulz coefficient type for the Muon optimizer. Valid types are discovered from the installed emerging_optimizers package (e.g. simple, quintic, polar_express, aol). Validated at optimizer creation time.", + choices=None, + nargs=None, + element_type=None, + ), + "muon_extra_scale_factor": MegatronArgMetadata( + arg_type=float, + default=1.0, + help="Additional scale factor for the muon update", + choices=None, + nargs=None, + element_type=None, + ), + "muon_fp32_matmul_prec": MegatronArgMetadata( + arg_type=str, + default="medium", + help="FP32 matmul precision for Newton-Schulz iteration", + choices=("low", "medium", "high"), + nargs=None, + element_type=None, + ), + "muon_momentum": MegatronArgMetadata( + arg_type=float, + default=0.9, + help="Momentum factor for Muon optimizer", + choices=None, + nargs=None, + element_type=None, + ), + "muon_nesterov": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to use Nesterov-style momentum in the internal SGD", + choices=None, + nargs=0, + element_type=None, + ), + "muon_num_ns_steps": MegatronArgMetadata( + arg_type=int, + default=5, + help="Number of Newton-Schulz steps for Muon optimizer", + choices=None, + nargs=None, + element_type=None, + ), + "muon_scalar_optimizer": MegatronArgMetadata( + arg_type=str, + default="adam", + help="Optimizer for scalar parameters (embeddings, biases, norms) when using muon. Defaults to adam.", + choices=("adam", "lion"), + nargs=None, + element_type=None, + ), + "muon_scale_mode": MegatronArgMetadata( + arg_type=str, + default="spectral", + help="Scale mode for Muon optimizer. With MuP, set --muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, or set --muon-scale-mode spectral to keep spectral scaling.", + choices=("spectral", "unit_rms_norm", "shape_scaling"), + nargs=None, + element_type=None, + ), + "muon_split_qkv": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Whether to split QKV parameters for Muon optimizer", choices=None, nargs=0, element_type=None, ), + "muon_tp_mode": MegatronArgMetadata( + arg_type=str, + default="blockwise", + help="How to perform NS calculation for tensor model parallel weights", + choices=("blockwise", "duplicated", "distributed"), + nargs=None, + element_type=None, + ), + "mup_attn_scale_power": MegatronArgMetadata( + arg_type=float, + default=1.0, + help="Power for attention scaling: softmax_scale = 1 / (kv_channels ** mup_attn_scale_power). 0.5 = standard attention (1/sqrt(d_head)), 1.0 = MuP attention (1/d_head). Default: 1.0 (MuP scaling when use_mup is True). Set to 0.5 for standard scaling.", + choices=None, + nargs=None, + element_type=None, + ), + "mup_base_head_dim": MegatronArgMetadata( + arg_type=float, + default=None, + help="Base head dimension for MuP attention scaling. When set, softmax_scale = sqrt(mup_base_head_dim) / (kv_channels ** mup_attn_scale_power). Set to base model's d_head (e.g., 64) to match standard 1/sqrt(d_head) scaling at the base width, ensuring non-MuP compatibility for that specific value.", + choices=None, + nargs=None, + element_type=None, + ), + "mup_base_hidden_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Base hidden size for MuP width scaling. This is the reference width from which scaling factors are computed. Defaults to hidden_size if not specified (base model case where width_mult=1.0). Set this to your base/proxy model's hidden size when scaling up.", + choices=None, + nargs=None, + element_type=None, + ), + "mup_embedding_mult": MegatronArgMetadata( + arg_type=float, + default=1.0, + help="Multiplier for embedding layer output. Applied after the embedding lookup. Default: 1.0 (no scaling).", + choices=None, + nargs=None, + element_type=None, + ), + "mup_output_mult": MegatronArgMetadata( + arg_type=float, + default=1.0, + help="Multiplier for output logits before softmax. When MuP is enabled and this is left at 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable across widths. Override to customize output scaling. Default: 1.0.", + choices=None, + nargs=None, + element_type=None, + ), + "mup_width_mult": MegatronArgMetadata( + arg_type=float, + default=1.0, + help="Width multiplier for MuP scaling, computed as hidden_size / mup_base_hidden_size. This value is automatically computed in __post_init__ when use_mup is enabled.", + choices=None, + nargs=None, + element_type=None, + ), "nccl_all_reduce_for_prefill": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "When using symmeric all reduce kernels this will use regular nccl kernels" - " for prefill. This can be more effecient when prefill is large as the nccl" - " kernels can be more bandwith optimized" - ), + help="If True, use NCCL all-reduce kernels when symmetric all-reduce is enabled.", choices=None, nargs=0, element_type=None, @@ -3486,12 +3943,7 @@ "nccl_communicator_config_path": MegatronArgMetadata( arg_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`." - ), + 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`.", choices=None, nargs=None, element_type=None, @@ -3499,11 +3951,7 @@ "nccl_ub": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use the userbuffer registration for DP/FSDP communication buffers.This" - " option will reduce GPU SM usage for the DP/FSDP communication,which is" - " improving the performance of the overlapped computation." - ), + help="Use the userbuffer registration for DP/FSDP communication buffers.This option will reduce GPU SM usage for the DP/FSDP communication,which is improving the performance of the overlapped computation.", choices=None, nargs=0, element_type=None, @@ -3527,11 +3975,7 @@ "no_persist_layer_norm": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -3539,15 +3983,7 @@ "no_rope_freq": MegatronArgMetadata( arg_type=None, default=None, - help=( - "Controls which layers to skip performing Rotary Position Embedding." - " Accepts either: - An integer N: Represents a 1:N ratio, meaning RoPE is" - " skipped every N-1 layers. - A string containing a Python list expression" - ' that defines a custom pattern, e.g.: "([0]*3+[1]*1)*3" evaluates to' - " [0,0,0,1,0,0,0,1,0,0,0,1] where 1 indicates no-rope layer. This patten is" - " equivalent to --no-rope-freq=4.By default this is disabled and set to" - " None, indicating RoPE will be performedon every layer." - ), + help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string containing a Python list expression that defines a custom pattern, e.g.: "([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.By default this is disabled and set to None, indicating RoPE will be performedon every layer.', choices=None, nargs=None, element_type=None, @@ -3568,16 +4004,19 @@ nargs=0, element_type=None, ), + "no_weight_decay_cond_type": MegatronArgMetadata( + arg_type=str, + default=None, + help='Type of no weight decay condition. Choices: None (default): apply weight decay to 1D weights and biases."apply_wd_to_qk_layernorm": additionally apply weight decay to qk layernorm as a special case.DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ', + choices=("apply_wd_to_qk_layernorm",), + nargs=None, + element_type=None, + ), "non_persistent_ckpt_type": MegatronArgMetadata( arg_type=str, default=None, - help=( - 'Type of non-persistent model checkpoints. "global" - Saved as a standard' - ' checkpoint (e.g., on Lustre) with old checkpoints being removed. "local"' - " - Each rank saves a portion of the checkpoint locally (e.g., on" - " SSD/ramdisk). None - No non-persistent checkpointing (default option)." - ), - 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).', + choices=("global", "local", "in_memory"), nargs=None, element_type=None, ), @@ -3613,19 +4052,27 @@ nargs=None, element_type=None, ), - "norm_epsilon": MegatronArgMetadata( - arg_type=float, - default=1e-05, - help="Epsilon for layer norm and RMS norm.", + "normalization": MegatronArgMetadata( + arg_type=str, + default="LayerNorm", + help="Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`.", + choices=("LayerNorm", "RMSNorm"), + nargs=None, + element_type=None, + ), + "null_tokenizer_eod_id": MegatronArgMetadata( + arg_type=int, + default=None, + help="EOD token id for NullTokenizer. Defaults to `vocab_size - 1`.", choices=None, nargs=None, element_type=None, ), - "normalization": MegatronArgMetadata( - arg_type=None, - default="LayerNorm", - help="Which normalization technique to use.", - choices=("LayerNorm", "RMSNorm"), + "null_tokenizer_pad_id": MegatronArgMetadata( + arg_type=int, + default=-1, + help="Pad token id for NullTokenizer. Defaults to -1 (no pad token). Set to a value outside the dataset to avoid masking real tokens.", + choices=None, nargs=None, element_type=None, ), @@ -3680,7 +4127,7 @@ "num_layers": MegatronArgMetadata( arg_type=int, default=None, - help="Number of transformer layers.", + help="Number of transformer layers in a transformer block.", choices=None, nargs=None, element_type=None, @@ -3688,9 +4135,7 @@ "num_layers_at_end_in_bf16": MegatronArgMetadata( arg_type=int, default=1, - help=( - "Number of layers at end to construct in bf16 when --first-last-layers-bf16 is enabled." - ), + help="Number of layers at the end of the model to keep in BF16 precision when first_last_layers_bf16 is True.", choices=None, nargs=None, element_type=None, @@ -3698,10 +4143,7 @@ "num_layers_at_start_in_bf16": MegatronArgMetadata( arg_type=int, default=1, - help=( - "Number of layers at start to construct in bf16 when" - " --first-last-layers-bf16 is enabled." - ), + help="Number of layers at the start of the model to keep in BF16 precision when first_last_layers_bf16 is True.", choices=None, nargs=None, element_type=None, @@ -3715,7 +4157,20 @@ element_type=None, ), "num_query_groups": MegatronArgMetadata( - arg_type=int, default=1, help=None, choices=None, nargs=None, element_type=None + arg_type=int, + default=1, + help="Number of query groups for group query attention. If None, normal attention is used.", + choices=None, + nargs=None, + element_type=None, + ), + "num_speculative_tokens": MegatronArgMetadata( + arg_type=int, + default=0, + help="Number of speculative tokens generated during decode", + choices=None, + nargs=None, + element_type=None, ), "num_virtual_stages_per_pipeline_rank": MegatronArgMetadata( arg_type=int, @@ -3733,6 +4188,14 @@ nargs=None, element_type=None, ), + "nvtx_ranges": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable NVTX range annotations for profiling. When enabled, inserts NVTX markers to categorize execution in profiler output.", + choices=None, + nargs=0, + element_type=None, + ), "object_storage_cache_path": MegatronArgMetadata( arg_type=str, default=None, @@ -3741,6 +4204,14 @@ nargs=None, element_type=None, ), + "offload_modules": MegatronArgMetadata( + arg_type=list, + default=[], + help='The submodules to offload its input. choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act". "attn_norm": offload the input of the normalization in the attention part. "qkv_linear": offload the input of the qkv linear part. "core_attn": offload the input of the core attention part. "attn_proj": offload the input of the attn linear projection part. "mlp_norm": offload the input of the normalization in the mlp part. "expert_fc1": offload the input of the expert fc1 part. "moe_act": offload the input of the moe act part.', + choices=None, + nargs="+", + element_type=str, + ), "one_logger_async": MegatronArgMetadata( arg_type=bool, default=False, @@ -3776,10 +4247,7 @@ "openai_gelu": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use OpenAIs GeLU implementation. This optionshould not be used unless for" - " backward compatibilityreasons." - ), + help="Use OpenAIs GeLU implementation. This optionshould not be used unless for backward compatibilityreasons.", choices=None, nargs=0, element_type=None, @@ -3787,8 +4255,8 @@ "optimizer": MegatronArgMetadata( arg_type=str, default="adam", - help="Optimizer function", - choices=("adam", "sgd"), + help="Optimizer function. Note: dist_muon is deprecated; use --optimizer muon with --use-distributed-optimizer instead.", + choices=("adam", "sgd", "muon", "dist_muon", "lion", "soap", "adaptive_muon"), nargs=None, element_type=None, ), @@ -3800,6 +4268,14 @@ nargs=0, element_type=None, ), + "optimizer_cuda_graph": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable CUDA graph for optimizer step", + choices=None, + nargs=0, + element_type=None, + ), "optimizer_offload_fraction": MegatronArgMetadata( arg_type=float, default=1.0, @@ -3808,13 +4284,18 @@ nargs=None, element_type=None, ), + "outer_dp_sharding_strategy": MegatronArgMetadata( + arg_type=str, + default="no_shard", + help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). Default: "no_shard".', + choices=("no_shard", "optim"), + nargs=None, + element_type=None, + ), "output_bert_embeddings": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Output Bert embeddings (via mean pooling) from model, rather than its" - " binary head output or entire hidden batch." - ), + help="Output Bert embeddings (via mean pooling) from model, rather than its binary head output or entire hidden batch.", choices=None, nargs=0, element_type=None, @@ -3827,6 +4308,14 @@ nargs=0, element_type=None, ), + "overlap_dispatch_backward_with_experts_wgrad": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Delay the weight gradient computation for TE Grouped GEMM MoE experts. When enabled with FSDP, the expert weight gradients are computed on a separate CUDA stream after the data gradients finish, allowing overlap of wgrad compute with EP A2A communication. The FSDP gradient reduce-scatter for expert parameters is deferred until the delayed wgrad computation completes. This requires transformer_engine with GroupedLinear support (TE >= 2.3.0).", + choices=None, + nargs=0, + element_type=None, + ), "overlap_grad_reduce": MegatronArgMetadata( arg_type=bool, default=False, @@ -3838,7 +4327,7 @@ "overlap_moe_expert_parallel_comm": MegatronArgMetadata( arg_type=bool, default=False, - help=("Overlap the EP A2A communication by batch-level overlapping in 1f1b stage."), + help="Overlap EP A2A communications with independent computations of different micro-batches in 1f1b phase of pipelining or non-pipelining schedule.", choices=None, nargs=0, element_type=None, @@ -3846,7 +4335,7 @@ "overlap_p2p_comm": MegatronArgMetadata( arg_type=bool, default=True, - help=("overlap pipeline parallel communication with forward and backward chunks in 1F1B"), + help="overlap pipeline parallel communication with forward and backward chunks in 1F1B", choices=None, nargs=0, element_type=None, @@ -3854,7 +4343,7 @@ "overlap_p2p_comm_warmup_flush": MegatronArgMetadata( arg_type=bool, default=False, - help="if set, overlap pipeline parallel communication in warmup and flush", + help="If true, overlap communication and computation in warm up and flush phase. Only valid when overlap_p2p_comm is True and batch_p2p_comm is False. Defaults to False.", choices=None, nargs=0, element_type=None, @@ -3878,12 +4367,15 @@ "override_opt_param_scheduler": MegatronArgMetadata( arg_type=bool, default=False, - 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. Notethat all the" - " above values will be reset." - ), + 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.", + choices=None, + nargs=0, + element_type=None, + ), + "pad_vocab_size": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Whether to pad vocab size of the model automatically if padded_vocab_size is not provided.", choices=None, nargs=0, element_type=None, @@ -3891,11 +4383,7 @@ "padded_vocab_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Vocabulary size of the model (padded to be divisible by tensor model" - " parallel size). If not provided, it will be automatically calculated from" - " vocab-size." - ), + help="Vocabulary size of the model (padded to be divisible by tensor model parallel size). If not provided, it will be automatically calculated from vocab-size.", choices=None, nargs=None, element_type=None, @@ -3908,15 +4396,18 @@ nargs=None, element_type=None, ), + "per_dataset_sequences_path": MegatronArgMetadata( + arg_type=None, + default=None, + help="Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.", + choices=None, + nargs=None, + element_type=None, + ), "per_split_data_args_path": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Path to per-split-data-args. Instead of feeding" - " `--(train|valid|test)-data-path` with weighted dataset, we pass in a file" - " path from which we read those arguments. This is useful when the list of" - " data is too big. Format is a json file with `train`, `valid, `test` keys" - ), + help="Path to per-split-data-args. Instead of feeding `--(train|valid|test)-data-path` with weighted dataset, we pass in a file path from which we read those arguments. This is useful when the list of data is too big. Format is a json file with `train`, `valid, `test` keys", choices=None, nargs=None, element_type=None, @@ -3924,10 +4415,7 @@ "perform_initialization": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "Do not perform initialization when building model, can reduce startup time" - " when definitely loading from a checkpoint" - ), + help="Controls weights initialization. This option can be useful when you know you are going to load values from a checkpoint.", choices=None, nargs=0, element_type=None, @@ -3940,6 +4428,14 @@ nargs=0, element_type=None, ), + "phase_transition_iterations": MegatronArgMetadata( + arg_type=str, + default=None, + help="Comma-separated list of iterations where phase transitions occur. Requires fixed global batch size across phases.", + choices=None, + nargs=None, + element_type=None, + ), "pin_cpu_grads": MegatronArgMetadata( arg_type=bool, default=True, @@ -3959,10 +4455,7 @@ "pipeline_model_parallel_comm_backend": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Select a communicator backend for pipeline parallel communication. If" - " None, the default backend will be used." - ), + help="Configuring backend option of pipeline parallel communication (e.g., nccl, ucc) If None, the default backend will be used.", choices=("nccl", "ucc"), nargs=None, element_type=None, @@ -3970,14 +4463,7 @@ "pipeline_model_parallel_layout": MegatronArgMetadata( arg_type=str, default=None, - help=( - "A string that describes a custom pipeline model parallel layout. e.g.," - ' "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer' - ' decoder layer, and mtp layer, respectively. Stages are split by "|".' - " Replicated stages or layers can be described with multiplication. Commas" - " can be used cosmetically. Default None is not using this argument to set" - " the layout." - ), + help='A string that describes a custom pipeline model parallel layout. e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer decoder layer, and mtp layer, respectively. Stages are split by "|". Replicated stages or layers can be described with multiplication. Commas can be used cosmetically. Default None is not using this argument to set the layout.', choices=None, nargs=None, element_type=None, @@ -3985,7 +4471,7 @@ "pipeline_model_parallel_size": MegatronArgMetadata( arg_type=int, default=1, - help="Degree of pipeline model parallelism.", + help="Inter-layer model parallelism. Splits transformer layers across GPU ranks.", choices=None, nargs=None, element_type=None, @@ -3994,7 +4480,7 @@ arg_type=str, default="learned_absolute", help="Position embedding type.", - choices=("learned_absolute", "rope", "mrope", "relative", "none"), + choices=("learned_absolute", "rope", "yarn", "mrope", "relative", "none"), nargs=None, element_type=None, ), @@ -4009,19 +4495,14 @@ "profile": MegatronArgMetadata( arg_type=bool, default=False, - 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 --force-overwrite true" - " --capture-range=cudaProfilerApi --capture-range-end=stop`." - ), + 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 --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop`.", choices=None, nargs=0, element_type=None, ), "profile_ranks": MegatronArgMetadata( arg_type=list, - default=[0], + default=[], help="Global ranks to profile.", choices=None, nargs="+", @@ -4043,6 +4524,30 @@ nargs=None, element_type=None, ), + "pytorch_profiler_collect_callstack": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Collect callstack in pytorch profiler.", + choices=None, + nargs=0, + element_type=None, + ), + "pytorch_profiler_collect_chakra": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Collect chakra trace in pytorch profiler.", + choices=None, + nargs=0, + element_type=None, + ), + "pytorch_profiler_collect_shapes": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Collect tensor shape in pytorch profiler.", + choices=None, + nargs=0, + element_type=None, + ), "q_lora_rank": MegatronArgMetadata( arg_type=int, default=None, @@ -4051,13 +4556,34 @@ nargs=None, element_type=None, ), + "qk_clip": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to clip the query and key weights. Needed for Muon MLA Model training.", + choices=None, + nargs=0, + element_type=None, + ), + "qk_clip_alpha": MegatronArgMetadata( + arg_type=float, + default=0.5, + help="The balancing alpha for qk-clip. Q = Q * (eta ** alpha)", + choices=None, + nargs=None, + element_type=None, + ), + "qk_clip_threshold": MegatronArgMetadata( + arg_type=float, + default=100, + help="The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0)", + choices=None, + nargs=None, + element_type=None, + ), "qk_head_dim": MegatronArgMetadata( arg_type=int, default=128, - help=( - "Dimension of the head in the QK projection. q_head_dim = qk_head_dim +" - " qk_pos_emb_head_dim" - ), + help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim", choices=None, nargs=None, element_type=None, @@ -4065,7 +4591,7 @@ "qk_l2_norm": MegatronArgMetadata( arg_type=bool, default=False, - help="Use llama 4 qk l2 norm", + help="Whether to apply llama 4-style qk L2 norm.", choices=None, nargs=0, element_type=None, @@ -4073,7 +4599,7 @@ "qk_layernorm": MegatronArgMetadata( arg_type=bool, default=False, - help="Whether to layer normalize the q and k attention embeddings.", + help="Whether to apply `normalization` type of normalization to the query and key embeddings.", choices=None, nargs=0, element_type=None, @@ -4103,28 +4629,17 @@ element_type=None, ), "rampup_batch_size": MegatronArgMetadata( - arg_type=list, + arg_type=int, default=None, - help=( - "Batch size ramp up with the following values: --rampup-batch-size " - " For example: --rampup-batch-size 16 8 300000 \\" - " --global-batch-size 1024will start with global batch size 16 and over " - " (1024 - 16) / 8 = 126 intervals will increasethe batch size linearly to" - " 1024. In each intervalwe will use approximately 300000 / 126 = 2380" - " samples." - ), + help="Batch size ramp up with the following values: , , 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. Deprecated. Use step_batch_size_schedule instead.", choices=None, - nargs="*", - element_type=str, + nargs=3, + element_type=None, ), "recompute_activations": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "recompute activation to allow for training with larger models, sequences," - " and batch sizes." - ), + help="recompute activation to allow for training with larger models, sequences, and batch sizes.", choices=None, nargs=0, element_type=None, @@ -4132,12 +4647,7 @@ "recompute_granularity": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Checkpoint activations to allow for training with larger models," - " sequences, and batch sizes. It is supported at two granularities 1) full:" - " whole transformer layer is recomputed, 2) selective: submodules set in" - " --recompute-modules are recomputed, default is core_attn." - ), + help="Determines which type of activation recompute to use. Megatron-core supports 'selective' activation checkpointing where the submodules set in --recompute-modules is checkpointed. The default is \"core_attn\" which is the memory intensive part of attention. These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+). See Reducing Activation Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198) for more details. 'full' will checkpoint the entire transformer layer. If None, no recompute is performed and all activations are saved. If set, must be 'selective' or 'full'. 'selective' always uses all layers.", choices=("full", "selective"), nargs=None, element_type=None, @@ -4145,14 +4655,7 @@ "recompute_method": MegatronArgMetadata( arg_type=str, default=None, - help=( - "1) uniform: uniformly divide the total number of Transformer layers and" - " recompute the input activation of each divided chunk at specified" - " granularity, 2) recompute the input activations of only a set number of" - " individual Transformer layers per pipeline stage and do the rest without" - " any recomputing at specified granularitydefault) do not apply activations" - " recompute to any layers" - ), + help="Determines which transformer layers will be recomputed. uniform will uniformly divide the total number of transformer layers in a transformer block and recompute the input activation of each divided chunk at the specified granularity. block will recompute the input activations for only a set number of transformer layers per pipeline stage. The rest of the layers in the pipeline stage will not have any activations recomputed. If None, and recompute is enabled, all layers will do recomputation. If set, must be 'uniform' or 'block'.", choices=("uniform", "block"), nargs=None, element_type=None, @@ -4160,31 +4663,15 @@ "recompute_modules": MegatronArgMetadata( arg_type=list, default=None, - help=( - 'The submodules to recompute. choices: "core_attn", "moe_act", "layernorm",' - ' "mla_up_proj", "mlp", "moe", "shared_experts". default:' - ' ["core_attn"]."core_attn": recompute the core attention part of the' - ' transformer layer. "moe_act": recompute the MoE MLP activation function.' - ' "layernorm": recompute the input_layernorm and pre_mlp_layernorm.' - ' "mla_up_proj": recompute the MLA up projection and RoPE applying' - ' parts."mlp": recompute the dense MLP layer."moe": recompute the MoE' - ' layer."shared_experts": recompute the shared experts in the MoE' - ' layer."moe_act", "layernorm", and "mla_up_proj" use output-discarding' - ' checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal' - " checkpointing." - ), + help='The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. "layernorm": recompute the input_layernorm and pre_mlp_layernorm. "mla_up_proj": recompute the MLA up projection and RoPE applying parts. "mlp": recompute the dense MLP submodule. "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing.', choices=None, - nargs="*", + nargs="+", element_type=str, ), "recompute_num_layers": MegatronArgMetadata( arg_type=int, default=None, - help=( - "1) uniform: the number of Transformer layers in each uniformly divided" - " recompute unit, 2) block: the number of individual Transformer layers to" - " recompute within each pipeline stage." - ), + help="When recompute_method is uniform, recompute_num_layers is the number of transformer layers in each uniformly divided recompute unit. When recompute_method is block, recompute_num_layers is the number of transformer layers to recompute within each pipeline stage. Must be None for 'selective' activation checkpointing.", choices=None, nargs=None, element_type=None, @@ -4197,6 +4684,22 @@ nargs=0, element_type=None, ), + "record_shapes": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Record shapes of tensors in `torch.autograd.profiler.emit_nvtx` for the Nsys profiler.", + choices=None, + nargs=0, + element_type=None, + ), + "refit_method": MegatronArgMetadata( + arg_type=str, + default="gloo", + help="Method to refit the model weights between training and inference models during RL. nccl: use NCCLCopyService to refit using NCCL; gloo: use GlooCopyService over CPU; nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM.", + choices=("nccl", "gloo", "nvshmem"), + nargs=None, + element_type=None, + ), "relative_attention_max_distance": MegatronArgMetadata( arg_type=int, default=128, @@ -4216,9 +4719,7 @@ "replication": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks." - ), + help="If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks.", choices=None, nargs=0, element_type=None, @@ -4234,12 +4735,7 @@ "replication_jump": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Specifies `J`, the spacing between ranks storing replicas of a given" - " rank's data. Replicas for rank `n` may be on ranks `n+J`, `n+2J`, ..., or" - " `n-J`, `n-2J`, etc. This flag has an effect only if --replication is" - " used. and must be consistent across all ranks." - ), + help="Specifies `J`, the spacing between ranks storing replicas of a given rank's data. Replicas for rank `n` may be on ranks `n+J`, `n+2J`, ..., or `n-J`, `n-2J`, etc. This flag has an effect only if --replication is used. and must be consistent across all ranks.", choices=None, nargs=None, element_type=None, @@ -4247,10 +4743,7 @@ "rerun_mode": MegatronArgMetadata( arg_type=str, default="validate_results", - help=( - "Use re-run engine to validate results (default) or to emit stats on" - " variability of computations due to non-deterministic algorithms." - ), + help="Use re-run engine to validate results (default) or to emit stats on variability of computations due to non-deterministic algorithms.", choices=("disabled", "validate_results", "report_stats"), nargs=None, element_type=None, @@ -4295,138 +4788,138 @@ nargs=0, element_type=None, ), - "retriever_seq_length": MegatronArgMetadata( - arg_type=int, - default=256, - help="Maximum sequence length for the biencoder model for retriever", - choices=None, - nargs=None, - element_type=None, - ), - "retro_add_retriever": MegatronArgMetadata( + "reuse_grad_buf_for_mxfp8_param_ag": MegatronArgMetadata( arg_type=bool, default=False, - help=("Add a retriever to the transformer, for use in pretraining a Retro model."), + help="If True, reuse the grad buffer for MXFP8 parameter all-gather.", choices=None, nargs=0, element_type=None, ), - "retro_attention_gate": MegatronArgMetadata( + "rl_default_temperature": MegatronArgMetadata( arg_type=float, - default=1, - help="Gated cross attention.", + default=1.0, + help="Default temperature for model inference.", choices=None, nargs=None, element_type=None, ), - "retro_cyclic_train_iters": MegatronArgMetadata( + "rl_default_top_k": MegatronArgMetadata( arg_type=int, - default=None, - help="Set number of training iterations for cyclic Retro training.", - choices=None, - nargs=None, - element_type=None, - ), - "retro_encoder_attention_dropout": MegatronArgMetadata( - arg_type=float, - default=0.1, - help="Attention dropout for retrieval encoder.", + default=-1, + help="Default top-k for model inference.", choices=None, nargs=None, element_type=None, ), - "retro_encoder_hidden_dropout": MegatronArgMetadata( + "rl_default_top_p": MegatronArgMetadata( arg_type=float, - default=0.1, - help="Hidden dropout for retrieval encoder.", + default=0, + help="Default top-p for model inference.", choices=None, nargs=None, element_type=None, ), - "retro_encoder_layers": MegatronArgMetadata( + "rl_generation_batch_size": MegatronArgMetadata( arg_type=int, - default=2, - help="Number of layers to use for the retrieval encoder.", + default=None, + help="Override the number of groups per generation batch. Defaults to grpo_prompts_per_step when --rl-num-parallel-generation-batches is set.", choices=None, nargs=None, element_type=None, ), - "retro_num_neighbors": MegatronArgMetadata( - arg_type=int, - default=2, - help="Number of neighbors to retrieve during pretraining.", + "rl_importance_sampling_truncation_coef": MegatronArgMetadata( + arg_type=float, + default=None, + help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.", choices=None, nargs=None, element_type=None, ), - "retro_num_retrieved_chunks": MegatronArgMetadata( + "rl_inference_expert_model_parallel_size": MegatronArgMetadata( arg_type=int, - default=2, - help="Number of chunks to retrieve from the retrieval database.", + default=None, + help="Degree of expert model parallelism for inference for RL.", choices=None, nargs=None, element_type=None, ), - "retro_project_dir": MegatronArgMetadata( - arg_type=None, + "rl_inference_expert_tensor_model_parallel_size": MegatronArgMetadata( + arg_type=int, default=None, - help=( - "Retro project directory, which contains the preprocessed data for" - " pretraining. This directory is built during preprocessing (see" - " tools/retro/README.md), and contains subdirectories for the chunk" - " database and pretraining neighbors." - ), + help="Degree of expert tensor model parallelism for inference for RL. For MoE models, this controls the TP size for expert layers specifically. Defaults to training expert_tensor_parallel_size if not specified.", choices=None, nargs=None, element_type=None, ), - "retro_verify_neighbor_count": MegatronArgMetadata( + "rl_inference_logprobs_is_correction": MegatronArgMetadata( arg_type=bool, - default=True, - help="Skip verifying that len(GPT dataset) == len(saved neighbors).", + default=False, + help="If set, use inference logprobs in importance sampling correction of the loss.", choices=None, nargs=0, element_type=None, ), - "reuse_grad_buf_for_mxfp8_param_ag": MegatronArgMetadata( - arg_type=bool, - default=False, - help="If True, reuse the grad buffer for MXFP8 parameter all-gather.", + "rl_inference_model_unified_memory_level": MegatronArgMetadata( + arg_type=int, + default=0, + help="Allocate the separate RL inference model parameters from a unified virtual memory (UVM) CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe device pointers.", + choices=(0, 1), + nargs=None, + element_type=None, + ), + "rl_inference_parsers": MegatronArgMetadata( + arg_type=list, + default=[], + help="List of response parsers to enable for RL inference (e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).", choices=None, - nargs=0, + nargs="*", + element_type=str, + ), + "rl_inference_pipeline_model_parallel_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Degree of pipeline model parallelism for inference for RL.", + choices=None, + nargs=None, element_type=None, ), - "rl_calculate_intra_group_similarity": MegatronArgMetadata( - arg_type=None, - default=False, - help="If set, calculate the intra-group similarity of rollouts.", + "rl_inference_tensor_model_parallel_size": MegatronArgMetadata( + arg_type=int, + default=None, + help="Degree of tensor model parallelism for inference for RL.", choices=None, - nargs=0, + nargs=None, element_type=None, ), - "rl_importance_sampling_truncation_coef": MegatronArgMetadata( - arg_type=float, + "rl_kv_cache_management_mode": MegatronArgMetadata( + arg_type=str, + default="persist", + help="KV cache management mode during RL training: persist: leave KV cache in GPU memory (default), offload: offload KV cache to CPU during training, recompute: deallocate KV cache and recompute from scratch each cycle", + choices=("persist", "offload", "recompute"), + nargs=None, + element_type=None, + ), + "rl_num_parallel_generation_batches": MegatronArgMetadata( + arg_type=int, default=None, - help=( - "If --inference-logprobs-is-correction is on and this coefficient is set," - " apply truncation for the IS correction at GRPO loss." - ), + help="Number of generation batches in flight. Set to L+1 to allow for L steps of staleness between the inference and training policies. Each batch contains grpo_prompts_per_step groups by default. Requires --rl-partial-rollouts. Mutually exclusive with --rl-num-parallel-generations.", choices=None, nargs=None, element_type=None, ), - "rl_inference_logprobs_is_correction": MegatronArgMetadata( - arg_type=bool, - default=False, - help=("If set, use inference logprobs in importance sampling correction of the loss."), + "rl_num_parallel_generations": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of rollouts being generated by the inference engine simultaneously. Internally divided by grpo_group_size. Requires --rl-partial-rollouts. Mutually exclusive with --rl-num-parallel-generation-batches.", choices=None, - nargs=0, + nargs=None, element_type=None, ), - "rl_offload_kv_cache_during_training": MegatronArgMetadata( + "rl_offload_inference_model_weights_when_idle": MegatronArgMetadata( arg_type=None, default=False, - help="Offload KV cache to CPU during training to save GPU memory", + help="When using a separate RL inference model, offload its weights to CPU when not doing rollout inference, and restore to GPU right before inference. Works with two backends: 1) UVM (when --rl-inference-model-unified-memory-level=1), or 2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).", choices=None, nargs=0, element_type=None, @@ -4434,67 +4927,75 @@ "rl_offload_optimizer_during_inference": MegatronArgMetadata( arg_type=bool, default=False, - help=("Offload optimizer state to CPU during inference/rollout to save GPU memory"), - choices=None, - nargs=0, - element_type=None, - ), - "rl_partial_rollouts": MegatronArgMetadata( - arg_type=None, - default=False, - help="If set, use partial rollouts.", + help="Offload optimizer state to CPU during inference/rollout to save GPU memory", choices=None, nargs=0, element_type=None, ), - "rl_prompts_per_eval": MegatronArgMetadata( + "rl_parallel_generation_tasks": MegatronArgMetadata( arg_type=int, - default=32, - help=( - "Number of prompts to evaluate for for each RL task.This evaluation can be" - " very expensive when using environmentsthat evaluate pass@k so we default" - " to a lower number." - ), + default=None, + help="Deprecated: use --rl-num-parallel-generations instead.", choices=None, nargs=None, element_type=None, ), - "rl_remove_kv_cache_during_training": MegatronArgMetadata( + "rl_partial_rollouts": MegatronArgMetadata( arg_type=None, default=False, - help="Remove KV cache during training to save GPU memory", + help="Allow inference to continue generating rollouts while training updates the policy weights. This enables off-policy training where rollouts may be generated with a stale version of the policy. Use --rl-num-parallel-generations or --rl-num-parallel-generation-batches to control the degree of staleness.", choices=None, nargs=0, element_type=None, ), - "rl_reset_cuda_graphs": MegatronArgMetadata( + "rl_persist_cuda_graphs": MegatronArgMetadata( arg_type=bool, default=False, - help="Reset CUDA graphs between inference/training to save GPU memory", + help="Persist CUDA graphs when the inference engine is suspended. If False, CUDA graphs are deleted on suspend and re-captured on resume.", choices=None, nargs=0, element_type=None, ), + "rl_prompts_per_eval": MegatronArgMetadata( + arg_type=int, + default=32, + help="Number of prompts to evaluate for for each RL task.This evaluation can be very expensive when using environmentsthat evaluate pass@k so we default to a lower number.", + choices=None, + nargs=None, + element_type=None, + ), "rl_sequence_packing_algo": MegatronArgMetadata( arg_type=str, default="fifo", - help=( - "Algorithm for distributing packed bins across ranks. fifo:" - " first-in-first-out sequential distribution, round-robin: distribute bins" - " cyclically across ranks for better load balancing" - ), + help="Algorithm for distributing packed bins across ranks. fifo: first-in-first-out sequential distribution, round-robin: distribute bins cyclically across ranks for better load balancing", choices=("fifo", "round-robin"), nargs=None, element_type=None, ), - "rl_sequence_packing_bin_size": MegatronArgMetadata( + "rl_sequence_packing_max_sequences_per_bin": MegatronArgMetadata( arg_type=int, - default=8192, - help="Override bin size for sequence packing.", + default=50, + help="Maximum number of sequences that can be packed into a single bin. ", choices=None, nargs=None, element_type=None, ), + "rl_skip_bos_token": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Skip BOS token at the beginning of the sequences. Default is False.", + choices=None, + nargs=0, + element_type=None, + ), + "rl_training_cuda_graphs": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If set, do not toggle CUDA graphs on/off between inference and training phases.", + choices=None, + nargs=0, + element_type=None, + ), "rl_use_sequence_packing": MegatronArgMetadata( arg_type=bool, default=False, @@ -4503,6 +5004,14 @@ nargs=0, element_type=None, ), + "rl_verify_model_weights_swap": MegatronArgMetadata( + arg_type=None, + default=False, + help="If set, verify that the model weights were correctly transferred by comparing forward pass outputs onthe first swap of model weights.", + choices=None, + nargs=0, + element_type=None, + ), "rope_scaling_factor": MegatronArgMetadata( arg_type=float, default=8.0, @@ -4514,10 +5023,7 @@ "rope_type": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Type of rope to use. Note that MLA takes yarn by default, and common" - " attention takes rope by default." - ), + help="Type of rope to use. Note that MLA takes yarn by default, and common attention takes rope by default.", choices=("rope", "yarn"), nargs=None, element_type=None, @@ -4533,7 +5039,7 @@ "rotary_interleaved": MegatronArgMetadata( arg_type=bool, default=False, - help="Use interleaved rotary embedding.", + help="True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of first half and second half (LLaMa style). Default to False.", choices=None, nargs=0, element_type=None, @@ -4586,13 +5092,29 @@ nargs=None, element_type=None, ), - "save_extra_steps": MegatronArgMetadata( - arg_type=list, - default=[], - help="Specific training iterations at which to save checkpoints", + "save_activations_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of iterations between act.name->act.data mapping saves.", choices=None, - nargs="*", - element_type=int, + nargs=None, + element_type=None, + ), + "save_dgrads_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of iterations between dgrad saves.", + choices=None, + nargs=None, + element_type=None, + ), + "save_extra_steps": MegatronArgMetadata( + arg_type=list, + default=[], + help="Specific training iterations at which to save checkpoints", + choices=None, + nargs="*", + element_type=int, ), "save_interval": MegatronArgMetadata( arg_type=int, @@ -4602,23 +5124,36 @@ nargs=None, element_type=None, ), + "save_params_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of iterations between param.name->param.data mapping saves.", + choices=None, + nargs=None, + element_type=None, + ), "save_retain_interval": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Number of iterations between retained checkpoints (othercheckpoints" - " _except the last checkpoint_ are automatically deleted)." - ), + help="Number of iterations between retained checkpoints (other checkpoints except the last checkpoint are automatically deleted).", choices=None, nargs=None, element_type=None, ), - "scatter_gather_tensors_in_pipeline": MegatronArgMetadata( - arg_type=bool, - default=True, - help=("If not set, use scatter/gather to optimize communication of tensors in pipeline."), + "save_tokens_per_expert_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of iterations between tokens-per-expert routing metadata saves.", choices=None, - nargs=0, + nargs=None, + element_type=None, + ), + "save_wgrads_interval": MegatronArgMetadata( + arg_type=int, + default=None, + help="Number of iterations between wgrad (main_grad) saves.", + choices=None, + nargs=None, element_type=None, ), "seed": MegatronArgMetadata( @@ -4640,7 +5175,7 @@ "sequence_parallel": MegatronArgMetadata( arg_type=bool, default=False, - help="Enable sequence parallel optimization.", + help="Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially. See Reducing Activation Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198) for more details.", choices=None, nargs=0, element_type=None, @@ -4672,10 +5207,7 @@ "sharp_enabled_group": MegatronArgMetadata( arg_type=str, default=None, - help=( - "IB SHARP can be enabled from only one communication group. By default, it" - " is enabled from dp group. Available options: [dp, dp_replica]" - ), + help="IB SHARP can be enabled from only one communication group. By default, it is enabled from dp group if not specified and use_sharp=True. Available options: [dp, dp_replica]", choices=("dp", "dp_replica"), nargs=None, element_type=None, @@ -4691,10 +5223,7 @@ "skip_train": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, bypass the training loop, optionally do evaluation for" - " validation/test, and exit." - ), + help="If set, bypass the training loop, perform evaluation for validation/test, and exit.", choices=None, nargs=0, element_type=None, @@ -4702,36 +5231,31 @@ "softmax_type": MegatronArgMetadata( arg_type=str, default="vanilla", - help=( - "Type of softmax to use for the attention. Supports both a fixed offset and" - " learnable offset." - ), - choices=("learnable", "vanilla", "off-by-one"), + help="Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and and learnable offset.", + choices=("vanilla", "off-by-one", "learnable"), nargs=None, element_type=None, ), "spec": MegatronArgMetadata( arg_type=list, default=None, - help=( - "Specify the 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`" - ), + help="Specify the 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`", choices=None, nargs="*", element_type=str, ), + "special_tokens": MegatronArgMetadata( + arg_type=list, + default=None, + help='List of special tokens. For TikTokenizer needs to have ["", "", "", "", "", "", ""]', + choices=None, + nargs="+", + element_type=str, + ), "split": MegatronArgMetadata( arg_type=str, default=None, - help=( - "Comma-separated list of proportions for training, validation, and test" - " split. For example the split `90,5,5` will use 90%% of data for training," - " 5%% for validation and 5%% for test." - ), + help="Comma-separated list of proportions for training, validation, and test split. For example the split `90,5,5` will use 90%% of data for training, 5%% for validation and 5%% for test.", choices=None, nargs=None, element_type=None, @@ -4744,6 +5268,14 @@ nargs=0, element_type=None, ), + "start_eval_at_iter": MegatronArgMetadata( + arg_type=int, + default=None, + help="If set, evaluation will only start after this iteration number. Useful for skipping evaluation during early training iterations when the model is not yet meaningful. If not set, evaluation starts from the first eval_interval.", + choices=None, + nargs=None, + element_type=None, + ), "start_weight_decay": MegatronArgMetadata( arg_type=float, default=None, @@ -4752,6 +5284,14 @@ nargs=None, element_type=None, ), + "step_batch_size_schedule": MegatronArgMetadata( + arg_type=str, + default=None, + help='Step-wise batch size schedule in format "THRESHOLD:BS THRESHOLD:BS ...". Thresholds support suffixes: K (1e3), M (1e6), B (1e9), T (1e12). If sequence length is provided, thresholds are interpreted as tokens; otherwise as samples. Example: step_batch_size_schedule = "0:768 250B:1536 500B:3072 750B:6144" Cannot be used together with decrease_batch_size_if_needed.', + choices=None, + nargs=None, + element_type=None, + ), "straggler_ctrlr_port": MegatronArgMetadata( arg_type=int, default=65535, @@ -4771,7 +5311,7 @@ "strict_fsdp_dtensor_load": MegatronArgMetadata( arg_type=bool, default=True, - help="Do not strict loading for fsdp_dtensor checkpoint format.", + help="Whether to enforce strict loading for FSDP DTensor checkpoints. When False, allows partial loading.", choices=None, nargs=0, element_type=None, @@ -4779,14 +5319,7 @@ "suggested_communication_unit_size": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Specifies the number of elements to communicate at once during FSDP (Fully" - " Sharded Data Parallel) operations. This flag also affects FSDP all-gather" - " prefetch behavior. Setting a larger value increases the communication" - " buffer size, while a smaller value disables prefetching and may degrade" - " performance. Adjust this value based on your system's memory and" - " performance requirements." - ), + help="Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system's memory and performance requirements.", choices=None, nargs=None, element_type=None, @@ -4810,21 +5343,23 @@ "symmetric_ar_type": MegatronArgMetadata( arg_type=str, default=None, - help=( - "What type of symmetric all reduce to use. The default is none which is no" - " use of symetric memory" - ), - choices=("two_shot", "one_shot", "multimem_all_reduce", None), + help="What type of symmetric all reduce to use. The default is None which is no use of symmetric memory.", + choices=("two_shot", "one_shot", "multimem_all_reduce"), + nargs=None, + element_type=None, + ), + "te_precision_config_file": MegatronArgMetadata( + arg_type=None, + default=None, + help="Configuration file to select per-module precision overrides. See TransformerEngineMixedPrecision.md", + choices=None, nargs=None, element_type=None, ), "te_rng_tracker": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use the Transformer Engine version of the random number generator." - " Required for CUDA graphs support." - ), + help="Use the Transformer Engine version of the random number generator. Required for CUDA graphs support.", choices=None, nargs=0, element_type=None, @@ -4832,7 +5367,7 @@ "tensor_model_parallel_size": MegatronArgMetadata( arg_type=int, default=1, - help="Degree of tensor model parallelism.", + help="Intra-layer model parallelism. Splits tensors across GPU ranks.", choices=None, nargs=None, element_type=None, @@ -4856,10 +5391,7 @@ "tensorboard_queue_size": MegatronArgMetadata( arg_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.' - ), + help="Size of the tensorboard queue for pending events and summaries before one of the 'add' calls forces a flush to disk.", choices=None, nargs=None, element_type=None, @@ -4867,10 +5399,7 @@ "test_data_path": MegatronArgMetadata( arg_type=list, default=None, - help=( - "The weight and prefix list for an independent test dataset. Follows the" - " same pattern rules as --data-path." - ), + help="The weight and prefix list for an independent test dataset. Follows the same pattern rules as --data-path.", choices=None, nargs="*", element_type=str, @@ -4886,7 +5415,7 @@ "tiktoken_num_special_tokens": MegatronArgMetadata( arg_type=int, default=1000, - help="Number of special tokens in tiktoken tokenizer", + help="Number of special tokens in tiktoken tokenizer.", choices=None, nargs=None, element_type=None, @@ -4895,33 +5424,14 @@ arg_type=str, default=None, help="Which tiktoken pattern to use. Options: [v1, v2]", - choices=None, + choices=("v1", "v2"), nargs=None, element_type=None, ), - "tiktoken_special_tokens": MegatronArgMetadata( - arg_type=list, - default=None, - help=( - 'List of tiktoken special tokens, needs to have ["", "", "",' - ' "", "", "", ""]' - ), - choices=None, - nargs="+", - element_type=str, - ), "timing_log_level": MegatronArgMetadata( arg_type=int, default=0, - 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." - ), + 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.", choices=(0, 1, 2), nargs=None, element_type=None, @@ -4929,11 +5439,7 @@ "timing_log_option": MegatronArgMetadata( arg_type=str, default="minmax", - 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." - ), + 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.", choices=("max", "minmax", "all"), nargs=None, element_type=None, @@ -4946,22 +5452,46 @@ nargs=None, element_type=None, ), - "tokenizer_metadata": MegatronArgMetadata( - arg_type=str, - default=None, - help="Path to tokenizer metadata in json format.", + "tokenizer_hf_no_include_special_tokens": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Converting text to ids will not include special for HuggingFace tokenizer.", choices=None, - nargs=None, + nargs=0, + element_type=None, + ), + "tokenizer_hf_no_use_fast": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to use fast HuggingFace tokenizer.", + choices=None, + nargs=0, element_type=None, ), "tokenizer_model": MegatronArgMetadata( arg_type=str, default=None, - help="Sentencepiece tokenizer model.", + help="Path to the tokenizer model.", choices=None, nargs=None, element_type=None, ), + "tokenizer_sentencepiece_ignore_extra_whitespaces": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Whether to ignore extra whitespaces in the input text while encoding.", + choices=None, + nargs=0, + element_type=None, + ), + "tokenizer_sentencepiece_legacy": MegatronArgMetadata( + arg_type=bool, + default=False, + help="SentencePiece tokenizer wrapper legacy behavior. Allows special tokens usage.", + choices=None, + nargs=0, + element_type=None, + ), "tokenizer_type": MegatronArgMetadata( arg_type=str, default=None, @@ -4986,10 +5516,7 @@ "torch_fsdp2_reshard_after_forward": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "Whether to reshard weights after forward pass when using PyTorch FSDP2." - " Set to enable FSDP ZeRO-2." - ), + help="Whether to reshard weights after forward pass when using PyTorch FSDP2. Set to enable FSDP ZeRO-2.", choices=None, nargs=0, element_type=None, @@ -5005,7 +5532,7 @@ "tp_comm_bulk_dgrad": MegatronArgMetadata( arg_type=bool, default=True, - help="Disables the All-Gather overlap with bprop activation gradient GEMM.", + help="Controls Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5013,7 +5540,7 @@ "tp_comm_bulk_wgrad": MegatronArgMetadata( arg_type=bool, default=True, - help="Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.", + help="Controls All-Gather overlap with Bprop activation gradient GEMM. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5021,7 +5548,7 @@ "tp_comm_overlap": MegatronArgMetadata( arg_type=bool, default=False, - help="Enables the overlap of Tensor parallel communication and GEMM kernels.", + help="If true, allows overlapping of Linear layer execution with tensor parallel communication collectives like AllGather/ReduceScatter. Overlapping is done for the linear layers wherever possible during the forward and the backward pass.", choices=None, nargs=0, element_type=None, @@ -5029,7 +5556,7 @@ "tp_comm_overlap_ag": MegatronArgMetadata( arg_type=bool, default=True, - help=("Disables the All-Gather overlap with GEMM by pipelining the GEMM and All-Gather."), + help="Controls All-Gather overlap with GEMM by pipelining the GEMM and All-Gather. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5045,10 +5572,7 @@ "tp_comm_overlap_rs": MegatronArgMetadata( arg_type=bool, default=True, - help=( - "Disables the Reduce-Scatter overlap with GEMM by pipelining the GEMM and" - " Reduce-Scatter." - ), + help="Controls Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5056,7 +5580,7 @@ "tp_comm_overlap_rs_dgrad": MegatronArgMetadata( arg_type=bool, default=False, - help="Enables the Reduce-Scatter overlap with dgrad GEMM.", + help="If true, allows Reduce-Scatter overlap with DGRAD GEMM by pipelining the GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5064,7 +5588,7 @@ "tp_comm_split_ag": MegatronArgMetadata( arg_type=bool, default=True, - help="Disables the All-Gather overlap with fprop GEMM.", + help="Deprecated from TransformerEngine v1.6.0. Controls All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather splits. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5072,7 +5596,7 @@ "tp_comm_split_rs": MegatronArgMetadata( arg_type=bool, default=True, - help="Disables the Reduce-Scatter overlap with fprop GEMM.", + help="Deprecated from TransformerEngine v1.6.0. Controls Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is False.", choices=None, nargs=0, element_type=None, @@ -5080,10 +5604,7 @@ "train_data_path": MegatronArgMetadata( arg_type=list, default=None, - help=( - "The weight and prefix list for an independent train dataset. Follows the" - " same pattern rules as --data-path." - ), + help="The weight and prefix list for an independent train dataset. Follows the same pattern rules as --data-path.", choices=None, nargs="*", element_type=str, @@ -5091,10 +5612,7 @@ "train_iters": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Total number of iterations to train over all training runs. Note that" - " either train-iters or train-samples should be provided." - ), + help="Total number of iterations to train over all training runs. Note that either train_iters or train_samples should be provided.", choices=None, nargs=None, element_type=None, @@ -5102,10 +5620,7 @@ "train_samples": MegatronArgMetadata( arg_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." - ), + help="Total number of samples to train over all training runs. Note that either train_iters or train_samples should be provided.", choices=None, nargs=None, element_type=None, @@ -5113,18 +5628,15 @@ "train_sync_interval": MegatronArgMetadata( arg_type=int, default=None, - help=( - "Training CPU-GPU synchronization interval, to ensure that CPU is not" - " running too far ahead of GPU." - ), + help="Training CPU-GPU synchronization interval, to ensure that CPU is not running too far ahead of GPU.", choices=None, nargs=None, element_type=None, ), "transformer_impl": MegatronArgMetadata( - arg_type=None, + arg_type=str, default="transformer_engine", - help="Which Transformer implementation to use.", + help="Transformer implementation to use. Options are 'transformer_engine' for Transformer Engine and 'local' for MCore.", choices=("local", "transformer_engine", "inference_optimized"), nargs=None, element_type=None, @@ -5132,7 +5644,7 @@ "trust_remote_code": MegatronArgMetadata( arg_type=bool, default=False, - help="Whether or not to allow PreTrainedTokenizer to execute remote code", + help="Whether or not to allow PreTrainedTokenizer to execute remote code.", choices=None, nargs=0, element_type=None, @@ -5148,7 +5660,7 @@ "use_checkpoint_args": MegatronArgMetadata( arg_type=bool, default=False, - help=("Override model-related command-line arguments with arguments from checkpoint"), + help="Override model-related command-line arguments with arguments from checkpoint", choices=None, nargs=0, element_type=None, @@ -5156,11 +5668,7 @@ "use_checkpoint_opt_param_scheduler": MegatronArgMetadata( arg_type=bool, default=False, - 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." - ), + 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.", choices=None, nargs=0, element_type=None, @@ -5168,10 +5676,7 @@ "use_cpu_initialization": MegatronArgMetadata( arg_type=bool, default=None, - help=( - "If set, initialize weights on the CPU. This eliminates init differences" - " based on tensor parallelism." - ), + help="When set to False, we initialize the weights directly on the GPU. CPU initialization is the same regardless of tensor model parallelism, but GPU initialization is not. Transferring weights from CPU to GPU can take a significant amount of time for large models.", choices=None, nargs=0, element_type=None, @@ -5195,7 +5700,7 @@ "use_flash_attn": MegatronArgMetadata( arg_type=bool, default=False, - help=("use FlashAttention implementation of attention. https://arxiv.org/abs/2205.14135"), + help="use FlashAttention implementation of attention. https://arxiv.org/abs/2205.14135", choices=None, nargs=0, element_type=None, @@ -5203,15 +5708,39 @@ "use_fused_weighted_squared_relu": MegatronArgMetadata( arg_type=bool, default=False, - help="Use fused weighted squared relu when using MoE.", + help="If True, uses fused weighted squared relu kernel when using MoE.", + choices=None, + nargs=0, + element_type=None, + ), + "use_gloo_process_groups": MegatronArgMetadata( + arg_type=bool, + default=True, + help="If enabled, create Gloo process groups for communications.", + choices=None, + nargs=0, + element_type=None, + ), + "use_grouped_gemm_for_dense_mlp": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Use GroupedLinear(num_groups=1) for dense MLP to trigger the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. Requires ``use_te_op_fuser=True`` and SwiGLU activation.", choices=None, nargs=0, element_type=None, ), - "use_legacy_models": MegatronArgMetadata( + "use_kitchen_attention": MegatronArgMetadata( arg_type=bool, default=False, - help="Use the legacy Megatron models, not Megatron-Core models.", + help="Use the kitchen extension for attention (instead of TE's attention).", + choices=None, + nargs=0, + element_type=None, + ), + "use_layer_wise_param_layout": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Opt out of the precomputed LayerWise param layout. When set, falls back to the legacy LayerWise ping-pong path: all params (including non-Muon embeddings, biases, layernorm) live in a single LayerWise buffer and the optimizer uses the allgather_params() codepath. The default (precomputed layout) routes non-Muon params through a separate DistributedOptimizer with byte-level sharding, which is faster and uses less padding but produces different bf16 reduction ordering and so will not match legacy-path loss curves bit-for-bit.", choices=None, nargs=0, element_type=None, @@ -5219,9 +5748,15 @@ "use_legacy_static_engine": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use legacy static engine. (Current static engine uses dynamic engine under the hood)" - ), + help="Use legacy static engine. (Current static engine uses dynamic engine under the hood)", + choices=None, + nargs=0, + element_type=None, + ), + "use_mamba_mem_eff_path": MegatronArgMetadata( + arg_type=bool, + default=True, + help="Controls usage of the memory efficient path for Mamba layers.", choices=None, nargs=0, element_type=None, @@ -5229,7 +5764,7 @@ "use_megatron_fsdp": MegatronArgMetadata( arg_type=bool, default=False, - help="Use the Megatron FSDP code path in DDP.", + help="Use Megatron's Fully Sharded Data Parallel. Cannot be used together with use_torch_fsdp2.", choices=None, nargs=0, element_type=None, @@ -5242,6 +5777,14 @@ nargs=0, element_type=None, ), + "use_mup": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Enable Maximal Update Parameterization (MuP) for hyperparameter transfer across model widths. When enabled, learning rates and initialization are scaled according to the width multiplier to ensure consistent training dynamics.", + choices=None, + nargs=0, + element_type=None, + ), "use_one_sent_docs": MegatronArgMetadata( arg_type=bool, default=False, @@ -5253,7 +5796,7 @@ "use_persistent_ckpt_worker": MegatronArgMetadata( arg_type=bool, default=False, - help="Enables a persitent checkpoint worker for async save", + help="Use a persistent background worker for async checkpoint saves. When enabled, creates a dedicated worker thread/process for handling async saves. When disabled, uses temporal workers that are created and destroyed for each save operation.", choices=None, nargs=0, element_type=None, @@ -5261,11 +5804,7 @@ "use_precision_aware_optimizer": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use the precision-aware optimizer in TransformerEngine, which allows" - " setting the main params and optimizer states to lower precision, such as" - " fp16, bf16 and fp8." - ), + help="Use the precision-aware optimizer in TransformerEngine, which allows setting the main params and optimizer states to lower precision, such as fp16, bf16 and fp8.", choices=None, nargs=0, element_type=None, @@ -5273,9 +5812,7 @@ "use_pytorch_profiler": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard." - ), + help="Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard.", choices=None, nargs=0, element_type=None, @@ -5283,11 +5820,7 @@ "use_ring_exchange_p2p": MegatronArgMetadata( arg_type=bool, 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." - ), + help="Use custom ring_exchange kernel instead of torch.distributed.batch_isend_irecv(). Requires custom built torch with torch.distributed.ring_exchange.", choices=None, nargs=0, element_type=None, @@ -5303,7 +5836,7 @@ "use_rotary_position_embeddings": MegatronArgMetadata( arg_type=bool, default=False, - help=("Use rotary positional embeddings or not. Deprecated: use --position-embedding-type"), + help="Use rotary positional embeddings or not. Deprecated: use --position-embedding-type", choices=None, nargs=0, element_type=None, @@ -5311,7 +5844,7 @@ "use_sharp": MegatronArgMetadata( arg_type=bool, default=False, - help="Required to enable SHARP communication.", + help="Set the use of SHARP for the collective communications of data-parallel process groups. When `True`, run barrier within each data-parallel process group, which specifies the SHARP application target groups.", choices=None, nargs=0, element_type=None, @@ -5319,7 +5852,7 @@ "use_te_activation_func": MegatronArgMetadata( arg_type=bool, default=False, - help="Use activation function kernel from Transformer Engine in MLP module.", + help="Whether to use ffn activation functions implemented by TransformerEngine", choices=None, nargs=0, element_type=None, @@ -5335,10 +5868,7 @@ "use_torch_fsdp2": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use the torch FSDP2 implementation. FSDP2 has not been tested with" - " pipeline parallelism, and may contain bugs." - ), + help="Use the torch FSDP2 implementation. FSDP2 is not currently working with Pipeline Parallel. It is still not in a stable release stage, and may therefore contain bugs or other potential issues.", choices=None, nargs=0, element_type=None, @@ -5346,10 +5876,7 @@ "use_torch_optimizer_for_cpu_offload": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu" - " offload mode." - ), + help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.", choices=None, nargs=0, element_type=None, @@ -5357,10 +5884,15 @@ "use_tp_pp_dp_mapping": MegatronArgMetadata( arg_type=bool, default=False, - help=( - "If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp" - " to tp-cp-ep-pp-dp." - ), + help="If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep-pp-dp.", + choices=None, + nargs=0, + element_type=None, + ), + "use_transformer_engine_op_fuser": MegatronArgMetadata( + arg_type=bool, + default=False, + help="If True, submodules may use Transformer Engine's operation fuser API to enable advanced fusions.", choices=None, nargs=0, element_type=None, @@ -5376,14 +5908,27 @@ "valid_data_path": MegatronArgMetadata( arg_type=list, default=None, - help=( - "The weight and prefix list for an independent validation dataset. Follows" - " the same pattern rules as --data-path." - ), + help="The weight and prefix list for an independent validation dataset. Follows the same pattern rules as --data-path.", choices=None, nargs="*", element_type=str, ), + "validation_set_names": MegatronArgMetadata( + arg_type=list, + default=None, + help="Optional list of names for multiple validation sets. When provided with --multiple-validation-sets, these names are used instead of numeric indices (e.g. 'validation-wiki' instead of 'validation-0'). The number of names must match the number of validation datasets.", + choices=None, + nargs="+", + element_type=str, + ), + "verify_integrity": MegatronArgMetadata( + arg_type=bool, + default=False, + help="Whether to hash checkpointing files during save and validate their integrity during load.", + choices=None, + nargs=0, + element_type=None, + ), "vision_backbone_type": MegatronArgMetadata( arg_type=str, default="vit", @@ -5411,9 +5956,7 @@ "vocab_extra_ids": MegatronArgMetadata( arg_type=int, default=0, - help=( - "Number of additional vocabulary tokens. They are used for span masking in the T5 model" - ), + help="Number of additional vocabulary tokens. They are used for span masking in the T5 model.", choices=None, nargs=None, element_type=None, @@ -5436,20 +5979,15 @@ ), "wandb_entity": MegatronArgMetadata( arg_type=str, - default="", - help=( - "The wandb entity name. It is useful when there are multiple sub-projects" - " in a project." - " https://community.wandb.ai/t/how-do-i-decide-which-account-private-or-team-to-upload-the-run-to/5704" - " Ignore wandb by default." - ), + default=None, + help="The wandb entity name. It is useful when there are multiple sub-projects in a project.", choices=None, nargs=None, element_type=None, ), "wandb_exp_name": MegatronArgMetadata( arg_type=str, - default="", + default=None, help="The wandb experiment name.", choices=None, nargs=None, @@ -5457,7 +5995,7 @@ ), "wandb_project": MegatronArgMetadata( arg_type=str, - default="", + default=None, help="The wandb project name. Ignore wandb by default.", choices=None, nargs=None, @@ -5465,7 +6003,7 @@ ), "wandb_save_dir": MegatronArgMetadata( arg_type=str, - default="", + default=None, help="Path to save the wandb results locally.", choices=None, nargs=None, @@ -5474,7 +6012,7 @@ "warmup": MegatronArgMetadata( arg_type=int, default=None, - help=("Old lr warmup argument, do not use. Use one of the--lr-warmup-* arguments above"), + help="Old lr warmup argument, do not use. Use one of the--lr-warmup-* arguments above", choices=None, nargs=None, element_type=None, @@ -5498,11 +6036,7 @@ "wgrad_deferral_limit": MegatronArgMetadata( arg_type=int, default=0, - help=( - "Number of micro-batches for whichweight gradient computation of vocabulary" - " projection is deferred, defaults to 0 whichmeans all the micro-batches" - " are deferred. Invalid if `defer-embedding-wgrad-compute`is not set" - ), + help="This value tunes the number of micro-batches for which the embedding weight gradient compute needs to be deferred to pipeline flush, this argument is invalid if `defer_embedding_wgrad_compute` is False. Defaults to 0, which means all micro-batches are deferred.", choices=None, nargs=None, element_type=None, @@ -5510,14 +6044,7 @@ "window_attn_skip_freq": MegatronArgMetadata( arg_type=None, default=None, - help=( - "Frequency of layers to skip window attention. Accepts either: - An integer" - " N: Represents a (N-1):1 ratio, meaning one full attention layer after" - " (N-1) SWA layers. - A string containing a Python list expression that" - ' defines a custom pattern, e.g.: "[1,1,1,0]*3" evaluates to' - " [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates SWA and 0 indicates full" - " attention. " - ), + help='Frequency of layers to skip window attention. Accepts either: - An integer N: Represents a (N-1):1 ratio, meaning one full attention layer after (N-1) SWA layers. - A string containing a Python list expression that defines a custom pattern, e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates SWA and 0 indicates full attention. ', choices=None, nargs=None, element_type=None, @@ -5525,9 +6052,7 @@ "window_size": MegatronArgMetadata( arg_type=None, default=None, - help=( - "Window size for window attention. If not provided, window attention will be disabled." - ), + help="Window size for window attention. If not provided, window attention will be disabled.", choices=None, nargs=None, element_type=None, @@ -5572,25 +6097,13 @@ default=None, ), "adam_beta1": MegatronActionSpec( - option_strings=("--adam-beta1",), - action_type="store", - nargs=None, - const=None, - default=0.9, + option_strings=("--adam-beta1",), action_type="store", nargs=None, const=None, default=0.9 ), "adam_beta2": MegatronActionSpec( - option_strings=("--adam-beta2",), - action_type="store", - nargs=None, - const=None, - default=0.999, + option_strings=("--adam-beta2",), action_type="store", nargs=None, const=None, default=0.999 ), "adam_eps": MegatronActionSpec( - option_strings=("--adam-eps",), - action_type="store", - nargs=None, - const=None, - default=1e-08, + option_strings=("--adam-eps",), action_type="store", nargs=None, const=None, default=1e-08 ), "add_bias_linear": MegatronActionSpec( option_strings=("--disable-bias-linear",), @@ -5628,7 +6141,7 @@ default=1000, ), "align_grad_reduce": MegatronActionSpec( - option_strings=("--no-align-grad-reduce",), + option_strings=("--no-align-grad-reduce", "--disable-align-grad-reduce"), action_type="store_false", nargs=0, const=False, @@ -5662,13 +6175,6 @@ const=None, default="0.0.0", ), - "apply_layernorm_1p": MegatronActionSpec( - option_strings=("--apply-layernorm-1p",), - action_type="store_true", - nargs=0, - const=True, - default=False, - ), "apply_query_key_layer_scaling": MegatronActionSpec( option_strings=("--apply-query-key-layer-scaling",), action_type="store_true", @@ -5690,19 +6196,47 @@ const=False, default=True, ), + "apply_wd_to_qk_layernorm": MegatronActionSpec( + option_strings=("--apply-wd-to-qk-layernorm",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "async_ckpt_cpu_priority": MegatronActionSpec( + option_strings=("--async-ckpt-cpu-priority",), + action_type="store", + nargs=None, + const=None, + default=10, + ), + "async_ckpt_io_priority": MegatronActionSpec( + option_strings=("--async-ckpt-io-priority",), + action_type="store", + nargs=None, + const=None, + default=3, + ), + "async_ckpt_use_cpu_shm": MegatronActionSpec( + option_strings=("--async-ckpt-use-cpu-shm",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "async_save": MegatronActionSpec( option_strings=("--async-save",), action_type="store_true", nargs=0, const=True, - default=None, + default=False, ), - "async_tensor_model_parallel_allreduce": MegatronActionSpec( - option_strings=("--no-async-tensor-model-parallel-allreduce",), - action_type="store_false", - nargs=0, - const=False, - default=True, + "async_strategy": MegatronActionSpec( + option_strings=("--async-strategy",), + action_type="store", + nargs=None, + const=None, + default="nvrx", ), "attention_backend": MegatronActionSpec( option_strings=("--attention-backend",), @@ -5718,6 +6252,13 @@ const=None, default=0.1, ), + "attention_output_gate": MegatronActionSpec( + option_strings=("--attention-output-gate",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "attention_softmax_in_fp32": MegatronActionSpec( option_strings=("--attention-softmax-in-fp32",), action_type="store_true", @@ -5739,12 +6280,15 @@ const=False, default=True, ), + "batch_invariant_mode": MegatronActionSpec( + option_strings=("--batch-invariant-mode",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "batch_size": MegatronActionSpec( - option_strings=("--batch-size",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--batch-size",), action_type="store", nargs=None, const=None, default=None ), "bert_binary_head": MegatronActionSpec( option_strings=("--bert-no-binary-head",), @@ -5761,18 +6305,10 @@ default="megatron", ), "bert_load": MegatronActionSpec( - option_strings=("--bert-load",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--bert-load",), action_type="store", nargs=None, const=None, default=None ), "bf16": MegatronActionSpec( - option_strings=("--bf16",), - action_type="store_true", - nargs=0, - const=True, - default=False, + option_strings=("--bf16",), action_type="store_true", nargs=0, const=True, default=False ), "bias_dropout_fusion": MegatronActionSpec( option_strings=("--no-bias-dropout-fusion",), @@ -5837,6 +6373,13 @@ const=True, default=False, ), + "chat_template": MegatronActionSpec( + option_strings=("--chat-template",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "check_for_large_grads": MegatronActionSpec( option_strings=("--check-for-large-grads",), action_type="store_true", @@ -5914,6 +6457,20 @@ const=True, default=False, ), + "ckpt_fully_parallel_load_exchange_algo": MegatronActionSpec( + option_strings=("--ckpt-fully-parallel-load-exchange-algo",), + action_type="store", + nargs=None, + const=None, + default="broadcast", + ), + "ckpt_fully_parallel_load_process_group": MegatronActionSpec( + option_strings=("--ckpt-fully-parallel-load-process-group",), + action_type="store", + nargs=None, + const=None, + default="dp", + ), "ckpt_fully_parallel_save": MegatronActionSpec( option_strings=("--no-ckpt-fully-parallel-save",), action_type="store_false", @@ -5928,12 +6485,25 @@ const=True, default=False, ), - "ckpt_step": MegatronActionSpec( - option_strings=("--ckpt-step",), + "ckpt_fully_parallel_save_process_group": MegatronActionSpec( + option_strings=("--ckpt-fully-parallel-save-process-group",), action_type="store", nargs=None, const=None, - default=None, + default="dp", + ), + "ckpt_load_validate_sharding_integrity": MegatronActionSpec( + option_strings=( + "--no-ckpt-load-validate-sharding-integrity", + "--disable-ckpt-load-validate-sharding-integrity", + ), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "ckpt_step": MegatronActionSpec( + option_strings=("--ckpt-step",), action_type="store", nargs=None, const=None, default=None ), "classes_fraction": MegatronActionSpec( option_strings=("--classes-fraction",), @@ -5943,14 +6513,13 @@ default=1.0, ), "clip_grad": MegatronActionSpec( - option_strings=("--clip-grad",), - action_type="store", - nargs=None, - const=None, - default=1.0, + option_strings=("--clip-grad",), action_type="store", nargs=None, const=None, default=1.0 ), "clone_scatter_output_in_embedding": MegatronActionSpec( - option_strings=("--no-clone-scatter-output-in-embedding",), + option_strings=( + "--no-clone-scatter-output-in-embedding", + "--disable-clone-scatter-output-in-embedding", + ), action_type="store_false", nargs=0, const=False, @@ -5984,6 +6553,20 @@ const=None, default=0, ), + "cpu_offloading_retain_pinned_cpu_buffers": MegatronActionSpec( + option_strings=("--cpu-offloading-retain-pinned-cpu-buffers",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "create_all_gather_group": MegatronActionSpec( + option_strings=("--create-all-gather-group",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "create_attention_mask_in_dataloader": MegatronActionSpec( option_strings=("--no-create-attention-mask-in-dataloader",), action_type="store_false", @@ -6012,12 +6595,19 @@ const=None, default="none", ), - "cuda_graph_scope": MegatronActionSpec( + "cuda_graph_modules": MegatronActionSpec( + option_strings=("--cuda-graph-modules",), + action_type="store", + nargs="+", + const=None, + default=[], + ), + "cuda_graph_scope_deprecated": MegatronActionSpec( option_strings=("--cuda-graph-scope",), action_type="store", - nargs=None, + nargs="+", const=None, - default="full", + default=None, ), "cuda_graph_warmup_steps": MegatronActionSpec( option_strings=("--cuda-graph-warmup-steps",), @@ -6052,14 +6642,10 @@ action_type="store", nargs=None, const=None, - default="no_shard", + default="optim_grads_params", ), "data_path": MegatronActionSpec( - option_strings=("--data-path",), - action_type="store", - nargs="*", - const=None, - default=None, + option_strings=("--data-path",), action_type="store", nargs="*", const=None, default=None ), "data_per_class_fraction": MegatronActionSpec( option_strings=("--data-per-class-fraction",), @@ -6075,6 +6661,20 @@ const=False, default=True, ), + "dataloader_defer_npy_index_mmap": MegatronActionSpec( + option_strings=("--dataloader-defer-npy-index-mmap",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "dataloader_fast_cache_load": MegatronActionSpec( + option_strings=("--dataloader-fast-cache-load",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "dataloader_type": MegatronActionSpec( option_strings=("--dataloader-type",), action_type="store", @@ -6110,6 +6710,13 @@ const=True, default=False, ), + "ddp_param_name_patterns_for_fp32_local_accumulation": MegatronActionSpec( + option_strings=("--ddp-param-name-patterns-for-fp32-local-accumulation",), + action_type="store", + nargs="+", + const=None, + default=[], + ), "ddp_reduce_scatter_with_fp32_accumulation": MegatronActionSpec( option_strings=("--ddp-reduce-scatter-with-fp32-accumulation",), action_type="store_true", @@ -6271,15 +6878,15 @@ const=True, default=False, ), - "disable_chunked_prefill": MegatronActionSpec( - option_strings=("--disable-chunked-prefill",), + "disable_jit_fuser": MegatronActionSpec( + option_strings=("--disable-jit-fuser",), action_type="store_true", nargs=0, const=True, default=False, ), - "disable_mamba_mem_eff_path": MegatronActionSpec( - option_strings=("--disable-mamba-mem-eff-path",), + "disable_msc_deprecated": MegatronActionSpec( + option_strings=("--disable-msc",), action_type="store_true", nargs=0, const=True, @@ -6327,6 +6934,13 @@ const=None, default="assume_ok_unexpected", ), + "dist_ckpt_workers": MegatronActionSpec( + option_strings=("--dist-ckpt-workers",), + action_type="store", + nargs=None, + const=None, + default=1, + ), "distrib_optim_fully_reshardable_mem_efficient": MegatronActionSpec( option_strings=("--distrib-optim-fully-reshardable-mem-efficient",), action_type="store_true", @@ -6362,6 +6976,41 @@ const=None, default=None, ), + "dsa_indexer_head_dim": MegatronActionSpec( + option_strings=("--dsa-indexer-head-dim",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "dsa_indexer_loss_coeff": MegatronActionSpec( + option_strings=("--dsa-indexer-loss-coeff",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "dsa_indexer_n_heads": MegatronActionSpec( + option_strings=("--dsa-indexer-n-heads",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "dsa_indexer_topk": MegatronActionSpec( + option_strings=("--dsa-indexer-topk",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "dsa_indexer_use_sparse_loss": MegatronActionSpec( + option_strings=("--dsa-indexer-use-sparse-loss",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "dump_param_to_param_group_map": MegatronActionSpec( option_strings=("--dump-param-to-param-group-map",), action_type="store", @@ -6390,6 +7039,13 @@ const=None, default=0, ), + "enable_chunked_prefill": MegatronActionSpec( + option_strings=("--enable-chunked-prefill",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "enable_cuda_graph": MegatronActionSpec( option_strings=("--enable-cuda-graph",), action_type="store_true", @@ -6418,19 +7074,12 @@ const=True, default=False, ), - "enable_gloo_process_groups": MegatronActionSpec( - option_strings=("--disable-gloo-process-groups",), - action_type="store_false", - nargs=0, - const=False, - default=True, - ), "enable_msc": MegatronActionSpec( - option_strings=("--disable-msc",), - action_type="store_false", + option_strings=("--enable-msc",), + action_type="store_true", nargs=0, - const=False, - default=True, + const=True, + default=False, ), "enable_one_logger": MegatronActionSpec( option_strings=("--no-one-logger",), @@ -6467,6 +7116,13 @@ const=True, default=False, ), + "ep_overlap_early_attn_memory_release": MegatronActionSpec( + option_strings=("--ep-overlap-early-attn-memory-release",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "error_injection_rate": MegatronActionSpec( option_strings=("--error-injection-rate",), action_type="store", @@ -6481,19 +7137,29 @@ const=None, default="transient_error", ), + "eval_global_batch_size": MegatronActionSpec( + option_strings=("--eval-global-batch-size",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "eval_interval": MegatronActionSpec( option_strings=("--eval-interval",), action_type="store", nargs=None, const=None, - default=1000, + default=None, ), "eval_iters": MegatronActionSpec( - option_strings=("--eval-iters",), + option_strings=("--eval-iters",), action_type="store", nargs=None, const=None, default=100 + ), + "eval_micro_batch_size": MegatronActionSpec( + option_strings=("--eval-micro-batch-size",), action_type="store", nargs=None, const=None, - default=100, + default=None, ), "evidence_data_path": MegatronActionSpec( option_strings=("--evidence-data-path",), @@ -6524,11 +7190,7 @@ default=False, ), "exit_signal": MegatronActionSpec( - option_strings=("--exit-signal",), - action_type="store", - nargs=None, - const=None, - default="SIGTERM", + option_strings=("--exit-signal",), action_type="store", nargs=None, const=None, default="15" ), "exit_signal_handler": MegatronActionSpec( option_strings=("--exit-signal-handler",), @@ -6537,6 +7199,20 @@ const=True, default=False, ), + "exit_signal_handler_for_dataloader": MegatronActionSpec( + option_strings=("--exit-signal-handler-for-dataloader",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "exit_signal_handler_for_training": MegatronActionSpec( + option_strings=("--exit-signal-handler-for-training",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "exp_avg_dtype": MegatronActionSpec( option_strings=("--exp-avg-dtype",), action_type="store", @@ -6551,6 +7227,13 @@ const=None, default="fp32", ), + "experimental_attention_variant": MegatronActionSpec( + option_strings=("--experimental-attention-variant",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "expert_model_parallel_size": MegatronActionSpec( option_strings=("--expert-model-parallel-size",), action_type="store", @@ -6579,6 +7262,69 @@ const=True, default=False, ), + "fault_injector_delay_start_iteration": MegatronActionSpec( + option_strings=("--fault-injector-delay-start-iteration",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_fault_delay": MegatronActionSpec( + option_strings=("--fault-injector-fault-delay",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_fault_probabilities": MegatronActionSpec( + option_strings=("--fault-injector-fault-probabilities",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_fault_types": MegatronActionSpec( + option_strings=("--fault-injector-fault-types",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_mtti_seconds": MegatronActionSpec( + option_strings=("--fault-injector-mtti-seconds",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_num_ranks": MegatronActionSpec( + option_strings=("--fault-injector-num-ranks",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_offset_seconds": MegatronActionSpec( + option_strings=("--fault-injector-offset-seconds",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_ranks": MegatronActionSpec( + option_strings=("--fault-injector-ranks",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "fault_injector_seed": MegatronActionSpec( + option_strings=("--fault-injector-seed",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "ffn_hidden_size": MegatronActionSpec( option_strings=("--ffn-hidden-size",), action_type="store", @@ -6587,11 +7333,7 @@ default=None, ), "fim_data": MegatronActionSpec( - option_strings=("--fim-data",), - action_type="store_true", - nargs=0, - const=True, - default=False, + option_strings=("--fim-data",), action_type="store_true", nargs=0, const=True, default=False ), "fim_eod_token": MegatronActionSpec( option_strings=("--fim-eod-token",), @@ -6636,11 +7378,7 @@ default="", ), "fim_rate": MegatronActionSpec( - option_strings=("--fim-rate",), - action_type="store", - nargs=None, - const=None, - default=0.5, + option_strings=("--fim-rate",), action_type="store", nargs=None, const=None, default=0.5 ), "fim_split_sample": MegatronActionSpec( option_strings=("--fim-split-sample",), @@ -6650,11 +7388,7 @@ default=None, ), "fim_spm_rate": MegatronActionSpec( - option_strings=("--fim-spm-rate",), - action_type="store", - nargs=None, - const=None, - default=0.5, + option_strings=("--fim-spm-rate",), action_type="store", nargs=None, const=None, default=0.5 ), "fim_suffix_token": MegatronActionSpec( option_strings=("--fim-suffix-token",), @@ -6663,13 +7397,23 @@ const=None, default="", ), - "finetune": MegatronActionSpec( - option_strings=("--finetune",), + "fine_grained_activation_offloading": MegatronActionSpec( + option_strings=("--fine-grained-activation-offloading",), action_type="store_true", nargs=0, const=True, default=False, ), + "fine_grained_offloading_max_inflight_offloads": MegatronActionSpec( + option_strings=("--fine-grained-offloading-max-inflight-offloads",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "finetune": MegatronActionSpec( + option_strings=("--finetune",), action_type="store_true", nargs=0, const=True, default=False + ), "first_last_layers_bf16": MegatronActionSpec( option_strings=("--first-last-layers-bf16",), action_type="store_true", @@ -6684,13 +7428,60 @@ const=True, default=False, ), - "fp16": MegatronActionSpec( - option_strings=("--fp16",), + "flight_recorder_dump_on_timeout": MegatronActionSpec( + option_strings=( + "--no-flight-recorder-dump-on-timeout", + "--disable-flight-recorder-dump-on-timeout", + ), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "flight_recorder_dump_path": MegatronActionSpec( + option_strings=("--flight-recorder-dump-path",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "flight_recorder_extra_dump_on_exec": MegatronActionSpec( + option_strings=( + "--no-flight-recorder-extra-dump-on-exec", + "--disable-flight-recorder-extra-dump-on-exec", + ), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "flight_recorder_include_only_active": MegatronActionSpec( + option_strings=( + "--no-flight-recorder-include-only-active", + "--disable-flight-recorder-include-only-active", + ), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "flight_recorder_include_stack_trace": MegatronActionSpec( + option_strings=("--flight-recorder-include-stack-trace",), action_type="store_true", nargs=0, const=True, default=False, ), + "flight_recorder_trace_buffer_size": MegatronActionSpec( + option_strings=("--flight-recorder-trace-buffer-size",), + action_type="store", + nargs=None, + const=None, + default=2000, + ), + "fp16": MegatronActionSpec( + option_strings=("--fp16",), action_type="store_true", nargs=0, const=True, default=False + ), "fp16_lm_cross_entropy": MegatronActionSpec( option_strings=("--fp16-lm-cross-entropy",), action_type="store_true", @@ -6706,13 +7497,9 @@ default=False, ), "fp4": MegatronActionSpec( - option_strings=("--fp4-format",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--fp4-format",), action_type="store", nargs=None, const=None, default=None ), - "fp4_param": MegatronActionSpec( + "fp4_param_gather": MegatronActionSpec( option_strings=("--fp4-param-gather",), action_type="store_true", nargs=0, @@ -6734,11 +7521,7 @@ default="nvfp4", ), "fp8": MegatronActionSpec( - option_strings=("--fp8-format",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--fp8-format",), action_type="store", nargs=None, const=None, default=None ), "fp8_amax_compute_algo": MegatronActionSpec( option_strings=("--fp8-amax-compute-algo",), @@ -6755,18 +7538,17 @@ default=1, ), "fp8_interval": MegatronActionSpec( - option_strings=("--fp8-interval",), - action_type="store", - nargs=None, - const=None, - default=1, + option_strings=("--fp8-interval",), action_type="store", nargs=None, const=None, default=1 ), "fp8_margin": MegatronActionSpec( - option_strings=("--fp8-margin",), - action_type="store", - nargs=None, - const=None, - default=0, + option_strings=("--fp8-margin",), action_type="store", nargs=None, const=None, default=0 + ), + "fp8_output_proj": MegatronActionSpec( + option_strings=("--fp8-output-proj",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "fp8_param_gather": MegatronActionSpec( option_strings=("--fp8-param-gather",), @@ -6790,7 +7572,7 @@ default="delayed", ), "fp8_wgrad": MegatronActionSpec( - option_strings=("--no-fp8-wgrad",), + option_strings=("--no-fp8-wgrad", "--disable-fp8-wgrad"), action_type="store_false", nargs=0, const=False, @@ -6803,6 +7585,20 @@ const=True, default=False, ), + "fsdp_manual_registration": MegatronActionSpec( + option_strings=("--fsdp-manual-registration",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "ft_num_warmup_iters": MegatronActionSpec( + option_strings=("--ft-num-warmup-iters",), + action_type="store", + nargs=None, + const=None, + default=5, + ), "full_validation": MegatronActionSpec( option_strings=("--full-validation",), action_type="store_true", @@ -6810,6 +7606,13 @@ const=True, default=False, ), + "fused_residual_rmsnorm": MegatronActionSpec( + option_strings=("--fused-residual-rmsnorm",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "global_batch_size": MegatronActionSpec( option_strings=("--global-batch-size",), action_type="store", @@ -6824,6 +7627,13 @@ const=None, default=0.0, ), + "gpu_sniff_test_interval": MegatronActionSpec( + option_strings=("--gpu-sniff-test-interval",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "grad_reduce_in_bf16": MegatronActionSpec( option_strings=("--grad-reduce-in-bf16",), action_type="store_true", @@ -6866,20 +7676,6 @@ const=None, default=0.01, ), - "grpo_default_temperature": MegatronActionSpec( - option_strings=("--grpo-default-temperature",), - action_type="store", - nargs=None, - const=None, - default=1.0, - ), - "grpo_default_top_p": MegatronActionSpec( - option_strings=("--grpo-default-top-p",), - action_type="store", - nargs=None, - const=None, - default=0, - ), "grpo_entropy_term_weight": MegatronActionSpec( option_strings=("--grpo-entropy-term-weight",), action_type="store", @@ -6923,11 +7719,7 @@ default=32, ), "head_lr_mult": MegatronActionSpec( - option_strings=("--head-lr-mult",), - action_type="store", - nargs=None, - const=None, - default=1.0, + option_strings=("--head-lr-mult",), action_type="store", nargs=None, const=None, default=1.0 ), "heterogeneous_layers_config_encoded_json": MegatronActionSpec( option_strings=("--heterogeneous-layers-config-encoded-json",), @@ -6951,11 +7743,7 @@ default=0.1, ), "hidden_size": MegatronActionSpec( - option_strings=("--hidden-size",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--hidden-size",), action_type="store", nargs=None, const=None, default=None ), "hierarchical_context_parallel_sizes": MegatronActionSpec( option_strings=("--hierarchical-context-parallel-sizes",), @@ -6964,26 +7752,33 @@ const=None, default=None, ), + "high_priority_a2a_comm_stream": MegatronActionSpec( + option_strings=("--high-priority-a2a-comm-stream",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "high_priority_stream_groups": MegatronActionSpec( option_strings=("--high-priority-stream-groups",), action_type="store", - nargs="*", + nargs="+", const=None, default=[], ), - "hybrid_attention_ratio": MegatronActionSpec( - option_strings=("--hybrid-attention-ratio",), - action_type="store", - nargs=None, - const=None, - default=0.0, + "hybrid_context_parallel": MegatronActionSpec( + option_strings=("--hybrid-context-parallel",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), - "hybrid_mlp_ratio": MegatronActionSpec( - option_strings=("--hybrid-mlp-ratio",), + "hybrid_layer_pattern": MegatronActionSpec( + option_strings=("--hybrid-layer-pattern",), action_type="store", nargs=None, const=None, - default=0.0, + default=None, ), "hybrid_override_pattern": MegatronActionSpec( option_strings=("--hybrid-override-pattern",), @@ -6993,11 +7788,7 @@ default=None, ), "hysteresis": MegatronActionSpec( - option_strings=("--hysteresis",), - action_type="store", - nargs=None, - const=None, - default=2, + option_strings=("--hysteresis",), action_type="store", nargs=None, const=None, default=2 ), "ict_head_size": MegatronActionSpec( option_strings=("--ict-head-size",), @@ -7007,25 +7798,13 @@ default=None, ), "ict_load": MegatronActionSpec( - option_strings=("--ict-load",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--ict-load",), action_type="store", nargs=None, const=None, default=None ), "img_h": MegatronActionSpec( - option_strings=("--img-h",), - action_type="store", - nargs=None, - const=None, - default=224, + option_strings=("--img-h",), action_type="store", nargs=None, const=None, default=224 ), "img_w": MegatronActionSpec( - option_strings=("--img-w",), - action_type="store", - nargs=None, - const=None, - default=224, + option_strings=("--img-w",), action_type="store", nargs=None, const=None, default=224 ), "indexer_batch_size": MegatronActionSpec( option_strings=("--indexer-batch-size",), @@ -7053,7 +7832,35 @@ action_type="store", nargs=None, const=None, - default=12346, + default=None, + ), + "inference_cuda_graph_all_prefills": MegatronActionSpec( + option_strings=("--inference-cuda-graph-all-prefills",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "inference_cuda_graph_scope": MegatronActionSpec( + option_strings=("--inference-cuda-graph-scope",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "inference_disable_ep_consensus": MegatronActionSpec( + option_strings=("--inference-disable-ep-consensus", "--no-inference-disable-ep-consensus"), + action_type="store", + nargs=0, + const=None, + default=False, + ), + "inference_disable_triton_nvls_kernels": MegatronActionSpec( + option_strings=("--inference-disable-triton-nvls-kernels",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "inference_dynamic_batching": MegatronActionSpec( option_strings=("--inference-dynamic-batching",), @@ -7076,19 +7883,36 @@ const=None, default=40.0, ), - "inference_dynamic_batching_cuda_graph_max_tokens": MegatronActionSpec( - option_strings=("--inference-dynamic-batching-cuda-graph-max-tokens",), + "inference_dynamic_batching_cuda_graph_mixed_prefill_count": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-cuda-graph-mixed-prefill-count",), action_type="store", nargs=None, const=None, - default=16384, + default=16, ), - "inference_dynamic_batching_cuda_graph_mixed_prefill_count": MegatronActionSpec( - option_strings=("--inference-dynamic-batching-cuda-graph-mixed-prefill-count",), + "inference_dynamic_batching_enable_prefix_caching": MegatronActionSpec( + option_strings=( + "--inference-dynamic-batching-prefix-caching", + "--no-inference-dynamic-batching-prefix-caching", + ), + action_type="store", + nargs=0, + const=None, + default=False, + ), + "inference_dynamic_batching_mamba_memory_ratio": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-mamba-memory-ratio",), action_type="store", nargs=None, const=None, - default=16, + default=None, + ), + "inference_dynamic_batching_max_requests": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-max-requests",), + action_type="store", + nargs=None, + const=None, + default=None, ), "inference_dynamic_batching_max_tokens": MegatronActionSpec( option_strings=("--inference-dynamic-batching-max-tokens",), @@ -7104,6 +7928,55 @@ const=None, default=16, ), + "inference_dynamic_batching_paused_buffer_size_gb": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-paused-buffer-size-gb",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "inference_dynamic_batching_prefix_caching_coordinator_policy": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-prefix-caching-coordinator-policy",), + action_type="store", + nargs=None, + const=None, + default="first_prefix_block", + ), + "inference_dynamic_batching_prefix_caching_eviction_policy": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-prefix-caching-eviction-policy",), + action_type="store", + nargs=None, + const=None, + default="ref_zero", + ), + "inference_dynamic_batching_prefix_caching_mamba_gb": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-prefix-caching-mamba-gb",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "inference_dynamic_batching_prefix_caching_routing_alpha": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-prefix-caching-routing-alpha",), + action_type="store", + nargs=None, + const=None, + default=0.5, + ), + "inference_dynamic_batching_sampling_backend": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-sampling-backend",), + action_type="store", + nargs=None, + const=None, + default="torch", + ), + "inference_dynamic_batching_track_generated_token_events": MegatronActionSpec( + option_strings=("--inference-dynamic-batching-track-generated-token-events",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "inference_dynamic_batching_track_paused_request_events": MegatronActionSpec( option_strings=("--inference-dynamic-batching-track-paused-request-events",), action_type="store_true", @@ -7116,9 +7989,30 @@ action_type="store", nargs=None, const=None, - default=1, + default=0, + ), + "inference_fuse_tp_communication": MegatronActionSpec( + option_strings=("--inference-fuse-tp-communication",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "inference_grouped_gemm_backend": MegatronActionSpec( + option_strings=("--inference-grouped-gemm-backend",), + action_type="store", + nargs=None, + const=None, + default="vllm", + ), + "inference_logging_step_interval": MegatronActionSpec( + option_strings=("--inference-logging-step-interval",), + action_type="store", + nargs=None, + const=None, + default=0, ), - "inference_max_batch_size": MegatronActionSpec( + "inference_max_requests": MegatronActionSpec( option_strings=("--inference-max-requests",), action_type="store", nargs=None, @@ -7132,6 +8026,20 @@ const=None, default=2560, ), + "inference_moe_disable_fused_quant_kernels": MegatronActionSpec( + option_strings=("--inference-moe-disable-fused-quant-kernels",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "inference_moe_token_dispatcher_type": MegatronActionSpec( + option_strings=("--inference-moe-token-dispatcher-type",), + action_type="store", + nargs=None, + const=None, + default="nvls", + ), "inference_rng_tracker": MegatronActionSpec( option_strings=("--inference-rng-tracker",), action_type="store_true", @@ -7139,12 +8047,32 @@ const=True, default=False, ), - "inference_wandb_logging_step_interval": MegatronActionSpec( - option_strings=("--inference-wandb-logging-step-interval",), + "inference_text_gen_server_logging": MegatronActionSpec( + option_strings=( + "--inference-text-gen-server-logging", + "--no-inference-text-gen-server-logging", + ), action_type="store", - nargs=None, + nargs=0, const=None, - default=0, + default=False, + ), + "inference_use_synchronous_zmq_collectives": MegatronActionSpec( + option_strings=( + "--inference-use-synchronous-zmq-collectives", + "--no-inference-use-synchronous-zmq-collectives", + ), + action_type="store", + nargs=0, + const=None, + default=False, + ), + "inference_wandb_logging": MegatronActionSpec( + option_strings=("--inference-wandb-logging", "--no-inference-wandb-logging"), + action_type="store", + nargs=0, + const=None, + default=False, ), "init_method_std": MegatronActionSpec( option_strings=("--init-method-std",), @@ -7314,96 +8242,114 @@ const=True, default=False, ), - "kitchen_config_file": MegatronActionSpec( - option_strings=("--kitchen-config-file",), + "kitchen_attention_backend": MegatronActionSpec( + option_strings=("--kitchen-attention-backend",), action_type="store", nargs=None, const=None, - default=None, + default="sdpa", + ), + "kv_channels": MegatronActionSpec( + option_strings=("--kv-channels",), action_type="store", nargs=None, const=None, default=None + ), + "kv_lora_rank": MegatronActionSpec( + option_strings=("--kv-lora-rank",), action_type="store", nargs=None, const=None, default=32 ), - "kitchen_recipe_number": MegatronActionSpec( - option_strings=("--kitchen-recipe-number",), + "langrl_env_config": MegatronActionSpec( + option_strings=("--langrl-env-config",), action_type="store", nargs=None, const=None, default=None, ), - "kv_channels": MegatronActionSpec( - option_strings=("--kv-channels",), + "layernorm_epsilon": MegatronActionSpec( + option_strings=("--norm-epsilon",), action_type="store", nargs=None, const=None, - default=None, + default=1e-05, ), - "kv_lora_rank": MegatronActionSpec( - option_strings=("--kv-lora-rank",), + "layernorm_zero_centered_gamma": MegatronActionSpec( + option_strings=("--apply-layernorm-1p",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "lazy_mpu_init": MegatronActionSpec( + option_strings=("--lazy-mpu-init",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "linear_attention_freq": MegatronActionSpec( + option_strings=("--linear-attention-freq",), action_type="store", nargs=None, const=None, - default=32, + default=None, ), - "langrl_env_config": MegatronActionSpec( - option_strings=("--langrl-env-config",), + "linear_conv_kernel_dim": MegatronActionSpec( + option_strings=("--linear-conv-kernel-dim",), action_type="store", nargs=None, const=None, - default=None, + default=4, ), - "langrl_external_server": MegatronActionSpec( - option_strings=("--langrl-external-server", "--no-langrl-external-server"), + "linear_key_head_dim": MegatronActionSpec( + option_strings=("--linear-key-head-dim",), action_type="store", - nargs=0, + nargs=None, const=None, - default=False, + default=128, ), - "langrl_inference_server_conversation_template": MegatronActionSpec( - option_strings=("--langrl-inference-server-conversation-template",), + "linear_num_key_heads": MegatronActionSpec( + option_strings=("--linear-num-key-heads",), action_type="store", nargs=None, const=None, - default=None, + default=16, ), - "langrl_inference_server_type": MegatronActionSpec( - option_strings=("--langrl-inference-server-type",), + "linear_num_value_heads": MegatronActionSpec( + option_strings=("--linear-num-value-heads",), action_type="store", nargs=None, const=None, - default="inplace_megatron", + default=32, ), - "lazy_mpu_init": MegatronActionSpec( - option_strings=("--lazy-mpu-init",), + "linear_value_head_dim": MegatronActionSpec( + option_strings=("--linear-value-head-dim",), action_type="store", nargs=None, const=None, - default=None, + default=128, ), - "legacy_tokenizer": MegatronActionSpec( - option_strings=("--legacy-tokenizer",), - action_type="store_true", - nargs=0, - const=True, - default=False, + "lion_beta1": MegatronActionSpec( + option_strings=("--lion-beta1",), action_type="store", nargs=None, const=None, default=0.95 + ), + "lion_beta2": MegatronActionSpec( + option_strings=("--lion-beta2",), action_type="store", nargs=None, const=None, default=0.98 ), "load": MegatronActionSpec( - option_strings=("--load",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--load",), action_type="store", nargs=None, const=None, default=None ), "load_main_params_from_ckpt": MegatronActionSpec( option_strings=("--load-main-params-from-ckpt",), action_type="store_true", nargs=0, const=True, - default=None, + default=False, ), "local_rank": MegatronActionSpec( - option_strings=("--local-rank",), - action_type="store", - nargs=None, - const=None, - default=0, + option_strings=("--local-rank",), action_type="store", nargs=None, const=None, default=0 + ), + "log_device_memory_used": MegatronActionSpec( + option_strings=("--log-device-memory-used",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "log_energy": MegatronActionSpec( option_strings=("--log-energy",), @@ -7413,19 +8359,32 @@ default=False, ), "log_interval": MegatronActionSpec( - option_strings=("--log-interval",), - action_type="store", - nargs=None, - const=None, - default=100, + option_strings=("--log-interval",), action_type="store", nargs=None, const=None, default=100 ), "log_loss_scale_to_tensorboard": MegatronActionSpec( - option_strings=("--no-log-loss-scale-to-tensorboard",), + option_strings=( + "--no-log-loss-scale-to-tensorboard", + "--disable-log-loss-scale-to-tensorboard", + ), action_type="store_false", nargs=0, const=False, default=True, ), + "log_max_attention_logit": MegatronActionSpec( + option_strings=("--log-max-attention-logit",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "log_memory_interval": MegatronActionSpec( + option_strings=("--log-memory-interval",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "log_memory_to_tensorboard": MegatronActionSpec( option_strings=("--log-memory-to-tensorboard",), action_type="store_true", @@ -7497,11 +8456,7 @@ default=None, ), "loss_scale": MegatronActionSpec( - option_strings=("--loss-scale",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--loss-scale",), action_type="store", nargs=None, const=None, default=None ), "loss_scale_window": MegatronActionSpec( option_strings=("--loss-scale-window",), @@ -7511,11 +8466,7 @@ default=1000, ), "lr": MegatronActionSpec( - option_strings=("--lr",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--lr",), action_type="store", nargs=None, const=None, default=None ), "lr_decay_iters": MegatronActionSpec( option_strings=("--lr-decay-iters",), @@ -7615,6 +8566,20 @@ const=None, default=64, ), + "mamba_inference_conv_states_dtype": MegatronActionSpec( + option_strings=("--mamba-inference-conv-states-dtype",), + action_type="store", + nargs=None, + const=None, + default="bf16", + ), + "mamba_inference_ssm_states_dtype": MegatronActionSpec( + option_strings=("--mamba-inference-ssm-states-dtype",), + action_type="store", + nargs=None, + const=None, + default="bf16", + ), "mamba_num_groups": MegatronActionSpec( option_strings=("--mamba-num-groups",), action_type="store", @@ -7644,7 +8609,7 @@ default=False, ), "manual_gc_eval": MegatronActionSpec( - option_strings=("--no-manual-gc-eval",), + option_strings=("--no-manual-gc-eval", "--disable-manual-gc-eval"), action_type="store_false", nargs=0, const=False, @@ -7658,18 +8623,10 @@ default=0, ), "mask_factor": MegatronActionSpec( - option_strings=("--mask-factor",), - action_type="store", - nargs=None, - const=None, - default=1.0, + option_strings=("--mask-factor",), action_type="store", nargs=None, const=None, default=1.0 ), "mask_prob": MegatronActionSpec( - option_strings=("--mask-prob",), - action_type="store", - nargs=None, - const=None, - default=0.15, + option_strings=("--mask-prob",), action_type="store", nargs=None, const=None, default=0.15 ), "mask_type": MegatronActionSpec( option_strings=("--mask-type",), @@ -7692,6 +8649,13 @@ const=None, default=None, ), + "max_seqlen_per_dp_cp_rank": MegatronActionSpec( + option_strings=("--max-seqlen-per-dp-cp-rank",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "max_tokens_to_oom": MegatronActionSpec( option_strings=("--max-tokens-to-oom",), action_type="store", @@ -7699,6 +8663,34 @@ const=None, default=12000, ), + "megatron_fsdp_enable_fine_grained_param_gather": MegatronActionSpec( + option_strings=("--megatron-fsdp-enable-fine-grained-param-gather",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "megatron_fsdp_grad_comm_dtype": MegatronActionSpec( + option_strings=("--megatron-fsdp-grad-comm-dtype",), + action_type="store", + nargs=None, + const=None, + default="auto", + ), + "megatron_fsdp_main_grads_dtype": MegatronActionSpec( + option_strings=("--megatron-fsdp-main-grads-dtype",), + action_type="store", + nargs=None, + const=None, + default="auto", + ), + "megatron_fsdp_main_params_dtype": MegatronActionSpec( + option_strings=("--megatron-fsdp-main-params-dtype",), + action_type="store", + nargs=None, + const=None, + default="fp32", + ), "memory_snapshot_path": MegatronActionSpec( option_strings=("--memory-snapshot-path",), action_type="store", @@ -7707,7 +8699,10 @@ default="snapshot.pickle", ), "merge_file": MegatronActionSpec( - option_strings=("--merge-file",), + option_strings=("--merge-file",), action_type="store", nargs=None, const=None, default=None + ), + "metadata_path": MegatronActionSpec( + option_strings=("--tokenizer-metadata",), action_type="store", nargs=None, const=None, @@ -7742,11 +8737,21 @@ default=1.0, ), "min_lr": MegatronActionSpec( - option_strings=("--min-lr",), + option_strings=("--min-lr",), action_type="store", nargs=None, const=None, default=0.0 + ), + "min_offloaded_tensor_size": MegatronActionSpec( + option_strings=("--min-offloaded-tensor-size",), action_type="store", nargs=None, const=None, - default=0.0, + default=1048576, + ), + "mla_down_proj_fusion": MegatronActionSpec( + option_strings=("--mla-down-proj-fusion",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "mlp_chunks_for_prefill": MegatronActionSpec( option_strings=("--mlp-chunks-for-prefill",), @@ -7755,6 +8760,13 @@ const=None, default=1, ), + "mlp_chunks_for_training": MegatronActionSpec( + option_strings=("--mlp-chunks-for-training",), + action_type="store", + nargs=None, + const=None, + default=1, + ), "mmap_bin_files": MegatronActionSpec( option_strings=("--no-mmap-bin-files",), action_type="store_false", @@ -7804,6 +8816,13 @@ const=True, default=False, ), + "moe_enable_routing_replay": MegatronActionSpec( + option_strings=("--moe-enable-routing-replay",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "moe_expert_capacity_factor": MegatronActionSpec( option_strings=("--moe-expert-capacity-factor",), action_type="store", @@ -7811,12 +8830,12 @@ const=None, default=None, ), - "moe_extended_tp": MegatronActionSpec( - option_strings=("--moe-extended-tp",), - action_type="store_true", - nargs=0, - const=True, - default=False, + "moe_expert_rank_capacity_factor": MegatronActionSpec( + option_strings=("--moe-expert-rank-capacity-factor",), + action_type="store", + nargs=None, + const=None, + default=None, ), "moe_ffn_hidden_size": MegatronActionSpec( option_strings=("--moe-ffn-hidden-size",), @@ -7839,12 +8858,33 @@ const=True, default=False, ), + "moe_hybridep_num_blocks_permute": MegatronActionSpec( + option_strings=("--moe-hybridep-num-blocks-permute",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "moe_hybridep_num_blocks_unpermute": MegatronActionSpec( + option_strings=("--moe-hybridep-num-blocks-unpermute",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "moe_hybridep_num_sms": MegatronActionSpec( option_strings=("--moe-hybridep-num-sms",), action_type="store", nargs=None, const=None, - default=16, + default=None, + ), + "moe_hybridep_num_sms_preprocessing": MegatronActionSpec( + option_strings=("--moe-hybridep-num-sms-preprocessing",), + action_type="store", + nargs=None, + const=None, + default=108, ), "moe_input_jitter_eps": MegatronActionSpec( option_strings=("--moe-input-jitter-eps",), @@ -7853,12 +8893,15 @@ const=None, default=None, ), - "moe_layer_freq": MegatronActionSpec( - option_strings=("--moe-layer-freq",), + "moe_latent_size": MegatronActionSpec( + option_strings=("--moe-latent-size",), action_type="store", nargs=None, const=None, - default=1, + default=None, + ), + "moe_layer_freq": MegatronActionSpec( + option_strings=("--moe-layer-freq",), action_type="store", nargs=None, const=None, default=1 ), "moe_layer_recompute": MegatronActionSpec( option_strings=("--moe-layer-recompute",), @@ -7867,6 +8910,13 @@ const=True, default=False, ), + "moe_mlp_glu_interleave_size": MegatronActionSpec( + option_strings=("--moe-mlp-glu-interleave-size",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "moe_pad_expert_input_to_capacity": MegatronActionSpec( option_strings=("--moe-pad-expert-input-to-capacity",), action_type="store_true", @@ -7881,6 +8931,34 @@ const=True, default=False, ), + "moe_paged_stash": MegatronActionSpec( + option_strings=("--moe-paged-stash",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_paged_stash_buffer_size_factor_cpu": MegatronActionSpec( + option_strings=("--moe-paged-stash-buffer-size-factor-cpu",), + action_type="store", + nargs=None, + const=None, + default=0.0, + ), + "moe_paged_stash_buffer_size_factor_cuda": MegatronActionSpec( + option_strings=("--moe-paged-stash-buffer-size-factor-cuda",), + action_type="store", + nargs=None, + const=None, + default=1.1, + ), + "moe_paged_stash_page_size": MegatronActionSpec( + option_strings=("--moe-paged-stash-page-size",), + action_type="store", + nargs=None, + const=None, + default=64, + ), "moe_per_layer_logging": MegatronActionSpec( option_strings=("--moe-per-layer-logging",), action_type="store_true", @@ -7895,6 +8973,13 @@ const=True, default=False, ), + "moe_permute_fusion_into_hybridep": MegatronActionSpec( + option_strings=("--moe-permute-fusion-into-hybridep",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "moe_router_bias_update_rate": MegatronActionSpec( option_strings=("--moe-router-bias-update-rate",), action_type="store", @@ -7916,6 +9001,13 @@ const=True, default=False, ), + "moe_router_force_biased": MegatronActionSpec( + option_strings=("--moe-router-force-biased",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "moe_router_force_load_balancing": MegatronActionSpec( option_strings=("--moe-router-force-load-balancing",), action_type="store_true", @@ -7951,152 +9043,295 @@ const=None, default=None, ), - "moe_router_padding_for_fp8": MegatronActionSpec( - option_strings=("--moe-router-padding-for-fp8",), + "moe_router_padding_for_fp8": MegatronActionSpec( + option_strings=("--moe-router-padding-for-fp8",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_router_padding_for_quantization": MegatronActionSpec( + option_strings=("--moe-router-padding-for-quantization",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_router_pre_softmax": MegatronActionSpec( + option_strings=("--moe-router-pre-softmax",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_router_score_function": MegatronActionSpec( + option_strings=("--moe-router-score-function",), + action_type="store", + nargs=None, + const=None, + default="softmax", + ), + "moe_router_topk": MegatronActionSpec( + option_strings=("--moe-router-topk",), + action_type="store", + nargs=None, + const=None, + default=2, + ), + "moe_router_topk_scaling_factor": MegatronActionSpec( + option_strings=("--moe-router-topk-scaling-factor",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "moe_shared_expert_gate": MegatronActionSpec( + option_strings=("--moe-shared-expert-gate",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_shared_expert_intermediate_size": MegatronActionSpec( + option_strings=("--moe-shared-expert-intermediate-size",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "moe_shared_expert_overlap": MegatronActionSpec( + option_strings=("--moe-shared-expert-overlap",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_single_grouped_bias": MegatronActionSpec( + option_strings=("--moe-single-grouped-bias",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_single_grouped_weight": MegatronActionSpec( + option_strings=("--moe-single-grouped-weight",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_token_dispatcher_type": MegatronActionSpec( + option_strings=("--moe-token-dispatcher-type",), + action_type="store", + nargs=None, + const=None, + default="allgather", + ), + "moe_token_drop_policy": MegatronActionSpec( + option_strings=("--moe-token-drop-policy",), + action_type="store", + nargs=None, + const=None, + default="probs", + ), + "moe_upcycling_granularity": MegatronActionSpec( + option_strings=("--moe-upcycling-granularity",), + action_type="store", + nargs=None, + const=None, + default=1, + ), + "moe_use_upcycling": MegatronActionSpec( + option_strings=("--moe-use-upcycling",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "moe_z_loss_coeff": MegatronActionSpec( + option_strings=("--moe-z-loss-coeff",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "mrope_section": MegatronActionSpec( + option_strings=("--mrope-section",), + action_type="store", + nargs="+", + const=None, + default=None, + ), + "mscale": MegatronActionSpec( + option_strings=("--mscale",), action_type="store", nargs=None, const=None, default=1.0 + ), + "mscale_all_dim": MegatronActionSpec( + option_strings=("--mscale-all-dim",), + action_type="store", + nargs=None, + const=None, + default=0.0, + ), + "mtp_hybrid_override_pattern": MegatronActionSpec( + option_strings=("--mtp-hybrid-override-pattern",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "mtp_loss_scaling_factor": MegatronActionSpec( + option_strings=("--mtp-loss-scaling-factor",), + action_type="store", + nargs=None, + const=None, + default=0.1, + ), + "mtp_num_layers": MegatronActionSpec( + option_strings=("--mtp-num-layers",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "mtp_standalone": MegatronActionSpec( + option_strings=("--mtp-standalone",), action_type="store_true", nargs=0, const=True, default=False, ), - "moe_router_padding_for_quantization": MegatronActionSpec( - option_strings=("--moe-router-padding-for-quantization",), + "mtp_use_repeated_layer": MegatronActionSpec( + option_strings=("--mtp-use-repeated-layer",), action_type="store_true", nargs=0, const=True, default=False, ), - "moe_router_pre_softmax": MegatronActionSpec( - option_strings=("--moe-router-pre-softmax",), + "multi_latent_attention": MegatronActionSpec( + option_strings=("--multi-latent-attention",), action_type="store_true", nargs=0, const=True, default=False, ), - "moe_router_score_function": MegatronActionSpec( - option_strings=("--moe-router-score-function",), + "multiple_validation_sets": MegatronActionSpec( + option_strings=("--multiple-validation-sets",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "muon_coefficient_type": MegatronActionSpec( + option_strings=("--muon-coefficient-type",), action_type="store", nargs=None, const=None, - default="softmax", + default="quintic", ), - "moe_router_topk": MegatronActionSpec( - option_strings=("--moe-router-topk",), + "muon_extra_scale_factor": MegatronActionSpec( + option_strings=("--muon-extra-scale-factor",), action_type="store", nargs=None, const=None, - default=2, + default=1.0, ), - "moe_router_topk_scaling_factor": MegatronActionSpec( - option_strings=("--moe-router-topk-scaling-factor",), + "muon_fp32_matmul_prec": MegatronActionSpec( + option_strings=("--muon-fp32-matmul-prec",), action_type="store", nargs=None, const=None, - default=None, + default="medium", ), - "moe_shared_expert_intermediate_size": MegatronActionSpec( - option_strings=("--moe-shared-expert-intermediate-size",), + "muon_momentum": MegatronActionSpec( + option_strings=("--muon-momentum",), action_type="store", nargs=None, const=None, - default=None, + default=0.9, ), - "moe_shared_expert_overlap": MegatronActionSpec( - option_strings=("--moe-shared-expert-overlap",), + "muon_nesterov": MegatronActionSpec( + option_strings=("--muon-nesterov",), action_type="store_true", nargs=0, const=True, default=False, ), - "moe_token_dispatcher_type": MegatronActionSpec( - option_strings=("--moe-token-dispatcher-type",), + "muon_num_ns_steps": MegatronActionSpec( + option_strings=("--muon-num-ns-steps",), action_type="store", nargs=None, const=None, - default="allgather", + default=5, ), - "moe_token_drop_policy": MegatronActionSpec( - option_strings=("--moe-token-drop-policy",), + "muon_scalar_optimizer": MegatronActionSpec( + option_strings=("--muon-scalar-optimizer",), action_type="store", nargs=None, const=None, - default="probs", + default="adam", ), - "moe_upcycling_granularity": MegatronActionSpec( - option_strings=("--moe-upcycling-granularity",), + "muon_scale_mode": MegatronActionSpec( + option_strings=("--muon-scale-mode",), action_type="store", nargs=None, const=None, - default=1, + default="spectral", ), - "moe_use_legacy_grouped_gemm": MegatronActionSpec( - option_strings=("--moe-use-legacy-grouped-gemm",), - action_type="store_true", + "muon_split_qkv": MegatronActionSpec( + option_strings=("--muon-no-split-qkv",), + action_type="store_false", nargs=0, - const=True, - default=False, + const=False, + default=True, ), - "moe_use_upcycling": MegatronActionSpec( - option_strings=("--moe-use-upcycling",), - action_type="store_true", - nargs=0, - const=True, - default=False, + "muon_tp_mode": MegatronActionSpec( + option_strings=("--muon-tp-mode",), + action_type="store", + nargs=None, + const=None, + default="blockwise", ), - "moe_z_loss_coeff": MegatronActionSpec( - option_strings=("--moe-z-loss-coeff",), + "mup_attn_scale_power": MegatronActionSpec( + option_strings=("--mup-attn-scale-power",), action_type="store", nargs=None, const=None, - default=None, + default=1.0, ), - "mrope_section": MegatronActionSpec( - option_strings=("--mrope-section",), + "mup_base_head_dim": MegatronActionSpec( + option_strings=("--mup-base-head-dim",), action_type="store", - nargs="+", + nargs=None, const=None, default=None, ), - "mscale": MegatronActionSpec( - option_strings=("--mscale",), + "mup_base_hidden_size": MegatronActionSpec( + option_strings=("--mup-base-hidden-size",), action_type="store", nargs=None, const=None, - default=1.0, + default=None, ), - "mscale_all_dim": MegatronActionSpec( - option_strings=("--mscale-all-dim",), + "mup_embedding_mult": MegatronActionSpec( + option_strings=("--mup-embedding-mult",), action_type="store", nargs=None, const=None, - default=0.0, + default=1.0, ), - "mtp_loss_scaling_factor": MegatronActionSpec( - option_strings=("--mtp-loss-scaling-factor",), + "mup_output_mult": MegatronActionSpec( + option_strings=("--mup-output-mult",), action_type="store", nargs=None, const=None, - default=0.1, + default=1.0, ), - "mtp_num_layers": MegatronActionSpec( - option_strings=("--mtp-num-layers",), + "mup_width_mult": MegatronActionSpec( + option_strings=("--mup-width-mult",), action_type="store", nargs=None, const=None, - default=None, - ), - "multi_latent_attention": MegatronActionSpec( - option_strings=("--multi-latent-attention",), - action_type="store_true", - nargs=0, - const=True, - default=False, - ), - "multiple_validation_sets": MegatronActionSpec( - option_strings=("--multiple-validation-sets",), - action_type="store_true", - nargs=0, - const=True, - default=False, + default=1.0, ), "nccl_all_reduce_for_prefill": MegatronActionSpec( option_strings=("--nccl-all-reduce-for-prefill",), @@ -8161,6 +9396,13 @@ const=True, default=None, ), + "no_weight_decay_cond_type": MegatronActionSpec( + option_strings=("--no-weight-decay-cond-type",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "non_persistent_ckpt_type": MegatronActionSpec( option_strings=("--non-persistent-ckpt-type",), action_type="store", @@ -8196,13 +9438,6 @@ const=None, default=None, ), - "norm_epsilon": MegatronActionSpec( - option_strings=("--norm-epsilon",), - action_type="store", - nargs=None, - const=None, - default=1e-05, - ), "normalization": MegatronActionSpec( option_strings=("--normalization",), action_type="store", @@ -8210,26 +9445,32 @@ const=None, default="LayerNorm", ), - "num_attention_heads": MegatronActionSpec( - option_strings=("--num-attention-heads",), + "null_tokenizer_eod_id": MegatronActionSpec( + option_strings=("--null-tokenizer-eod-id",), action_type="store", nargs=None, const=None, default=None, ), - "num_channels": MegatronActionSpec( - option_strings=("--num-channels",), + "null_tokenizer_pad_id": MegatronActionSpec( + option_strings=("--null-tokenizer-pad-id",), action_type="store", nargs=None, const=None, - default=3, + default=-1, ), - "num_classes": MegatronActionSpec( - option_strings=("--num-classes",), + "num_attention_heads": MegatronActionSpec( + option_strings=("--num-attention-heads",), action_type="store", nargs=None, const=None, - default=1000, + default=None, + ), + "num_channels": MegatronActionSpec( + option_strings=("--num-channels",), action_type="store", nargs=None, const=None, default=3 + ), + "num_classes": MegatronActionSpec( + option_strings=("--num-classes",), action_type="store", nargs=None, const=None, default=1000 ), "num_dataset_builder_threads": MegatronActionSpec( option_strings=("--num-dataset-builder-threads",), @@ -8246,18 +9487,10 @@ default=1, ), "num_experts": MegatronActionSpec( - option_strings=("--num-experts",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--num-experts",), action_type="store", nargs=None, const=None, default=None ), "num_layers": MegatronActionSpec( - option_strings=("--num-layers",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--num-layers",), action_type="store", nargs=None, const=None, default=None ), "num_layers_at_end_in_bf16": MegatronActionSpec( option_strings=("--num-layers-at-end-in-bf16",), @@ -8287,6 +9520,13 @@ const=None, default=1, ), + "num_speculative_tokens": MegatronActionSpec( + option_strings=("--num-speculative-tokens",), + action_type="store", + nargs=None, + const=None, + default=0, + ), "num_virtual_stages_per_pipeline_rank": MegatronActionSpec( option_strings=("--num-virtual-stages-per-pipeline-rank",), action_type="store", @@ -8295,11 +9535,14 @@ default=None, ), "num_workers": MegatronActionSpec( - option_strings=("--num-workers",), - action_type="store", - nargs=None, - const=None, - default=2, + option_strings=("--num-workers",), action_type="store", nargs=None, const=None, default=2 + ), + "nvtx_ranges": MegatronActionSpec( + option_strings=("--nvtx-ranges",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "object_storage_cache_path": MegatronActionSpec( option_strings=("--object-storage-cache-path",), @@ -8308,6 +9551,13 @@ const=None, default=None, ), + "offload_modules": MegatronActionSpec( + option_strings=("--offload-modules",), + action_type="store", + nargs="+", + const=None, + default=[], + ), "one_logger_async": MegatronActionSpec( option_strings=("--one-logger-async",), action_type="store_true", @@ -8330,11 +9580,7 @@ default=None, ), "onnx_safe": MegatronActionSpec( - option_strings=("--onnx-safe",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--onnx-safe",), action_type="store", nargs=None, const=None, default=None ), "openai_gelu": MegatronActionSpec( option_strings=("--openai-gelu",), @@ -8344,11 +9590,7 @@ default=False, ), "optimizer": MegatronActionSpec( - option_strings=("--optimizer",), - action_type="store", - nargs=None, - const=None, - default="adam", + option_strings=("--optimizer",), action_type="store", nargs=None, const=None, default="adam" ), "optimizer_cpu_offload": MegatronActionSpec( option_strings=("--optimizer-cpu-offload",), @@ -8357,6 +9599,13 @@ const=True, default=False, ), + "optimizer_cuda_graph": MegatronActionSpec( + option_strings=("--optimizer-cuda-graph",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "optimizer_offload_fraction": MegatronActionSpec( option_strings=("--optimizer-offload-fraction",), action_type="store", @@ -8364,6 +9613,13 @@ const=None, default=1.0, ), + "outer_dp_sharding_strategy": MegatronActionSpec( + option_strings=("--outer-dp-sharding-strategy",), + action_type="store", + nargs=None, + const=None, + default="no_shard", + ), "output_bert_embeddings": MegatronActionSpec( option_strings=("--output-bert-embeddings",), action_type="store_true", @@ -8378,6 +9634,13 @@ const=True, default=False, ), + "overlap_dispatch_backward_with_experts_wgrad": MegatronActionSpec( + option_strings=("--overlap-dispatch-backward-with-experts-wgrad",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "overlap_grad_reduce": MegatronActionSpec( option_strings=("--overlap-grad-reduce",), action_type="store_true", @@ -8421,15 +9684,19 @@ default=False, ), "override_opt_param_scheduler": MegatronActionSpec( - option_strings=( - "--override-opt_param-scheduler", - "--override-opt-param-scheduler", - ), + option_strings=("--override-opt_param-scheduler", "--override-opt-param-scheduler"), action_type="store_true", nargs=0, const=True, default=False, ), + "pad_vocab_size": MegatronActionSpec( + option_strings=("--no-pad-vocab-size", "--disable-pad-vocab-size"), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), "padded_vocab_size": MegatronActionSpec( option_strings=("--padded-vocab-size",), action_type="store", @@ -8438,11 +9705,14 @@ default=None, ), "patch_dim": MegatronActionSpec( - option_strings=("--patch-dim",), + option_strings=("--patch-dim",), action_type="store", nargs=None, const=None, default=16 + ), + "per_dataset_sequences_path": MegatronActionSpec( + option_strings=("--per-dataset-sequences-path",), action_type="store", nargs=None, const=None, - default=16, + default=None, ), "per_split_data_args_path": MegatronActionSpec( option_strings=("--per-split-data-args-path",), @@ -8465,6 +9735,13 @@ const=True, default=False, ), + "phase_transition_iterations": MegatronActionSpec( + option_strings=("--phase-transition-iterations",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "pin_cpu_grads": MegatronActionSpec( option_strings=("--no-pin-cpu-grads",), action_type="store_false", @@ -8515,18 +9792,10 @@ default=None, ), "profile": MegatronActionSpec( - option_strings=("--profile",), - action_type="store_true", - nargs=0, - const=True, - default=False, + option_strings=("--profile",), action_type="store_true", nargs=0, const=True, default=False ), "profile_ranks": MegatronActionSpec( - option_strings=("--profile-ranks",), - action_type="store", - nargs="+", - const=None, - default=[0], + option_strings=("--profile-ranks",), action_type="store", nargs="+", const=None, default=[] ), "profile_step_end": MegatronActionSpec( option_strings=("--profile-step-end",), @@ -8542,19 +9811,49 @@ const=None, default=10, ), + "pytorch_profiler_collect_callstack": MegatronActionSpec( + option_strings=("--pytorch-profiler-collect-callstack",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "pytorch_profiler_collect_chakra": MegatronActionSpec( + option_strings=("--pytorch-profiler-collect-chakra",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "pytorch_profiler_collect_shapes": MegatronActionSpec( + option_strings=("--pytorch-profiler-collect-shapes",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "q_lora_rank": MegatronActionSpec( - option_strings=("--q-lora-rank",), + option_strings=("--q-lora-rank",), action_type="store", nargs=None, const=None, default=None + ), + "qk_clip": MegatronActionSpec( + option_strings=("--qk-clip",), action_type="store_true", nargs=0, const=True, default=False + ), + "qk_clip_alpha": MegatronActionSpec( + option_strings=("--qk-clip-alpha",), action_type="store", nargs=None, const=None, - default=None, + default=0.5, ), - "qk_head_dim": MegatronActionSpec( - option_strings=("--qk-head-dim",), + "qk_clip_threshold": MegatronActionSpec( + option_strings=("--qk-clip-threshold",), action_type="store", nargs=None, const=None, - default=128, + default=100, + ), + "qk_head_dim": MegatronActionSpec( + option_strings=("--qk-head-dim",), action_type="store", nargs=None, const=None, default=128 ), "qk_l2_norm": MegatronActionSpec( option_strings=("--qk-l2-norm",), @@ -8594,7 +9893,7 @@ "rampup_batch_size": MegatronActionSpec( option_strings=("--rampup-batch-size",), action_type="store", - nargs="*", + nargs=3, const=None, default=None, ), @@ -8622,7 +9921,7 @@ "recompute_modules": MegatronActionSpec( option_strings=("--recompute-modules",), action_type="store", - nargs="*", + nargs="+", const=None, default=None, ), @@ -8640,6 +9939,20 @@ const=True, default=False, ), + "record_shapes": MegatronActionSpec( + option_strings=("--record-shapes",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "refit_method": MegatronActionSpec( + option_strings=("--refit-method",), + action_type="store", + nargs=None, + const=None, + default="gloo", + ), "relative_attention_max_distance": MegatronActionSpec( option_strings=("--relative-attention-max-distance",), action_type="store", @@ -8717,121 +10030,125 @@ const=True, default=False, ), - "retriever_seq_length": MegatronActionSpec( - option_strings=("--retriever-seq-length",), - action_type="store", - nargs=None, - const=None, - default=256, - ), - "retro_add_retriever": MegatronActionSpec( - option_strings=("--retro-add-retriever",), + "reuse_grad_buf_for_mxfp8_param_ag": MegatronActionSpec( + option_strings=("--reuse-grad-buf-for-mxfp8-param-ag",), action_type="store_true", nargs=0, const=True, default=False, ), - "retro_attention_gate": MegatronActionSpec( - option_strings=("--retro-attention-gate",), + "rl_default_temperature": MegatronActionSpec( + option_strings=("--rl-default-temperature",), + action_type="store", + nargs=None, + const=None, + default=1.0, + ), + "rl_default_top_k": MegatronActionSpec( + option_strings=("--rl-default-top-k",), + action_type="store", + nargs=None, + const=None, + default=-1, + ), + "rl_default_top_p": MegatronActionSpec( + option_strings=("--rl-default-top-p",), + action_type="store", + nargs=None, + const=None, + default=0, + ), + "rl_generation_batch_size": MegatronActionSpec( + option_strings=("--rl-generation-batch-size",), + action_type="store", + nargs=None, + const=None, + default=None, + ), + "rl_importance_sampling_truncation_coef": MegatronActionSpec( + option_strings=("--rl-importance-sampling-truncation-coef",), action_type="store", nargs=None, const=None, - default=1, + default=None, ), - "retro_cyclic_train_iters": MegatronActionSpec( - option_strings=("--retro-cyclic-train-iters",), + "rl_inference_expert_model_parallel_size": MegatronActionSpec( + option_strings=("--rl-inference-expert-model-parallel-size",), action_type="store", nargs=None, const=None, default=None, ), - "retro_encoder_attention_dropout": MegatronActionSpec( - option_strings=("--retro-encoder-attention-dropout",), + "rl_inference_expert_tensor_model_parallel_size": MegatronActionSpec( + option_strings=("--rl-inference-expert-tensor-model-parallel-size",), action_type="store", nargs=None, const=None, - default=0.1, + default=None, ), - "retro_encoder_hidden_dropout": MegatronActionSpec( - option_strings=("--retro-encoder-hidden-dropout",), + "rl_inference_logprobs_is_correction": MegatronActionSpec( + option_strings=( + "--rl-inference-logprobs-is-correction", + "--no-rl-inference-logprobs-is-correction", + ), action_type="store", - nargs=None, + nargs=0, const=None, - default=0.1, + default=False, ), - "retro_encoder_layers": MegatronActionSpec( - option_strings=("--retro-encoder-layers",), + "rl_inference_model_unified_memory_level": MegatronActionSpec( + option_strings=("--rl-inference-model-unified-memory-level",), action_type="store", nargs=None, const=None, - default=2, + default=0, ), - "retro_num_neighbors": MegatronActionSpec( - option_strings=("--retro-num-neighbors",), + "rl_inference_parsers": MegatronActionSpec( + option_strings=("--rl-inference-parsers",), action_type="store", - nargs=None, + nargs="*", const=None, - default=2, + default=[], ), - "retro_num_retrieved_chunks": MegatronActionSpec( - option_strings=("--retro-num-retrieved-chunks",), + "rl_inference_pipeline_model_parallel_size": MegatronActionSpec( + option_strings=("--rl-inference-pipeline-model-parallel-size",), action_type="store", nargs=None, const=None, - default=2, + default=None, ), - "retro_project_dir": MegatronActionSpec( - option_strings=("--retro-project-dir",), + "rl_inference_tensor_model_parallel_size": MegatronActionSpec( + option_strings=("--rl-inference-tensor-model-parallel-size",), action_type="store", nargs=None, const=None, default=None, ), - "retro_verify_neighbor_count": MegatronActionSpec( - option_strings=("--retro-no-verify-neighbor-count",), - action_type="store_false", - nargs=0, - const=False, - default=True, - ), - "reuse_grad_buf_for_mxfp8_param_ag": MegatronActionSpec( - option_strings=("--reuse-grad-buf-for-mxfp8-param-ag",), - action_type="store_true", - nargs=0, - const=True, - default=False, - ), - "rl_calculate_intra_group_similarity": MegatronActionSpec( - option_strings=( - "--rl-calculate-intra-group-similarity", - "--no-rl-calculate-intra-group-similarity", - ), + "rl_kv_cache_management_mode": MegatronActionSpec( + option_strings=("--rl-kv-cache-management-mode",), action_type="store", - nargs=0, + nargs=None, const=None, - default=False, + default="persist", ), - "rl_importance_sampling_truncation_coef": MegatronActionSpec( - option_strings=("--rl-importance-sampling-truncation-coef",), + "rl_num_parallel_generation_batches": MegatronActionSpec( + option_strings=("--rl-num-parallel-generation-batches",), action_type="store", nargs=None, const=None, default=None, ), - "rl_inference_logprobs_is_correction": MegatronActionSpec( - option_strings=( - "--rl-inference-logprobs-is-correction", - "--no-rl-inference-logprobs-is-correction", - ), + "rl_num_parallel_generations": MegatronActionSpec( + option_strings=("--rl-num-parallel-generations",), action_type="store", - nargs=0, + nargs=None, const=None, - default=False, + default=None, ), - "rl_offload_kv_cache_during_training": MegatronActionSpec( + "rl_offload_inference_model_weights_when_idle": MegatronActionSpec( option_strings=( - "--rl-offload-kv-cache-during-training", - "--no-rl-offload-kv-cache-during-training", + "--rl-offload-inference-model-weights-when-idle", + "--no-rl-offload-inference-model-weights-when-idle", ), action_type="store", nargs=0, @@ -8845,6 +10162,13 @@ const=True, default=False, ), + "rl_parallel_generation_tasks": MegatronActionSpec( + option_strings=("--rl-parallel-generation-tasks",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "rl_partial_rollouts": MegatronActionSpec( option_strings=("--rl-partial-rollouts", "--no-rl-partial-rollouts"), action_type="store", @@ -8852,6 +10176,13 @@ const=None, default=False, ), + "rl_persist_cuda_graphs": MegatronActionSpec( + option_strings=("--rl-persist-cuda-graphs", "--no-rl-persist-cuda-graphs"), + action_type="store", + nargs=0, + const=None, + default=False, + ), "rl_prompts_per_eval": MegatronActionSpec( option_strings=("--rl-prompts-per-eval",), action_type="store", @@ -8859,42 +10190,46 @@ const=None, default=32, ), - "rl_remove_kv_cache_during_training": MegatronActionSpec( - option_strings=( - "--rl-remove-kv-cache-during-training", - "--no-rl-remove-kv-cache-during-training", - ), + "rl_sequence_packing_algo": MegatronActionSpec( + option_strings=("--rl-sequence-packing-algo",), action_type="store", - nargs=0, + nargs=None, const=None, - default=False, + default="fifo", + ), + "rl_sequence_packing_max_sequences_per_bin": MegatronActionSpec( + option_strings=("--rl-sequence-packing-max-sequences-per-bin",), + action_type="store", + nargs=None, + const=None, + default=50, ), - "rl_reset_cuda_graphs": MegatronActionSpec( - option_strings=("--rl-reset-cuda-graphs", "--no-rl-reset-cuda-graphs"), + "rl_skip_bos_token": MegatronActionSpec( + option_strings=("--rl-skip-bos-token", "--no-rl-skip-bos-token"), action_type="store", nargs=0, const=None, default=False, ), - "rl_sequence_packing_algo": MegatronActionSpec( - option_strings=("--rl-sequence-packing-algo",), + "rl_training_cuda_graphs": MegatronActionSpec( + option_strings=("--rl-training-cuda-graphs", "--no-rl-training-cuda-graphs"), action_type="store", - nargs=None, + nargs=0, const=None, - default="fifo", + default=False, ), - "rl_sequence_packing_bin_size": MegatronActionSpec( - option_strings=("--rl-sequence-packing-bin-size",), + "rl_use_sequence_packing": MegatronActionSpec( + option_strings=("--rl-use-sequence-packing", "--no-rl-use-sequence-packing"), action_type="store", - nargs=None, + nargs=0, const=None, - default=8192, + default=False, ), - "rl_use_sequence_packing": MegatronActionSpec( - option_strings=("--rl-use-sequence-packing",), - action_type="store_true", + "rl_verify_model_weights_swap": MegatronActionSpec( + option_strings=("--rl-verify-model-weights-swap", "--no-rl-verify-model-weights-swap"), + action_type="store", nargs=0, - const=True, + const=None, default=False, ), "rope_scaling_factor": MegatronActionSpec( @@ -8905,11 +10240,7 @@ default=8.0, ), "rope_type": MegatronActionSpec( - option_strings=("--rope-type",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--rope-type",), action_type="store", nargs=None, const=None, default=None ), "rotary_base": MegatronActionSpec( option_strings=("--rotary-base",), @@ -8954,14 +10285,20 @@ default=False, ), "sample_rate": MegatronActionSpec( - option_strings=("--sample-rate",), + option_strings=("--sample-rate",), action_type="store", nargs=None, const=None, default=1.0 + ), + "save": MegatronActionSpec( + option_strings=("--save",), action_type="store", nargs=None, const=None, default=None + ), + "save_activations_interval": MegatronActionSpec( + option_strings=("--save-activations-interval",), action_type="store", nargs=None, const=None, - default=1.0, + default=None, ), - "save": MegatronActionSpec( - option_strings=("--save",), + "save_dgrads_interval": MegatronActionSpec( + option_strings=("--save-dgrads-interval",), action_type="store", nargs=None, const=None, @@ -8981,6 +10318,13 @@ const=None, default=None, ), + "save_params_interval": MegatronActionSpec( + option_strings=("--save-params-interval",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "save_retain_interval": MegatronActionSpec( option_strings=("--save-retain-interval",), action_type="store", @@ -8988,27 +10332,26 @@ const=None, default=None, ), - "scatter_gather_tensors_in_pipeline": MegatronActionSpec( - option_strings=("--no-scatter-gather-tensors-in-pipeline",), - action_type="store_false", - nargs=0, - const=False, - default=True, - ), - "seed": MegatronActionSpec( - option_strings=("--seed",), + "save_tokens_per_expert_interval": MegatronActionSpec( + option_strings=("--save-tokens-per-expert-interval",), action_type="store", nargs=None, const=None, - default=1234, + default=None, ), - "seq_length": MegatronActionSpec( - option_strings=("--seq-length",), + "save_wgrads_interval": MegatronActionSpec( + option_strings=("--save-wgrads-interval",), action_type="store", nargs=None, const=None, default=None, ), + "seed": MegatronActionSpec( + option_strings=("--seed",), action_type="store", nargs=None, const=None, default=1234 + ), + "seq_length": MegatronActionSpec( + option_strings=("--seq-length",), action_type="store", nargs=None, const=None, default=None + ), "sequence_parallel": MegatronActionSpec( option_strings=("--sequence-parallel",), action_type="store_true", @@ -9017,11 +10360,7 @@ default=False, ), "sft": MegatronActionSpec( - option_strings=("--sft",), - action_type="store_true", - nargs=0, - const=True, - default=False, + option_strings=("--sft",), action_type="store_true", nargs=0, const=True, default=False ), "sft_tokenizer_prompt_format": MegatronActionSpec( option_strings=("--sft-tokenizer-prompt-format",), @@ -9031,11 +10370,7 @@ default="nemotron-h-aligned", ), "sgd_momentum": MegatronActionSpec( - option_strings=("--sgd-momentum",), - action_type="store", - nargs=None, - const=None, - default=0.9, + option_strings=("--sgd-momentum",), action_type="store", nargs=None, const=None, default=0.9 ), "sharp_enabled_group": MegatronActionSpec( option_strings=("--sharp-enabled-group",), @@ -9066,18 +10401,17 @@ default="vanilla", ), "spec": MegatronActionSpec( - option_strings=("--spec",), + option_strings=("--spec",), action_type="store", nargs="*", const=None, default=None + ), + "special_tokens": MegatronActionSpec( + option_strings=("--tokenizer-special-tokens",), action_type="store", - nargs="*", + nargs="+", const=None, default=None, ), "split": MegatronActionSpec( - option_strings=("--split",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--split",), action_type="store", nargs=None, const=None, default=None ), "squared_relu": MegatronActionSpec( option_strings=("--squared-relu",), @@ -9086,6 +10420,13 @@ const=True, default=False, ), + "start_eval_at_iter": MegatronActionSpec( + option_strings=("--start-eval-at-iter",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "start_weight_decay": MegatronActionSpec( option_strings=("--start-weight-decay",), action_type="store", @@ -9093,6 +10434,13 @@ const=None, default=None, ), + "step_batch_size_schedule": MegatronActionSpec( + option_strings=("--step-batch-size-schedule",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "straggler_ctrlr_port": MegatronActionSpec( option_strings=("--straggler-ctrlr-port",), action_type="store", @@ -9108,7 +10456,7 @@ default=1, ), "strict_fsdp_dtensor_load": MegatronActionSpec( - option_strings=("--no-strict-fsdp-dtensor-load",), + option_strings=("--no-strict-fsdp-dtensor-load", "--disable-strict-fsdp-dtensor-load"), action_type="store_false", nargs=0, const=False, @@ -9122,11 +10470,7 @@ default=None, ), "swiglu": MegatronActionSpec( - option_strings=("--swiglu",), - action_type="store_true", - nargs=0, - const=True, - default=False, + option_strings=("--swiglu",), action_type="store_true", nargs=0, const=True, default=False ), "swin_backbone_type": MegatronActionSpec( option_strings=("--swin-backbone-type",), @@ -9142,6 +10486,13 @@ const=None, default=None, ), + "te_precision_config_file": MegatronActionSpec( + option_strings=("--te-precision-config-file",), + action_type="store", + nargs=None, + const=None, + default=None, + ), "te_rng_tracker": MegatronActionSpec( option_strings=("--te-rng-tracker",), action_type="store_true", @@ -9205,13 +10556,6 @@ const=None, default=None, ), - "tiktoken_special_tokens": MegatronActionSpec( - option_strings=("--tiktoken-special-tokens",), - action_type="store", - nargs="+", - const=None, - default=None, - ), "timing_log_level": MegatronActionSpec( option_strings=("--timing-log-level",), action_type="store", @@ -9233,12 +10577,19 @@ const=None, default=None, ), - "tokenizer_metadata": MegatronActionSpec( - option_strings=("--tokenizer-metadata",), - action_type="store", - nargs=None, - const=None, - default=None, + "tokenizer_hf_no_include_special_tokens": MegatronActionSpec( + option_strings=("--tokenizer-hf-no-include-special-tokens",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "tokenizer_hf_no_use_fast": MegatronActionSpec( + option_strings=("--tokenizer-hf-no-use-fast",), + action_type="store_true", + nargs=0, + const=True, + default=False, ), "tokenizer_model": MegatronActionSpec( option_strings=("--tokenizer-model",), @@ -9247,6 +10598,23 @@ const=None, default=None, ), + "tokenizer_sentencepiece_ignore_extra_whitespaces": MegatronActionSpec( + option_strings=( + "--no-tokenizer-sentencepiece-ignore-extra-whitespaces", + "--disable-tokenizer-sentencepiece-ignore-extra-whitespaces", + ), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "tokenizer_sentencepiece_legacy": MegatronActionSpec( + option_strings=("--tokenizer-sentencepiece-legacy",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "tokenizer_type": MegatronActionSpec( option_strings=("--tokenizer-type",), action_type="store", @@ -9269,14 +10637,14 @@ default="nccl", ), "tp_comm_bulk_dgrad": MegatronActionSpec( - option_strings=("--disable-tp-comm-bulk-dgrad",), + option_strings=("--no-tp-comm-bulk-dgrad", "--disable-tp-comm-bulk-dgrad"), action_type="store_false", nargs=0, const=False, default=True, ), "tp_comm_bulk_wgrad": MegatronActionSpec( - option_strings=("--disable-tp-comm-bulk-wgrad",), + option_strings=("--no-tp-comm-bulk-wgrad", "--disable-tp-comm-bulk-wgrad"), action_type="store_false", nargs=0, const=False, @@ -9290,7 +10658,7 @@ default=False, ), "tp_comm_overlap_ag": MegatronActionSpec( - option_strings=("--disable-tp-comm-overlap-ag",), + option_strings=("--no-tp-comm-overlap-ag", "--disable-tp-comm-overlap-ag"), action_type="store_false", nargs=0, const=False, @@ -9304,7 +10672,7 @@ default=None, ), "tp_comm_overlap_rs": MegatronActionSpec( - option_strings=("--disable-tp-comm-overlap-rs",), + option_strings=("--no-tp-comm-overlap-rs", "--disable-tp-comm-overlap-rs"), action_type="store_false", nargs=0, const=False, @@ -9318,14 +10686,14 @@ default=False, ), "tp_comm_split_ag": MegatronActionSpec( - option_strings=("--disable-tp-comm-split-ag",), + option_strings=("--no-tp-comm-split-ag", "--disable-tp-comm-split-ag"), action_type="store_false", nargs=0, const=False, default=True, ), "tp_comm_split_rs": MegatronActionSpec( - option_strings=("--disable-tp-comm-split-rs",), + option_strings=("--no-tp-comm-split-rs", "--disable-tp-comm-split-rs"), action_type="store_false", nargs=0, const=False, @@ -9339,11 +10707,7 @@ default=None, ), "train_iters": MegatronActionSpec( - option_strings=("--train-iters",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--train-iters",), action_type="store", nargs=None, const=None, default=None ), "train_samples": MegatronActionSpec( option_strings=("--train-samples",), @@ -9432,13 +10796,34 @@ const=True, default=False, ), - "use_legacy_models": MegatronActionSpec( - option_strings=("--use-legacy-models",), + "use_gloo_process_groups": MegatronActionSpec( + option_strings=("--disable-gloo-process-groups",), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), + "use_grouped_gemm_for_dense_mlp": MegatronActionSpec( + option_strings=("--use-grouped-gemm-for-dense-mlp",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), + "use_kitchen_attention": MegatronActionSpec( + option_strings=("--use-kitchen-attention",), action_type="store_true", nargs=0, const=True, default=False, ), + "use_layer_wise_param_layout": MegatronActionSpec( + option_strings=("--no-use-layer-wise-param-layout",), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), "use_legacy_static_engine": MegatronActionSpec( option_strings=("--use-legacy-static-engine",), action_type="store_true", @@ -9446,6 +10831,13 @@ const=True, default=False, ), + "use_mamba_mem_eff_path": MegatronActionSpec( + option_strings=("--disable-mamba-mem-eff-path",), + action_type="store_false", + nargs=0, + const=False, + default=True, + ), "use_megatron_fsdp": MegatronActionSpec( option_strings=("--use-megatron-fsdp",), action_type="store_true", @@ -9460,6 +10852,9 @@ const=True, default=False, ), + "use_mup": MegatronActionSpec( + option_strings=("--use-mup",), action_type="store_true", nargs=0, const=True, default=False + ), "use_one_sent_docs": MegatronActionSpec( option_strings=("--use-one-sent-docs",), action_type="store_true", @@ -9524,7 +10919,10 @@ default=False, ), "use_tokenizer_model_from_checkpoint_args": MegatronActionSpec( - option_strings=("--no-use-tokenizer-model-from-checkpoint-args",), + option_strings=( + "--no-use-tokenizer-model-from-checkpoint-args", + "--disable-use-tokenizer-model-from-checkpoint-args", + ), action_type="store_false", nargs=0, const=False, @@ -9551,12 +10949,15 @@ const=True, default=False, ), + "use_transformer_engine_op_fuser": MegatronActionSpec( + option_strings=("--use-transformer-engine-op-fuser",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "v_head_dim": MegatronActionSpec( - option_strings=("--v-head-dim",), - action_type="store", - nargs=None, - const=None, - default=128, + option_strings=("--v-head-dim",), action_type="store", nargs=None, const=None, default=128 ), "valid_data_path": MegatronActionSpec( option_strings=("--valid-data-path",), @@ -9565,6 +10966,20 @@ const=None, default=None, ), + "validation_set_names": MegatronActionSpec( + option_strings=("--validation-set-names",), + action_type="store", + nargs="+", + const=None, + default=None, + ), + "verify_integrity": MegatronActionSpec( + option_strings=("--verify-integrity",), + action_type="store_true", + nargs=0, + const=True, + default=False, + ), "vision_backbone_type": MegatronActionSpec( option_strings=("--vision-backbone-type",), action_type="store", @@ -9594,53 +11009,41 @@ default=0, ), "vocab_file": MegatronActionSpec( - option_strings=("--vocab-file",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--vocab-file",), action_type="store", nargs=None, const=None, default=None ), "vocab_size": MegatronActionSpec( - option_strings=("--vocab-size",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--vocab-size",), action_type="store", nargs=None, const=None, default=None ), "wandb_entity": MegatronActionSpec( option_strings=("--wandb-entity",), action_type="store", nargs=None, const=None, - default="", + default=None, ), "wandb_exp_name": MegatronActionSpec( option_strings=("--wandb-exp-name",), action_type="store", nargs=None, const=None, - default="", + default=None, ), "wandb_project": MegatronActionSpec( option_strings=("--wandb-project",), action_type="store", nargs=None, const=None, - default="", + default=None, ), "wandb_save_dir": MegatronActionSpec( option_strings=("--wandb-save-dir",), action_type="store", nargs=None, const=None, - default="", + default=None, ), "warmup": MegatronActionSpec( - option_strings=("--warmup",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--warmup",), action_type="store", nargs=None, const=None, default=None ), "weight_decay": MegatronActionSpec( option_strings=("--weight-decay",), @@ -9671,18 +11074,10 @@ default=None, ), "window_size": MegatronActionSpec( - option_strings=("--window-size",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--window-size",), action_type="store", nargs=None, const=None, default=None ), "yaml_cfg": MegatronActionSpec( - option_strings=("--yaml-cfg",), - action_type="store", - nargs=None, - const=None, - default=None, + option_strings=("--yaml-cfg",), action_type="store", nargs=None, const=None, default=None ), } diff --git a/oellm_autoexp/backends/megatron/config_schema.py b/oellm_autoexp/backends/megatron/config_schema.py index afcfce87..9b205785 100644 --- a/oellm_autoexp/backends/megatron/config_schema.py +++ b/oellm_autoexp/backends/megatron/config_schema.py @@ -10,40 +10,884 @@ class MegatronConfig(ConfigInterface): """Typed projection of Megatron-LM CLI arguments.""" # Generated by scripts/generate_megatron_dataclass.py - # Number of transformer layers. + # Intra-layer model parallelism. Splits tensors across GPU ranks. + tensor_model_parallel_size: int = 1 + + # Configuring backend option of pipeline parallel communication (e.g., nccl, ucc) If None, + # the default backend will be used. + pipeline_model_parallel_comm_backend: Literal["nccl", "ucc"] | None = None + + # Inter-layer model parallelism. Splits transformer layers across GPU ranks. + pipeline_model_parallel_size: int = 1 + + # Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer + # norms and dropout sequentially. See Reducing Activation Recomputation in Large + # Transformer Models (https://arxiv.org/abs/2205.05198) for more details. + sequence_parallel: bool = False + + # Splits network input along sequence dimension across GPU ranks. + context_parallel_size: int = 1 + + # Degrees of the hierarchical context parallelism. Users should provide a list to specify + # the sizes for different levels. Taking the a2a+p2p cp comm type as example, it contains + # groups of two levels, so the first value of the list indicates the group size of the a2a + # communication type, and the second value indicates the group size of the p2p + # communication type. + hierarchical_context_parallel_sizes: list[int] | None = None + + # Maximum sequence length per DPxCP rank. This is the maximum sequence length each rank + # can handle without overflowing the memory. Typically, a good starting point is to set + # this to maximum sequence length / context parallel size. This is used to calculate the + # number and length of sub-samples assigned to each rank when using + # hybrid_context_parallel. + max_seqlen_per_dp_cp_rank: int | None = None + + # If true, enables hybrid context parallel. This is used to balance the workload of each + # CP rank when we use packed samples with variable sequence lengths. Please set + # max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. + hybrid_context_parallel: bool = False + + # Distributes Moe Experts across sub data parallel dimension. + expert_model_parallel_size: int = 1 + + # Intra-layer tensor model parallelism for expert layer. Splits tensors across GPU ranks. + # Default is None, which will be set to the value of tensor_model_parallel_size. + expert_tensor_parallel_size: int | None = None + + # Controls weights initialization. This option can be useful when you know you are going + # to load values from a checkpoint. + perform_initialization: bool = True + + # When set to False, we initialize the weights directly on the GPU. CPU initialization is + # the same regardless of tensor model parallelism, but GPU initialization is not. + # Transferring weights from CPU to GPU can take a significant amount of time for large + # models. + use_cpu_initialization: bool | None = None + + # If true, train with fp16 mixed precision training. + fp16: bool = False + + # If true, train with bf16 mixed precision training. + bf16: bool = False + + # If true, code that has deterministic execution will be chosen. This usually means slower + # execution, but is good for debugging and testing. Defaults to False. + deterministic_mode: bool = False + + # If true, allows overlapping of Linear layer execution with tensor parallel communication + # collectives like AllGather/ReduceScatter. Overlapping is done for the linear layers + # wherever possible during the forward and the backward pass. + tp_comm_overlap: bool = False + + # Controls All-Gather overlap with Bprop activation gradient GEMM. Don't care if + # tp_comm_overlap is False. + tp_comm_bulk_wgrad: bool = True + + # Controls Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if + # tp_comm_overlap is False. + tp_comm_bulk_dgrad: bool = True + + # Controls All-Gather overlap with GEMM by pipelining the GEMM and All-Gather. Don't care + # if tp_comm_overlap is False. + tp_comm_overlap_ag: bool = True + + # Controls Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter. + # Don't care if tp_comm_overlap is False. + tp_comm_overlap_rs: bool = True + + # If true, allows Reduce-Scatter overlap with DGRAD GEMM by pipelining the GEMM and + # Reduce-Scatter splits. Don't care if tp_comm_overlap is False. + tp_comm_overlap_rs_dgrad: bool = False + + # Deprecated from TransformerEngine v1.6.0. Controls All-Gather overlap with Fprop GEMM by + # pipelining the GEMM and All-Gather splits. Don't care if tp_comm_overlap is False. + tp_comm_split_ag: bool = True + + # Deprecated from TransformerEngine v1.6.0. Controls Reduce-Scatter overlap with Fprop + # GEMM by pipelining the GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is + # False. + tp_comm_split_rs: bool = True + + # If this is enabled, the fused cross entropy implementation would be used. Defaults to + # False. + cross_entropy_loss_fusion: bool = False + + # If 'native', MCore based CE loss fusion is used, if 'te', Parallel CE loss from + # Transformer Engine library is used. Defaults to 'native'. + cross_entropy_fusion_impl: Literal["native", "te"] = "native" + + # Set the bootstrapping backend of Tensor parallel communications. + tp_comm_bootstrap_backend: Literal["nccl", "mpi", "gloo"] = "nccl" + + # Overlap EP A2A communications with independent computations of different micro-batches + # in 1f1b phase of pipelining or non-pipelining schedule. + overlap_moe_expert_parallel_comm: bool = False + + # Delay the weight gradient computation to improve batch-level communication overlapping + delay_wgrad_compute: bool = False + + # Delay the weight gradient computation for TE Grouped GEMM MoE experts. When enabled with + # FSDP, the expert weight gradients are computed on a separate CUDA stream after the data + # gradients finish, allowing overlap of wgrad compute with EP A2A communication. The FSDP + # gradient reduce-scatter for expert parameters is deferred until the delayed wgrad + # computation completes. This requires transformer_engine with GroupedLinear support (TE + # >= 2.3.0). + overlap_dispatch_backward_with_experts_wgrad: bool = False + + # Enable early memory release of attention activations during EP overlap. EP overlap can + # increase peak memory usage when the overlapped forward module allocates more memory than + # what is freed by the backward module. This flag addresses this by reordering the + # attention backward pass to occur earlier in the schedule. Specifically: - Without this + # flag: attn_bwd executes after moe_combine_fwd - With this flag: attn_bwd executes before + # mlp_fwd The earlier execution releases attention activations sooner, reducing peak + # memory. Note: This may impact performance as moe_combine_fwd and moe_dispatch_bwd become + # exposed (not overlapped with other computation). + ep_overlap_early_attn_memory_release: bool = False + + # Use custom ring_exchange kernel instead of torch.distributed.batch_isend_irecv(). + # Requires custom built torch with torch.distributed.ring_exchange. + use_ring_exchange_p2p: bool = False + + # If true, defers the embedding WGRAD GEMMs while pipeline flush is taking place enabling + # us to hide pipeline flush latency. Defaults to False. + defer_embedding_wgrad_compute: bool = False + + # This value tunes the number of micro-batches for which the embedding weight gradient + # compute needs to be deferred to pipeline flush, this argument is invalid if + # `defer_embedding_wgrad_compute` is False. Defaults to 0, which means all micro-batches + # are deferred. + wgrad_deferral_limit: int = 0 + + # If true, overlap communication and computation in warm up and flush phase. Only valid + # when overlap_p2p_comm is True and batch_p2p_comm is False. Defaults to False. + overlap_p2p_comm_warmup_flush: bool = False + + # This value specifies the number of micro-batches that are executed at a time for a given + # virtual stage (both forward and backward). Default (in __post_init__() method below) to + # pipeline_parallel_size which specifies a depth-first schedule. Example: for PP=2 VP=2, + # when microbatch_group_size_per_vp_stage=2, num_microbatches = 4, we have rank 0 | 0 1 0 + # 1 2 3 2 3 rank 1 | 0 1 0 1 2 3 2 3 When microbatch_group_size_per_vp_stage=3, + # num_microbatches = 5, we have rank 0 | 0 1 2 0 1 2 3 4 3 4 rank 1 | 0 1 2 0 1 2 3 4 3 4 + microbatch_group_size_per_vp_stage: int | None = None + + # This will be set automatically according to the pipeline layout, and will be set to True + # if MTP is in a separate vpp stage. + mtp_standalone: bool = False + + # Tells the number of transformer layers for which activations has to be offloaded. + cpu_offloading_num_layers: int = 0 + + # If True, the pinned CPU buffers are retained after offloading and reused for the next + # iteration. It is useful for cuda graphs capture. + cpu_offloading_retain_pinned_cpu_buffers: bool = False + + # Number of transformer layers in a transformer block. num_layers: int | None = None - # Number of encoder transformer layers. - encoder_num_layers: int | None = None + # Number of Multi-Token Prediction (MTP) Layers. MTP extends the prediction scope to + # multiple future tokens at each position. This MTP implementation sequentially predict + # additional tokens by using D sequential modules to predict D additional tokens. + mtp_num_layers: int | None = None + + # Weighting factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP + # losses across all depths, and multiply it the scaling factor to obtain the overall MTP + # loss, which serves as an additional training objective. + mtp_loss_scaling_factor: float = 0.1 + + # Use a single MTP layer repeatedly instead of multiple separate layers. + mtp_use_repeated_layer: bool = False + + # DEPRECATED: Use unified hybrid_layer_pattern instead. Legacy argument for loading old + # checkpoints. Force a specific hybrid layer pattern for MTP layers. + mtp_hybrid_override_pattern: str | None = None + + # If set, the embedding layer will be treated as a standard transformer layer in the + # context of partition and placement for pipeline parallelism. + account_for_embedding_in_pipeline_split: bool = False + + # If set, the loss layer will be treated as a standard transformer layer in the context of + # partition and placement for pipeline parallelism. + account_for_loss_in_pipeline_split: bool = False + + # Transformer hidden size. + hidden_size: int | None = None + + # Number of transformer attention heads. + num_attention_heads: int | None = None + + # Attention backend to run. By default we let transformer engine decide the best backend + # to run (except in the case of local). If attention backend is local we use the local + # pytorch implementation in mcore. Users can specify exact backend by changing this + # config. + attention_backend: str = "auto" + + # Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. + # Supports both TE FusedAttention and local unfused attention. Supports both a fixed + # offset and and learnable offset. + softmax_type: Literal["vanilla", "off-by-one", "learnable"] = "vanilla" + + # Number of query groups for group query attention. If None, normal attention is used. + num_query_groups: int = 1 + + # Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size if not + # provided. + ffn_hidden_size: int | None = None + + # Projection weights dimension in multi-head attention. This is set to hidden_size // + # num_attention_heads if not provided. + kv_channels: int | None = None + + # Dropout probability for transformer hidden state. + hidden_dropout: float = 0.1 + + # Post attention dropout probability. + attention_dropout: float = 0.1 + + # If true, move residual connections to fp32. + fp32_residual_connection: bool = False + + # If True, uses the original BERT residule connection ordering. + apply_residual_connection_post_layernorm: bool = False + + # Epsilon value for any LayerNorm/RMSNorm operations. + layernorm_epsilon: float = 1e-05 + + # If set to True, the LayerNorm is adjusted to center the gamma values around 0. This + # improves numerical stability. + layernorm_zero_centered_gamma: bool = False + + # Include/exclude a bias term in all linear layers (QKV projections, after core attention, + # and two in MLP layer). + add_bias_linear: bool = True + + # Add a bias term only for QKV projections. + add_qkv_bias: bool = False + + # Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). + # Only used when gated_linear_unit is True + glu_linear_offset: float = 0.0 + + # Clamp the output of the linear_fc1 in the activation function. Only used when + # activation_func is quick_gelu. + activation_func_clamp_value: float | None = None + + # True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs + # of first half and second half (LLaMa style). Default to False. + rotary_interleaved: bool = False + + # Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`. + normalization: Literal["LayerNorm", "RMSNorm"] = "LayerNorm" + + # Whether to apply `normalization` type of normalization to the query and key embeddings. + qk_layernorm: bool = False + + # Whether to apply llama 4-style qk L2 norm. + qk_l2_norm: bool = False + + # Whether to clip the query and key weights. Needed for Muon MLA Model training. + qk_clip: bool = False + + # The balancing alpha for qk-clip. Q = Q * (eta ** alpha) + qk_clip_alpha: float = 0.5 + + # The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0) + qk_clip_threshold: float = 100 + + # Whether to apply output gate to the attention layers. + attention_output_gate: bool = False + + # Whether cross entropy loss is calculated over the actual number of non-padded tokens in + # the global batch, versus the default behavior of assuming all tokens are non-padded. + calculate_per_token_loss: bool = False + + # Whether to use multi-latent attention. + multi_latent_attention: bool = False + + # Type of attention variant to use. Currently support gated_delta_net and dsa. + experimental_attention_variant: Literal["gated_delta_net", "dsa"] | None = None + + # Number of DSA indexer heads. + dsa_indexer_n_heads: int | None = None + + # Dimension per DSA indexer head. + dsa_indexer_head_dim: int | None = None + + # Number of top-k tokens to select in DSA indexer. + dsa_indexer_topk: int | None = None + + # Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss. + dsa_indexer_loss_coeff: float | None = None + + # Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using + # the top-k indices. + dsa_indexer_use_sparse_loss: bool = False + + # Conv kernel dimension for the gated delta net. + linear_conv_kernel_dim: int = 4 + + # Query and key head dimension for the gated delta net. + linear_key_head_dim: int = 128 + + # Value and gate head dimension for the gated delta net. + linear_value_head_dim: int = 128 + + # Number of query and key heads for the gated delta net. + linear_num_key_heads: int = 16 + + # Number of value and gate heads for the gated delta net. + linear_num_value_heads: int = 32 + + # Standard deviation of the zero mean normal for the default initialization method, not + # used if init_method and output_layer_init_method are provided. + init_method_std: float = 0.02 + + # Standard deviation of the zero mean normal for the default initialization method for the + # embedding layer. If None, will be set to init_method_std. Setting this to a value around + # 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying + # weight decay on embedding weights to avoid shrinkage towards zero. See + # https://arxiv.org/abs/2312.16903 for more details. + embedding_init_method_std: float | None = None + + # If True, initializes the model with the meta device. This is helpful for training of + # very large models. This feature is only works when megatron fsdp is turned on. + init_model_with_meta_device: bool = False + + # Enable Maximal Update Parameterization (MuP) for hyperparameter transfer across model + # widths. When enabled, learning rates and initialization are scaled according to the + # width multiplier to ensure consistent training dynamics. + use_mup: bool = False + + # Width multiplier for MuP scaling, computed as hidden_size / mup_base_hidden_size. This + # value is automatically computed in __post_init__ when use_mup is enabled. + mup_width_mult: float = 1.0 + + # Base hidden size for MuP width scaling. This is the reference width from which scaling + # factors are computed. Defaults to hidden_size if not specified (base model case where + # width_mult=1.0). Set this to your base/proxy model's hidden size when scaling up. + mup_base_hidden_size: int | None = None + + # Multiplier for embedding layer output. Applied after the embedding lookup. Default: 1.0 + # (no scaling). + mup_embedding_mult: float = 1.0 + + # Multiplier for output logits before softmax. When MuP is enabled and this is left at + # 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable across widths. + # Override to customize output scaling. Default: 1.0. + mup_output_mult: float = 1.0 + + # Base head dimension for MuP attention scaling. When set, softmax_scale = + # sqrt(mup_base_head_dim) / (kv_channels ** mup_attn_scale_power). Set to base model's + # d_head (e.g., 64) to match standard 1/sqrt(d_head) scaling at the base width, ensuring + # non-MuP compatibility for that specific value. + mup_base_head_dim: float | None = None + + # Power for attention scaling: softmax_scale = 1 / (kv_channels ** mup_attn_scale_power). + # 0.5 = standard attention (1/sqrt(d_head)), 1.0 = MuP attention (1/d_head). Default: 1.0 + # (MuP scaling when use_mup is True). Set to 0.5 for standard scaling. + mup_attn_scale_power: float = 1.0 + + # If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training + # with fp16. Also sets `attention_softmax_in_fp32` to True. + apply_query_key_layer_scaling: bool = False + + # If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to + # prevent matmul from using reduced precision accumulation when using BF16. + disable_bf16_reduced_precision_matmul: bool = False + + # If True, uses fused weighted squared relu kernel when using MoE. + use_fused_weighted_squared_relu: bool = False + + # If True, fuses residual connection and RMSNorm backward pass when TE is used. + fused_residual_rmsnorm: bool = False + + # If True, submodules may use Transformer Engine's operation fuser API to enable advanced + # fusions. + use_transformer_engine_op_fuser: bool = False + + # Determines which type of activation recompute to use. Megatron-core supports 'selective' + # activation checkpointing where the submodules set in --recompute-modules is + # checkpointed. The default is "core_attn" which is the memory intensive part of + # attention. These memory intensive activations are also less compute intensive which + # makes activation checkpointing more efficient for LLMs (20B+). See Reducing Activation + # Recomputation in Large Transformer Models (https://arxiv.org/abs/2205.05198) for more + # details. 'full' will checkpoint the entire transformer layer. If None, no recompute is + # performed and all activations are saved. If set, must be 'selective' or 'full'. + # 'selective' always uses all layers. + recompute_granularity: Literal["full", "selective"] | None = None + + # Determines which transformer layers will be recomputed. uniform will uniformly divide + # the total number of transformer layers in a transformer block and recompute the input + # activation of each divided chunk at the specified granularity. block will recompute the + # input activations for only a set number of transformer layers per pipeline stage. The + # rest of the layers in the pipeline stage will not have any activations recomputed. If + # None, and recompute is enabled, all layers will do recomputation. If set, must be + # 'uniform' or 'block'. + recompute_method: Literal["uniform", "block"] | None = None + + # When recompute_method is uniform, recompute_num_layers is the number of transformer + # layers in each uniformly divided recompute unit. When recompute_method is block, + # recompute_num_layers is the number of transformer layers to recompute within each + # pipeline stage. Must be None for 'selective' activation checkpointing. + recompute_num_layers: int | None = None + + # If True, distribute recomputed activations across the model parallel group. + distribute_saved_activations: bool = False + + # The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", + # "mla_up_proj", "mlp", "moe", "shared_experts". default: ["core_attn"]. "core_attn": + # recompute the core attention part of the transformer layer. "moe_act": recompute the MoE + # MLP activation function. "layernorm": recompute the input_layernorm and + # pre_mlp_layernorm. "mla_up_proj": recompute the MLA up projection and RoPE applying + # parts. "mlp": recompute the dense MLP submodule. "moe": recompute the MoE layer. + # "shared_experts": recompute the shared experts in the MoE layer. "moe_act", "layernorm", + # and "mla_up_proj" use output-discarding checkpointing, "core_attn", "mlp", "moe", and + # "shared_experts" use normal checkpointing. + recompute_modules: list[str] | None = None + + # If set, enables the use of FP8 precision through Transformer Engine. There are 2 + # predefined choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses + # e4m3 for all FP8 activation and weight tensors and e5m2 for all FP8 output activation + # gradient tensors. + fp8: Literal["e4m3", "hybrid"] | None = None + + # If set, enables the use of FP8 precision through Transformer Engine. There are 5 + # predefined choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) + # 'delayed' uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, 4) + # 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe. + fp8_recipe: Literal["tensorwise", "delayed", "mxfp8", "blockwise", "custom"] = "delayed" + + # Python import path to a callable quantizer factory, e.g., + # package.module.quantizer_factory. Required when fp8_recipe is custom. + fp8_quantizer_factory: str | None = None + + # Margin for the scaling factor computation. + fp8_margin: int = 0 + + # DEPRECATED from TransformerEngine v1.8.0. This flag is ignored. Controls how often the + # scaling factor is recomputed. + fp8_interval: int = 1 + + # The length of the amax history window used for scaling factor computation. + fp8_amax_history_len: int = 1 + + # Algorithm used for choosing the `amax` value for the scaling factor computation. There + # are 2 predefined choices: `max` chooses the largest `amax` in the history window, while + # `most_recent` always chooses the most recently seen value. + fp8_amax_compute_algo: Literal["most_recent", "max"] = "most_recent" + + # When set to False, override FP8 config options and do the wgrad computation in higher + # precision. + fp8_wgrad: bool = True + + # If True, run the LM-head output projection with a TE ColumnParallelLinear under the + # MXFP8 autocast context. Only active when fp8=True and fp8_recipe='mxfp8'. + fp8_output_proj: bool = False + + # If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8. + first_last_layers_bf16: bool = False + + # Number of layers at the start of the model to keep in BF16 precision when + # first_last_layers_bf16 is True. + num_layers_at_start_in_bf16: int = 1 + + # Number of layers at the end of the model to keep in BF16 precision when + # first_last_layers_bf16 is True. + num_layers_at_end_in_bf16: int = 1 + + # Use the kitchen extension for attention (instead of TE's attention). + use_kitchen_attention: bool = False + + # Which kitchen attention backend to use when use_kitchen_attention=True. "sdpa" for + # KitchenDotProductAttention, "fa" for KitchenFlashAttention. + kitchen_attention_backend: Literal["sdpa", "fa"] = "sdpa" + + # If set, enables the use of FP4 precision through Transformer Engine. Currently only + # supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0). + fp4: Literal["e2m1"] | None = None + + # If set, enables the use of FP4 precision through Transformer Engine. Currently only + # 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture. + fp4_recipe: Literal["nvfp4", "custom"] = "nvfp4" + + # Python import path to a callable quantizer factory, e.g., + # package.module.quantizer_factory. Required when fp4_recipe is custom. + fp4_quantizer_factory: str | None = None + + # If True, the communication stream created by set_streams for combined 1f1b a2a overlap + # is created with CUDA high priority. + high_priority_a2a_comm_stream: bool = False + + # Shared expert total ffn hidden size. It should be equal to 'num_shared_experts * + # ffn_size_of_each_shared_expert' if there are multiple shared experts. None means no + # shared expert. By default, the shared experts execute before the router. However, when + # moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set, the shared experts + # execute after the router, before the routed experts. This makes the gradients from the + # router and the shared experts added in different orders to the hidden_states, causing + # minor numerical differences in the hidden_states gradient. + moe_shared_expert_intermediate_size: int | None = None + + # Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size + # is set. + moe_shared_expert_gate: bool = False + + # Enable overlapping between shared expert computations and dispatcher communications. + # Without this, the shared experts execute before the router. Only effective when moe- + # shared-expert-intermediate-size is set. + moe_shared_expert_overlap: bool = False + + # MoE Feed-Forward Network hidden size. If not specified, defaults to the ffn_hidden_size. + moe_ffn_hidden_size: int | None = None + + # Number of experts to route to for each token. + moe_router_topk: int = 2 + + # If True, enable the routing replay feature for MoE layers. + moe_enable_routing_replay: bool = False + + # Whether to pad the routing_map to make sure the number of tokens each expert receives is + # a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the + # explicit padding in the GroupedMLP layer. + moe_router_padding_for_quantization: bool = False + + # [Compatibility alias for moe_router_padding_for_quantization] Enabling this will also + # enable moe_router_padding_for_quantization. + moe_router_padding_for_fp8: bool = False + + # Number of groups to divide experts into for group-limited routing. When using group- + # limited routing: 1. Experts are divided into 'moe_router_num_groups' equal-sized groups + # 2. For each token, 'moe_router_group_topk' groups are selected based on sum of + # top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group 3. From + # these selected groups, 'moe_router_topk' individual experts are chosen Two common use + # cases: - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel + # size (EP) to limit each token to experts on a subset of devices (See DeepSeek-V2: + # https://arxiv.org/pdf/2405.04434) - Node-limited routing: Set 'moe_router_num_groups' + # equal to number of nodes in EP group to limit each token to experts on a subset of nodes + # (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437) + moe_router_num_groups: int | None = None + + # Number of selected groups for group-limited routing. + moe_router_group_topk: int | None = None + + # Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the top-k + # selection. By default, softmax is done after top-k. + moe_router_pre_softmax: bool = False + + # Scaling factor for routing score in top-k selection, only works when + # moe_router_pre_softmax enabled. Defaults to None, which means no scaling. + moe_router_topk_scaling_factor: float | None = None + + # Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus". + moe_router_score_function: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "softmax" + + # Data type for routing and expert output weighted averaging. Using fp32 or fp64 can + # improve stability especially when the number of experts is large (e.g. finegrained-moe). + # None means no changes for dtype. + moe_router_dtype: Literal["fp32", "fp64"] | None = None + + # TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. + # The routing decision is based on the sum of the routing scores and the expert bias. See + # https://arxiv.org/abs/2408.15664 for details. + moe_router_enable_expert_bias: bool = False + + # The expert bias is updated based on the number of assigned tokens to each expert in a + # global batch, where the bias is increased for the experts with less assigned tokens and + # decreased for the experts with more assigned tokens. The default value 1e-3 is same as + # that used in DeepSeekV3. + moe_router_bias_update_rate: float = 0.001 + + # [Experimental] Force load balancing with random logits for MoE router, supports naive + # topk and group-limited topk. This is an experimental feature and only for benchmark. + moe_router_force_load_balancing: bool = False + + # Apply random expert bias in normal distribution with specified std to router logits. + # Shared seed across all ranks ensures identical bias. If positive, generates new random + # bias each forward pass. If negative, generates bias once per layer and reuses it (abs + # value is std). This is an experimental feature for benchmarking purposes. + moe_router_force_biased: float | None = None + + # Use GroupedLinear(num_groups=1) for dense MLP to trigger the + # ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. Requires + # ``use_te_op_fuser=True`` and SwiGLU activation. + use_grouped_gemm_for_dense_mlp: bool = False + + # When there are multiple experts per rank, compress multiple local (potentially small) + # gemms in a single kernel launch to improve the utilization and performance by leveraging + # the Grouped GEMM feature introduced since CUTLASS 2.8 + # (https://github.com/fanshiqing/grouped_gemm). + moe_grouped_gemm: bool = False + + # When using TE GroupedLinear for MoE experts, store expert weights as a single grouped + # parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``. + moe_single_grouped_weight: bool = False + + # When using TE GroupedLinear for MoE experts, store expert biases as a single grouped + # parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` + # and ``add_bias_linear=True``. + moe_single_grouped_bias: bool = False + + # Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended. + moe_z_loss_coeff: float | None = None + + # Add noise to the input tensor by applying jitter with a specified epsilon value. + moe_input_jitter_eps: float | None = None + + # The type of token dispatcher to use. The default is 'allgather'. Options are + # 'allgather','alltoall' and 'flex'. + moe_token_dispatcher_type: Literal["allgather", "alltoall", "flex"] = "allgather" + + # [Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models. + moe_enable_deepep: bool = False + + # [Experimental] The backend to use for flex token dispatcher. The default is "deepep". + # Options are "deepep" and "hybridep". Currently only "hybridep" backend supports the + # MNNVL case. + moe_flex_dispatcher_backend: Literal["deepep", "hybridep"] = "deepep" + + # Fuse token rearrangement ops during token dispatching for HybridEP. + moe_permute_fusion_into_hybridep: bool = False + + # Enable per-layer logging for MoE, currently supports auxiliary loss and z loss. + moe_per_layer_logging: bool = False + + # moe_expert_capacity_factor (float): The capacity factor for each expert, None means no + # token will be dropped. The default is None. + moe_expert_capacity_factor: float | None = None + + # moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to + # match the expert capacity length, effective only after the moe_expert_capacity_factor is + # set. The default setting is False. + moe_pad_expert_input_to_capacity: bool = False + + # moe_pad_experts_for_cuda_graph_inference (bool): If True, the router will switch to + # dropping and padding during decode time which does not have a D2H sync. The capacity + # factor is set to the max that an expert could see during inference so no tokens are + # actually dropped. The default setting is False. + moe_pad_experts_for_cuda_graph_inference: bool = False + + # 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. + moe_token_drop_policy: Literal["probs", "position"] = "probs" + + # Memory optimization: checkpointing moe_layer to save actiavtion memory. + moe_layer_recompute: bool = False + + # Fuse token rearrangement ops during token dispatching. + moe_permute_fusion: bool = False + + # Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in + # TransformerEngine 2.7.0 and above. + moe_router_fusion: bool = False + + # Apply probs on input of experts instead of applying after activation and glu. + moe_apply_probs_on_input: bool = False + + # Latent projection dimension for MoE. If None, MoE latent projections are not used. + moe_latent_size: int | None = None + + # Number of SMs to use for DeepEP. + moe_deepep_num_sms: int = 20 + + # Number of SMs to use for HybridEP. None uses the default from DeepEP. In pure NVL + # scenarios, 16 SMs can generally achieve good bandwidth. + moe_hybridep_num_sms: int | None = None + + # Number of CUDA thread blocks for the permute part in HybridEP. When + # permute_fusion_into_hybridep is True, this sets the number of SMs for the permute part + # (only 1 block per SM). + moe_hybridep_num_blocks_permute: int | None = None + + # Number of CUDA thread blocks for the unpermute part in HybridEP. When + # permute_fusion_into_hybridep is True, this sets the number of SMs for the unpermute part + # (only 1 block per SM). + moe_hybridep_num_blocks_unpermute: int | None = None + + # Number of SMs to use for HybridEP preprocessing (metadata scan kernel). + moe_hybridep_num_sms_preprocessing: int = 108 + + # When set, GLU activations in the MoE grouped MLP layer will use a block interleaved + # format. Instead of interpreting the input tensor as a concatenation of gates and linear + # units, it will be interpreted as alternating blocks of gates and linear units. This data + # format is experimental and primarily intended to enable advanced fused kernels. + moe_mlp_glu_interleave_size: int | None = None + + # moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. + # Tokens exceeding this budget will be dropped. None means no token will be dropped. The + # default is None. + moe_expert_rank_capacity_factor: float | None = None + + # DEPRECATED and replaced by cuda_graph_impl. When set to true, either partial CUDA graph + # (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole + # iteration excluding optimizer) is enabled. cuda_graph_modules determines the scope of + # graph capture. + enable_cuda_graph: bool = False + + # Number of warmup steps for CUDA graphs + cuda_graph_warmup_steps: int = 3 + + # DEPRECATED and replaced by cuda_graph_impl. When set to true, TransformerLayer layers + # are swapped with user provided CUDA graphs. + external_cuda_graph: bool = False + + # Determines the CUDA graph capture implementation. "none": no CUDA graph. "local": MCore + # CUDA graph implementation. During training, graphable modules own per-layer CUDA graphs + # controlled by cuda_graph_modules. During inference, graph ownership is controlled + # separately by inference_cuda_graph_scope. "transformer_engine": Transformer Engine CUDA + # graph implementation. During training, TE make_graphed_callables() creates per-layer + # CUDA graphs controlled by cuda_graph_modules. Inference CUDA graphs are not supported; + # inference_cuda_graph_scope must be "none". "full_iteration": full-iteration CUDA graph + # implementation for the training iteration (1 CUDA graph for the whole forward-backward + # path excluding the optimizer step). Inference CUDA graphs are not supported; + # inference_cuda_graph_scope must be "none". cuda_graph_modules has no effect when + # cuda_graph_impl="none" and must be empty when cuda_graph_impl="full_iteration". + cuda_graph_impl: Literal["none", "local", "transformer_engine", "full_iteration"] = "none" + + # Controls the CUDA graph scope during inference. When unset, the effective default is + # derived from cuda_graph_impl: "local" -> "layer", all other impls -> "none". "none": + # inference runs in eager mode (no CUDA graphs). "layer": inference graphs are owned at + # the module/layer boundary, e.g. TransformerLayer or MambaLayer. "block": inference + # graphs are owned by the enclosing block, e.g. TransformerBlock or HybridBlock. Currently + # supported combinations are: cuda_graph_impl="local" -> "layer" or "block"; all other + # cuda_graph_impl values -> "none". + inference_cuda_graph_scope: Literal["none", "layer", "block"] | None = None + + # When set to True, clone the output of scatter_to_sequence_parallel_region in embedding + # layer to facilitate garbage collection of input. + clone_scatter_output_in_embedding: bool = True + + # When non-empty, dumps entry-point configs to config_logger_dir + config_logger_dir: str = "" + + # Use the optimized flash decoding kernel during inference. + flash_decode: bool = False + + # If true, uses batch-invariant kernels that provide deterministic forward execution + # regardless of batch size. This ensures bitwise identical results when the same inputs + # are processed in different batch configurations. This will significantly affect speed of + # training and inference as the kernels are not full optimized. Defaults to False. + batch_invariant_mode: bool = False + + # Whether to use ffn activation functions implemented by TransformerEngine + use_te_activation_func: bool = False + + # What type of symmetric all reduce to use. The default is None which is no use of + # symmetric memory. + symmetric_ar_type: Literal["two_shot", "one_shot", "multimem_all_reduce"] | None = None - # Number of decoder transformer layers. - decoder_num_layers: int | None = None + # If True, use NCCL all-reduce kernels when symmetric all-reduce is enabled. + nccl_all_reduce_for_prefill: bool = False - # Transformer hidden size. - hidden_size: int | None = None + # If true, uses a fused reduce-scatter-residual-norm-allgather kernel during inference. + inference_fuse_tp_communication: bool = False + + # If true, disables the use of Triton NVLS kernels during inference. + inference_disable_triton_nvls_kernels: bool = False + + # Specifies the backend to use for grouped GEMM operations during inference. Options: - + # 'flashinfer': Uses FlashInfer cutlass_fused_moe. Not compatible with MXFP8. - 'torch': + # Uses torch.nn.functional.grouped_mm (mcore_fused_moe with Triton kernels). Supports both + # BF16 and MXFP8. - 'vllm': Uses vLLM's Triton fused MoE kernel (BF16). Avoids physical + # token permutation via indirect addressing. + inference_grouped_gemm_backend: Literal["flashinfer", "torch", "vllm"] = "vllm" + + # When False (default), use fused kernels that combine permute/activation with MXFP8 + # quantization + swizzle into a single kernel launch. Only applies when + # fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel launches + # (useful for debugging). + inference_moe_disable_fused_quant_kernels: bool = False + + # Token dispatcher to use for MoE expert parallelism during inference. - 'nccl': + # AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires decode-only CUDA + # graphs (forced automatically). - 'nvls': Variable-count AllGather-V/ReduceScatter-V via + # NVLS multimem kernels. Requires Hopper+ GPUs with NVLink and symmetric memory. Default. + # Only applies when transformer_impl='inference_optimized' and EP > 1. + inference_moe_token_dispatcher_type: Literal["nccl", "nvls"] = "nvls" + + # Multimodal rope section is for channel dimension of temporal, height and width in rope + # calculation. + mrope_section: list[int] | None = None - # Transformer Feed-Forward Network hidden size. This is set to 4*hidden-size if not - # provided - ffn_hidden_size: int | None = None + # Indicates whether this is a hybrid model. + is_hybrid_model: bool = False - # Number of transformer attention heads. - num_attention_heads: int | None = None + # The dimensionality of the state representation in Mamba layers. + mamba_state_dim: int = 128 - # Attention backend to use (flash,fused,unfused,local,auto). Defaults to auto - attention_backend: str = "auto" + # The dimensionality of the heads in the Mamba layers. + mamba_head_dim: int = 64 - # Projection weights dimension in multi-head attention. This is set to args.hidden_size - # // args.num_attention_heads if not provided. - kv_channels: int | None = None + # The number of groups used in Mamba layers. + mamba_num_groups: int = 8 - # Use group-query attention. - group_query_attention: bool = False + # The number of heads used in Mamba layers. If None, the number of heads will be + # hidden_size * expand // mamba_head_dim. + mamba_num_heads: int | None = None - num_query_groups: int = 1 + # Controls usage of the memory efficient path for Mamba layers. + use_mamba_mem_eff_path: bool = True + + # The number of chunks along the sequence dimension to use for MLP computation during + # prefill. + mlp_chunks_for_prefill: int = 1 + + # The number of chunks along the sequence dimension to use for MLP computation during + # training. + mlp_chunks_for_training: int = 1 + + # Transformer implementation to use. Options are 'transformer_engine' for Transformer + # Engine and 'local' for MCore. + transformer_impl: Literal["local", "transformer_engine", "inference_optimized"] = ( + "transformer_engine" + ) + + # If True, offload the input of the specified modules to the CPU. Fine-grained activation + # offloading is a module-level offloading method instead of a layer-level offloading + # method like cpu_offloading. + fine_grained_activation_offloading: bool = False + + # The submodules to offload its input. choices: "attn_norm", "qkv_linear", "core_attn", + # "attn_proj", "mlp_norm", "expert_fc1", "moe_act". "attn_norm": offload the input of the + # normalization in the attention part. "qkv_linear": offload the input of the qkv linear + # part. "core_attn": offload the input of the core attention part. "attn_proj": offload + # the input of the attn linear projection part. "mlp_norm": offload the input of the + # normalization in the mlp part. "expert_fc1": offload the input of the expert fc1 part. + # "moe_act": offload the input of the moe act part. + offload_modules: list[str] = field(default_factory=lambda: []) + + # The minimum size of the tensor to be offloaded. + min_offloaded_tensor_size: int = 1048576 + + # If True, enable paged stash for all routed-expert activations needed for backward + moe_paged_stash: bool = False + + # Number of tokens per page for paged stash memory management. + moe_paged_stash_page_size: int = 64 + + # Scale factor for paged stash CUDA buffer allocation. Sign selects sizing: positive = + # avg-based, negative = actual-max. Magnitude is headroom (e.g. 1.10 = 10%). + moe_paged_stash_buffer_size_factor_cuda: float = 1.1 + + # Scale factor for paged stash host buffer. 0 disables host buffer. Same sign convention + # as moe_paged_stash_buffer_size_factor_cuda: positive = avg-based, negative = actual-max; + # scale = abs(factor). + moe_paged_stash_buffer_size_factor_cpu: float = 0.0 + + # Per fine-grained offloading group name, max number of inflight offloads for that name + # not yet joined on the main stream (wait_event on D2H). The same cap applies to every + # name (e.g., ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait + # after every offload for that name. 1 = at most one not-yet-waited offload per name, etc. + # None = do not insert these joins. This feature is particularly useful when using with + # full-iteration CUDA graphs + fine_grained_offloading_max_inflight_offloads: int | None = None + + # Number of encoder transformer layers. + encoder_num_layers: int | None = None + + # Number of decoder transformer layers. + decoder_num_layers: int | None = None - # Type of softmax to use for the attention. Supports both a fixed offset and learnable - # offset. - softmax_type: Literal["learnable", "vanilla", "off-by-one"] = "vanilla" + # Use group-query attention. + group_query_attention: bool = False # Window size for window attention. If not provided, window attention will be disabled. window_size: Any | None = None @@ -59,9 +903,9 @@ class MegatronConfig(ConfigInterface): max_position_embeddings: int | None = None # Position embedding type. - position_embedding_type: Literal["learned_absolute", "rope", "mrope", "relative", "none"] = ( - "learned_absolute" - ) + position_embedding_type: Literal[ + "learned_absolute", "rope", "yarn", "mrope", "relative", "none" + ] = "learned_absolute" # Number of buckets for relative position embeddings. relative_attention_num_buckets: int = 32 @@ -78,9 +922,6 @@ class MegatronConfig(ConfigInterface): # Percent of rotary dimension to use, default 100%% rotary_percent: float = 1.0 - # Use interleaved rotary embedding. - rotary_interleaved: bool = False - # Sequence length interpolation factor for rotary embeddings. rotary_seq_len_interpolation_factor: int | None = None @@ -101,26 +942,10 @@ class MegatronConfig(ConfigInterface): # Disable position embedding. Deprecated: use --position-embedding-type add_position_embedding: bool = True - # Multimodal rope section is for channel dimension, empty by default. - mrope_section: list[int] | None = None - # Pad the vocab size to be divisible by this value.This is added for computational # efficieny reasons. make_vocab_size_divisible_by: int = 128 - # Which normalization technique to use. - normalization: Literal["LayerNorm", "RMSNorm"] = "LayerNorm" - - # Epsilon for layer norm and RMS norm. - norm_epsilon: float = 1e-05 - - # Adjust LayerNorm weights such that they are centered around zero. This improves - # numerical stability. - apply_layernorm_1p: bool = False - - # If set, use original BERT residula connection ordering. - apply_residual_connection_post_layernorm: bool = False - # Use OpenAIs GeLU implementation. This optionshould not be used unless for backward # compatibilityreasons. openai_gelu: bool = False @@ -134,14 +959,6 @@ class MegatronConfig(ConfigInterface): # Use quick geglu activation instead of default gelu quick_geglu: bool = False - # Clamp the output of the linear_fc1 in the activation function. Only used when - # activation_func is quick_gelu. - activation_func_clamp_value: float | None = None - - # Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). - # Only used when gated_linear_unit is True - glu_linear_offset: float = 0.0 - # Use workarounds for known problems with Torch ONNX exporter onnx_safe: bool | None = None @@ -151,36 +968,11 @@ class MegatronConfig(ConfigInterface): # Untie embeddings and output weights. untie_embeddings_and_output_weights: bool = False - # Use multi-latent attention for model. - multi_latent_attention: bool = False - - # Number of Multi-Token Prediction (MTP) Layers.MTP extends the prediction scope to - # multiple future tokens at each position.This MTP implementation sequentially predict - # additional tokens by using D sequential modules to predict D additional tokens. - mtp_num_layers: int | None = None - - # Scaling factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP - # losses across all depths, and multiply it the scaling factor to obtain the overall MTP - # loss, which serves as an additional training objective. - mtp_loss_scaling_factor: float = 0.1 - - # Post attention dropout probability. - attention_dropout: float = 0.1 - - # Dropout probability for hidden state transformer. - hidden_dropout: float = 0.1 - # Weight decay coefficient for L2 regularization. weight_decay: float = 0.01 - # Initial weight decay coefficient for L2 regularization. - start_weight_decay: float | None = None - - # End of run weight decay coefficient for L2 regularization. - end_weight_decay: float | None = None - - # Weight decay increment function. - weight_decay_incr_style: Literal["constant", "linear", "cosine"] = "constant" + # Apply weight decay to qk layernorm as a special case. + apply_wd_to_qk_layernorm: bool = False # Gradient clipping based on global L2 norm. clip_grad: float = 1.0 @@ -197,82 +989,50 @@ class MegatronConfig(ConfigInterface): # Momentum factor for sgd sgd_momentum: float = 0.9 - # Batch size per model instance (local batch size). Global batch size is local batch size - # times data parallel size times number of micro batches. - micro_batch_size: int | None = None - - # Old batch size parameter, do not use. Use --micro-batch-size instead - batch_size: int | None = None - - # Training batch size. If set, it should be a multiple of micro-batch-size times data- - # parallel-size. If this value is None, then use micro-batch-size * data-parallel-size as - # the global batch size. This choice will result in 1 for number of micro-batches. - global_batch_size: int | None = None + # Momentum factor for Muon optimizer + muon_momentum: float = 0.9 - # Batch size ramp up with the following values: --rampup-batch-size - # For example: --rampup- - # batch-size 16 8 300000 \ --global-batch-size 1024will start with global batch size 16 - # and over (1024 - 16) / 8 = 126 intervals will increasethe batch size linearly to 1024. - # In each intervalwe will use approximately 300000 / 126 = 2380 samples. - rampup_batch_size: list[str] | None = None - - # If set, decrease batch size if microbatch_size * dp_sizedoes not divide batch_size. - # Useful for KSO (Keep Soldiering On)to continue making progress if number of healthy GPUs - # (andcorresponding dp_size) does not support current batch_size.Old batch_size will be - # restored if training is re-started withdp_size that divides batch_size // - # microbatch_size. - decrease_batch_size_if_needed: bool = False + # Whether to split QKV parameters for Muon optimizer + muon_split_qkv: bool = True - # recompute activation to allow for training with larger models, sequences, and batch - # sizes. - recompute_activations: bool = False + # Whether to use Nesterov-style momentum in the internal SGD + muon_nesterov: bool = False - # Checkpoint activations to allow for training with larger models, sequences, and batch - # sizes. It is supported at two granularities 1) full: whole transformer layer is - # recomputed, 2) selective: submodules set in --recompute-modules are recomputed, default - # is core_attn. - recompute_granularity: Literal["full", "selective"] | None = None + # Scale mode for Muon optimizer. With MuP, set --muon-scale-mode unit_rms_norm to use + # unit_rms_norm scaling, or set --muon-scale-mode spectral to keep spectral scaling. + muon_scale_mode: Literal["spectral", "unit_rms_norm", "shape_scaling"] = "spectral" - # Check for NaNs in loss and grad - check_for_nan_in_loss_and_grad: bool = True + # FP32 matmul precision for Newton-Schulz iteration + muon_fp32_matmul_prec: Literal["low", "medium", "high"] = "medium" - # Check for spiky loss - check_for_spiky_loss: bool = False + # Newton-Schulz coefficient type for the Muon optimizer. Valid types are discovered from + # the installed emerging_optimizers package (e.g. simple, quintic, polar_express, aol). + # Validated at optimizer creation time. + muon_coefficient_type: str = "quintic" - # Check for unexpectedly large grads - check_for_large_grads: bool = False + # Number of Newton-Schulz steps for Muon optimizer + muon_num_ns_steps: int = 5 - # If set, distribute recomputed activations across model parallel group. - distribute_saved_activations: bool = False + # How to perform NS calculation for tensor model parallel weights + muon_tp_mode: Literal["blockwise", "duplicated", "distributed"] = "blockwise" - # 1) uniform: uniformly divide the total number of Transformer layers and recompute the - # input activation of each divided chunk at specified granularity, 2) recompute the input - # activations of only a set number of individual Transformer layers per pipeline stage and - # do the rest without any recomputing at specified granularitydefault) do not apply - # activations recompute to any layers - recompute_method: Literal["uniform", "block"] | None = None + # Additional scale factor for the muon update + muon_extra_scale_factor: float = 1.0 - # 1) uniform: the number of Transformer layers in each uniformly divided recompute unit, - # 2) block: the number of individual Transformer layers to recompute within each pipeline - # stage. - recompute_num_layers: int | None = None + # Optimizer for scalar parameters (embeddings, biases, norms) when using muon. Defaults to + # adam. + muon_scalar_optimizer: Literal["adam", "lion"] = "adam" - # The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", - # "mla_up_proj", "mlp", "moe", "shared_experts". default: - # ["core_attn"]."core_attn": recompute the core attention part of the transformer layer. - # "moe_act": recompute the MoE MLP activation function. "layernorm": recompute the - # input_layernorm and pre_mlp_layernorm. "mla_up_proj": recompute the MLA up projection - # and RoPE applying parts."mlp": recompute the dense MLP layer."moe": recompute the MoE - # layer."shared_experts": recompute the shared experts in the MoE layer."moe_act", - # "layernorm", and "mla_up_proj" use output-discarding checkpointing, "core_attn", "mlp", - # "moe", and "shared_experts" use normal checkpointing. - recompute_modules: list[str] | None = None + # First beta coefficient for Lion optimizer (used in sign update). Default: 0.95. + lion_beta1: float = 0.95 - # The number of Transformer layers to offload to CPU. - cpu_offloading_num_layers: int = 0 + # Second beta coefficient for Lion optimizer (used in momentum EMA update). Default: 0.98. + lion_beta2: float = 0.98 - # If not set, clone the output of the scatter in embedding layer to GC original tensor. - clone_scatter_output_in_embedding: bool = True + # Type of no weight decay condition. Choices: None (default): apply weight decay to 1D + # weights and biases."apply_wd_to_qk_layernorm": additionally apply weight decay to qk + # layernorm as a special case.DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. + no_weight_decay_cond_type: Literal["apply_wd_to_qk_layernorm"] | None = None # 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 @@ -286,20 +1046,20 @@ class MegatronConfig(ConfigInterface): # Global step to stop profiling. profile_step_end: int = 12 - # List of iterations to skip, empty by default. - iterations_to_skip: list[int] = field(default_factory=lambda: []) + # Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard. + use_pytorch_profiler: bool = False - # Optional name of file tracking `result_rejected` events. - result_rejected_tracker_filename: str | None = None + # Collect tensor shape in pytorch profiler. + pytorch_profiler_collect_shapes: bool = False - # Disables creation and usage of Gloo process groups. - enable_gloo_process_groups: bool = True + # Collect callstack in pytorch profiler. + pytorch_profiler_collect_callstack: bool = False - # Use the built-in pytorch profiler. Useful if you wish to view profiles in tensorboard. - use_pytorch_profiler: bool = False + # Collect chakra trace in pytorch profiler. + pytorch_profiler_collect_chakra: bool = False # Global ranks to profile. - profile_ranks: list[int] = field(default_factory=lambda: [0]) + profile_ranks: list[int] = field(default_factory=lambda: []) # Record memory history in last rank. record_memory_history: bool = False @@ -307,83 +1067,125 @@ class MegatronConfig(ConfigInterface): # Specifies where to dump the memory history pickle. memory_snapshot_path: str = "snapshot.pickle" - # Enables the overlap of Tensor parallel communication and GEMM kernels. - tp_comm_overlap: bool = False - - # Config file when tp_comm_overlap is enabled. - tp_comm_overlap_cfg: str | None = None - - # Disables the All-Gather overlap with GEMM by pipelining the GEMM and All-Gather. - tp_comm_overlap_ag: bool = True - - # Disables the Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter. - tp_comm_overlap_rs: bool = True - - # Enables the Reduce-Scatter overlap with dgrad GEMM. - tp_comm_overlap_rs_dgrad: bool = False + # Record shapes of tensors in `torch.autograd.profiler.emit_nvtx` for the Nsys profiler. + record_shapes: bool = False - # Disables the All-Gather overlap with bprop activation gradient GEMM. - tp_comm_bulk_dgrad: bool = True + # Enable NVTX range annotations for profiling. When enabled, inserts NVTX markers to + # categorize execution in profiler output. + nvtx_ranges: bool = False - # Disables the Reduce-Scatter overlap with bprop weight gradient GEMM. - tp_comm_bulk_wgrad: bool = True + # Batch size per model instance (local batch size). Global batch size is local batch size + # times data parallel size times number of micro batches. + micro_batch_size: int | None = None - # Set the bootstrapping backend of Tensor parallel communications. - tp_comm_bootstrap_backend: Literal["nccl", "mpi", "gloo"] = "nccl" + # Training batch size. If set, it should be a multiple of micro-batch-size times data- + # parallel-size. If this value is None, then use micro-batch-size * data-parallel-size as + # the global batch size. This choice will result in 1 for number of micro-batches. + global_batch_size: int | None = None - # If set, initialize weights on the CPU. This eliminates init differences based on tensor - # parallelism. - use_cpu_initialization: bool | None = None + # Batch size ramp up with the following values: , , 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. Deprecated. Use step_batch_size_schedule + # instead. + rampup_batch_size: int | None = None + + # Step-wise batch size schedule in format "THRESHOLD:BS THRESHOLD:BS ...". Thresholds + # support suffixes: K (1e3), M (1e6), B (1e9), T (1e12). If sequence length is provided, + # thresholds are interpreted as tokens; otherwise as samples. Example: + # step_batch_size_schedule = "0:768 250B:1536 500B:3072 750B:6144" Cannot be used together + # with decrease_batch_size_if_needed. + step_batch_size_schedule: str | None = None + + # If set, decrease batch size if microbatch_size * dp_size does not divide batch_size. Old + # batch_size will be restored if training is re-started with dp_size that divides + # batch_size // microbatch_size. Not supported with step-batch-size-schedule. + decrease_batch_size_if_needed: bool = False # Call torch.cuda.empty_cache() each iteration (training and eval), to reduce - # fragmentation.0=off, 1=moderate, 2=aggressive. + # fragmentation. 0=off, 1=moderate, 2=aggressive. empty_unused_memory_level: Literal[0, 1, 2] = 0 - # Choose code that has deterministic execution. This usually means slower execution, but - # is good for debugging and testing. - deterministic_mode: bool = False - # Interval to check weight hashes are same across DP replicas. If not specified, weight # hashes not checked. check_weight_hash_across_dp_replicas_interval: int | None = None - # 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. - calculate_per_token_loss: bool = False + # Interval (in iterations) to run GPU performance sniff tests (GEMMs, all-to-all, + # send/recv). Flags any rank whose throughput differs from the mean by more than one + # standard deviation. If not specified, sniff tests are not run. + gpu_sniff_test_interval: int | None = None # Training CPU-GPU synchronization interval, to ensure that CPU is not running too far # ahead of GPU. train_sync_interval: int | None = None - # Checkpoint activation to allow for training with larger models, sequences, and batch - # sizes. - checkpoint_activations: bool = False - - # Total number of iterations to train over all training runs. Note that either train-iters - # or train-samples should be provided. + # Total number of iterations to train over all training runs. Note that either train_iters + # or train_samples should be provided. train_iters: int | None = None - # Total number of samples to train over all training runs. Note that either train-iters or - # train-samples should be provided. + # Total number of samples to train over all training runs. Note that either train_iters or + # train_samples should be provided. train_samples: int | None = None - # Report loss and timing interval. - log_interval: int = 100 - # Exit the program after the iteration is divisible by this value. exit_interval: int | None = None # Exit the program after this many minutes. exit_duration_in_mins: int | None = None - # Dynamically save the checkpoint and shutdown the training if signal is received + # Dynamically save the checkpoint and shutdown the training if SIGTERM is received exit_signal_handler: bool = False - # Signal to use for exit signal handler. If not specified, defaults to SIGTERM. - exit_signal: Literal["SIGTERM", "SIGINT", "SIGUSR1", "SIGUSR2"] = "SIGTERM" + # Signal for the signal handler to detect. + exit_signal: str = "15" - # Write TensorBoard logs to this directory. - tensorboard_dir: str | None = None + # Use signal handler for dataloader workers + exit_signal_handler_for_dataloader: bool = False + + # Shutdown the training when SIGINT or SIGTERM received to avoid unclear traceback + exit_signal_handler_for_training: bool = False + + # 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. + manual_gc: bool = False + + # Training step interval to trigger manual garbage collection. Values > 0 will trigger + # garbage collections between training steps. + manual_gc_interval: int = 0 + + # When using manual garbage collection, this controls garbage collection at the start and + # the end of each evaluation run. + manual_gc_eval: bool = True + + # List of 1-indexed iterations to skip during training, empty by default. + iterations_to_skip: list[int] = field(default_factory=lambda: []) + + # Old batch size parameter, do not use. Use --micro-batch-size instead + batch_size: int | None = None + + # recompute activation to allow for training with larger models, sequences, and batch + # sizes. + recompute_activations: bool = False + + # Check for NaNs in loss and grad + check_for_nan_in_loss_and_grad: bool = True + + # Check for unexpectedly large grads + check_for_large_grads: bool = False + + # Optional name of file tracking `result_rejected` events. + result_rejected_tracker_filename: str | None = None + + # Config file when tp_comm_overlap is enabled. + tp_comm_overlap_cfg: str | None = None + + # Checkpoint activation to allow for training with larger models, sequences, and batch + # sizes. + checkpoint_activations: bool = False # Disable fusion of query_key_value scaling, masking, and softmax. masked_softmax_fusion: bool = True @@ -394,9 +1196,6 @@ class MegatronConfig(ConfigInterface): # Disable bias and swiglu fusion, the fusion is available only when using megatron-core. bias_swiglu_fusion: bool = True - # Use fused weighted squared relu when using MoE. - use_fused_weighted_squared_relu: bool = False - # Disable bias and dropout fusion. bias_dropout_fusion: bool = True @@ -407,27 +1206,19 @@ class MegatronConfig(ConfigInterface): # rope by default. rope_type: Literal["rope", "yarn"] | None = None - # Enabled fusion of cross entropy loss calculation. - cross_entropy_loss_fusion: bool = False - - # Implementation of cross entropy loss calculation. - cross_entropy_fusion_impl: Literal["native", "te"] = "native" - # use FlashAttention implementation of attention. https://arxiv.org/abs/2205.14135 use_flash_attn: bool = False - # Disable bias in the linear layers - add_bias_linear: bool = True - - # Enable bias only in the QKV linear layers - add_qkv_bias: bool = False - - # Optimizer function - optimizer: Literal["adam", "sgd"] = "adam" + # Optimizer function. Note: dist_muon is deprecated; use --optimizer muon with --use- + # distributed-optimizer instead. + optimizer: Literal["adam", "sgd", "muon", "dist_muon", "lion", "soap", "adaptive_muon"] = "adam" # Offload optimizer state to CPU optimizer_cpu_offload: bool = False + # Enable CUDA graph for optimizer step + optimizer_cuda_graph: bool = False + # Ratio of optimizer state to offload to CPU optimizer_offload_fraction: float = 1.0 @@ -450,57 +1241,17 @@ class MegatronConfig(ConfigInterface): # Single pass vs multiple pass data loader dataloader_type: Literal["single", "cyclic", "external"] | None = None - # DEPRECATED. This flag is ignored. - async_tensor_model_parallel_allreduce: bool = True - # 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. no_persist_layer_norm: bool = False - # Enable sequence parallel optimization. - sequence_parallel: bool = False - # Disable fusing gradient accumulation to weight gradient computation of linear layers gradient_accumulation_fusion: bool = True # DEPRECATED. Use the implementation from megatron core.Now ignored and mcore models are - # the default, use --use-legacy-models to not use core models. + # the default. deprecated_use_mcore_models: bool = False - # Use the legacy Megatron models, not Megatron-Core models. - use_legacy_models: bool = False - - # 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. - manual_gc: bool = False - - # Training step interval to trigger manual garbage collection. When the value is set to 0, - # garbage collection is not triggered between training steps. - manual_gc_interval: int = 0 - - # When using manual garbage collection, disable garbage collection at the start and the - # end of each evaluation run. - manual_gc_eval: bool = True - - # Disables the All-Gather overlap with fprop GEMM. - tp_comm_split_ag: bool = True - - # Disables the Reduce-Scatter overlap with fprop GEMM. - tp_comm_split_rs: bool = True - - # Select a communicator backend for pipeline parallel communication. If None, the default - # backend will be used. - pipeline_model_parallel_comm_backend: Literal["nccl", "ucc"] | None = None - - # The communicator group names to use high priority streams. - high_priority_stream_groups: list[str] = field(default_factory=lambda: []) - - # Use activation function kernel from Transformer Engine in MLP module. - use_te_activation_func: bool = False - # Use the RL training step. perform_rl_step: bool = False @@ -514,6 +1265,21 @@ class MegatronConfig(ConfigInterface): # Number of samples per a GRPO group. grpo_group_size: int = 2 + # Number of rollouts being generated by the inference engine simultaneously. Internally + # divided by grpo_group_size. Requires --rl-partial-rollouts. Mutually exclusive with + # --rl-num-parallel-generation-batches. + rl_num_parallel_generations: int | None = None + + # Number of generation batches in flight. Set to L+1 to allow for L steps of staleness + # between the inference and training policies. Each batch contains grpo_prompts_per_step + # groups by default. Requires --rl-partial-rollouts. Mutually exclusive with --rl-num- + # parallel-generations. + rl_num_parallel_generation_batches: int | None = None + + # Override the number of groups per generation batch. Defaults to grpo_prompts_per_step + # when --rl-num-parallel-generation-batches is set. + rl_generation_batch_size: int | None = None + # Number of iterations per a GRPO implementation. grpo_iterations: int = 2 @@ -532,38 +1298,34 @@ class MegatronConfig(ConfigInterface): # Filter groups with same reward. grpo_filter_groups_with_same_reward: bool = False + # Path to YAML config file for RL environment configuration. + langrl_env_config: str | None = None + # Default temperature for model inference. - grpo_default_temperature: float = 1.0 + rl_default_temperature: float = 1.0 # Default top-p for model inference. - grpo_default_top_p: float = 0 - - # Type of inference server to use. - langrl_inference_server_type: Literal["inplace_megatron", "inplace_megatron_chat"] = ( - "inplace_megatron" - ) - - # Conversation template, if using a chat server. - langrl_inference_server_conversation_template: str | None = None - - langrl_external_server: bool = False + rl_default_top_p: float = 0 - # Path to YAML config file for RL environment configuration. - langrl_env_config: str | None = None + # Default top-k for model inference. + rl_default_top_k: int = -1 # Offload optimizer state to CPU during inference/rollout to save GPU memory rl_offload_optimizer_during_inference: bool = False - # Offload KV cache to CPU during training to save GPU memory - rl_offload_kv_cache_during_training: bool = False + # KV cache management mode during RL training: persist: leave KV cache in GPU memory + # (default), offload: offload KV cache to CPU during training, recompute: deallocate KV + # cache and recompute from scratch each cycle + rl_kv_cache_management_mode: Literal["persist", "offload", "recompute"] = "persist" - # Remove KV cache during training to save GPU memory - rl_remove_kv_cache_during_training: bool = False + # Persist CUDA graphs when the inference engine is suspended. If False, CUDA graphs are + # deleted on suspend and re-captured on resume. + rl_persist_cuda_graphs: bool = False - # Reset CUDA graphs between inference/training to save GPU memory - rl_reset_cuda_graphs: bool = False - - # If set, use partial rollouts. + # Allow inference to continue generating rollouts while training updates the policy + # weights. This enables off-policy training where rollouts may be generated with a stale + # version of the policy. Use --rl-num-parallel-generations or --rl-num-parallel- + # generation-batches to control the degree of staleness. rl_partial_rollouts: bool = False # If set, use inference logprobs in importance sampling correction of the loss. @@ -573,61 +1335,99 @@ class MegatronConfig(ConfigInterface): # truncation for the IS correction at GRPO loss. rl_importance_sampling_truncation_coef: float | None = None - # If set, calculate the intra-group similarity of rollouts. - rl_calculate_intra_group_similarity: bool = False - # Enable sequence packing rl_use_sequence_packing: bool = False - # Override bin size for sequence packing. - rl_sequence_packing_bin_size: int = 8192 + # Maximum number of sequences that can be packed into a single bin. + rl_sequence_packing_max_sequences_per_bin: int = 50 # Algorithm for distributing packed bins across ranks. fifo: first-in-first-out sequential # distribution, round-robin: distribute bins cyclically across ranks for better load # balancing rl_sequence_packing_algo: Literal["fifo", "round-robin"] = "fifo" + # If set, do not toggle CUDA graphs on/off between inference and training phases. + rl_training_cuda_graphs: bool = False + + # Degree of tensor model parallelism for inference for RL. + rl_inference_tensor_model_parallel_size: int | None = None + + # Degree of pipeline model parallelism for inference for RL. + rl_inference_pipeline_model_parallel_size: int | None = None + + # Degree of expert model parallelism for inference for RL. + rl_inference_expert_model_parallel_size: int | None = None + + # Degree of expert tensor model parallelism for inference for RL. For MoE models, this + # controls the TP size for expert layers specifically. Defaults to training + # expert_tensor_parallel_size if not specified. + rl_inference_expert_tensor_model_parallel_size: int | None = None + + # Allocate the separate RL inference model parameters from a unified virtual memory (UVM) + # CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the + # inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe + # device pointers. + rl_inference_model_unified_memory_level: Literal[0, 1] = 0 + + # When using a separate RL inference model, offload its weights to CPU when not doing + # rollout inference, and restore to GPU right before inference. Works with two backends: + # 1) UVM (when --rl-inference-model-unified-memory-level=1), or 2) torch_memory_saver + # (when UVM is not enabled; requires torch_memory_saver to be installed). + rl_offload_inference_model_weights_when_idle: bool = False + + # Method to refit the model weights between training and inference models during RL. nccl: + # use NCCLCopyService to refit using NCCL; gloo: use GlooCopyService over CPU; nvshmem: + # use NVSHMEMCopyService to refit using the NVSHMEM. + refit_method: Literal["nccl", "gloo", "nvshmem"] = "gloo" + + # If set, verify that the model weights were correctly transferred by comparing forward + # pass outputs onthe first swap of model weights. + rl_verify_model_weights_swap: bool = False + + # Deprecated: use --rl-num-parallel-generations instead. + rl_parallel_generation_tasks: int | None = None + + # Skip BOS token at the beginning of the sequences. Default is False. + rl_skip_bos_token: bool = False + + # List of response parsers to enable for RL inference (e.g. --rl-inference-parsers + # deepseek-r1-reasoning qwen3-coder-tool). + rl_inference_parsers: list[str] = field(default_factory=lambda: []) + # Random seed used for python, numpy, pytorch, and cuda. seed: int = 1234 - # Enable random initialization of params across data parallel ranks - data_parallel_random_init: bool = False + # Use the Transformer Engine version of the random number generator. Required for CUDA + # graphs support. + te_rng_tracker: bool = False - # Standard deviation of the zero mean normal distribution used for weight initialization. - init_method_std: float = 0.02 + # Use a random number generator configured for inference. + inference_rng_tracker: bool = False - # Standard deviation of the zero mean normal distribution used for embedding weight - # initialization. If unset, embeddings will be initialized the same way as other weights. - # Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to - # any value will also skip applying weight decay on embedding weights to avoid shrinkage - # towards zero. See https://arxiv.org/abs/2312.16903 for more details. - embedding_init_method_std: float | None = None + # Enable random initialization of params across data parallel ranks + data_parallel_random_init: bool = False # Enable Xavier uniform parameter initialization init_method_xavier_uniform: bool = False - # Initial learning rate. Depending on decay style and initial warmup, the learning rate at - # each iteration would be different. - lr: float | None = None - # Learning rate decay function. lr_decay_style: Literal["constant", "linear", "cosine", "inverse-square-root", "WSD"] = "linear" # Decay style for the annealing phase of WSD lr_wsd_decay_style: Literal["exponential", "linear", "cosine", "minus_sqrt"] = "exponential" - # number of iterations to decay learning rate over, If None defaults to `--train-iters` + # number of iterations to decay learning rate over, If None defaults to train iters lr_decay_iters: int | None = None - # number of samples to decay learning rate over, If None defaults to `--train-samples` + # number of samples to decay learning rate over, If None defaults to train samples lr_decay_samples: int | None = None - # number of samples for the annealing phase in the wsd schedule - lr_wsd_decay_samples: int | None = None - # number of iterations for the annealing phase in the wsd schedule lr_wsd_decay_iters: int | None = None + # number of samples for the annealing phase in the wsd schedule + lr_wsd_decay_samples: int | None = None + # fraction of lr-warmup-(iters/samples) to use for warmup (as a float) lr_warmup_fraction: float | None = None @@ -640,22 +1440,35 @@ class MegatronConfig(ConfigInterface): # Initial value for learning rate warmup. The scheduler starts warmup from this value. lr_warmup_init: float = 0.0 + # 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. + override_opt_param_scheduler: bool = False + + # 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. + use_checkpoint_opt_param_scheduler: bool = False + + # Initial weight decay coefficient for L2 regularization. + start_weight_decay: float | None = None + + # End of run weight decay coefficient for L2 regularization. + end_weight_decay: float | None = None + + # Weight decay increment function. + weight_decay_incr_style: Literal["constant", "linear", "cosine"] = "constant" + + # Initial learning rate. Depending on decay style and initial warmup, the learning rate at + # each iteration would be different. + lr: float | None = None + # Old lr warmup argument, do not use. Use one of the--lr-warmup-* arguments above warmup: int | None = None # Minimum value for learning rate. The schedulerclip values below this threshold. min_lr: float = 0.0 - # 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. Notethat all the above values will be reset. - override_opt_param_scheduler: bool = False - - # 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. - use_checkpoint_opt_param_scheduler: bool = False - # Separate learning rate for the input and output layer decoupled_lr: float | None = None @@ -669,42 +1482,42 @@ class MegatronConfig(ConfigInterface): # Number of iterations between persistent checkpoint saves. save_interval: int | None = None - # Specific training iterations at which to save checkpoints - save_extra_steps: list[int] = field(default_factory=lambda: []) - - # Number of iterations between retained checkpoints (othercheckpoints _except the last - # checkpoint_ are automatically deleted). - save_retain_interval: int | None = None + # Number of iterations between param.name->param.data mapping saves. + save_params_interval: int | None = None - # Do not save current optimizer. - no_save_optim: bool | None = None + # Number of iterations between act.name->act.data mapping saves. + save_activations_interval: int | None = None - # Do not save current rng state. - no_save_rng: bool | None = None + # Number of iterations between tokens-per-expert routing metadata saves. + save_tokens_per_expert_interval: int | None = None - # Directory containing a model checkpoint. - load: str | None = None + # Number of iterations between wgrad (main_grad) saves. + save_wgrads_interval: int | None = None - # Do not load optimizer when loading checkpoint. - no_load_optim: bool | None = None + # Number of iterations between dgrad saves. + save_dgrads_interval: int | None = None - # Load main parameters from checkpoint directly. - load_main_params_from_ckpt: bool | None = None + # Number of iterations between retained checkpoints (other checkpoints except the last + # checkpoint are automatically deleted). + save_retain_interval: int | None = None - # Do not load rng state when loading checkpoint. - no_load_rng: bool | None = None + # Directory containing a model checkpoint. + load: str | None = None - # Do not strict loading for fsdp_dtensor checkpoint format. - strict_fsdp_dtensor_load: bool = True + # Load main parameters from checkpoint. When loading a model from a checkpoint without + # loading the optimizer, the model parameters are updated but for fp16 optimizer with main + # parameters, the main parameters need to also be updated. + load_main_params_from_ckpt: bool = False # Number of iterations between non-persistent saves. non_persistent_save_interval: int | None = None # Type of non-persistent model checkpoints. "global" - Saved as a standard checkpoint - # (e.g., on Lustre) with old checkpoints being removed. "local" - Each rank saves a - # portion of the checkpoint locally (e.g., on SSD/ramdisk). None - No non-persistent + # (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). - non_persistent_ckpt_type: Literal["global", "local", "in_memory", None] = None + non_persistent_ckpt_type: Literal["global", "local", "in_memory"] | None = None # Directory containing global non-persistent model checkpoints. non_persistent_global_ckpt_dir: str | None = None @@ -725,10 +1538,6 @@ class MegatronConfig(ConfigInterface): # Checkpoint step to load model from. ckpt_step: int | None = None - # Do not perform initialization when building model, can reduce startup time when - # definitely loading from a checkpoint - perform_initialization: bool = True - # Override model-related command-line arguments with arguments from checkpoint use_checkpoint_args: bool = False @@ -738,60 +1547,95 @@ class MegatronConfig(ConfigInterface): # If set, do not use tokenizer model path from checkpoint use_tokenizer_model_from_checkpoint_args: bool = True - # If '--load' is set, but checkpoint is not found (e.g., path typo), then exit instead of + # If 'load' is set, but checkpoint is not found (e.g., path typo), then exit instead of # random initialization. exit_on_missing_checkpoint: bool = False - # Deprecated: see --ckpt-format. - use_dist_ckpt_deprecated: bool = False - - # Enables a persitent checkpoint worker for async save - use_persistent_ckpt_worker: bool = False + # Checkpoint format to use. torch is the format used by torch.save/load. torch_dist is a + # megatron built-in distributed checkpointing format. torch_dcp is the + # torch.distributed.checkpoint format. fsdp_dtensor is a torch DCP native, Megatron FSDP + # training-specific checkpoint format. + ckpt_format: Literal["torch", "torch_dist", "torch_dcp", "fsdp_dtensor"] = "torch_dist" # 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). auto_detect_ckpt_format: bool = False - # Deprecated: see --ckpt-format. - dist_ckpt_format_deprecated: Any | None = None - - # Checkpoint format to use. torch is the format used by torch.save/load. torch_dist is a - # megatron built-in distributed checkpointing format. torch_dcp is the - # torch.distributed.checkpoint format. fsdp_dtensor is a torch DCP native, Megatron FSDP - # training-specific checkpoint format. - ckpt_format: Literal["torch", "torch_dist", "zarr", "torch_dcp", "fsdp_dtensor"] = "torch_dist" - # Checkpoint format for conversion. - ckpt_convert_format: Literal["torch", "torch_dist", "zarr"] | None = None + ckpt_convert_format: Literal["torch", "torch_dist"] | None = None # Save directory for converted checkpoint. - ckpt_convert_save: Any | None = None + ckpt_convert_save: str | None = None # 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. ckpt_convert_update_legacy_dist_opt_format: bool = False - # Deprecated: see --no-ckpt-fully-parallel-save. - ckpt_fully_parallel_save_deprecated: bool = False - # 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. ckpt_fully_parallel_save: bool = True - # Apply async checkpointing save. Currently works only with`torch_dist` distributed + # Apply async checkpointing save. Currently works only with `torch_dist` distributed # checkpoint format. - async_save: bool | None = None + async_save: bool = False + + # Which async save strategy to use. Available strategies: nvrx, mcore. + async_strategy: Literal["nvrx", "mcore"] = "nvrx" + + # Use a persistent background worker for async checkpoint saves. When enabled, creates a + # dedicated worker thread/process for handling async saves. When disabled, uses temporal + # workers that are created and destroyed for each save operation. + use_persistent_ckpt_worker: bool = False + + # CPU nice value target (0-19, higher = lower priority) for the async checkpoint writer + # process. If it exceeds 19, it will be set to 19. If the current nice value is greater + # than the target, it will be left unchanged. Only applies when using persistent ckpt + # worker. + async_ckpt_cpu_priority: int = 10 + + # I/O scheduling class (0-3, 3=idle) for the async checkpoint writer process. + async_ckpt_io_priority: int = 3 + + # Copy GPU tensors to CPU shared-memory in the training process before handing off to the + # async checkpoint worker. Avoids CUDA IPC / NVLink fabric handles in the worker + # subprocess. Useful on MNNVL systems where fabric resources are exhausted. Only applies + # with the nvrx async strategy. + async_ckpt_use_cpu_shm: bool = False # Apply full load parallelization across DP for distributed checkpoints. ckpt_fully_parallel_load: bool = False - # If the model and optimizer state dict structure isconstant throughout a *single training - # job*, it allows fordifferent checkpointing performance optimizations. + # Algorithm for fully parallel load of distributed checkpoints. "broadcast"(default): + # Broadcast the checkpoint from rank 0 to all other ranks. "gather_rounds": Gather the + # checkpoint from all ranks in rounds. "gather_object": Gather the checkpoint from all + # ranks in a single operation. + ckpt_fully_parallel_load_exchange_algo: Literal[ + "broadcast", "gather_rounds", "gather_object" + ] = "broadcast" + + # Process group for fully parallel save of distributed checkpoints. "dp"(default): Data + # parallel process group. "ep_dp": Expert data parallel process group. + ckpt_fully_parallel_save_process_group: Literal["dp", "ep_dp"] = "dp" + + # Process group for fully parallel load of distributed checkpoints. "dp"(default): Data + # parallel process group. "ep_dp": Expert data parallel process group. + ckpt_fully_parallel_load_process_group: Literal["dp", "ep_dp"] = "dp" + + # Assume the checkpoint structure is constant across saves to enable optimizations. ckpt_assume_constant_structure: bool = False + # Whether to validate sharding access integrity when loading a distributed checkpoint. + # When True (default), each tensor shard is checked to be accessed exactly once as main + # replica by some rank. Disabling skips this validation + ckpt_load_validate_sharding_integrity: bool = True + + # Whether to enforce strict loading for FSDP DTensor checkpoints. When False, allows + # partial loading. + strict_fsdp_dtensor_load: bool = True + # 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. @@ -817,15 +1661,53 @@ class MegatronConfig(ConfigInterface): # During distributed optimizer checkpoint save and load tries to use as little memory as # possible by using Gloo (instead of NCCL) and only one rank for saving. Turn on only if - # experiencing host or device memory issues. Has affect only with `--dist-ckpt-optim- - # fully-reshardable` flag. + # experiencing host or device memory issues. Has affect only with + # `dist_ckpt_optim_fully_reshardable` flag. distrib_optim_fully_reshardable_mem_efficient: bool = False - # Run model in fp16 mode. - fp16: bool = False + # If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks. + replication: bool = False - # Run model in bfloat16 mode. - bf16: bool = False + # Specifies `J`, the spacing between ranks storing replicas of a given rank's data. + # Replicas for rank `n` may be on ranks `n+J`, `n+2J`, ..., or `n-J`, `n-2J`, etc. This + # flag has an effect only if --replication is used. and must be consistent across all + # ranks. + replication_jump: int | None = None + + # Number of machines storing the replica of a given rank's data. + replication_factor: int = 2 + + # Whether to hash checkpointing files during save and validate their integrity during + # load. + verify_integrity: bool = False + + # Specific training iterations at which to save checkpoints + save_extra_steps: list[int] = field(default_factory=lambda: []) + + # Do not save current optimizer. + no_save_optim: bool | None = None + + # Do not save current rng state. + no_save_rng: bool | None = None + + # Do not load optimizer when loading checkpoint. + no_load_optim: bool | None = None + + # Do not load rng state when loading checkpoint. + no_load_rng: bool | None = None + + # Deprecated: see --ckpt-format. + use_dist_ckpt_deprecated: bool = False + + # Deprecated: see --ckpt-format. + dist_ckpt_format_deprecated: Any | None = None + + # Number of workers for distributed checkpointing. Only used for async save. If set to 1, + # the checkpointing is performed in a single process. + dist_ckpt_workers: int = 1 + + # Deprecated: see --no-ckpt-fully-parallel-save. + ckpt_fully_parallel_save_deprecated: bool = False # Reduce gradients in bfloat16. grad_reduce_in_bf16: bool = False @@ -846,13 +1728,6 @@ class MegatronConfig(ConfigInterface): # hysteresis for dynamic loss scaling hysteresis: int = 2 - # Move residual connections to fp32. - fp32_residual_connection: bool = False - - # Scale Q * K^T by 1 / layer-number. Useful for fp16 training. Also sets - # `attention_softmax_in_fp32` to True. - apply_query_key_layer_scaling: bool = False - # Run attention masking and softmax in fp32. attention_softmax_in_fp32: bool = False @@ -862,18 +1737,88 @@ class MegatronConfig(ConfigInterface): # Move the cross entropy unreduced loss calculationfor lm head to fp16. fp16_lm_cross_entropy: bool = False - # If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to - # prevent matmul from using reduced precision accumulation when using BF16. - disable_bf16_reduced_precision_matmul: bool = False - # If True, reuse the grad buffer for MXFP8 parameter all-gather. reuse_grad_buf_for_mxfp8_param_ag: bool = False - # Degree of tensor model parallelism. - tensor_model_parallel_size: int = 1 + # Which backend to use for distributed training. + distributed_backend: Literal["nccl", "gloo"] = "nccl" - # Degree of pipeline model parallelism. - pipeline_model_parallel_size: int = 1 + # Timeout minutes for torch.distributed. + distributed_timeout_minutes: int = 10 + + # If not set, all PP stages will launch gradient reduces simultaneously. Otherwise, each + # PP stage will independently launch as needed. + align_grad_reduce: bool = True + + # local rank passed from distributed launcher. + local_rank: int = 0 + + # 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. + lazy_mpu_init: bool = False + + # Use Megatron's Fully Sharded Data Parallel. Cannot be used together with + # use_torch_fsdp2. + use_megatron_fsdp: bool = False + + # Use the torch FSDP2 implementation. FSDP2 is not currently working with Pipeline + # Parallel. It is still not in a stable release stage, and may therefore contain bugs or + # other potential issues. + use_torch_fsdp2: bool = False + + # 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`. + nccl_communicator_config_path: str | None = None + + # If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep- + # pp-dp. + use_tp_pp_dp_mapping: bool = False + + # If enabled, create Gloo process groups for communications. + use_gloo_process_groups: bool = True + + # Set the use of SHARP for the collective communications of data-parallel process groups. + # When `True`, run barrier within each data-parallel process group, which specifies the + # SHARP application target groups. + use_sharp: bool = False + + # IB SHARP can be enabled from only one communication group. By default, it is enabled + # from dp group if not specified and use_sharp=True. Available options: [dp, dp_replica] + sharp_enabled_group: Literal["dp", "dp_replica"] | None = None + + # Specify which communicator groups should use high priority streams during creation. + # Assigning high priority to communication streams ensures that communication kernels are + # scheduled with higher priority, minimizing the exposed communication when it is + # overlapped with other computation kernels. + high_priority_stream_groups: list[str] = field(default_factory=lambda: []) + + # Timeout in seconds for process groups after initialization. This timeout is applied to + # all process groups after initialization and the first iteration completes. + distributed_timeout_seconds_after_init: int | None = None + + # Path for NCCL flight recorder trace dumps. Sets TORCH_FR_DUMP_TEMP_FILE and + # TORCH_NCCL_DEBUG_INFO_TEMP_FILE env variables before distributed init. + flight_recorder_dump_path: str | None = None + + # Size of the NCCL flight recorder trace buffer (TORCH_NCCL_TRACE_BUFFER_SIZE). + flight_recorder_trace_buffer_size: int = 2000 + + # Dump flight recorder traces on NCCL timeout (TORCH_NCCL_DUMP_ON_TIMEOUT). + flight_recorder_dump_on_timeout: bool = True + + # Include stack traces in flight recorder dumps (TORCH_INCLUDE_STACK_TRACE). + flight_recorder_include_stack_trace: bool = False + + # Include only active operations in flight recorder dumps (TORCH_INCLUDE_ONLY_ACTIVE). + flight_recorder_include_only_active: bool = True + + # Enable extra flight recorder dump on execution (TORCH_NCCL_EXTRA_DUMP_ON_EXEC). + flight_recorder_extra_dump_on_exec: bool = True + + # Disable the JIT fuser. + disable_jit_fuser: bool = False # 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 @@ -899,41 +1844,12 @@ class MegatronConfig(ConfigInterface): # Number of virtual pipeline stages per pipeline parallelism rank num_virtual_stages_per_pipeline_rank: int | None = None - # Number of contiguous microbatches per virtual pipeline stage - microbatch_group_size_per_vp_stage: int | None = None - # overlap pipeline parallel communication with forward and backward chunks in 1F1B overlap_p2p_comm: bool = True - # if set, overlap pipeline parallel communication in warmup and flush - overlap_p2p_comm_warmup_flush: bool = False - - # Which backend to use for distributed training. - distributed_backend: Literal["nccl", "gloo"] = "nccl" - - # Default timeout minutes for torch.distributed. - distributed_timeout_minutes: int = 10 - - # Timeout seconds for process groups after initialization.This timeout is applied to all - # process groups after initialization. - distributed_timeout_seconds_after_init: int | None = None - # If set, overlap DDP grad reduce. overlap_grad_reduce: bool = False - # If set, defers the vocabulary projection linear layer weightgradient compute to pipeline - # flush. - defer_embedding_wgrad_compute: bool = False - - # Number of micro-batches for whichweight gradient computation of vocabulary projection is - # deferred, defaults to 0 whichmeans all the micro-batches are deferred. Invalid if - # `defer-embedding-wgrad-compute`is not set - wgrad_deferral_limit: int = 0 - - # If not set, all PP stages will launch gradient reduces simultaneously. Otherwise, each - # PP stage will independently launch as needed. - align_grad_reduce: bool = True - # Number of buckets for data-parallel communication ddp_num_buckets: int | None = None @@ -951,6 +1867,12 @@ class MegatronConfig(ConfigInterface): # ring implementation) but performs accumulation locally in FP32. ddp_reduce_scatter_with_fp32_accumulation: bool = False + # List of param_name patterns (in Python's fnmatch format) to match against to do local + # gradient accumulation in FP32. The special pattern 'all' matches every parameter. + ddp_param_name_patterns_for_fp32_local_accumulation: list[str] = field( + default_factory=lambda: [] + ) + # If set, average directly in data-parallel communication collective. ddp_average_in_collective: bool = False @@ -964,32 +1886,18 @@ class MegatronConfig(ConfigInterface): # PP stage will independently launch as needed. align_param_gather: bool = True - # If not set, use scatter/gather to optimize communication of tensors in pipeline. - scatter_gather_tensors_in_pipeline: bool = True - - # 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. - use_ring_exchange_p2p: bool = False - - # local rank passed from distributed launcher. - local_rank: int = 0 - - # 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. - lazy_mpu_init: bool | None = None - - # If set, *input* embedding layer will be treated as a standard transformerlayer in the - # context of partition and placement for pipeline parallelism. - account_for_embedding_in_pipeline_split: bool = False - - # If set, loss layer will be treated as a standard transformerlayer in the context of - # partition and placement for pipeline parallelism. - account_for_loss_in_pipeline_split: bool = False - # Use distributed optimizer. use_distributed_optimizer: bool = False + # Opt out of the precomputed LayerWise param layout. When set, falls back to the legacy + # LayerWise ping-pong path: all params (including non-Muon embeddings, biases, layernorm) + # live in a single LayerWise buffer and the optimizer uses the allgather_params() + # codepath. The default (precomputed layout) routes non-Muon params through a separate + # DistributedOptimizer with byte-level sharding, which is faster and uses less padding but + # produces different bf16 reduction ordering and so will not match legacy-path loss curves + # bit-for-bit. + use_layer_wise_param_layout: bool = True + # Use the userbuffer registration for DP/FSDP communication buffers.This option will # reduce GPU SM usage for the DP/FSDP communication,which is improving the performance of # the overlapped computation. @@ -999,22 +1907,25 @@ class MegatronConfig(ConfigInterface): # will force to use conventional (local) userbuffer registration when use-nccl-ub is set. disable_symmetric_registration: bool = False - # Required to enable SHARP communication. - use_sharp: bool = False - - # IB SHARP can be enabled from only one communication group. By default, it is enabled - # from dp group. Available options: [dp, dp_replica] - sharp_enabled_group: Literal["dp", "dp_replica"] | None = None - - # Use the Megatron FSDP code path in DDP. - use_megatron_fsdp: bool = False + # Manually register the FSDP communication buffers to NCCL user buffer.This option is only + # effective when use-megatron-fsdp and use-nccl-ub is set. + fsdp_manual_registration: bool = False - init_model_with_meta_device: bool = False + # Create a separate process group for all-gather operations to overlap reduce-scatter and + # all-gather operations. + create_all_gather_group: bool = False # Sharding strategy of data parallelism. data_parallel_sharding_strategy: Literal[ "no_shard", "optim", "optim_grads", "optim_grads_params" - ] = "no_shard" + ] = "optim_grads_params" + + # Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) + # mode. Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid + # Sharding). The "optim" option is only supported when --data-parallel-sharding-strategy + # is "optim_grads_params". This option is only effective when Hybrid FSDP is enabled + # (i.e., when dp_outer_dim is not None). Default: "no_shard". + outer_dp_sharding_strategy: Literal["no_shard", "optim"] = "no_shard" # If not set, fuse the division in gradient reduce. gradient_reduce_div_fusion: bool = True @@ -1042,49 +1953,15 @@ class MegatronConfig(ConfigInterface): # Number of Distributed Optimizer copies across Data Parallel domain. num_distributed_optimizer_instances: int = 1 - # Use the torch FSDP2 implementation. FSDP2 has not been tested with pipeline parallelism, - # and may contain bugs. - use_torch_fsdp2: bool = False - # Whether to reshard weights after forward pass when using PyTorch FSDP2. Set to enable # FSDP ZeRO-2. torch_fsdp2_reshard_after_forward: bool = True - # Degree of context parallelism. - context_parallel_size: int = 1 - - # Inter-gpu communication type for context parallelism: p2p, a2a, allgather or a2a+p2p. If - # a single string is provided, all layers will share the same communication type. Users - # can also specify separated types for each layer like --cp-comm-type p2p p2p a2a a2a - # a2a+p2p a2a+p2p - cp_comm_type: list[str] = field(default_factory=lambda: ["p2p"]) - - # Degrees of the hierarchical context parallelism. Users should provide a list to specify - # the sizes for different levels. --hierarchical-context-parallel-sizes 2 4 indicates - # every two adjacent gpus forms the first level of cp groups and the cp ranks with the - # same odevity forms the second level of cp groups. - hierarchical_context_parallel_sizes: list[int] | None = None - - # 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`. - nccl_communicator_config_path: str | None = None - - # If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep- - # pp-dp. - use_tp_pp_dp_mapping: bool = False - - # If set, replication of local checkpoints is enabled. Needs to be enabled on all ranks. - replication: bool = False - - # Specifies `J`, the spacing between ranks storing replicas of a given rank's data. - # Replicas for rank `n` may be on ranks `n+J`, `n+2J`, ..., or `n-J`, `n-2J`, etc. This - # flag has an effect only if --replication is used. and must be consistent across all - # ranks. - replication_jump: int | None = None - - # Number of machines storing the replica of a given rank's data. - replication_factor: int = 2 + # Inter-gpu communication type for context parallelism: p2p, a2a, allgather or a2a+p2p. If + # a single string is provided, all layers will share the same communication type. Users + # can also specify separated types for each layer like --cp-comm-type p2p p2p a2a a2a + # a2a+p2p a2a+p2p + cp_comm_type: list[str] = field(default_factory=lambda: ["p2p"]) # If set, initialize with fake distributed process group and all distributed communication # operations will be skipped. This is quite useful for profiling @@ -1092,27 +1969,48 @@ class MegatronConfig(ConfigInterface): # WORLD_SIZE and RANK to the specific values for target distribtued scale. fake_process_group: bool = False + # Number of iterations to run for evaluation. Used for both validation and test. If not + # set, evaluation will not run. + eval_iters: int = 100 + + # Interval between running evaluation on validation set. If not set, evaluation will not + # run during training. + eval_interval: int | None = None + + # If set, evaluation will only start after this iteration number. Useful for skipping + # evaluation during early training iterations when the model is not yet meaningful. If not + # set, evaluation starts from the first eval_interval. + start_eval_at_iter: int | None = None + + # Global batch size to use during evaluation. If not set, defaults to global_batch_size. + # Must be divisible by (eval_micro_batch_size * data_parallel_size). + eval_global_batch_size: int | None = None + + # Micro batch size to use during evaluation. If not set, defaults to micro_batch_size. + # Changing this affects per-device memory usage during eval and the number of microbatches + # per eval step. + eval_micro_batch_size: int | None = None + + # If set, bypass the training loop, perform evaluation for validation/test, and exit. + skip_train: bool = False + + # Run all real-time test alongside the experiment. + test_mode: bool = False + # If set, each time validation occurs it uses the full validation dataset(s). This # currently only works for GPT datasets! full_validation: bool = False # If set, multiple datasets listed in the validation split are evaluated independently # with a separate loss for each dataset in the list. This argument requires that no - # weights are included in the list + # weights are included in the list. multiple_validation_sets: bool = False - # Number of iterations to run for evaluationvalidation/test for. - eval_iters: int = 100 - - # Interval between running evaluation on validation set. - eval_interval: int = 1000 - - # Run all real-time test alongside the experiment. - test_mode: bool = False - - # If set, bypass the training loop, optionally do evaluation for validation/test, and - # exit. - skip_train: bool = False + # Optional list of names for multiple validation sets. When provided with --multiple- + # validation-sets, these names are used instead of numeric indices (e.g. 'validation-wiki' + # instead of 'validation-0'). The number of names must match the number of validation + # datasets. + validation_set_names: list[str] | None = None # The weight and prefix list for a set of train, validation, and testdatasets which split # according to --split. The accepted formats are: (1) a single prefix, (2) a list of @@ -1121,6 +2019,10 @@ class MegatronConfig(ConfigInterface): # datasets. This argument is exclusive to the other independent --*-data-path arguments. data_path: list[str] | None = None + # Comma-separated list of iterations where phase transitions occur. Requires fixed global + # batch size across phases. + phase_transition_iterations: str | None = None + # Comma-separated list of proportions for training, validation, and test split. For # example the split `90,5,5` will use 90%% of data for training, 5%% for validation and # 5%% for test. @@ -1149,6 +2051,18 @@ class MegatronConfig(ConfigInterface): # `test` keys per_split_data_args_path: str | None = None + # Path to a json file with the sequences per dataset. Check the + # tools/build_sequences_per_dataset.py script to build this file. + per_dataset_sequences_path: Any | None = None + + # Option to use the fast cache loading path when building the datasets. Requires all the + # dataset caches to be built and stored in --data-cache-path. + dataloader_fast_cache_load: bool = False + + # Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all + # the dataset caches to be built and stored in --data-cache-path. + dataloader_defer_npy_index_mmap: bool = False + # Path to a directory to hold cached index files. data_cache_path: Any | None = None @@ -1168,9 +2082,6 @@ class MegatronConfig(ConfigInterface): # Maximum decoder sequence length to process. decoder_seq_length: int | None = None - # Maximum sequence length for the biencoder model for retriever - retriever_seq_length: int = 256 - # sample rate for training data. Supposed to be 0 < sample_rate < 1 sample_rate: float = 1.0 @@ -1256,13 +2167,17 @@ class MegatronConfig(ConfigInterface): # not provided, it will be automatically calculated from vocab-size. padded_vocab_size: int | None = None + # Whether to pad vocab size of the model automatically if padded_vocab_size is not + # provided. + pad_vocab_size: bool = True + # Path to the vocab file. vocab_file: str | None = None # Path to the BPE merge file. merge_file: str | None = None - # Number of additional vocabulary tokens. They are used for span masking in the T5 model + # Number of additional vocabulary tokens. They are used for span masking in the T5 model. vocab_extra_ids: int = 0 # What type of tokenizer to use. @@ -1284,28 +2199,47 @@ class MegatronConfig(ConfigInterface): | None ) = None - # Sentencepiece tokenizer model. + # Path to the tokenizer model. tokenizer_model: str | None = None - # Path to tokenizer metadata in json format. - tokenizer_metadata: str | None = None + # Path to the tokenizer metadata file in json format. + metadata_path: str | None = None + + # List of special tokens. For TikTokenizer needs to have ["", "", "", + # "", "", "", ""] + special_tokens: list[str] | None = None # Which tiktoken pattern to use. Options: [v1, v2] - tiktoken_pattern: str | None = None + tiktoken_pattern: Literal["v1", "v2"] | None = None - # Number of special tokens in tiktoken tokenizer + # Number of special tokens in tiktoken tokenizer. tiktoken_num_special_tokens: int = 1000 - # List of tiktoken special tokens, needs to have ["", "", "", "", - # "", "", ""] - tiktoken_special_tokens: list[str] | None = None + # SentencePiece tokenizer wrapper legacy behavior. Allows special tokens usage. + tokenizer_sentencepiece_legacy: bool = False - # To use legacy tokenizer system. - legacy_tokenizer: bool = False + # Whether to ignore extra whitespaces in the input text while encoding. + tokenizer_sentencepiece_ignore_extra_whitespaces: bool = True - # Whether or not to allow PreTrainedTokenizer to execute remote code + # Whether to use fast HuggingFace tokenizer. + tokenizer_hf_no_use_fast: bool = False + + # Converting text to ids will not include special for HuggingFace tokenizer. + tokenizer_hf_no_include_special_tokens: bool = False + + # Whether or not to allow PreTrainedTokenizer to execute remote code. trust_remote_code: bool = False + # EOD token id for NullTokenizer. Defaults to `vocab_size - 1`. + null_tokenizer_eod_id: int | None = None + + # Pad token id for NullTokenizer. Defaults to -1 (no pad token). Set to a value outside + # the dataset to avoid masking real tokens. + null_tokenizer_pad_id: int = -1 + + # Custom chat template in jinja format for conversation formatting. + chat_template: str | None = None + # Enable autoresume on adlr cluster. adlr_autoresume: bool = False @@ -1432,19 +2366,6 @@ class MegatronConfig(ConfigInterface): # warmup teacher temperaure epochs dino_warmup_teacher_temp_epochs: int = 30 - # Whether to layer normalize the q and k attention embeddings. - qk_layernorm: bool = False - - # Use llama 4 qk l2 norm - qk_l2_norm: bool = False - - # Degree of expert model parallelism. - expert_model_parallel_size: int = 1 - - # Degree of expert model parallelism. Default is None, which will be set to the value of - # --tensor-model-paralle-size. - expert_tensor_parallel_size: int | None = None - # Number of Experts in MoE (None means no MoE) num_experts: int | None = None @@ -1453,44 +2374,10 @@ class MegatronConfig(ConfigInterface): # containing a Python list expression that defines a custom pattern, e.g.: # "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] where 1 indicates an expert # layer and 0 indicates a dense layer. Examples: "([0]+[1]*23)": 1 dense layer followed by - # 23 experts layers, "([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, + # 23 expert layers, "([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, # repeated twice. moe_layer_freq: Any = 1 - # The hidden size of each expert's feed-forward network (ffn). If not specified, defaults - # to the ffn_hidden_size. - moe_ffn_hidden_size: int | None = None - - # Shared expert total ffn hidden size. It should be equal to "num_shared_experts * - # ffn_size_of_each_shared_expert" if there are multiple shared experts. None means no - # shared expert. By default, the shared experts execute before the router. However, when - # --moe-shared-expert-overlap or --overlap-moe-expert-parallel-comm is set, the shared - # experts execute after the router, before the routed experts. This makes the gradients - # from the router and the shared experts added in different orders to the hidden_states, - # causing minor numerical differences in the hidden_states gradient. - moe_shared_expert_intermediate_size: int | None = None - - # Enable overlapping between shared expert computations and dispatcher communications. - # Without this, the shared experts execute before the router. Only effective when moe- - # shared-expert-intermediate-size is set. - moe_shared_expert_overlap: bool = False - - # 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. - moe_grouped_gemm: bool = False - - # Use legacy GroupedMLP rather than TEGroupedMLP. Note: The legacy one will be deprecated - # soon. - moe_use_legacy_grouped_gemm: bool = False - - # Enable checkpointing for moe_layer, should be used when memory is not sufficient. - # Deprecated. Use "--recompute-granularity selective --recompute-modules moe" instead. - moe_layer_recompute: bool = False - - # Deprecated. Use --expert-tensor-parallel-size instead. - moe_extended_tp: bool = False - # 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. @@ -1505,136 +2392,15 @@ class MegatronConfig(ConfigInterface): "aux_loss", "seq_aux_loss", "global_aux_loss", "sinkhorn", "none" ] = "aux_loss" - # Data type for routing computation and expert output weighted averaging. Fp32/fp64 - # enhances numerical stability, especially with numerous experts. The perf impact should - # be negligible when used with permute fusion. None means no changes for dtype. - moe_router_dtype: Literal["fp32", "fp64"] | None = None - - # Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in - # TransformerEngine 2.7.0 and above. - moe_router_fusion: bool = False - - # Score function for MoE TopK routing. Can be "softmax" or "sigmoid". - moe_router_score_function: Literal["softmax", "sigmoid"] = "softmax" - - # Number of experts to route to for each token. The default is 2. - moe_router_topk: int = 2 - - # Enable pre-softmax routing for MoE, which means softmax is before the top-k selection. - # By default, softmax is done after top-k. - moe_router_pre_softmax: bool = False - - # Number of groups to divide experts into for group-limited routing. When using group- - # limited routing: 1) Experts are divided into equal-sized groups, 2) For each token, a - # subset of groups are selected based on routing scores (sum of top-2 expert scores within - # each group), 3) From these selected groups, moe_router_topk experts are chosen.Two - # common use cases: 1) Device-limited routing: Set equal to expert parallel size (EP) to - # limit each token to experts on a subset of devices (See DeepSeek-V2: - # https://arxiv.org/pdf/2405.04434) 2) Node-limited routing: Set equal to number of nodes - # in EP group to limit each token to experts on a subset of nodes (See DeepSeek-V3: - # https://arxiv.org/pdf/2412.19437) - moe_router_num_groups: int | None = None - - # Number of selected groups for group-limited routing. - moe_router_group_topk: int | None = None - - # Scaling factor for routing score in top-k selection, only works when --moe-router-pre- - # softmax enabled. Defaults to None, which means no scaling. - moe_router_topk_scaling_factor: float | None = None - - # TopK routing with dynamic expert bias in the aux-loss-free load balancing strategy. The - # routing decision is based on the sum of the routing scores and the expert bias. See - # https://arxiv.org/abs/2408.15664 for details. - moe_router_enable_expert_bias: bool = False - - # Expert bias update rate in the aux-loss-free load balancing strategy. The expert bias is - # updated based on the number of assigned tokens to each expert in a global batch, where - # the bias is increased for the experts with less assigned tokens and decreased for the - # experts with more assigned tokens. The default value 1e-3 is same as that used in - # DeepSeekV3. - moe_router_bias_update_rate: float = 0.001 - - # [Experimental] Force override routing to balance token distribution using random logits - # for MoE routers, supporting naive top-k and group-limited top-k. This experimental - # feature is for benchmarking purposes only! - moe_router_force_load_balancing: bool = False - - # Pad the routing_map to make sure the number of tokens each expert received is a multiple - # of 16/32 for FP8/FP4 precision. It is suggested to enable this for dropless training - # with FP8/FP4 precision when num_local_experts > 1. This is a more efficient way to pad - # for FP8/FP4 which eliminates the explicit padding in the GroupedMLP layer. - moe_router_padding_for_quantization: bool = False - - # [Compatibility alias for --moe-router-padding-for-quantization] Enabling this will also - # enable --moe-router-padding-for-quantization. - moe_router_padding_for_fp8: bool = False - # Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended. moe_aux_loss_coeff: list[float] | float = 0.0 - # Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended. - moe_z_loss_coeff: float | None = None - - # Add noise to the input tensor by applying jitter with a specified epsilon value. - moe_input_jitter_eps: float | None = None - - # Enable per-layer logging for MoE, currently supports auxiliary loss and z loss. - moe_per_layer_logging: bool = False - - # The type of token dispatcher to use. The default is 'allgather'. Options are - # 'allgather', 'alltoall'. We recommend using 'alltoall' when applying expert parallelism. - # For more information, please refer to the documentation in core/moe/README. - moe_token_dispatcher_type: Literal["allgather", "alltoall", "flex"] = "allgather" - - # DEPRECATED: Please use --moe-flex-dispatcher-backend=deepep instead. - moe_enable_deepep: bool = False - - # The backend to use for flex token dispatcher. The default is "deepep". Options are - # "deepep" and "hybridep". - moe_flex_dispatcher_backend: Literal["deepep", "hybridep"] = "deepep" - - # Number of SMs to use for DeepEP. - moe_deepep_num_sms: int = 20 - - # Number of SMs to use for HybridEP. - moe_hybridep_num_sms: int = 16 - - # Fuse token rearrangement ops during token dispatching. - moe_permute_fusion: bool = False - - # The capacity factor for each expert, None means no token will be dropped. - moe_expert_capacity_factor: float | None = None - - # Pads the input for each expert to match the expert capacity length, effective only after - # the --moe-expert-capacity-factor is set. - moe_pad_expert_input_to_capacity: bool = False - - # 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. - moe_token_drop_policy: Literal["probs", "position"] = "probs" - - # Apply probs before mlp activation for moe routing. - moe_apply_probs_on_input: bool = False - - # Overlap the EP A2A communication by batch-level overlapping in 1f1b stage. - overlap_moe_expert_parallel_comm: bool = False - - # Delay the wgrad compute for batch-level overlapping - delay_wgrad_compute: bool = False - # This param sepecifics how many times smaller is the expert hidden size compared with the # original dense FFN hidden size. For using granular upcycling strategy, please set this # param as a positive integer. If this param is set to 1, it means using the default # upcycling strategy. moe_upcycling_granularity: int = 1 - # some MoE routers have a D2H sync that will break cuda graphs. If this flag is set the - # router will switch to dropping and padding during decode time which does not have a D2H - # sync. The capacity factor is set to the max that an expert could see during inference so - # no tokens are actually dropped. - moe_pad_experts_for_cuda_graph_inference: bool = False - # Rank of Query tensor's low rank representation. q_lora_rank: int | None = None @@ -1663,6 +2429,19 @@ class MegatronConfig(ConfigInterface): # If set caches the mla down projected latents with mla flash decode. cache_mla_latents: bool = False + # Enable fused q/kv down-projection and fused input layernorm when backend supports. + # Otherwise fall back to the unfused MLA. + mla_down_proj_fusion: bool = False + + # Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) + # layers. Accepts either: - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA + # layers for every 1 SDPA layer - A string containing a Python list expression that + # defines a custom pattern, e.g.: "([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] + # where 1 indicates an LA layer and 0 indicates a SDPA layer. Examples: "([0]+[1]*23)": 1 + # SDPA layer followed by 23 LA layers, "([1]*3+[0]*2)*2": Three LA layers followed by two + # SDPA layers, repeated twice. + linear_attention_freq: Any | None = None + # Path to json file containing heterogeneous model configuration. Use the format of the # HuggingFace config files in llama nemotron models, e.g. # https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json. @@ -1675,12 +2454,12 @@ class MegatronConfig(ConfigInterface): # Super-49B-v1/resolve/main/config.json. heterogeneous_layers_config_encoded_json: str | None = None + # Report loss and timing interval. + log_interval: int = 100 + # If set, calculate and log parameters norm. log_params_norm: bool = False - # If set, calculate and log the number of zeros in gradient. - log_num_zeros_in_grad: bool = False - # If set, calculate and log throughput per GPU. log_throughput: bool = False @@ -1688,30 +2467,25 @@ class MegatronConfig(ConfigInterface): # point operations) to progress.txt file in checkpoint directory. log_progress: bool = False - # 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. + # 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. timing_log_level: Literal[0, 1, 2] = 0 - # If set, log energy consumption (in Joules) - log_energy: bool = False - - # 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. - barrier_with_L1_time: bool = True - - # 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. + # 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. timing_log_option: Literal["max", "minmax", "all"] = "minmax" + # Write TensorBoard logs to this directory. + tensorboard_dir: str | None = None + # Report to tensorboard interval. tensorboard_log_interval: int = 1 - # Size of the tensorboard queue for pending events and summaries before one of the "add" + # Size of the tensorboard queue for pending events and summaries before one of the 'add' # calls forces a flush to disk. tensorboard_queue_size: int = 1000 @@ -1727,38 +2501,56 @@ class MegatronConfig(ConfigInterface): # Enable memory logging to tensorboard. log_memory_to_tensorboard: bool = False + # Report memory interval. + log_memory_interval: int | None = None + + # Log device memory used (as reported by nvidia-smi). + log_device_memory_used: bool = False + + # If set, calculate and log the number of zeros in gradient. + log_num_zeros_in_grad: bool = False + + # Enable max attention logit logging to tensorboard. + log_max_attention_logit: bool = False + + # If not disabled, 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. + barrier_with_L1_time: bool = True + # Enable world size logging to tensorboard. log_world_size_to_tensorboard: bool = False # The wandb project name. Ignore wandb by default. - wandb_project: str = "" - - # The wandb entity name. It is useful when there are multiple sub-projects in a project. - # https://community.wandb.ai/t/how-do-i-decide-which-account-private-or-team-to-upload- - # the-run-to/5704 Ignore wandb by default. - wandb_entity: str = "" + wandb_project: str | None = None # The wandb experiment name. - wandb_exp_name: str = "" + wandb_exp_name: str | None = None # Path to save the wandb results locally. - wandb_save_dir: str = "" + wandb_save_dir: str | None = None + + # The wandb entity name. It is useful when there are multiple sub-projects in a project. + wandb_entity: str | None = None # Set default logging level logging_level: int | None = None + # If set, log energy consumption (in Joules). + log_energy: bool = False + # If set, tracks and logs straggler per GPU. log_straggler: bool = False - # If set, StragglerDetector is disabled on startup. - disable_straggler_on_startup: bool = False - # Port number to toggle StragglerDetector on/off at runtime straggler_ctrlr_port: int = 65535 # Number of ranks to report with high/low estimated throughput straggler_minmax_count: int = 1 + # If set, StragglerDetector is disabled on startup. + disable_straggler_on_startup: bool = False + # If set, enables workload inspector server for on-demand profiling. run_workload_inspector_server: bool = False @@ -1778,46 +2570,35 @@ class MegatronConfig(ConfigInterface): # Select either Megatron or Huggingface as the Bert embedder. bert_embedder_type: Literal["megatron", "huggingface"] = "megatron" - # Whether to use the flash decoding kernel. - flash_decode: bool = False - - # Deprecated. Use --cuda-graph-impl=local instead. Use local implementation of CUDA graph - # capture and replay. --cuda-graph-scope="full_iteration" enables whole iteration CUDA - # graph. - enable_cuda_graph: bool = False - - # Number of CUDA graph warmup steps - cuda_graph_warmup_steps: int = 3 - - # Deprecated. Use --cuda-graph-impl=transformer_engine instead. Use TE - # make_graphed_callables() to capture the CUDA graph. - external_cuda_graph: bool = False - - # Determines the CUDA graph capture implementation. "none": no CUDA graph. "local": - # capture the CUDA graph using MCore local implementation. --cuda-graph- - # scope="full_iteration" enables whole iteration CUDA graph. "transformer_engine": capture - # the CUDA graph using TE make_graphed_callables(). - cuda_graph_impl: Literal["none", "local", "transformer_engine"] = "none" - - # Determines the CUDA graphs capturing scope. Valid values are "full", "attn" and - # "full_iteration". "Full" scope captures a whole Transformer layer. "Attn" scope only - # captures operations in TransformerLayer._forward_attention(). "ful_iteration" scope - # captures a whole iteration. full_iteration scope is only supported with --cuda-graph- - # impl=local, attn scope is only supported with --cuda-graph-impl=transformer_engine. - cuda_graph_scope: Literal["full", "attn", "full_iteration"] = "full" + # ==SUPPRESS== + cuda_graph_scope_deprecated: list[Any] | None = None + + # Selects training capture coverage within per-layer CUDA graphs (local and + # transformer_engine implementations). Valid values are "attn", "mlp", "moe", + # "moe_router", "moe_preprocess", and "mamba": "attn": captures operations in + # TransformerLayer._forward_attention(). "mlp": captures operations in + # TransformerLayer._forward_mlp() for a dense layer. "moe": captures operations in + # TransformerLayer._forward_mlp() for a MoE layer. "moe_router": captures operations in + # TransformerLayer._forward_mlp() up to MoELayer.router(), including the shared experts if + # they are not overlapped with EP comm. "moe_preprocess": captures operations in + # MoELayer.preprocess(). Must be used together with "moe_router". "mamba": captures the + # mamba layer. An empty list means capturing the whole Transformer layer. This field is + # meaningless when --cuda-graph-impl=full_iteration and must be empty. Backward + # compatibility: "full" is deprecated but kept for backward compatibility; it is + # transformed to an empty list in validate_args. The deprecated values "full_iteration" + # and "full_iteration_inference" are also accepted and migrated to the new API in + # validate_args. + cuda_graph_modules: list[Any] = field(default_factory=lambda: []) # Use legacy static engine. (Current static engine uses dynamic engine under the hood) use_legacy_static_engine: bool = False - # Maximum batch size for inference. - inference_max_batch_size: int = 8 + # Maximum number of requests for inference. + inference_max_requests: int = 8 # Maximum sequence length expected for inference (prefill + decode). inference_max_seq_length: int = 2560 - # Maximum batch size for inference. - inference_max_batch_size: int = 8 - # Enable dynamic batching mode. inference_dynamic_batching: bool = False @@ -1826,156 +2607,134 @@ class MegatronConfig(ConfigInterface): # memory (UVM) level (via `--inference-dynamic-batching-unified-memory-level`).If the UVM # level is 0, then only GPU memory is used and the total memory equals `buffer_size_gb`. # If the UVM level is 1, then additional memory is utilized on the CPU and the total - # memory equals `2 * buffer_size_gb`. + # memory equals `buffer_size_gb + paused_buffer_size_gb`. inference_dynamic_batching_buffer_size_gb: float = 40.0 + # Amount of memory reserved for paused requests in the dynamic inference context. Active + # requests are paused when there are not enough active blocks available to continue + # generating a request. + inference_dynamic_batching_paused_buffer_size_gb: float | None = None + + # Percentage of memory buffer to allocate for Mamba states. If not specified, allocates + # Mamba state tensors for each KV cache block. Only used for hybrid models. + inference_dynamic_batching_mamba_memory_ratio: float | None = None + # KV cache block size. It should be a multiple of 256 inference_dynamic_batching_block_size: int = 256 + # Override the inference context's `max_requests`. By default, `max_requests` is set to + # the number of blocks in the context's memory buffer. + inference_dynamic_batching_max_requests: int | None = None + # Override the inference context's default `max_tokens`. inference_dynamic_batching_max_tokens: int | None = None # Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 # to `max_requests`. (See `dynamic_context.py` for details on how `max_requests` is # computed). Due to rounding, the actual number of cuda graphs may not equal this - # argument. + # argument.The user can also pass -1, in which case we automatically determine the number + # of graphs to capture based on the `max_requests`. inference_dynamic_batching_num_cuda_graphs: int = 16 # Track paused request ids by adding 'paused' events to each request's event history. This # has a very minor impact on latency. inference_dynamic_batching_track_paused_request_events: bool = False + # Track per-token events with timestamps for each generated token. When enabled, each + # generated token creates a GENERATED_TOKEN event with a timestamp, useful for per-token + # latency analysis. + inference_dynamic_batching_track_generated_token_events: bool = False + # Only use cuda graphs for decode-only steps, not prefill and mixed steps. decode_only_cuda_graphs: bool = False + # Extend prefill/mixed CUDA graph capture up to `max_tokens`. By default, all graphs are + # limited by the decode limit of `max_requests * (num_speculative_tokens + 1)`. + inference_cuda_graph_all_prefills: bool = False + # Set unified memory usage within the dynamic inference context. The levels are: 0) no # unified memory, 1) allocate `memory_buffer` in unified memory. Eventually, additional # levels will be included to control other tensors within the context. - inference_dynamic_batching_unified_memory_level: Literal[0, 1] = 1 + inference_dynamic_batching_unified_memory_level: Literal[0, 1] = 0 - # What type of symmetric all reduce to use. The default is none which is no use of - # symetric memory - symmetric_ar_type: Literal["two_shot", "one_shot", "multimem_all_reduce", None] = None + # Enable chunked prefill (disabled by default) + enable_chunked_prefill: bool = False - # When using symmeric all reduce kernels this will use regular nccl kernels for prefill. - # This can be more effecient when prefill is large as the nccl kernels can be more - # bandwith optimized - nccl_all_reduce_for_prefill: bool = False + # Number of speculative tokens generated during decode + num_speculative_tokens: int = 0 - # Number of chunks along sequence dimension for MLP computation during prefill - mlp_chunks_for_prefill: int = 1 + # Enable/disable prefix caching for dynamic batching inference. When disabled, KV cache + # blocks cannot be shared between requests with identical prompt prefixes. + inference_dynamic_batching_enable_prefix_caching: bool = False + + # Eviction policy for prefix caching blocks. "ref_zero" (default) immediately returns + # blocks to the free pool when ref_count hits 0. "lru" keeps blocks cached and evicts via + # LRU only when space is needed. + inference_dynamic_batching_prefix_caching_eviction_policy: Literal["ref_zero", "lru"] = ( + "ref_zero" + ) - # Disable chunked prefill (chunked prefill is enabled by default). - disable_chunked_prefill: bool = False + # Coordinator routing policy for prefix caching. "first_prefix_block" (default) routes + # based on the first block hash only. "longest_prefix" routes to the rank with the longest + # matching prefix. "round_robin" ignores prefix affinity and cycles through ranks. + inference_dynamic_batching_prefix_caching_coordinator_policy: Literal[ + "longest_prefix", "first_prefix_block", "round_robin" + ] = "first_prefix_block" - # Maximum number of tokens to capture in a cuda graph. - inference_dynamic_batching_cuda_graph_max_tokens: int = 16384 + # Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * + # normalized_load. Higher alpha favors prefix cache hits; lower alpha favors load balance. + # Default: 0.5. + inference_dynamic_batching_prefix_caching_routing_alpha: float = 0.5 + + # GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid + # models. When set, Mamba states at block boundaries are cached for reuse. + inference_dynamic_batching_prefix_caching_mamba_gb: float | None = None # Number of mixed prefill requests to capture in a cuda graph. inference_dynamic_batching_cuda_graph_mixed_prefill_count: int = 16 - # Step interval for logging inference metrics to wandb. Default to 0 to disable inference - # wandb logging. - inference_wandb_logging_step_interval: int = 0 - - # This port will be used to setup the inference coordinator on node-0 - inference_coordinator_port: int = 12346 + # Which sampling kernels to use during inference. Falls back to "torch" with a warning if + # "flashinfer" is requested but the package is not installed. + inference_dynamic_batching_sampling_backend: Literal["torch", "flashinfer"] = "torch" - # Which fp8 format scheme to use for FP8 tensors in the forward and backward pass - fp8: Literal["e4m3", "hybrid"] | None = None + # Step interval for logging inference metrics. Default to 0 to disable inference logging. + inference_logging_step_interval: int = 0 - # Which fp8 recipe to use for FP8 tensors in the forward and backward pass - fp8_recipe: Literal["tensorwise", "delayed", "mxfp8", "blockwise", "custom"] = "delayed" + # Enable per-request logging in the inference text generation server. + inference_text_gen_server_logging: bool = False - # Python import path to a callable quantizer factory, e.g., - # package.module.quantizer_factory. - fp8_quantizer_factory: Any | None = None + # Enable inference wandb logging. + inference_wandb_logging: bool = False - # Scaling margin for fp8 - fp8_margin: int = 0 + # This port will be used to setup the inference coordinator on node-0 + inference_coordinator_port: int | None = None - # DEPRECATED. This flag is ignored. Scaling update interval for fp8 - fp8_interval: int = 1 + # Dtype for the Mamba inference conv states tensor + mamba_inference_conv_states_dtype: Literal["bf16", "fp16", "fp32"] = "bf16" - # Number of steps for which amax history is recorded per tensor - fp8_amax_history_len: int = 1 + # Dtype for the Mamba inference SSM states tensor + mamba_inference_ssm_states_dtype: Literal["bf16", "fp16", "fp32"] = "bf16" - # Algorithm for computing amax from history - fp8_amax_compute_algo: Literal["most_recent", "max"] = "most_recent" + # Use synchronous ZMQ collectives for inference. Helps in reducing performance variability + # for MoEs. + inference_use_synchronous_zmq_collectives: bool = False - # Execute wgrad in higher precision even for FP8 runs - fp8_wgrad: bool = True + # Skip the EP-group consensus all-reduce in the inference engine control loop and step on + # local state only. Pause/unpause take effect as soon as the signal is delivered to a + # rank. Only safe when EP coordination is not required (e.g. ep_world_size == 1). + inference_disable_ep_consensus: bool = False - # Which Transformer implementation to use. - transformer_impl: Literal["local", "transformer_engine", "inference_optimized"] = ( - "transformer_engine" - ) + # Keep the compute param in fp4 (do not use any other intermediate dtype) and perform the + # param all-gather in fp4. + fp4_param_gather: bool = False # Keep the compute param in fp8 (do not use any other intermediate dtype) and perform the # param all-gather in fp8. fp8_param_gather: bool = False - # Construct first and last layers in bf16 when doing FP8 training. - first_last_layers_bf16: bool = False - - # Number of layers at start to construct in bf16 when --first-last-layers-bf16 is enabled. - num_layers_at_start_in_bf16: int = 1 - - # Number of layers at end to construct in bf16 when --first-last-layers-bf16 is enabled. - num_layers_at_end_in_bf16: int = 1 - - # Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass - fp4: Literal["e2m1"] | None = None - - # Which fp4 recipe to use for FP4 tensors in the forward and backward pass - fp4_recipe: Literal["nvfp4", "custom"] = "nvfp4" - - # Python import path to a callable quantizer factory, e.g., - # package.module.quantizer_factory. - fp4_quantizer_factory: Any | None = None - - # Keep the compute param in fp4 (do not use any other intermediate dtype) and perform the - # param all-gather in fp4. - fp4_param: bool = False - - # Use the Transformer Engine version of the random number generator. Required for CUDA - # graphs support. - te_rng_tracker: bool = False - - # Use a random number generator configured for inference. - inference_rng_tracker: bool = False - - # Retro project directory, which contains the preprocessed data for pretraining. This - # directory is built during preprocessing (see tools/retro/README.md), and contains - # subdirectories for the chunk database and pretraining neighbors. - retro_project_dir: Any | None = None - - # Add a retriever to the transformer, for use in pretraining a Retro model. - retro_add_retriever: bool = False - - # Set number of training iterations for cyclic Retro training. - retro_cyclic_train_iters: int | None = None - - # Number of layers to use for the retrieval encoder. - retro_encoder_layers: int = 2 - - # Hidden dropout for retrieval encoder. - retro_encoder_hidden_dropout: float = 0.1 - - # Attention dropout for retrieval encoder. - retro_encoder_attention_dropout: float = 0.1 - - # Number of neighbors to retrieve during pretraining. - retro_num_neighbors: int = 2 - - # Number of chunks to retrieve from the retrieval database. - retro_num_retrieved_chunks: int = 2 - - # Gated cross attention. - retro_attention_gate: float = 1 - - # Skip verifying that len(GPT dataset) == len(saved neighbors). - retro_verify_neighbor_count: bool = True + # Configuration file to select per-module precision overrides. See + # TransformerEngineMixedPrecision.md + te_precision_config_file: Any | None = None # Enable experimental features. enable_experimental: bool = False @@ -1986,37 +2745,17 @@ class MegatronConfig(ConfigInterface): # `transformer_block.py`, or `transformer_layer.py` spec: list[str] | None = None - # Ratio of attention layers to total layers, in the range [0.0, 1.0]. - hybrid_attention_ratio: float = 0.0 + # Specify a hybrid layer pattern using M (mamba), G (gdn), * (attention), D (dsa), - + # (mlp), E (moe). Use | to define pipeline stage boundaries for flexible virtual pipeline + # parallel (fVPP). Use / to separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or + # "M-M-|M-M*-/MM/MM". When this flag is used, it is the sole indicator that a hybrid model + # is being run. + hybrid_layer_pattern: str | None = None - # Ratio of mlp layers to total layers, in the range [0.0, 1.0]. - hybrid_mlp_ratio: float = 0.0 - - # Force a specific hybrid layer pattern. The valueshould be a string of characters chosen - # fromcore.ssm.mamba_hybrid_layer_allocation.Symbols.If a value greater than 0.0 is - # supplied to any of the hybrid ratio arguments, then the number of each typeof layer in - # the override pattern must match number inthe overidden pattern + # Deprecated. Use --hybrid-layer-pattern instead. If specified, its value will be + # forwarded to --hybrid-layer-pattern. hybrid_override_pattern: str | None = None - # State dimension for Mamba layers. - mamba_state_dim: int = 128 - - # Head dimension for Mamba layers. - mamba_head_dim: int = 64 - - # Number of groups for Mamba layers. - mamba_num_groups: int = 8 - - # Number of heads for Mamba layers.If not set, then the number of heads will be --hidden- - # size * expand // --mamba-head-dim - mamba_num_heads: int | None = None - - # Indicates whether the model is a hybrid model. - is_hybrid_model: bool = False - - # Disable Mamba efficient path. - disable_mamba_mem_eff_path: bool = False - # Config file to add additional arguments yaml_cfg: str | None = None @@ -2024,10 +2763,10 @@ class MegatronConfig(ConfigInterface): # params and optimizer states to lower precision, such as fp16, bf16 and fp8. use_precision_aware_optimizer: bool = False - # Dtype of main grads when enabling precision-aware-optimizer + # Dtype of main grads when enabling precision-aware-optimizer. main_grads_dtype: Literal["fp32", "bf16"] = "fp32" - # Dtype of main params when enabling precision-aware-optimizer + # Dtype of main params when enabling precision-aware-optimizer. main_params_dtype: Literal["fp32", "fp16"] = "fp32" # Dtype of exp_avg (1st moment in adam optimizer) when enabling precision-aware-optimizer. @@ -2040,6 +2779,30 @@ class MegatronConfig(ConfigInterface): # but does not affect the precision in the kernel computation. exp_avg_sq_dtype: Literal["fp32", "fp16", "bf16", "fp8"] = "fp32" + # Data type for the main weight buffer utilized for distributed optimization and + # quantization with Megatron-FSDP. If 'auto', then the native model parameter data-type + # will be used for the main weight data-type. Replaces --main-params-dtype. + megatron_fsdp_main_params_dtype: Literal["fp32", "bf16", "fp16", "auto"] = "fp32" + + # Data type for the main gradient buffer utilized for distributed optimization with + # Megatron-FSDP. If 'auto', then the native model gradient data-type will be used for the + # main gradient / accumulation data-type. Replaces --main-grads-dtype. + megatron_fsdp_main_grads_dtype: Literal["fp32", "bf16", "fp16", "auto"] = "auto" + + # When using Megatron-FSDP, this controls the data-type used when communicating model + # gradients during FSDP. If 'auto', then the main gradient data-type will be used for the + # gradient communication / reduction data-type. When using NCCL v2.27+, reduction is + # always computed in FP32 if using NCCL Symmetric kernels. + megatron_fsdp_grad_comm_dtype: Literal["fp32", "fp16", "bf16", "auto"] = "auto" + + # If set, enables fine-grained parameter gathering for Megatron-FSDP. This allows greater + # overlap between parameter all-gather operations and forward computation, at the cost of + # additional communication calls. For MXFP8, this helps save memory during fine-grained + # activation recomputation, because MXFP8 forward and backward passes use different + # parameter representations (rowwise data for forward, colwise data for backward). Only + # the rowwise parameters of modules involved in recomputation will be unsharded. + megatron_fsdp_enable_fine_grained_param_gather: bool = False + # If set, disable using one_logger to track E2E metricsNote that one_logger is an internal # tool and not available externally. For installation, please go to # https://confluence.nvidia.com/display/MLWFO/Package+Repositoriesfor more details @@ -2119,8 +2882,11 @@ class MegatronConfig(ConfigInterface): # for Nvidia internal use only. calc_ft_timeouts: bool = False - # If set, will dump all configs to --config-logger-dir - config_logger_dir: str = "" + # Number of warmup iterations before monitoring step section and out-of-section timeouts. + # The first N iterations are excluded from timeout monitoring as they can be significantly + # slower than steady-state. Default: 5. Note: This feature is for Nvidia internal use + # only. + ft_num_warmup_iters: int = 5 # Rate at which to inject unexpected results, e.g. 1000 means once every 1000 result # validations @@ -2135,14 +2901,16 @@ class MegatronConfig(ConfigInterface): # computations due to non-deterministic algorithms. rerun_mode: Literal["disabled", "validate_results", "report_stats"] = "validate_results" - # Disable the usage of Multi-Storage Client (MSC) in Megatron Core. - enable_msc: bool = True + # Check for spiky loss. + check_for_spiky_loss: bool = False - # Use the config .yaml file at the specified location to configure kitchen quantization. - kitchen_config_file: str | None = None + # Enable the usage of Multi-Storage Client (MSC) in Megatron Core. Disabled by default; + # pass this flag to opt in. + enable_msc: bool = False - # Use a default kitchen recipe for all layers as defined by QAT_PARAMS index - kitchen_recipe_number: int | None = None + # [DEPRECATED] MSC is disabled by default; this flag is a no-op and will be removed in a + # future release. + disable_msc_deprecated: bool = False # Megatron SFT training sft: bool = False @@ -2150,6 +2918,34 @@ class MegatronConfig(ConfigInterface): # SFT prompt format. sft_tokenizer_prompt_format: str = "nemotron-h-aligned" + # Comma-separated list of ranks to inject faults on. + fault_injector_ranks: str | None = None + + # Number of ranks to inject faults on (random selection). + fault_injector_num_ranks: int | None = None + + # Comma-separated list of fault types to inject (e.g. 'hang,crash'). + fault_injector_fault_types: str | None = None + + # Comma-separated list of fault probabilities (normalized at runtime). + fault_injector_fault_probabilities: str | None = None + + # Force a specific fault delay in seconds from training start or delay_start_iteration. + fault_injector_fault_delay: float | None = None + + # Start the fault delay timer after iteration N completes. If unset, fault delay timing + # starts from the beginning of training. + fault_injector_delay_start_iteration: int | None = None + + # Mean time to inject (MTTI) in seconds; used when fault_delay is None. + fault_injector_mtti_seconds: float | None = None + + # Offset seconds added to the sampled fault delay. + fault_injector_offset_seconds: float | None = None + + # RNG seed for the fault injector. + fault_injector_seed: int | None = None + aux: dict[str, Any] = field(default_factory=dict) diff --git a/submodules/Megatron-LM b/submodules/Megatron-LM index 055f7def..5f3865f7 160000 --- a/submodules/Megatron-LM +++ b/submodules/Megatron-LM @@ -1 +1 @@ -Subproject commit 055f7defc2284738b207c81c5cdc9f1fdc9dc897 +Subproject commit 5f3865f7c32b6dee08792e1637f71777f95d3e9b