# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron arguments.""" import argparse import dataclasses import json import os import re import types from pathlib import Path import torch from megatron.core.config import set_experimental_flag from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.rerun_state_machine import RerunStateMachine from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, get_deprecated_cuda_graph_modules_migration, normalize_cuda_graph_modules, normalize_inference_cuda_graph_scope, validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.tuning import AutotunePolicy from megatron.core.utils import ( get_torch_version, is_flashinfer_min_version, is_te_min_version, is_torch_min_version, ) from megatron.training import global_vars from megatron.training.argument_utils import ( # noqa: F401 # pylint: disable=unused-import ArgumentGroupFactory, _default_config_from_args, _wide_residual_config_from_args, core_transformer_config_from_args, ) from megatron.training.utils import ( get_device_arch_version, print_rank_0, update_use_dist_ckpt, warn_rank_0, ) def add_megatron_arguments(parser: argparse.ArgumentParser): """"Add Megatron-LM arguments to the given parser.""" # Standard arguments. parser = _add_network_size_args(parser) parser = _add_wide_residual_args(parser) parser = _add_triton_autotune_args(parser) parser = _add_regularization_args(parser) parser = _add_training_args(parser) parser = _add_rl_args(parser) parser = _add_initialization_args(parser) parser = _add_learning_rate_args(parser) parser = _add_checkpointing_args(parser) parser = _add_mixed_precision_args(parser) parser = _add_distributed_args(parser) parser = _add_validation_args(parser) parser = _add_data_args(parser) parser = _add_tokenizer_args(parser) parser = _add_autoresume_args(parser) parser = _add_vision_args(parser) parser = _add_moe_args(parser) parser = _add_mla_args(parser) parser = _add_experimental_attention_variant_args(parser) parser = _add_heterogeneous_args(parser) parser = _add_logging_args(parser) parser = _add_logits_distillation_args(parser) parser = _add_straggler_detector_args(parser) parser = _add_workload_inspector_server_args(parser) parser = _add_inference_args(parser) parser = _add_transformer_engine_args(parser) parser = _add_experimental_args(parser) parser = _add_inprocess_restart_args(parser) parser = _add_ft_package_args(parser) parser = _add_rerun_machine_args(parser) parser = _add_msc_args(parser) parser = _add_kitchen_quantization_arguments(parser) parser = _add_sft_args(parser) parser = _add_varlen_dataset_args(parser) parser = _add_fault_injector_args(parser) return parser def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, args_defaults={}): """Prepare and register CLI inputs without constructing runtime services. Checkpoint overrides and validation precede config construction. Callers initialize runtime services explicitly after preparing their configuration. """ args = parse_args(extra_args_provider, ignore_unknown_args) if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): from megatron.training.checkpointing import load_args_from_checkpoint assert args.load is not None or args.pretrained_checkpoint is not None, "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" assert args.non_persistent_ckpt_type != "local", ( "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. " "Two-stage checkpoint loading is not implemented, and all arguments must be defined " "before initializing LocalCheckpointManager." ) load_args_from_checkpoint(args, load_arg='pretrained_checkpoint') load_args_from_checkpoint(args) if args.yaml_cfg is not None: from megatron.training.yaml_arguments import validate_yaml args = validate_yaml(args, args_defaults) else: validate_args(args, args_defaults) global_vars._ensure_var_is_not_initialized(global_vars._GLOBAL_ARGS, 'args') global_vars.set_args(args) # Model config construction can use experimental features. Enabling the # feature gate does not construct any runtime services. if args.enable_experimental: set_experimental_flag(True) return args def parse_args(extra_args_provider=None, ignore_unknown_args=False): """Parse all arguments.""" parser = argparse.ArgumentParser(description='Megatron-LM Arguments', allow_abbrev=False) parser = add_megatron_arguments(parser) # Custom arguments. if extra_args_provider is not None: parser = extra_args_provider(parser) # Parse. if ignore_unknown_args: args, _ = parser.parse_known_args() else: args = parser.parse_args() args._is_global_batch_size_explicitly_specified = args.global_batch_size is not None # Experimental yaml if args.yaml_cfg is not None: cli_policy = _default_config_from_args(AutotunePolicy, args, prefix='triton_autotune') if cli_policy != AutotunePolicy(): raise ValueError( 'Triton autotune CLI arguments cannot be combined with --yaml-cfg; ' 'set triton_autotune in the YAML configuration instead.' ) if _wide_residual_config_from_args(args) is not None: raise ValueError( 'Wide-residual CLI arguments cannot be combined with --yaml-cfg because ' 'YAML model configuration replaces argparse model arguments.' ) from .yaml_arguments import load_yaml args = load_yaml(args.yaml_cfg) # Args from environment args.rank = int(os.getenv('RANK', '0')) args.world_size = int(os.getenv("WORLD_SIZE", '1')) # Args to enable MSC (opt-in: disabled by default) if args.disable_msc_deprecated: warn_rank_0( '--disable-msc is deprecated and will be removed in a future release. ' 'MSC is now disabled by default; pass --enable-msc to opt in.' ) # Preserve legacy semantics: --disable-msc forces MSC off, even if # --enable-msc was also passed. args.enable_msc = False if args.enable_msc: MultiStorageClientFeature.enable() if not MultiStorageClientFeature.is_enabled(): raise RuntimeError( "--enable-msc was passed but the multistorageclient package is not " "installed. Install it with `pip install multi-storage-client`." ) warn_rank_0('The MSC feature is enabled.') else: MultiStorageClientFeature.disable() assert MultiStorageClientFeature.is_enabled() is False return args def validate_model_config_args_from_heterogeneous_config(args): """Validate model config arguments from heterogeneous config. This function takes model arguments and validates them based on a heterogeneous layer configuration. The heterogeneous config can be provided either as a path to a JSON file or as an encoded JSON string. The function enforces certain model architecture choices like SiLU activation, RMSNorm, grouped query attention, and RoPE positional embeddings. It also sets model dimensions like number of layers, hidden size, and attention heads based on the heterogeneous config. Args: args: Model configuration arguments to be overridden. Expected to have attributes: - heterogeneous_layers_config_path (str): Path to JSON config file - heterogeneous_layers_config_encoded_json (str): Encoded JSON config string Returns: None """ if ( args.heterogeneous_layers_config_path is None and args.heterogeneous_layers_config_encoded_json is None ): return if args.heterogeneous_layers_config_encoded_json is None: args.heterogeneous_layers_config_encoded_json = Path( args.heterogeneous_layers_config_path ).read_text() hf_config_dict = types.SimpleNamespace(**json.loads(args.heterogeneous_layers_config_encoded_json)) assert hf_config_dict.hidden_act == "silu", ( f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" ) n_kv_heads_in_group = [ config["attention"]["n_heads_in_group"] for config in hf_config_dict.block_configs if config["attention"]["n_heads_in_group"] is not None ] assert all(num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group), "num query head must be consistent across all layers" args_to_validate = { "swiglu": True, "normalization": "RMSNorm", "group_query_attention": True, "position_embedding_type": "rope", "rotary_percent": 1.0, "use_rope_scaling": True, "use_rotary_position_embeddings": True, "num_layers": hf_config_dict.num_hidden_layers, "hidden_size": hf_config_dict.hidden_size, "num_attention_heads": hf_config_dict.num_attention_heads, "untie_embeddings_and_output_weights": not hf_config_dict.tie_word_embeddings, "rotary_base": hf_config_dict.rope_theta, "rope_scaling_factor": hf_config_dict.rope_scaling["factor"], "num_query_groups": hf_config_dict.num_attention_heads // n_kv_heads_in_group[0], } incompatible_args = {} for key, value in args_to_validate.items(): provided_value = getattr(args, key, None) if provided_value != value: incompatible_args[key] = (provided_value, value) if incompatible_args: incompatible_args_str = ', '.join([ f"{k}: {provided_value} (provided) != {value} (expected)" for k, (provided_value, value) in incompatible_args.items() ]) raise ValueError( f"Arguments differ from heterogeneous config: {incompatible_args_str}" ) def _eval_pattern(pattern): """ Validate and evaluate a string containing a Python list expression """ assert isinstance(pattern, str) # validate input, only allow comma, digits, [, ], (, ), +, and * if bool(re.compile(r'[^,\d\[\]\(\)\+\*]').search(pattern)): raise ValueError(f"Invalid pattern: {pattern}") return eval(pattern) def _parse_cuda_graph_modules_arg(scope): """Parse CUDA graph module CLI values while preserving deprecated spellings for migration.""" if scope in {"full", "full_iteration", "full_iteration_inference"}: return scope return CudaGraphModule[scope] def _normalize_cuda_graph_modules_args(args): """Normalize cuda_graph_modules to enums and apply deprecated scope migrations.""" normalized_scopes, deprecated_scopes, used_full_scope = normalize_cuda_graph_modules( args.cuda_graph_modules ) validate_deprecated_cuda_graph_modules_migration_inputs( deprecated_scopes, args.cuda_graph_impl, args.inference_cuda_graph_scope, ) if used_full_scope: warn_rank_0('full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.') for scope, attr, value in deprecated_scopes: migration = get_deprecated_cuda_graph_modules_migration( scope, attr, value, args.cuda_graph_impl ) if migration is None: warn_rank_0( f"--cuda-graph-modules '{scope}' is deprecated and has no effect when " "--cuda-graph-impl=none. Use --cuda-graph-impl=local with " "--inference-cuda-graph-scope=block to enable inference CUDA graphs." ) continue migration_attr, migration_value = migration warn_rank_0( f"--cuda-graph-modules '{scope}' is deprecated. " f"Setting --{migration_attr.replace('_', '-')}={migration_value} instead." ) setattr(args, migration_attr, migration_value) args.cuda_graph_modules = normalized_scopes def _normalize_inference_cuda_graph_scope_arg(args): """Normalize inference_cuda_graph_scope and apply the impl-derived default.""" args.inference_cuda_graph_scope = normalize_inference_cuda_graph_scope( args.inference_cuda_graph_scope, args.cuda_graph_impl ) def no_rope_freq_type(x): """ Controls which layers to skip performing Rotary Position Embedding. - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string "N": Same as above, but provided as a string - 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 rope is skipped on the layer. This allows defining arbitrary patterns of rope skipping. The pattern length must match the total number of transformer layers. Examples: "([1]+[0]*23)": Only first layer has rope skipped for a 24-layer network. "([0]*3+[1]*1)*2": Every 4 layers the rope is skipped on the last layer. Repeat twice. """ if x is None or isinstance(x, int): return x assert isinstance(x, str) if '[' in x: # it's a custom pattern return _eval_pattern(x) else: # it's a single int but in str return int(x) def compress_ratios_type(x): """Parse per-layer compression ratios for compressed sparse attention.""" if isinstance(x, list): return x assert isinstance(x, str) return _eval_pattern(x) def moe_freq_type(x): """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 "N": Same as above, but provided as a string - 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. This allows defining arbitrary patterns of expert and dense layers. The pattern length must match the total number of transformer layers. 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. """ if isinstance(x, int): return x assert isinstance(x, str) if '[' in x: # it's a custom pattern return _eval_pattern(x) else: # it's a single int but in str return int(x) def la_freq_type(x): """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 "N": Same as above, but provided as a string - 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. This allows defining arbitrary patterns of LA and SDPA layers. The pattern length must match the total number of transformer layers. 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. """ if x is None or isinstance(x, int): return x assert isinstance(x, str) if '[' in x: # it's a custom pattern return _eval_pattern(x) else: # it's a single int but in str return int(x) def tuple_type(x): """ Convert a string to a tuple of integers. Examples: "1,2,3" -> (1, 2, 3) "(1,2,3)" -> (1, 2, 3) """ if x is None or isinstance(x, tuple): return x assert isinstance(x, str) return tuple(int(i) for i in x.strip('()').split(',')) def validate_args(args, defaults={}): # Prep for checkpoint conversion. if args.ckpt_convert_format is not None: assert args.ckpt_convert_save is not None assert args.load is not None args.exit_on_missing_checkpoint = True # Temporary assert args.non_persistent_ckpt_type in ['global', 'local', None], \ 'Currently only global and local checkpoints are supported' if args.non_persistent_ckpt_type == 'local': try: from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( LocalCheckpointManager, ) except ModuleNotFoundError as e: raise RuntimeError('nvidia_resiliency_ext is required for local checkpointing') from e # validate model config args from heterogeneous config (if provided). validate_model_config_args_from_heterogeneous_config(args) update_use_dist_ckpt(args) # Resolve the actual weight degree before deriving the rank layout. Independent # sample count is computed separately below, including when CP overlaps GTP. from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.tensor_parallel.gtp_parallel_layout import ( resolve_tensor_parallel_sequence_shards, ) (args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size) = ( resolve_tensor_parallel_weight_shards( args.tensor_model_parallel_size, args.tensor_parallel_num_weight_shards, getattr(args, "gtp_weight_remat_size", 1), ) ) (args.tensor_parallel_num_sequence_shards, args.gtp_remat_num_sequence_shards) = ( resolve_tensor_parallel_sequence_shards( args.tensor_model_parallel_size, args.tensor_parallel_num_sequence_shards, args.gtp_weight_remat_size, args.sequence_parallel, ) ) total_model_size = ( args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size * args.gtp_weight_remat_size ) # Normalize CLI CP once: downstream args/config consumers expect the full runtime group size. args.context_parallel_size *= args.gtp_remat_num_sequence_shards if args.gtp_remat_num_sequence_shards > 1: if args.ckpt_format != "torch_dist": raise ValueError("GTP sequence sharding requires --ckpt-format torch_dist") if args.use_torch_fsdp2 or args.use_megatron_fsdp: raise ValueError("GTP sequence sharding uses Megatron DDP, not FSDP") if args.use_tp_pp_dp_mapping: raise ValueError("GTP sequence sharding does not support --use-tp-pp-dp-mapping") # Total model size. assert args.world_size % total_model_size == 0, ( f"world size ({args.world_size}) is not divisible by total_model_size ({total_model_size=})" ) if args.attention_backend == AttnBackend.local: assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local' # Pipeline model parallel size. args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size args.data_parallel_size = args.world_size // total_model_size args.batch_parallel_size = args.world_size // ( args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size ) from megatron.training.config import InferenceSetupConfig, RLConfig rl_cfg = _default_config_from_args(RLConfig, args) _default_config_from_args(InferenceSetupConfig, args).validate() print_rank_0('using world size: {}, data-parallel size: {}, ' 'resolved context-parallel group size: {}, ' 'hierarchical context-parallel sizes: {}, ' 'tensor-model-parallel size: {}, ' 'pipeline-model-parallel size: {}'.format( args.world_size, args.data_parallel_size, args.context_parallel_size, args.hierarchical_context_parallel_sizes, args.tensor_model_parallel_size, args.pipeline_model_parallel_size)) # Checks. if args.hierarchical_context_parallel_sizes: from numpy import prod assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes) if "a2a+p2p" in args.cp_comm_type: assert args.hierarchical_context_parallel_sizes is not None, \ "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" if args.expert_tensor_parallel_size is None: args.expert_tensor_parallel_size = args.tensor_model_parallel_size # Deprecated arguments. assert args.batch_size is None, '--batch-size argument is no longer ' \ 'valid, use --micro-batch-size instead' del args.batch_size assert args.warmup is None, '--warmup argument is no longer valid, use ' \ '--lr-warmup-fraction instead' del args.warmup assert args.model_parallel_size is None, '--model-parallel-size is no ' \ 'longer valid, use --tensor-model-parallel-size instead' del args.model_parallel_size if args.checkpoint_activations: print_rank_0('--checkpoint-activations is no longer valid, use --recompute-activations, ' 'or, for more control, --recompute-granularity and --recompute-method.') exit() del args.checkpoint_activations if args.recompute_activations: args.recompute_granularity = 'selective' del args.recompute_activations if args.enable_cuda_graph or args.external_cuda_graph: assert ( args.cuda_graph_impl == "none" ), "Do not use --enable-cuda-graph or --external-cuda-graph with --cuda-graph-impl." assert ( not args.enable_cuda_graph or not args.external_cuda_graph ), "--enable-cuda-graph and --external-cuda-graph cannot be enabled at the same time." if args.enable_cuda_graph: print_rank_0( '--enable-cuda-graph is deprecated, use --cuda-graph-impl=local instead.', args.rank ) args.cuda_graph_impl = "local" del args.enable_cuda_graph if args.external_cuda_graph: print_rank_0( '--external-cuda-graph is deprecated, use --cuda-graph-impl=transformer_engine instead.', args.rank, ) args.cuda_graph_impl = "transformer_engine" del args.external_cuda_graph if getattr(args, 'cuda_graph_scope_deprecated', None) is not None: assert not args.cuda_graph_modules, ( "--cuda-graph-scope and --cuda-graph-modules cannot be used together." ) warn_rank_0( '--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.' ) args.cuda_graph_modules = args.cuda_graph_scope_deprecated del args.cuda_graph_scope_deprecated # Normalize cuda_graph_modules and inference_cuda_graph_scope early so that # all subsequent validation sees fully-typed enum values. _normalize_cuda_graph_modules_args(args) _normalize_inference_cuda_graph_scope_arg(args) assert ( args.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl] ), ( "Invalid inference CUDA graph scope " f"{args.inference_cuda_graph_scope.name!r} for " f"--cuda-graph-impl={args.cuda_graph_impl!r}." ) # Set input defaults. for key in defaults: # For default to be valid, it should not be provided in the # arguments that are passed to the program. We check this by # ensuring the arg is set to None. if getattr(args, key, None) is not None: warn_rank_0('Overriding default arguments for {key}:{v} ' 'with {key}:{v2}'.format(key=key, v=defaults[key], v2=getattr(args, key))) else: setattr(args, key, defaults[key]) if args.data_path is not None and args.split is None: legacy_default_split_value = '969, 30, 1' warn_rank_0('Please specify --split when using --data-path. Using legacy default value ' f'of "{legacy_default_split_value}"') args.split = legacy_default_split_value use_data_path = (args.data_path is not None) or (args.data_args_path is not None) if use_data_path: # Exactly one of the two has to be None if we use it. assert (args.data_path is None) or (args.data_args_path is None) use_per_split_data_path = any( elt is not None for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) or \ args.per_split_data_args_path is not None if use_per_split_data_path: # Exactly one of the two has to be None if we use it. assert any(elt is not None for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) is False or \ args.per_split_data_args_path is None if args.phase_transition_iterations: args.phase_transition_iterations = sorted( int(x.strip()) for x in args.phase_transition_iterations.split(",") ) # Batch size. assert args.micro_batch_size is not None assert args.micro_batch_size > 0 is_global_batch_size_explicitly_specified = getattr( args, '_is_global_batch_size_explicitly_specified', args.global_batch_size is not None ) if args.step_batch_size_schedule is not None and is_global_batch_size_explicitly_specified: raise ValueError( 'Cannot specify both --step-batch-size-schedule and --global-batch-size' ) if args.global_batch_size is None: args.global_batch_size = args.micro_batch_size * args.batch_parallel_size print_rank_0('setting global batch size to {}'.format(args.global_batch_size)) assert args.global_batch_size > 0 # Eval batch size defaults and validation. if args.eval_global_batch_size is None: args.eval_global_batch_size = args.global_batch_size if args.eval_micro_batch_size is None: args.eval_micro_batch_size = args.micro_batch_size # Evaluation microbatches depend on independent sequences, not weight shards. assert ( args.eval_global_batch_size % (args.eval_micro_batch_size * args.batch_parallel_size) == 0 ), ( f"eval_global_batch_size ({args.eval_global_batch_size}) must be divisible by " f"eval_micro_batch_size ({args.eval_micro_batch_size}) " f"* batch_parallel_size ({args.batch_parallel_size})" ) # === Hybrid layer pattern: deprecation handling and validation === # Backward compat: --hybrid-override-pattern is deprecated in favor of --hybrid-layer-pattern used_hybrid_override_pattern = False if args.hybrid_override_pattern is not None: assert args.hybrid_layer_pattern is None, ( '--hybrid-override-pattern and --hybrid-layer-pattern cannot both be specified. ' '--hybrid-override-pattern is deprecated; use --hybrid-layer-pattern instead.' ) warn_rank_0( "--hybrid-override-pattern is deprecated. Use --hybrid-layer-pattern instead.", args.rank, ) args.hybrid_layer_pattern = args.hybrid_override_pattern used_hybrid_override_pattern = True if args.mtp_hybrid_override_pattern is not None: warn_rank_0( "--mtp-hybrid-override-pattern is deprecated. " "For new hybrid models with MTP, use unified --hybrid-layer-pattern instead. " "Example: 'M*M*/MM/MM' means main='M*M*', MTP pattern='MM' with 2 depths. " "This argument is kept only for loading old checkpoints.", args.rank, ) from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_total_layer_count, get_hybrid_total_pipeline_segment_count, parse_hybrid_pattern, ) sep = Symbols.MTP_SEPARATOR # Backward compat: convert legacy mtp_hybrid_override_pattern to unified format if ( args.mtp_hybrid_override_pattern is not None and args.mtp_num_layers is not None and args.mtp_num_layers > 0 and (args.hybrid_layer_pattern is None or sep not in args.hybrid_layer_pattern) ): main_pattern = args.hybrid_layer_pattern or '' mtp_pattern = args.mtp_hybrid_override_pattern args.hybrid_layer_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) args.mtp_hybrid_override_pattern = None print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_layer_pattern}") if args.hybrid_layer_pattern is not None: # Derive num_layers from pattern; hybrid_layer_pattern always overrides --num-layers when # both are present (e.g. when loading from checkpoint with --use-checkpoint-args). num_layers_in_pattern = get_hybrid_total_layer_count(args.hybrid_layer_pattern) if args.num_layers is not None and args.num_layers != num_layers_in_pattern: warn_rank_0( f'--hybrid-layer-pattern is set; ignoring --num-layers ({args.num_layers}) and ' f'using the layer count derived from the pattern ({num_layers_in_pattern}).', args.rank, ) args.num_layers = num_layers_in_pattern # first/last pipeline num layers are incompatible with pipe-separated patterns # (the pipe separators already define the pipeline layout explicitly), but are # allowed for pipe-free patterns where they control uneven PP splitting. has_pipes = Symbols.PIPE in args.hybrid_layer_pattern.split(sep)[0] if has_pipes: assert args.decoder_first_pipeline_num_layers is None, ( 'If --hybrid-layer-pattern contains pipe separators, ' '--decoder-first-pipeline-num-layers should not be specified ' 'as the pipeline layout is explicitly defined.' ) assert args.decoder_last_pipeline_num_layers is None, ( 'If --hybrid-layer-pattern contains pipe separators, ' '--decoder-last-pipeline-num-layers should not be specified ' 'as the pipeline layout is explicitly defined.' ) assert args.num_layers_per_virtual_pipeline_stage is None, ( '--num-layers-per-virtual-pipeline-stage should not be used with ' '--hybrid-layer-pattern. To specify virtual pipelining, describe a number of ' 'pipeline segments in --hybrid-layer-pattern that is a multiple of ' '--pipeline-model-parallel-size greater than 1' ) assert args.num_virtual_stages_per_pipeline_rank is None, ( '--num-virtual-stages-per-pipeline-rank should not be used with ' '--hybrid-layer-pattern. Virtual pipeline stages are derived from the ' 'number of | segments in the pattern.' ) assert args.pipeline_model_parallel_layout is None, ( '--pipeline-model-parallel-layout should not be used with --hybrid-layer-pattern. ' 'Pipeline stage layout is defined by | separators in the pattern.' ) assert not args.account_for_embedding_in_pipeline_split, ( '--account-for-embedding-in-pipeline-split should not be used with ' '--hybrid-layer-pattern. Pipeline stage layout is defined by | separators ' 'in the pattern.' ) assert not args.account_for_loss_in_pipeline_split, ( '--account-for-loss-in-pipeline-split should not be used with ' '--hybrid-layer-pattern. Pipeline stage layout is defined by | separators ' 'in the pattern.' ) # Derive VPP from pipe segments in the pattern hybrid_pipeline_segments = get_hybrid_total_pipeline_segment_count( args.hybrid_layer_pattern ) if hybrid_pipeline_segments == 1 and args.transformer_pipeline_model_parallel_size > 1: # No pipes in pattern -- PP will be handled by select_pipeline_segment # at model init time (for backwards compatibility). args.virtual_pipeline_model_parallel_size = None else: assert hybrid_pipeline_segments % args.transformer_pipeline_model_parallel_size == 0, ( 'The number of hybrid pipeline segments described by --hybrid-layer-pattern must ' 'be evenly divisible by --pipeline-model-parallel-size. ' f'Got {hybrid_pipeline_segments} segments and ' f'{args.transformer_pipeline_model_parallel_size} pipeline parallel size.' ) if hybrid_pipeline_segments > args.transformer_pipeline_model_parallel_size: # Must be set here in order to assign virtual parallel ranks in # training.py/get_model args.virtual_pipeline_model_parallel_size = ( hybrid_pipeline_segments // args.transformer_pipeline_model_parallel_size ) else: args.virtual_pipeline_model_parallel_size = None # Infer mtp_num_layers from unified pattern if args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern: parsed = parse_hybrid_pattern(args.hybrid_layer_pattern) if parsed.mtp_pattern and parsed.mtp_num_depths > 0: inferred_mtp_num_layers = parsed.mtp_num_depths if args.mtp_num_layers is None: args.mtp_num_layers = inferred_mtp_num_layers elif args.mtp_num_layers != inferred_mtp_num_layers: warn_rank_0( f"--mtp-num-layers ({args.mtp_num_layers}) conflicts with " f"MTP depth count ({inferred_mtp_num_layers}) in pattern " f"'{args.hybrid_layer_pattern}'. " f"Using the inferred value ({inferred_mtp_num_layers}).", args.rank ) args.mtp_num_layers = inferred_mtp_num_layers # MTP validation if args.mtp_num_layers: assert args.position_embedding_type == "rope" or args.position_embedding_type == "none", ( f"Multi-Token Prediction (MTP) is not supported with {args.position_embedding_type} position embedding type." + f"The supported position embedding types are rope and none." ) if args.freeze_base_model_for_mtp: assert not args.freeze_all_layers, ( "--freeze-base-model-for-mtp cannot be combined with --freeze-all-layers." ) # All MLA-based hybrid attention symbols use MLA projections. if args.hybrid_layer_pattern and any( symbol in args.hybrid_layer_pattern for symbol in Symbols.MLA_ATTENTION ): args.multi_latent_attention = True # === End of hybrid layer pattern: deprecation handling and validation === # Uneven virtual pipeline parallelism assert ( int(args.num_layers_per_virtual_pipeline_stage is not None) + int(args.num_virtual_stages_per_pipeline_rank is not None) + int(args.pipeline_model_parallel_layout is not None) ) <= 1, ( 'No more than one of the following arguments can be set at the same time: ' '--num-layers-per-virtual-pipeline-stage, --num-virtual-stages-per-pipeline-rank,' '--pipeline-model-parallel-layout. ' f'{args.num_layers_per_virtual_pipeline_stage=}, ' f'{args.num_virtual_stages_per_pipeline_rank=}, ' f'{args.pipeline_model_parallel_layout=}.' ) if args.pipeline_model_parallel_layout is not None: # Parse the input flattened layout to a list and get the vpp size. # We will validate the layout more carefully in the TransformerConfig constructor. num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(args.pipeline_model_parallel_layout) assert num_stages % args.pipeline_model_parallel_size == 0, ( f"The length of pipeline_model_parallel_layout must be divisible" f" by pipeline_model_parallel_size ({num_stages=}," f" {args.pipeline_model_parallel_size=})" ) args.virtual_pipeline_model_parallel_size = num_stages // args.pipeline_model_parallel_size if args.virtual_pipeline_model_parallel_size == 1: args.virtual_pipeline_model_parallel_size = None elif args.num_layers_per_virtual_pipeline_stage is not None or args.num_virtual_stages_per_pipeline_rank is not None: if args.num_virtual_stages_per_pipeline_rank is None: assert args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None, \ 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' if args.num_layers is not None: num_layers = args.num_layers else: num_layers = args.decoder_num_layers if args.account_for_embedding_in_pipeline_split: num_layers += 1 if args.account_for_loss_in_pipeline_split: num_layers += 1 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ 'number of layers of the model must be divisible pipeline model parallel size' num_layers_per_pipeline_stage = num_layers // args.transformer_pipeline_model_parallel_size assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ args.num_layers_per_virtual_pipeline_stage else: args.virtual_pipeline_model_parallel_size = args.num_virtual_stages_per_pipeline_rank if args.virtual_pipeline_model_parallel_size == 1: args.virtual_pipeline_model_parallel_size = None else: # Only set VPP to None if it wasn't already derived from --hybrid-layer-pattern if args.hybrid_layer_pattern is None: args.virtual_pipeline_model_parallel_size = None if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None: # Divisibility check not applicable for T5 models which specify encoder_num_layers # and decoder_num_layers, or for hybrid models using --hybrid-layer-pattern. if args.num_layers is not None and args.hybrid_layer_pattern is None: num_layers = args.num_layers if args.account_for_embedding_in_pipeline_split: num_layers += 1 if args.account_for_loss_in_pipeline_split: num_layers += 1 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ 'Number of layers should be divisible by the pipeline-model-parallel size' if args.virtual_pipeline_model_parallel_size is not None: if args.overlap_p2p_comm: assert args.pipeline_model_parallel_size > 1, \ 'When interleaved schedule is used, pipeline-model-parallel size '\ 'should be greater than 1' else: assert args.pipeline_model_parallel_size > 2, \ 'When interleaved schedule is used and p2p communication overlap is disabled, '\ 'pipeline-model-parallel size should be greater than 2 to avoid having multiple '\ 'p2p sends and recvs between same 2 ranks per communication batch' else: # Overlap P2P communication is disabled if not using the interleaved schedule. args.overlap_p2p_comm = False args.align_param_gather = False # Only print warning if PP size > 1. if args.rank == 0 and args.pipeline_model_parallel_size > 1: print('WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' 'since non-interleaved schedule does not support overlapping p2p communication ' 'and aligned param AG') print_rank_0( f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}" ) # A shortcut block calls its paired layers' sub-methods directly rather than their forward, so # the FSDP parameter all-gather hooks registered on the TransformerLayer/MambaLayer FSDP units # never fire and those parameters stay sharded. The expert-parallel overlap schedule hit the # same problem and needed explicit release hooks that only cover TransformerLayer, HybridStack # and MTP layers, none of which a shortcut block is. assert not ( args.moe_shortcut_connection and (args.use_torch_fsdp2 or args.use_megatron_fsdp) ), ( "FSDP is not supported with --moe-shortcut-connection: the shortcut block bypasses the " "per-layer FSDP parameter all-gather hooks, leaving the paired attention and MoE layer " "parameters sharded. Use DDP or --use-distributed-optimizer instead." ) if args.use_torch_fsdp2: assert is_torch_min_version("2.4.0"), \ 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' assert args.pipeline_model_parallel_size == 1, \ '--use-torch-fsdp2 is not supported with pipeline parallelism' assert args.expert_model_parallel_size == 1, \ '--use-torch-fsdp2 is not supported with expert parallelism' assert not args.gradient_accumulation_fusion, \ '--use-torch-fsdp2 is not supported with gradient accumulation fusion' assert args.ckpt_format in ('torch_dist', 'torch_dcp'), \ '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' assert args.untie_embeddings_and_output_weights, \ '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' assert not args.fp16, \ '--use-torch-fsdp2 not supported with fp16 yet' assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' if args.fp8_param_gather and is_te_min_version("2.0.0"): args.fp8_param_gather = False warn_rank_0( 'FSDP2 FP8 param gather is not supported yet in TE 2.0, will fallback to bf16' 'all_gather instead, turning off fp8_param_gather', args.rank, ) if args.fp4_param_gather and not is_te_min_version("2.7.0.dev0"): args.fp4_param_gather = False warn_rank_0( 'FSDP2 FP4 param gather is not supported yet in TE 2.0, will fallback to bf16' 'all_gather instead, turning off fp4_param_gather', args.rank, ) # Map string data-type to torch.dtype. dtype_map = { 'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16, 'fp8': torch.uint8, 'auto': None, None: None, } map_dtype = lambda d: d if isinstance(d, torch.dtype) else dtype_map[d] args.mamba_inference_conv_states_dtype = map_dtype(args.mamba_inference_conv_states_dtype) args.mamba_inference_ssm_states_dtype = map_dtype(args.mamba_inference_ssm_states_dtype) args.mamba_training_ssm_states_dtype = map_dtype(args.mamba_training_ssm_states_dtype) args.logit_dtype = map_dtype(getattr(args, 'logit_dtype', None)) args.megatron_fsdp_main_params_dtype = map_dtype(args.megatron_fsdp_main_params_dtype) args.megatron_fsdp_main_grads_dtype = map_dtype(args.megatron_fsdp_main_grads_dtype) args.megatron_fsdp_grad_comm_dtype = map_dtype(args.megatron_fsdp_grad_comm_dtype) if args.fp8_param_gather: assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' if args.fp4_param_gather: assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ '--fp4-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") # FP4 param requires FP4 mode if args.fp4_param_gather and not args.fp4: raise ValueError("--fp4-param-gather must be used together with --fp4-format.") # FP4 requires TE >= 2.7.0.dev0 if args.fp4 and not is_te_min_version("2.7.0.dev0"): raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.") if ( args.fp8_recipe == 'mxfp8' and args.transformer_impl == 'inference_optimized' and not is_flashinfer_min_version("0.6.4") ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") if args.moe_megakernel_backend == "mok" and (args.use_megatron_fsdp or args.use_torch_fsdp2): raise ValueError("MOK has not been validated with Megatron-FSDP or Torch FSDP2") if args.use_megatron_fsdp: # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. # Please use `use_megatron_fsdp` instead, as all functionality will be migrated there. # Future updates will drop support for `use_custom_fsdp` to avoid confusion. args.use_custom_fsdp = True if args.megatron_fsdp_version == 2: assert not args.use_distributed_optimizer, \ '--megatron-fsdp-version 2 is not compatible with --use-distributed-optimizer' else: # Megatron-FSDP v1 requires the DistributedOptimizer. if not args.use_distributed_optimizer: warn_rank_0( 'Megatron-FSDP v1 is only compatible with --use-distributed-optimizer. ' 'Using DistributedOptimizer...', args.rank, ) args.use_distributed_optimizer = True # Optimizer step MXFP8 buffer operation that is not relevant or supported for Megatron-FSDP. args.reuse_grad_buf_for_mxfp8_param_ag = False if args.moe_single_grouped_weight or args.moe_single_grouped_bias: # Megatron-FSDP currently remaps module parameters through plain Tensor and TE # Float8Tensor/MXFP8Tensor storage paths. TE GroupedTensor parameters need their # grouped backing storage remapped instead; quantized grouped tensors also need # grouped scale/amax handling. DDP has a separate GroupedTensor-aware path. raise ValueError( "Megatron-FSDP does not currently support moe_single_grouped_weight or " "moe_single_grouped_bias. Disable single grouped MoE parameters or use the " "regular DDP/distributed optimizer path until Megatron-FSDP supports TE " "GroupedTensor param buffers." ) # Optimizer compatibility check. assert args.optimizer in ('sgd', 'adam'), \ f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." # Expert parameters may be sharded differently from non-expert parameters, in which # case both strategies have to be considered by model-wide checks. sharding_strategies = {args.data_parallel_sharding_strategy} if args.expert_data_parallel_sharding_strategy is not None: sharding_strategies.add(args.expert_data_parallel_sharding_strategy) if ( sharding_strategies & {"optim_grads_params", "optim_grads"} and args.gradient_accumulation_fusion ): warn_rank_0( 'Verify that fused gradient accumulation is supported by TransformerEngine for Megatron-FSDP.', args.rank, ) if "optim_grads_params" in sharding_strategies: assert args.check_weight_hash_across_dp_replicas_interval is None, \ 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ 'FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 or unset.' assert args.ckpt_format == "fsdp_dtensor", \ "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." if args.nccl_ub: # In Megatron-LM, required implementation for manual registration is already provided. # So we enable the manual registration by default when nccl-ub and use_megatron_fsdp is set. args.fsdp_manual_registration = True args.fsdp_double_buffer = True warn_rank_0('FSDP double buffer and manual registration is enabled by default when --nccl-ub is enabled!') if args.megatron_fsdp_max_pool_double_buffer: # MaxPoolAllocator is a type of FSDP double buffer. args.fsdp_double_buffer = True if args.init_model_with_meta_device and sharding_strategies == {"no_shard"}: raise ValueError( "Meta device initialization (init_model_with_meta_device=True) is not " "supported or necessary for the 'no_shard' / 0 sharding strategy." ) if args.fsdp_manual_registration: assert args.use_megatron_fsdp, "FSDP manual registration is only supported with Megatron FSDP." assert args.nccl_ub, "FSDP manual registration is only supported with --nccl-ub argument." # Parameters dtype. args.params_dtype = torch.float if args.fp16: assert not args.bf16 args.params_dtype = torch.half # Turn off checking for NaNs in loss and grads if using dynamic loss scaling, # where NaNs in grads / loss are signal to the loss scaler. if not args.loss_scale: args.check_for_nan_in_loss_and_grad = False warn_rank_0('Setting args.check_for_nan_in_loss_and_grad to False since ' 'dynamic loss scaling is being used') # main_grads_dtype is still the CLI string here; it is mapped to a torch.dtype later, in # OptimizerConfig.finalize_optimizer_config. main_grads_is_fp32 = args.main_grads_dtype in ('fp32', torch.float32) if args.bf16: assert not args.fp16 args.params_dtype = torch.bfloat16 # bfloat16 requires gradient accumulation and all-reduce to # be done in fp32. if args.accumulate_allreduce_grads_in_fp32: assert main_grads_is_fp32, \ "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" if args.grad_reduce_in_bf16: args.accumulate_allreduce_grads_in_fp32 = False elif not args.accumulate_allreduce_grads_in_fp32 and main_grads_is_fp32: args.accumulate_allreduce_grads_in_fp32 = True print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.') if args.accumulate_allreduce_grads_in_fp32: # The FP32-accumulation reduce-scatters only exist to recover FP32 accumulation from a # lower-precision wire dtype. With FP32 main_grads the plain ring reduce-scatter already # sums in FP32 on both axes, so they buy no precision and cost one extra # unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter. for arg_name in ( 'ddp_reduce_scatter_with_fp32_accumulation', 'gtp_remat_reduce_scatter_with_fp32_accumulation', ): if getattr(args, arg_name, False): setattr(args, arg_name, False) warn_rank_0( f"Setting args.{arg_name} to False since " "--accumulate-allreduce-grads-in-fp32 already reduces in fp32" ) if args.inference_cuda_graph_scope == InferenceCudaGraphScope.block: if args.fp8 is not None: assert args.transformer_impl == "inference_optimized", \ "fp8 with --inference-cuda-graph-scope=block is only supported with " \ "--transformer-impl=inference_optimized" assert args.fp8_recipe == "mxfp8", \ "Only --fp8-recipe=mxfp8 is supported with --inference-cuda-graph-scope=block" if args.cuda_graph_impl == 'local': assert args.inference_dynamic_batching_num_cuda_graphs > 0 or args.inference_dynamic_batching_num_cuda_graphs == -1, \ 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' \ '-1 means that we will automatically determine the number of CUDA graphs to capture based on the `max_requests` value.' if args.dataloader_type is None: args.dataloader_type = 'single' # data assert args.num_dataset_builder_threads > 0 if args.num_layers is not None: assert args.encoder_num_layers is None, \ 'cannot have both num-layers and encoder-num-layers specified' args.encoder_num_layers = args.num_layers else: assert args.encoder_num_layers is not None, \ 'either num-layers or encoder-num-layers should be specified' args.num_layers = args.encoder_num_layers # Check required arguments. required_args = ['num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'] for req_arg in required_args: _check_arg_is_not_none(args, req_arg) # Checks. if args.ffn_hidden_size is None: if args.swiglu: # reduce the dimnesion for MLP since projections happens on # two linear layers. this keeps the number of paramters in # the same ballpark as the counterpart with 4*h size # we keep it a multiple of 64, which means the actual tensor size # will be a multiple of 64 / tp_size args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64 else: args.ffn_hidden_size = 4 * args.hidden_size if args.kv_channels is None: assert args.hidden_size % args.num_attention_heads == 0 args.kv_channels = args.hidden_size // args.num_attention_heads if args.seq_length is not None and args.context_parallel_size > 1: assert args.seq_length % (args.context_parallel_size * 2) == 0, \ f"seq-length ({args.seq_length}) must be a multiple of 2 * the resolved " \ f"context-parallel group size ({args.context_parallel_size}), " \ "including GTP sequence shards." if getattr(args, 'dataloader_inter_document_masking', False): # The dataset omits attention_mask when inter-document masking is # enabled; disable the flag to avoid a TP broadcast mismatch. if args.create_attention_mask_in_dataloader: args.create_attention_mask_in_dataloader = False if args.seq_length is not None: assert args.encoder_seq_length is None args.encoder_seq_length = args.seq_length else: assert args.encoder_seq_length is not None args.seq_length = args.encoder_seq_length if args.seq_length is not None: assert args.max_position_embeddings >= args.seq_length, \ f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " \ f"or equal to seq_length ({args.seq_length})." if args.decoder_seq_length is not None: assert args.max_position_embeddings >= args.decoder_seq_length if args.save is not None: assert args.save_interval is not None assert args.save_interval > 0 if args.save_retain_interval is not None: assert args.save_retain_interval > 0 assert args.save_retain_interval % args.save_interval == 0 if args.save_params_interval is not None: assert not args.overlap_param_gather # Mixed precision checks. if args.fp16_lm_cross_entropy: assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' if args.fp32_residual_connection: assert args.fp16 or args.bf16, \ 'residual connection in fp32 only supported when using fp16 or bf16.' if args.moe_grouped_gemm: dc = torch.cuda.get_device_capability() assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." # Persistent fused layer norm. if not is_torch_min_version("1.11.0a0"): args.no_persist_layer_norm = True print_rank_0('Persistent fused layer norm kernel is supported from ' 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' 'Defaulting to no_persist_layer_norm=True') # Activation recomputing. if args.distribute_saved_activations: assert args.tensor_model_parallel_size > 1, 'can distribute ' \ 'recomputed activations only across tensor model ' \ 'parallel groups' assert args.recompute_granularity == 'full', \ 'distributed recompute activations is only '\ 'application to full recompute granularity' assert args.recompute_method is not None, \ 'for distributed recompute activations to work you '\ 'need to use a recompute method ' assert is_torch_min_version("1.10.0a0"), \ 'distributed recompute activations are supported for pytorch ' \ 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ f'pytorch version is v{get_torch_version()}.' if args.recompute_granularity == 'selective': assert args.recompute_method is None, \ 'recompute method is not yet supported for ' \ 'selective recomputing granularity' # disable sequence parallelism when tp=1 # to avoid change in numerics when # sequence_parallelism is enabled. if args.tensor_model_parallel_size == 1: if args.sequence_parallel: warn_rank_0( "Disabling sequence parallelism because tensor model parallelism is disabled", args.rank, ) args.sequence_parallel = False if args.tp_comm_overlap: assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' if args.hybrid_context_parallel: warn_rank_0( "--hybrid-context-parallel is deprecated. Use --dynamic-context-parallel instead.", args.rank, ) if args.dynamic_context_parallel: raise ValueError( "Cannot set both --hybrid-context-parallel and --dynamic-context-parallel" ) args.dynamic_context_parallel = True args.hybrid_context_parallel = False if args.dynamic_context_parallel: assert not args.use_megatron_fsdp, 'Dynamic context parallelism not supported with Megatron FSDP' assert args.dataloader_type == 'single', 'Dynamic context parallelism only supported with single dataloader type' if args.sequence_packing_scheduler is None: args.sequence_packing_scheduler = 'default_dynamic_cp' if args.sequence_packing_scheduler != 'default_dynamic_cp': raise ValueError( "Dynamic context parallelism requires " "--sequence-packing-scheduler default_dynamic_cp" ) dp_cp_size = args.data_parallel_size * args.context_parallel_size from megatron.core.parallel_state import ( get_valid_dynamic_context_parallel_group_sizes, ) valid_group_sizes = get_valid_dynamic_context_parallel_group_sizes(dp_cp_size) if args.min_dynamic_context_parallel_size not in valid_group_sizes: raise ValueError( "--min-dynamic-context-parallel-size=" f"{args.min_dynamic_context_parallel_size} is invalid for DPxCP size " f"{dp_cp_size}; expected one of {valid_group_sizes}" ) # Support for variable sequence lengths across batches/microbatches. # set it if the dataloader supports generation of variable sequence lengths # across batches/microbatches. Due to additional communication overhead # during pipeline parallelism, it should not be set if sequence length # is constant during training. args.variable_seq_lengths = False if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None: args.sft_mock_dataset_config_json = json.dumps( { "mode": "distribution", "type": "lognormal", "min_seq_len": args.seq_length // 2, "max_seq_len": args.seq_length, "mean_seq_len": args.seq_length // 4 * 3, "lognormal_sigma": 1.1, } ) # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ and get_device_arch_version() < 10: # CUDA_DEVICE_MAX_CONNECTIONS requirement no longer exists since the Blackwell architecture if args.use_torch_fsdp2 or args.use_megatron_fsdp: fsdp_impl = "Torch-FSDP2" if args.use_torch_fsdp2 else "Megatron-FSDP" warn_rank_0( f"Using tensor model parallelism or context parallelism with {fsdp_impl} together. " "Try not to using them together since they require different CUDA_MAX_CONNECTIONS " "settings for best performance. sequence parallelism requires setting the " f"environment variable CUDA_DEVICE_MAX_CONNECTIONS to 1 while {fsdp_impl} " "requires not setting CUDA_DEVICE_MAX_CONNECTIONS=1 for better parallelization.", args.rank, ) elif args.overlap_moe_expert_parallel_comm: warn_rank_0( "For Hopper and before, try not to use tensor model parallelism or context parallelism with overlap_moe_expert_parallel_comm. " "Using tensor/context model parallelism requires setting the environment " "variable CUDA_DEVICE_MAX_CONNECTIONS to 1 to maximize the performance. " "While overlap_moe_expert_parallel_comm requires setting a larger CUDA_DEVICE_MAX_CONNECTIONS " "for better parallelization. If you want to use both, you can set CUDA_DEVICE_MAX_CONNECTIONS to 1 or 32, " "which depends on which parallelization you want to prioritize.", args.rank, ) else: assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \ "Using tensor model parallelism or context parallelism require setting the environment variable " \ "CUDA_DEVICE_MAX_CONNECTIONS to 1" # Setting FSDP communication groups for high priority streams for Blackwell and later architectures # 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. if args.use_torch_fsdp2 or args.use_megatron_fsdp and get_device_arch_version() >= 10: if 'dp_cp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('dp_cp') if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('ep_dp') # Derive the internal expert_gtp_weight_remat_size from the # user-facing --expert-tensor-parallel-num-weight-shards (expert_tensor_parallel_size is # defaulted earlier in validate_args). expert_gtp_weight_remat_size has no CLI flag. (args.expert_tensor_parallel_num_weight_shards, args.expert_gtp_weight_remat_size) = ( resolve_tensor_parallel_weight_shards( args.expert_tensor_parallel_size, args.expert_tensor_parallel_num_weight_shards, getattr(args, "expert_gtp_weight_remat_size", 1), ) ) if args.gtp_weight_remat_size > 1 or args.expert_gtp_weight_remat_size > 1: if args.fp4 and not args.fp4_param_gather: raise ValueError( "GTP (--tensor-parallel-num-weight-shards / " "--expert-tensor-parallel-num-weight-shards > 1) with --fp4-format requires " "--fp4-param-gather so NVFP4 weights are all-gathered as native NVFP4." ) gtp_weight_remat_size = args.gtp_weight_remat_size egtp_weight_remat_size = args.expert_gtp_weight_remat_size if get_device_arch_version() >= 10: # Setting GTP communication groups for high priority streams for Blackwell and later # architectures. 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. if 'gtp_remat' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('gtp_remat') warn_rank_0("Setting 'gtp_remat' group for high priority streams.") if ( egtp_weight_remat_size > 1 and 'expt_gtp_remat' not in args.high_priority_stream_groups ): args.high_priority_stream_groups.append('expt_gtp_remat') warn_rank_0("Setting 'expt_gtp_remat' group for high priority streams.") # Sanity check for 'CUDA_GRAPHS_USE_NODE_PRIORITY'. if args.cuda_graph_impl != "none": assert os.environ.get('CUDA_GRAPHS_USE_NODE_PRIORITY') == "1", \ 'GTP requires CUDA_GRAPHS_USE_NODE_PRIORITY=1 to make sure fine-grained GTP ' \ 'comms can be well overlapped with GEMMs when CudaGraph is enabled for ' \ 'Blackwell and later architecture.' # Sanity check for 'NCCL_PROTO'. if os.environ.get('NCCL_PROTO', '').lower() == "simple": warn_rank_0( "Generally GTP prefers 'NCCL_PROTO=LL128 or LL' while get 'NCCL_PROTO=simple', " "force setting NCCL_PROTO=Simple might introduce bad perf." ) assert not args.ddp_average_in_collective, ( "GTP requires --ddp-average-in-collective off (the default); averaged collectives " "would need per-buffer 1/gtp_remat scaling." ) assert args.ckpt_format in ('torch', 'torch_dist'), ( f"GTP supports only --ckpt-format 'torch' (legacy) or 'torch_dist', got " f"'{args.ckpt_format}'." ) assert not ( getattr(args, 'dist_ckpt_optim_fully_reshardable', False) and getattr(args, 'distrib_optim_fully_reshardable_mem_efficient', False) ), ( "GTP does not support the distributed-optimizer fully-reshardable + " "mem-efficient checkpoint mode. Disable " "--distrib-optim-fully-reshardable-mem-efficient (or " "--dist-ckpt-optim-fully-reshardable)." ) # GTP with the mxfp8 recipe requires --fp8-param-gather: GTP keeps no bf16 weight and # relies on the optimizer maintaining the fp8 shard (the forward all-gathers fp8 and does # not re-quantize). Without fp8-param-gather the fp8 forward weight would never be updated. if getattr(args, 'fp8_recipe', None) == 'mxfp8': assert getattr(args, 'fp8_param_gather', False), ( "GTP + mxfp8 requires --fp8-param-gather (the optimizer maintains the fp8 shard; " "GTP does not keep or re-quantize a bf16 weight)." ) # MXFP8 params cannot be mapped into the contiguous param buffer (TE's # replace_raw_data does not support the MXFP8 tile-scaling layout), so the param # all-gather must reuse the grad buffer instead. assert getattr(args, 'reuse_grad_buf_for_mxfp8_param_ag', False), ( "GTP + mxfp8 + --fp8-param-gather requires --reuse-grad-buf-for-mxfp8-param-ag " "(MXFP8 params keep their own quantized storage; mapping them into the param " "buffer via replace_raw_data is unsupported)." ) # GTP symmetric memory registers pools with symmetric=True (NVLS needs symmetric # windows), which contradicts --disable-symmetric-registration. if getattr(args, 'gtp_remat_nccl_ub', False) or getattr(args, 'gtp_expert_remat_nccl_ub', False): assert not getattr(args, 'disable_symmetric_registration', False), ( "--gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub require symmetric window registration and " "cannot be combined with --disable-symmetric-registration." ) if getattr(args, 'gtp_remat_reduce_scatter_with_fp32_accumulation', False): print_rank_0( "WARNING: --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub take precedence over " "--gtp-remat-reduce-scatter-with-fp32-accumulation on their groups: NVLS " "symmetric reduce-scatters accumulate in fp32 in-switch " "(NCCL multimem.ld_reduce .acc::f32)." ) # Disable bias gelu fusion if we are disabling bias altogether if not args.add_bias_linear: args.bias_gelu_fusion = False # Keep the 'add bias' args in sync; add_qkv_bias is more targeted. if args.add_bias_linear: args.add_qkv_bias = True if args.qk_clip: assert is_te_min_version("2.9.0"), \ '--qk-clip is only supported with TE >= 2.9.0.' assert 0.0 < args.qk_clip_alpha < 1.0, \ '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' assert args.qk_clip_threshold > 0, \ '--qk-clip-threshold must be greater than 0 when using --qk-clip.' # decoupled log max attention logit check if args.log_max_attention_logit: assert is_te_min_version("2.9.0"), \ '--log-max-attention-logit is only supported with TE >= 2.9.0.' # Legacy RoPE arguments if args.use_rotary_position_embeddings: args.position_embedding_type = 'rope' if args.position_embedding_type != 'rope': args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now # don't allow it to keep things simple if not args.add_position_embedding and args.position_embedding_type != 'rope': raise RuntimeError('--no-position-embedding is deprecated, use --position-embedding-type') # Relative position embeddings arguments if args.position_embedding_type == 'relative': assert ( args.transformer_impl == "transformer_engine" ), 'Local transformer implementation currently does not support attention bias-based position embeddings.' # MultiModal rotary embeddings arguments if args.position_embedding_type == "mrope": assert args.mrope_section is not None, \ '--mrope-section should be set when using --position-embedding-type mrope.' # MoE Spec check if args.num_experts == 0: args.num_experts = None if args.num_experts is not None and args.moe_ffn_hidden_size is None: args.moe_ffn_hidden_size = args.ffn_hidden_size warn_rank_0("moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.") # Expert parallelism check if args.expert_model_parallel_size > 1: assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" assert args.num_experts % args.expert_model_parallel_size == 0, \ "Number of experts should be a multiple of expert model parallel_size." # MoE router check if isinstance(args.moe_router_load_balancing_type, list) and len(args.moe_router_load_balancing_type) == 1: args.moe_router_load_balancing_type = args.moe_router_load_balancing_type[0] if isinstance(args.moe_aux_loss_coeff, list) and len(args.moe_aux_loss_coeff) == 1: args.moe_aux_loss_coeff = args.moe_aux_loss_coeff[0] # torch_dcp (torch.distributed.checkpoint) checkpointing format checks. if args.ckpt_format == "torch_dcp": assert args.use_torch_fsdp2, "--ckpt-format torch_dcp is only tested with FSDP." assert args.tensor_model_parallel_size <= 1, \ "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." assert args.pipeline_model_parallel_size <= 1, \ "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." # fsdp_dtensor checkpointing format checks. if args.ckpt_format == "fsdp_dtensor": assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." # --use-varlen-dataset: independent of --sft. Cannot be combined with --sft # because they are mutually-exclusive top-level dataset selectors that both # drive the packed-sequence (THD) path. These stay in validate_args: the # selectors are CLI-level args, not core config fields. if args.use_varlen_dataset: assert not args.sft, ( "--use-varlen-dataset and --sft are mutually exclusive; both " "select the packed-sequence dataset family. Pick one." ) if args.varlen_sbhd_validation: assert args.sequence_packing_scheduler is None, ( "--varlen-sbhd-validation does not use a sequence packing " "scheduler; drop --sequence-packing-scheduler." ) # SBHD validation is a real-data numerical-reference path only; # MockVarlenDataset does not implement it. assert not args.mock_data, ( "--varlen-sbhd-validation is not supported with --mock-data; " "SBHD validation requires a real dataset." ) else: # VarlenDataset emits one unpacked sample per __getitem__; it # relies on an upstream packing scheduler to group variable-length # samples into THD batches. Auto-pick ``dp_balanced`` when the # user did not request one explicitly. if args.sequence_packing_scheduler is None: args.sequence_packing_scheduler = ( 'default_dynamic_cp' if args.dynamic_context_parallel else 'dp_balanced' ) # Runs after the varlen auto-select above so it sees the final resolved # scheduler. Scheduler-name and max-seqlen validation live in # ModelParallelConfig.__post_init__; only the buffer-size check stays here # because seq_length is not a core config field. The None case for # max_seqlen_per_dp_cp_rank is rejected by the config check. if ( args.sequence_packing_scheduler == 'default_dynamic_cp' and not args.dynamic_context_parallel ): raise ValueError( "--sequence-packing-scheduler default_dynamic_cp requires " "--dynamic-context-parallel" ) if args.sequence_packing_scheduler is not None: args.variable_seq_lengths = True # Packed microbatches carry different numbers of valid tokens, so the # default per-microbatch loss averaging would weight tokens unevenly # depending on how samples happened to be packed (same reasoning as # the hybrid-context-parallel check above). assert args.calculate_per_token_loss, ( 'Sequence packing must be used with --calculate-per-token-loss' ) if args.max_seqlen_per_dp_cp_rank is not None: total_cp_ranks = args.context_parallel_size if args.sequence_packing_scheduler == 'default_dynamic_cp': total_cp_ranks *= args.data_parallel_size assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' f'must be >= single sequence max length ({args.seq_length})' ) if args.sequence_packing_scheduler == "dp_balanced" and args.cuda_graph_impl == "transformer_engine": assert args.use_varlen_dataset and args.micro_batch_size == 1, ( "Scheduler CUDA graphs require one variable-length sample per dataset item " "and --micro-batch-size 1." ) assert args.step_batch_size_schedule is None, ( "Scheduler CUDA graphs require a fixed source global batch size." ) # Data blend checks assert args.mock_data + \ bool(args.data_path) + \ any([args.train_data_path, args.valid_data_path, args.test_data_path]) \ <= 1, "A single data source must be provided in training mode, else None" if args.fim_data: extra_tokens = [ args.fim_prefix_token, args.fim_middle_token, args.fim_suffix_token, args.fim_pad_token, args.fim_eod_token, ] assert not args.mock_data, "Mock dataset is not supported with FIM dataset." assert args.fim_rate, "--fim-rate should be specified." assert args.fim_spm_rate, "--fim-spm-rate should be specified." assert all(token is not None for token in extra_tokens), "FIM extra tokens should be specified." assert not ( args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te' and not is_te_min_version("2.19.0") ), ( "Transformer Engine cross entropy loss fusion is disabled due to stability issues. " "Use --cross-entropy-fusion-impl native, or omit --cross-entropy-loss-fusion." ) # Deterministic mode — env vars + config overrides + torch global state. # Implementation lives in ``megatron/training/determinism.py`` so the # same setup is reachable from tests / profiling scripts that don't go # through argparse. if args.deterministic_mode: from megatron.training.determinism import apply_determinism_to_args apply_determinism_to_args(args) # Update the printed args to reflect that `apply_query_key_layer_scaling` also controls `attention_softmax_in_fp32` if args.apply_query_key_layer_scaling: args.attention_softmax_in_fp32 = True if args.result_rejected_tracker_filename is not None: # Append to passed-in args.iterations_to_skip. iterations_to_skip_from_file = RerunStateMachine.get_skipped_iterations_from_tracker_file( args.result_rejected_tracker_filename ) args.iterations_to_skip.extend(iterations_to_skip_from_file) if args.optimizer not in ('sgd', 'adam'): assert not args.use_torch_fsdp2, "Emerging optimizer does not support Torch-FSDP2 for now." assert not args.use_megatron_fsdp, "Emerging optimizer does not support Megatron-FSDP for now." assert args.ckpt_format in ["torch", "torch_dist"], "Emerging optimizer supports torch and torch_dist checkpoint format." # Make sure all functionality that requires Gloo process groups is disabled. if not args.use_gloo_process_groups: if args.use_distributed_optimizer and args.optimizer in ('sgd', 'adam'): # If using distributed optimizer, must use distributed checkpointing. # Legacy checkpointing uses Gloo process groups to collect full distributed # optimizer state in the CPU memory of DP rank 0. assert args.use_dist_ckpt if args.dist_ckpt_optim_fully_reshardable: assert not args.distrib_optim_fully_reshardable_mem_efficient, \ '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' if args.fake_process_group: assert args.moe_token_dispatcher_type != "flex", "Fake process group is not supported with flex token dispatcher." # Disable nan check for fake process group args.check_for_nan_in_loss_and_grad = False warn_rank_0('check_for_nan_in_loss_and_grad is set to False for fake process group.') # Disable gloo process groups for fake process group args.use_gloo_process_groups = False warn_rank_0('use_gloo_process_groups is set to False for fake process group.') # Checkpointing if args.ckpt_fully_parallel_save_deprecated and args.rank == 0: print('--ckpt-fully-parallel-save flag is deprecated and has no effect.' ' Use --no-ckpt-fully-parallel-save to disable parallel save.') if ( args.use_dist_ckpt and not args.ckpt_fully_parallel_save and args.use_distributed_optimizer and args.rank == 0 ): print('Warning: With non-parallel ckpt save and DistributedOptimizer,' ' it will be impossible to resume training with different parallelism.' ' Consider removing flag --no-ckpt-fully-parallel-save.') if args.use_dist_ckpt_deprecated and args.rank == 0: print('--use-dist-ckpt is deprecated and has no effect.' ' Use --ckpt-format to select the checkpoint format.') if args.dist_ckpt_format_deprecated and args.rank == 0: print('--dist-ckpt-format is deprecated and has no effect.' ' Use --ckpt-format to select the checkpoint format.') if args.use_dist_ckpt and args.ckpt_fully_parallel_load: if args.ckpt_fully_parallel_load_exchange_algo != "broadcast": warn_rank_0( "Currently only the 'broadcast' exchange algorithm is supported for fully parallel load. " "Other algorithms cannot guarantee numerical stability yet." ) # The PG-distribution cache stores one file per parallelization group, keyed by # that group's ordered global ranks, and each file holds both the save and the # load distribution derived from one gather over that group. A single create # run (even one that only loads) therefore covers later saves only when save # and load use the same group; otherwise a read run finds no cache file for the # other side's group. Require both sides to use the same group. if args.ckpt_pg_tensors_cache_path is not None: assert ( args.ckpt_fully_parallel_save_process_group == args.ckpt_fully_parallel_load_process_group ), ( "--ckpt-pg-tensors-cache-path requires --ckpt-fully-parallel-save-process-group and " "--ckpt-fully-parallel-load-process-group to be identical, but got " f"save='{args.ckpt_fully_parallel_save_process_group}' and " f"load='{args.ckpt_fully_parallel_load_process_group}'. The cache is keyed by the " "parallelization group, so with mixed groups a run that creates the cache for one " "side leaves the other side without a cache file." ) if args.load_main_params_from_ckpt: assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' if args.stream_ckpt_dequant: # Inference-only load path: the loaded state dict holds the re-quantized params # themselves, so nothing that consumes the loaded values in high precision (optimizer # main params) may run. `load_checkpoint` additionally refuses any optimizer. assert args.no_load_optim and not args.finetune, ( '--stream-ckpt-dequant is an inference-only load path: it requires --no-load-optim ' 'and cannot be used with --finetune.' ) assert not args.load_main_params_from_ckpt, ( '--stream-ckpt-dequant cannot be used with --load-main-params-from-ckpt: the main ' 'params would be initialized from the re-quantized params instead of the ' 'high-precision checkpoint values.' ) assert args.fp8 is None or args.fp8_recipe != 'delayed', ( '--stream-ckpt-dequant does not support --fp8-recipe delayed (its scale is only ' 'restored by load_state_dict); use tensorwise, mxfp8, blockwise or nvfp4.' ) assert ( not args.ckpt_fully_parallel_load or args.ckpt_fully_parallel_load_exchange_algo == 'broadcast' ), ( '--stream-ckpt-dequant supports --ckpt-fully-parallel-load only with ' '--ckpt-fully-parallel-load-exchange-algo broadcast.' ) if args.use_dist_ckpt and args.async_save: if not args.use_persistent_ckpt_worker: warn_rank_0( '--async-save is not supported without --use-persistent-ckpt-worker. ' 'Disabling --async-save.' ) args.async_save = False if args.logits_save_dir is not None: assert args.async_save, ( '--logits-save-dir requires --async-save (and --use-persistent-ckpt-worker). ' 'Logits are flushed as an async request in the checkpoint queue.' ) if not args.freeze_all_layers: warn_rank_0( '--logits-save-dir without --freeze-all-layers: the LM loss is still computed and ' 'gradients will update the model while logits are dumped. This is intended only ' 'when dumping logits during active training; for a frozen-teacher dump pass ' '--freeze-all-layers.' ) assert args.rampup_batch_size is None, ( '--logits-save-dir does not support --rampup-batch-size: the sample<->iteration ' 'mapping assumes a fixed global batch size.' ) if args.freeze_all_layers: if args.use_distributed_optimizer: warn_rank_0( '--freeze-all-layers incompatible with use_distributed_optimizer. Disabling use_distributed_optimizer.' ) args.use_distributed_optimizer = False if args.overlap_param_gather: warn_rank_0( '--freeze-all-layers incompatible with overlap_param_gather. Disabling overlap_param_gather.' ) args.overlap_param_gather = False if args.override_ckpt_iteration is not None: assert not args.finetune, "Cannot override checkpoint iteration together with finetune flag." # Inference args if args.inference_batch_times_seqlen_threshold > -1: assert args.pipeline_model_parallel_size > 1, \ "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." assert ( args.cuda_graph_impl == "none" ), "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" # MoE upcycling check if args.moe_use_upcycling: assert args.save is not None, "When using upcycling, the --save option must be specified." if not args.no_load_optim: args.no_load_optim = True warn_rank_0('enabling --no-load-optim for upcycling.') if not args.no_load_rng: args.no_load_rng = True warn_rank_0('enabling --no-load-rng for upcycling.') # --skip-train checks. if args.skip_train and not args.no_load_optim: args.no_load_optim = True warn_rank_0('enabling --no-load-optim when skipping training.') rl_cfg.validate_run( cuda_graph_impl=args.cuda_graph_impl, micro_batch_size=args.micro_batch_size, global_batch_size=args.global_batch_size, train_iters=args.train_iters, train_samples=args.train_samples, skip_train=args.skip_train, load_optim=not args.no_load_optim, save_interval=args.save_interval, exit_interval=args.exit_interval, separate_inference_model=rl_cfg.uses_separate_inference_model( context_parallel_size=args.context_parallel_size, inference_shards=args.inference_shards ), ) if args.non_persistent_ckpt_type == "local": assert args.non_persistent_local_ckpt_dir is not None, "Tried to use local checkpointing without specifying --local-ckpt-dir!" if args.replication: assert args.replication_jump is not None, "--replication requires the value of --replication-jump!" assert args.non_persistent_ckpt_type == "local", f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" elif args.replication_jump: warn_rank_0("--replication-jump was specified despite not using replication. Ignoring.") args.replication_jump = None if args.delay_wgrad_compute: assert args.transformer_impl == 'transformer_engine', \ "Delaying wgrad compute is only supported with transformer_engine implementation" if args.overlap_grad_reduce: assert is_te_min_version("2.8.0"), ( "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" ) wgrad_in_graph_scope = CudaGraphModule.attn in args.cuda_graph_modules or ( CudaGraphModule.moe_router in args.cuda_graph_modules and args.moe_shared_expert_intermediate_size is not None and not args.moe_shared_expert_overlap ) if wgrad_in_graph_scope: assert is_te_min_version( "2.12.0" ), "CUDA graph with delay_wgrad_compute requires TE version >= 2.12.0." assert args.gradient_accumulation_fusion, ( 'CUDA graph with delay_wgrad_compute requires gradient_accumulation_fusion ' 'to be enabled. This is because the default gradient accumulation does not ' 'use static memory addresses, which breaks CUDA graph requirements.' ) if CudaGraphModule.attn in args.cuda_graph_modules: assert ( not args.add_bias_linear and not args.add_qkv_bias ), "CUDA graph with delay_wgrad_compute doesn't support attn bias for now." if not args.gradient_accumulation_fusion: assert is_te_min_version("2.7.0"), ( "disabling gradient_accumulation_fusion is only supported with TE >= 2.7.0 " "when enabling delay_wgrad_compute" ) if args.fine_grained_activation_offloading: assert args.transformer_impl == 'transformer_engine', \ "Fine-grained activation offloading is only supported with transformer_engine implementation" if is_te_min_version("2.10.0"): assert os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1", \ "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." if args.mtp_num_layers: # MTP is compatible with position embedding types that use position_ids. supported_position_types = ["learned_absolute", "rope", "mrope", "none"] assert args.position_embedding_type in supported_position_types, ( f"Multi-Token Prediction (MTP) is not supported with '{args.position_embedding_type}' position embedding type. " f"The supported position embedding types are: {', '.join(supported_position_types)}." ) if args.cpu_offloading_num_layers > 0: args.cpu_offloading = True # CUDA Graphs if args.cuda_graph_impl != "none": if ( "transformer_engine" in (args.transformer_impl, args.cuda_graph_impl) and not args.te_rng_tracker ): args.te_rng_tracker = True warn_rank_0("te_rng_tracker is not enabled, enabling it for CUDA graphs.", args.rank) if args.cuda_graph_impl == "transformer_engine": assert ( "expandable_segments:True" not in os.getenv("PYTORCH_CUDA_ALLOC_CONF", "") or os.getenv("NCCL_GRAPH_REGISTER", "") == "0" ), ( "Setting NCCL_GRAPH_REGISTER=0 to avoid illegal memory access when using " "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True." ) assert not ( args.cuda_graph_impl == "full_iteration" and args.cuda_graph_modules ), '--cuda-graph-modules must be empty when --cuda-graph-impl=full_iteration.' assert not ( args.moe_shortcut_connection and args.cuda_graph_impl != "none" and not ( args.cuda_graph_impl == "local" and args.inference_cuda_graph_scope == InferenceCudaGraphScope.block ) ), ( "--moe-shortcut-connection supports CUDA graphs only with " "--cuda-graph-impl=local and --inference-cuda-graph-scope=block." ) if args.multi_latent_attention: assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." if args.mla_down_proj_fusion: assert args.multi_latent_attention, "--mla-down-proj-fusion requires --multi-latent-attention" assert ( not args.moe_use_norm_before_up_proj or args.moe_latent_size is not None ), "--moe-use-norm-before-up-proj requires --moe-latent-size to be set." # MoE latent projections if args.moe_latent_size is not None: assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero." assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models." # Print arguments. _print_args("arguments", args) return args def _print_args(title, args): """Print arguments.""" from megatron.training.utils import is_rank0 if is_rank0(): print(f'------------------------ {title} ------------------------', flush=True) str_list = [] for arg in vars(args): dots = '.' * (48 - len(arg)) str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) for arg in sorted(str_list, key=lambda x: x.lower()): print(arg, flush=True) print(f'-------------------- end of {title} ---------------------', flush=True) def _check_arg_is_not_none(args, arg): assert getattr(args, arg) is not None, '{} argument is None'.format(arg) def _add_transformer_engine_args(parser): group = parser.add_argument_group(title='Transformer-Engine') # FP4 related arguments group.add_argument('--fp4-param-gather', action='store_true', help='Keep the compute param in fp4 (do not use any other intermediate ' 'dtype) and perform the param all-gather in fp4.') # FP8 related arguments group.add_argument('--fp8-param-gather', action='store_true', help='Keep the compute param in fp8 (do not use any other intermediate ' 'dtype) and perform the param all-gather in fp8.') # TE precision config file group.add_argument('--te-precision-config-file', default=None, help='Configuration file to select per-module precision overrides. ' 'See TransformerEngineMixedPrecision.md') return parser def _add_inference_args(parser): group = parser.add_argument_group(title='inference') group.add_argument('--inference-batch-times-seqlen-threshold', type=int, default=-1, help='If (batch-size * sequence-length) is smaller than this threshold' 'then 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.') group.add_argument('--max-tokens-to-oom', type=int, default=12000, help='Maximum number of tokens during inference' 'tokens here is # in prompt + # to generate' 'Allows us to throw an error before OOM crashes server') group.add_argument('--output-bert-embeddings', action='store_true', help='Output Bert embeddings (via mean pooling) from ' 'model, rather than its binary head output or entire ' 'hidden batch.') group.add_argument('--bert-embedder-type', default="megatron", choices=["megatron", "huggingface"], help='Select either Megatron or Huggingface as the ' 'Bert embedder.') group.add_argument('--cuda-graph-scope', nargs='+', type=_parse_cuda_graph_modules_arg, default=None, dest='cuda_graph_scope_deprecated', help=argparse.SUPPRESS) # hidden; use --cuda-graph-modules instead group.add_argument('--cuda-graph-modules', nargs='+', type=_parse_cuda_graph_modules_arg, 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.') group.add_argument('--use-legacy-static-engine', action='store_true', default=False, help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', dest='use_legacy_static_engine') group.add_argument('--inference-max-requests', type=int, default=8, help='Maximum number of requests for inference.', dest='inference_max_requests') group.add_argument('--inference-max-seq-length', type=int, default=2560, help='Maximum sequence length expected for inference (prefill + decode).', dest='inference_max_seq_length') group.add_argument('--inference-dynamic-batching', action='store_true', default=False, help='Enable dynamic batching mode.') group.add_argument('--inference-dynamic-batching-buffer-size-gb', type=float, default=40., help='On-GPU portion of the shared KV cache block pool. ' '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`.') group.add_argument('--inference-dynamic-batching-paused-buffer-size-gb', type=float, default=None, help='Memory used to derive the paused-request block retention ' 'budget. This does not reserve blocks from active requests: ' 'active requests may use the entire shared pool of usable KV ' 'cache blocks. Under allocation pressure, paused requests ' 'retain blocks only within this budget and excess paused ' 'requests may be evicted.') group.add_argument('--inference-dynamic-batching-mamba-memory-ratio', 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.') group.add_argument('--inference-dynamic-batching-block-size', type=int, default=256, help='KV cache block size. ' 'It should be a multiple of 256') group.add_argument('--inference-dynamic-batching-max-requests', type=int, 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.') group.add_argument('--inference-dynamic-batching-max-tokens', type=int, default=None, help='Override the inference context\'s default `max_tokens`.') group.add_argument('--inference-dynamic-batching-num-cuda-graphs', 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.' 'The user can also pass -1, in which case we automatically determine the number of graphs ' \ 'to capture based on the `max_requests`.') group.add_argument('--inference-dynamic-batching-track-paused-request-events', action='store_true', help='Track paused request ids by adding \'paused\' events ' 'to each request\'s event history. This has a very minor ' 'impact on latency.') group.add_argument('--inference-dynamic-batching-track-generated-token-events', action='store_true', 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.') group.add_argument('--decode-only-cuda-graphs', action='store_true', default=False, help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.') group.add_argument('--inference-cuda-graph-all-prefills', action='store_true', 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)`.') group.add_argument('--inference-cuda-graph-max-tokens', type=int, default=512, dest='inference_cuda_graph_max_tokens', help='Token ceiling for the largest captured prefill/mixed CUDA ' 'graph (default: 512). Clamped to at least the decode limit ' '`max_requests * (num_speculative_tokens + 1)` and at most ' '`max_tokens`. Ignored when --inference-cuda-graph-all-prefills ' 'is set (which extends capture to the full `max_tokens`).') group.add_argument('--inference-dynamic-batching-unified-memory-level', type=int, default=0, choices=[0, 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.') group.add_argument('--inference-kv-cache-management-mode', '--rl-kv-cache-management-mode', dest='kv_cache_management_mode', type=str, default='persist', choices=['persist', 'offload', 'recompute'], help='What the engine does with the KV cache while it is suspended: ' 'persist keeps it on GPU, offload moves it to CPU and restores it on ' 'resume, recompute frees it and recomputes it on resume. ' '--rl-kv-cache-management-mode is the deprecated spelling.') group.add_argument('--inference-static-kv-memory-pointers', '--rl-persist-cuda-graphs', dest='static_kv_memory_pointers', action=argparse.BooleanOptionalAction, default=False, help='Keep the KV cache at fixed device addresses across suspend/resume so ' 'CUDA graphs that reference it stay valid and are not recaptured. Requires ' 'UVM or torch_memory_saver unless the cache persists on GPU. ' '--rl-persist-cuda-graphs / --no-rl-persist-cuda-graphs are the deprecated ' 'spellings.') group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill', action='store_true', default=False, help="Enable chunked prefill (disabled by default)") group.add_argument('--num-speculative-tokens', type=int, default=0, help='Number of speculative tokens generated during decode') group.add_argument('--inference-dynamic-batching-prefix-caching', dest='inference_dynamic_batching_enable_prefix_caching', action=argparse.BooleanOptionalAction, 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.') group.add_argument('--inference-dynamic-batching-prefix-caching-eviction-policy', type=str, default='lru', choices=['ref_zero', 'lru'], dest='inference_dynamic_batching_prefix_caching_eviction_policy', help='Eviction policy for prefix caching blocks. ' '"ref_zero" immediately returns blocks to the ' 'free pool when ref_count hits 0. "lru" (default) keeps ' 'blocks cached and evicts via LRU only when space is needed.') group.add_argument('--inference-dynamic-batching-prefix-caching-coordinator-policy', type=str, default='longest_prefix', choices=['longest_prefix', 'first_prefix_block', 'load_balanced'], dest='inference_dynamic_batching_prefix_caching_coordinator_policy', help='Coordinator routing policy for prefix caching. ' '"load_balanced" routes to the rank with the fewest ' 'in-flight requests, ignoring prefix affinity. ' '"first_prefix_block" routes based on the first block hash only. ' '"longest_prefix" (default) routes to the rank with the longest ' 'matching prefix. "first_prefix_block" and "longest_prefix" both combine ' 'prefix affinity with load balancing and fall back to ' 'load-balanced routing when prefix caching is disabled or no ' 'prefix match exists.') group.add_argument('--inference-dynamic-batching-prefix-caching-routing-alpha', type=float, default=1.0, dest='inference_dynamic_batching_prefix_caching_routing_alpha', help='How hard to penalise load when routing on prefix ' 'affinity: score = cache_score - alpha * relative_load, where ' 'relative_load is a rank load measured against the fleet mean. ' '0 is pure prefix affinity; higher values divert to idle ranks ' 'more readily as the fleet becomes lopsided. Dimensionless and ' 'not capped at 1. Default: 0.5.') group.add_argument('--inference-dynamic-batching-prefix-cache-ttl-seconds', type=float, default=300.0, dest='inference_dynamic_batching_prefix_cache_ttl_seconds', help='How long the coordinator assumes an engine still holds a ' 'block it routed there. The coordinator never observes evictions, ' 'so entries untouched for this long are dropped rather than kept ' 'forever. Default: 300.0.') group.add_argument('--inference-dynamic-batching-media-cache-coordinator-policy', type=str, default='affinity', choices=['affinity', 'load_balanced'], dest='inference_dynamic_batching_media_cache_coordinator_policy', help='Coordinator routing policy for media caching. ' '"affinity" prefers a rank assigned the same media; ' '"load_balanced" ignores standalone media affinity.') group.add_argument('--inference-dynamic-batching-media-cache-routing-weight', type=float, default=1.0, dest='inference_dynamic_batching_media_cache_routing_weight', help='Media-cache hit weight in equivalent compact-prompt blocks. ' 'Default: 1.0.') group.add_argument('--inference-dynamic-batching-vision-embedding-cache-max-bytes', type=int, default=0, dest='inference_dynamic_batching_vision_embedding_cache_max_bytes', help='Maximum GPU bytes retained per engine for reusable vision ' 'embeddings. Zero disables the cache. Default: 0.') group.add_argument('--inference-dynamic-batching-allow-stale-multimodal-embeddings', action='store_true', dest='inference_dynamic_batching_allow_stale_multimodal_embeddings', help='Allow request-local and cached multimodal embeddings to survive ' 'suspend/resume and generation-epoch changes. Use only when model ' 'weights do not change across these boundaries.') group.add_argument('--inference-dynamic-batching-prefix-caching-mamba-gb', type=float, default=None, dest='inference_dynamic_batching_prefix_caching_mamba_gb', 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. This budget ' 'covers both the durable cache (the ssm_states/conv_states ' 'slots reused across requests) and the per-step extraction ' 'scratch (the intermediate_ssm_out/intermediate_conv_out ' 'buffers, sized to min(ceil(max_tokens / block_size), ' '3 * max_requests) slots); the scratch is reserved first, so a ' 'smaller max_tokens (or max_requests) shrinks the scratch and ' 'leaves more durable slots.') group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count', type=int, default=16, help='Number of mixed prefill requests to capture in a cuda graph.') group.add_argument('--inference-dynamic-batching-cuda-graph-sizing-distribution', type=str, default='hybrid', choices=['exponential', 'linear', 'hybrid'], dest='inference_dynamic_batching_cuda_graph_sizing_distribution', help='Spacing of CUDA graph token counts. "hybrid" (default) uses ' 'exponential spacing for prefill/mixed graphs and linear spacing ' 'for decode-only graphs, whose token counts are capped at ' 'max_requests and are far too small for halving to cover well. ' '"exponential" halves from cuda_graph_max_tokens down to tp_size, ' 'giving a log-spaced distribution with bounded relative padding. ' '"linear" uses varying linear strides across the range.') group.add_argument('--inference-dynamic-batching-sampling-backend', type=str, default='torch', choices=['torch', 'flashinfer'], 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.') group.add_argument('--use-same-sampling-seed-across-dp-ranks', action='store_false', dest='offset_sampling_seed_by_dp_rank', default=True, help='Use the same inference sampling seed on every data-parallel rank. ' '--deterministic-mode also uses the same seed on every DP rank.') group.add_argument('--inference-dynamic-batching-async-sched-mode', type=str, default='async', choices=['async', 'legacy'], help='Async scheduling mode for dynamic batching. ' '"async" (default) overlaps asynchronous scheduling phases by ' 'reordering them to prepare-before-resolve. Select "legacy" to ' 'disable async scheduling and use the resolve-before-prepare path.') group.add_argument('--inference-dynamic-batching-logprobs-mode', type=str, default='raw_logprobs', choices=['raw_logprobs', 'processed_logprobs'], help='How returned inference log-probs are computed engine-wide. ' '"raw_logprobs" (default) uses the unmodified model logits; ' '"processed_logprobs" applies temperature and top-k/top-p to ' 'sampled-token and top-N log-probs. Prompt log-probs, ' 'including prompt top-N, are always raw.') group.add_argument('--inference-logging-step-interval', type=int, default=0, help='Step interval for logging inference metrics. ' 'Default to 0 to disable inference logging.') group.add_argument('--inference-text-gen-server-logging', action=argparse.BooleanOptionalAction, required=False, default=False, help='Enable per-request logging in the inference text generation server.') group.add_argument('--inference-wandb-logging', action=argparse.BooleanOptionalAction, required=False, default=False, help='Enable inference wandb logging.') group.add_argument("--inference-coordinator-port", type=int, help="This port will be used to setup the inference coordinator on node-0") group.add_argument('--mamba-inference-conv-states-dtype', type=str, choices=['bf16', 'fp16', 'fp32'], default='bf16', help='Dtype for the Mamba inference conv states tensor') group.add_argument('--mamba-inference-ssm-states-dtype', type=str, choices=['bf16', 'fp16', 'fp32'], default='bf16', help='Dtype for the Mamba inference SSM states tensor') group.add_argument('--inference-use-synchronous-zmq-collectives', action=argparse.BooleanOptionalAction, required=False, default=False, help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.') group.add_argument('--inference-shards', type=str, default=None, metavar='SPEC', help='Partition the world into independent inference models, each with ' 'its own parallelism, e.g. "tp=2,role=prefill+tp=1,dp=2,role=decode". ' 'Shards are separated by "+" or ";"; per-shard keys are ' 'tp,pp,ep,expt_tp,dp (each defaults to 1) and must partition the full ' 'world. Tagging shards role=prefill|decode enables disaggregated ' 'inference (prefill hands KV to the decode pool); a dp>1 decode shard ' 'is several independent decode instances.') group.add_argument('--disagg-kv-transport-backend', type=str, default='nixl', choices=['nccl', 'nixl'], help='State transport for native disaggregated inference: ' '"nccl" uses two-sided transfers and "nixl" uses one-sided pulls.') return parser def _add_network_size_args(parser): exclude = [ # cannot provide callables over CLI "timers", "finalize_model_grads_func", "grad_scale_func", "moe_grad_scale_func", "mtp_grad_scale_func", "no_sync_func", "grad_sync_func", "param_sync_func", "_cpu_offloading_context", "init_method", "output_layer_init_method", "embedding_init_method", "activation_func", "experimental_attention_variant_loss_scale_func", # types affect docstring "pipeline_model_parallel_layout", "window_size", "window_attn_skip_freq", "no_rope_freq", "moe_layer_freq", "linear_attention_freq", "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "moe_router_quantile_balancing_estimation_scope", "moe_router_qb_num_bins", "cp_comm_type", "cuda_graph_modules", "cuda_graph_scope", # deprecated alias; handled manually by --cuda-graph-scope flag # no CLI argument exists for these "virtual_pipeline_model_parallel_size", "params_dtype", "enable_autocast", "autocast_dtype", "num_microbatches_with_partial_activation_checkpoints", "tp_comm_overlap_disable_qkv", "tp_comm_overlap_disable_fc1", "pipeline_dtype", "variable_seq_lengths", "batch_p2p_comm", "batch_p2p_sync", "deallocate_pipeline_outputs", "cpu_offloading", "cpu_offloading_activations", "cpu_offloading_weights", "cpu_offloading_double_buffering", "num_layers_in_first_pipeline_stage", "num_layers_in_last_pipeline_stage", "softmax_scale", "gated_linear_unit", "bias_activation_fusion", "activation_func_fp8_input_store", "memory_efficient_layer_norm", "fused_single_qkv_rope", "fp8_dot_product_attention", "fp8_multi_head_attention", "tp_only_amax_red", "use_kitchen", "moe_token_dropping", "cuda_graph_use_single_mempool", "cuda_graph_retain_backward_graph", "disable_parameter_transpose_cache", "inference_sampling_seed", "use_inference_optimized_layers", "heterogeneous_block_specs", "hetereogenous_dist_checkpoint", "quant_recipe", # deprecated and no CLI arg exists "tp_comm_atomic_ag", "tp_comm_atomic_rs", "moe_router_topk_limited_devices", # already generated by another config "inference_rng_tracker", "use_te_rng_tracker", # args uses same var with a different name "num_moe_experts", "hash_moe_vocab_size", "fp8_param", "fp4_param", # filled from the explicit --seq-length argument below "max_seq_length_per_sample", # incompatible defaults in dataclass "gradient_accumulation_fusion", "overlap_p2p_comm", "attention_softmax_in_fp32", "masked_softmax_fusion", "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", "mamba_training_ssm_states_dtype", "max_seqlen_per_dp_cp_rank", "pad_packed_seq_alignment", # integer/max CLI parsing is handled explicitly "hybrid_context_parallel", "dynamic_context_parallel", "min_dynamic_context_parallel_size", "sequence_packing_scheduler", # internal/derived: controlled only via --tensor-parallel-num-weight-shards "gtp_weight_remat_size", # internal/derived: controlled only via --expert-tensor-parallel-num-weight-shards "expert_gtp_weight_remat_size", # Constructed from the dedicated flat CLI arguments below. "wide_residual", # Parsed manually as a JSON object below. "moe_megakernel_backend_config", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") transformer_group.add_argument( "--moe-megakernel-backend-config", type=json.loads, default=None, metavar="JSON", help="Backend-specific MoE megakernel options as a JSON object.", ) group = parser.add_argument_group(title='network size') group.add_argument('--encoder-num-layers', type=int, default=None, help='Number of encoder transformer layers.') group.add_argument('--decoder-num-layers', type=int, default=None, help='Number of decoder transformer layers.') group.add_argument('--group-query-attention', action='store_true', help='Use group-query attention.') group.add_argument('--window-size', type=tuple_type, default=None, help='Window size for window attention. If not provided, ' 'window attention will be disabled.') group.add_argument('--window-attn-skip-freq', type=moe_freq_type, 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. ') group.add_argument('--max-position-embeddings', type=int, default=None, help='Maximum number of position embeddings to use. ' 'This is the size of position embedding.') group.add_argument('--position-embedding-type', type=str, default='learned_absolute', choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], help='Position embedding type.') group.add_argument('--yarn-original-max-position-embeddings', type=int, default=None, help='Original maximum position embeddings for YaRN RoPE frequency correction.') group.add_argument('--yarn-beta-fast', type=float, default=None, help='Beta fast for YaRN RoPE frequency correction.') group.add_argument('--yarn-beta-slow', type=float, default=None, help='Beta slow for YaRN RoPE frequency correction.') group.add_argument('--yarn-correction-range-round-to-int', action='store_true', default=None, help='Round YaRN correction range endpoints to integers.') group.add_argument('--no-yarn-correction-range-round-to-int', action='store_false', dest='yarn_correction_range_round_to_int', help='Do not round YaRN correction range endpoints to integers.') group.add_argument('--relative-attention-num-buckets', type=int, default=32, help='Number of buckets for relative position embeddings.') group.add_argument('--relative-attention-max-distance', type=int, default=128, help='Maximum distance for relative position embeddings calculation.') group.add_argument('--use-rotary-position-embeddings', action='store_true', help='Use rotary positional embeddings or not. ' 'Deprecated: use --position-embedding-type') group.add_argument('--rotary-base', type=int, default=10000, help='Base to use for rotary positional embeddings, default 10000') group.add_argument('--rotary-percent', type=float, default=1.0, help='Percent of rotary dimension to use, default 100%%') group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, help='Sequence length interpolation factor for rotary embeddings.') group.add_argument('--use-rope-scaling', action='store_true', help='Apply rope scaling as used in llama3.x') group.add_argument('--rope-scaling-factor', type=float, default=8.0, help='Rope scaling factor in llama3.x models') group.add_argument('--no-rope-freq', type=no_rope_freq_type, 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 performed' 'on every layer.' ) group.add_argument('--no-position-embedding', action='store_false', help='Disable position embedding. Deprecated: use --position-embedding-type', dest='add_position_embedding') group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, help='Pad the vocab size to be divisible by this value.' 'This is added for computational efficieny reasons.') group.add_argument('--openai-gelu', action='store_true', help='Use OpenAIs GeLU implementation. This option' 'should not be used unless for backward compatibility' 'reasons.') group.add_argument('--squared-relu', action='store_true', help='Use squared relu activation instead of default gelu') group.add_argument('--swiglu', action='store_true', help='Use gated linear units and SiLU activation instead of default gelu') group.add_argument('--quick-geglu', action='store_true', help='Use quick geglu activation instead of default gelu') group.add_argument('--onnx-safe', type=bool, required=False, help='Use workarounds for known problems with ' 'Torch ONNX exporter') group.add_argument('--bert-no-binary-head', action='store_false', help='Disable BERT binary head.', dest='bert_binary_head') group.add_argument('--untie-embeddings-and-output-weights', action='store_true', help='Untie embeddings and output weights.') return parser def _parse_triton_autotune_block_size(value): """Parse a named Triton block-size override for the nested policy.""" name, separator, size = value.partition('=') if separator and name.startswith('BLOCK') and name.isidentifier(): try: size = int(size) except ValueError: pass else: if size > 0: return name, size raise argparse.ArgumentTypeError( f'Invalid block size {value!r}; expected BLOCK_NAME=positive_integer.' ) def _add_triton_autotune_args(parser): """Add CLI controls for the process-wide Triton autotune policy.""" autotune_factory = ArgumentGroupFactory( AutotunePolicy, exclude=['block_sizes'], prefix='triton_autotune' ) group = autotune_factory.build_group(parser, title='Triton autotune policy') # NAME=VALUE pairs cannot be inferred from the field type. group.add_argument( '--triton-autotune-block-sizes', type=_parse_triton_autotune_block_size, nargs='+', default=(), metavar='BLOCK_NAME=VALUE', help='Block sizes to match when no tuned table entry applies, ' 'for example BLOCK_SIZE=128 BLOCK_M=64.', ) return parser def _add_wide_residual_args(parser): """Add CLI arguments used to construct ``WideResidualConfig``.""" group = parser.add_argument_group(title='wide residual') group.add_argument( '--wide-residual', dest='wide_residual_num_streams', type=int, default=None, help='Enable streamwise wide residuals with this many hidden-size streams.', ) group.add_argument( '--wide-residual-streamwise-sigmoid-init-scale', type=float, default=0.01, help='Symmetric initialization spread for streamwise write logits.', ) group.add_argument( '--wide-residual-learned-retention', action='store_true', help='Apply one bounded learned carry factor to every residual stream.', ) group.add_argument( '--wide-residual-retention-init', type=float, default=0.999, help='Initial retention factor for learned wide-residual retention.', ) group.add_argument( '--wide-residual-retention-max-forget', type=float, default=0.10, help='Maximum forget rate for learned wide-residual retention.', ) return parser def _add_straggler_detector_args(parser): from megatron.training.config import StragglerDetectionConfig straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig) group = straggler_factory.build_group(parser, "straggler") return parser def _add_workload_inspector_server_args(parser): group = parser.add_argument_group(title='workload inspector') group.add_argument('--run-workload-inspector-server', action='store_true', help='If set, enables workload inspector server for on-demand profiling.') return parser def _add_inprocess_restart_args(parser): group = parser.add_argument_group(title='In-process restart') group.add_argument('--inprocess-restart', action='store_true', help='Enables in-process restart.') group.add_argument('--inprocess-max-iterations', default=None, type=int, help='Maximum number of in-process restart iterations.') group.add_argument('--inprocess-monitor-thread-interval', default=1.0, type=float, help='Monitoring interval (in seconds) for the monitoring thread.') group.add_argument('--inprocess-monitor-process-interval', default=1.0, type=float, help='Monitoring interval (in seconds) for the monitoring process.') group.add_argument('--inprocess-progress-watchdog-interval', default=1.0, type=float, help='Interval (in seconds) for automatic progress watchdog timestamp ' 'updates.') group.add_argument('--inprocess-heartbeat-interval', default=30, type=float, help='Monitoring interval (in seconds) for detecting unresponsive ranks.') group.add_argument('--inprocess-soft-timeout', default=60, type=float, help='Soft progress timeout (in seconds).') group.add_argument('--inprocess-hard-timeout', default=90, type=float, help='Hard progress timeout (in seconds).') group.add_argument('--inprocess-heartbeat-timeout', default=60, type=float, help='Timeout (in seconds) for a missing rank detection heartbeat.') group.add_argument('--inprocess-barrier-timeout', default=120, type=float, help='Timeout (in seconds) for internal distributed barrier') group.add_argument('--inprocess-completion-timeout', default=120, type=float, help='Timeout (in seconds) for barrier on completion on all ranks') group.add_argument('--inprocess-last-call-wait', default=1, type=float, help='Time interval (in seconds) for other ranks to report concurrent ' 'terminal failures.') group.add_argument('--inprocess-termination-grace-time', default=1, type=float, help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' 'timeout') group.add_argument('--inprocess-granularity', default='node', type=str, choices=['node', 'rank'], help='Granularity for in-process restart.') group.add_argument('--inprocess-active-world-size', default=int(os.getenv('WORLD_SIZE', '1')), type=int, help='The number of ranks initially executing the workload. ' 'The remaining ranks from the allocation are set aside ' 'as warm reserve.') group.add_argument('--inprocess-empty-cuda-cache', action='store_true', help='Release all unoccupied cached GPU memory on every in-process restart.') return parser def _add_ft_package_args(parser): group = parser.add_argument_group(title='ft_package') group.add_argument('--enable-ft-package', action='store_true', help='If set, Fault Tolerance package is enabled. ' 'Note: This feature is for Nvidia internal use only.') group.add_argument('--calc-ft-timeouts', action='store_true', help='If set, FT package will try to automatically compute the timeouts. ' 'Note: This feature is for Nvidia internal use only.') group.add_argument('--ft-num-warmup-iters', 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.') return parser def _add_logging_args(parser): from megatron.training.config import LoggerConfig log_factory = ArgumentGroupFactory(LoggerConfig, exclude = ["log_throughput_to_tensorboard", "throughput_window_size", "memory_keys", "log_l2_norm_grad_to_tensorboard", "log_runtime_to_tensorboard", "runtime_time_unit", "filter_warnings", "modules_to_filter", "set_level_for_all_loggers", "save_config_filepath"]) group = log_factory.build_group(parser, title="logging") return parser def _add_regularization_args(parser): group = parser.add_argument_group(title='regularization') group.add_argument('--log-num-zeros-in-grad', action='store_true', help='If set, calculate and log the number of zeros in gradient.') group.add_argument('--weight-decay', type=float, default=0.01, help='Weight decay coefficient for L2 regularization.') group.add_argument( '--rmsnorm-weight-decay-mult', type=float, default=None, help='Multiplier on scheduled global weight decay for RMSNorm gamma ' 'only (e.g. 0.5 for half decay, 1.0 for full decay, 0.0 for none). ' 'Unset preserves the existing policy. Overrides ' '--apply-wd-to-qk-layernorm for RMSNorm gamma only.', ) group.add_argument("--embedding-min-norm", type=float, default=None, help="After each optimizer update, rescale input word-embedding rows " "with nonzero L2 norms below this threshold to twice the threshold. " "Requires the distributed optimizer and is disabled by default.") group.add_argument('--apply-wd-to-qk-layernorm', action='store_true', help='Apply weight decay to qk layernorm as a special case.') group.add_argument('--clip-grad', type=float, default=1.0, help='Gradient clipping based on global L2 norm.') group.add_argument('--adam-beta1', type=float, default=0.9, help='First coefficient for computing running averages ' 'of gradient and its square') group.add_argument('--adam-beta2', type=float, default=0.999, help='Second coefficient for computing running averages ' 'of gradient and its square') group.add_argument('--adam-eps', type=float, default=1e-08, help='Term added to the denominator to improve' 'numerical stability') group.add_argument('--sgd-momentum', type=float, default=0.9, help='Momentum factor for sgd') group.add_argument('--muon-momentum', type=float, default=0.9, help='Momentum factor for Muon optimizer') group.add_argument('--muon-no-split-qkv', action='store_false', default=True, dest='muon_split_qkv', help='Whether to split QKV parameters for Muon optimizer') group.add_argument('--muon-nesterov', action='store_true', help='Whether to use Nesterov-style momentum in the internal SGD') group.add_argument('--muon-scale-mode', type=str, default='spectral', choices=['spectral', 'unit_rms_norm', 'shape_scaling'], 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.') group.add_argument('--muon-fp32-matmul-prec', type=str, default='medium', choices=['low', 'medium', 'high'], help='FP32 matmul precision for Newton-Schulz iteration') group.add_argument('--muon-coefficient-type', 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.') group.add_argument('--muon-num-ns-steps', type=int, default=5, help='Number of Newton-Schulz steps for Muon optimizer') group.add_argument('--muon-tp-mode', type=str, default='duplicated', choices=['blockwise', 'duplicated', 'distributed', 'auto', 'layer_sharded'], help='How to perform NS calculation for tensor model parallel weights. ' 'blockwise orthogonalizes each shard on its own, so the update rule ' 'depends on the parallelism config; duplicated and distributed both ' 'orthogonalize the whole matrix and give TP-invariant results; auto ' 'select between duplicated and distributed mode per-weight for ' 'dense weights; layer_sharded assigns each 2D weight one NS home ' 'rank in the (gtp_remat x tp) domain and routes the shards there ' 'with all_to_all (same math as duplicated, no redundant NS); ' 'requires --use-distributed-optimizer and --muon-no-split-qkv. See ' 'OptimizerConfig.muon_tp_mode.') group.add_argument('--muon-ns-batch-size', type=int, default=1, help='Max number of same-shape matrices fused into one batched ' 'Newton-Schulz on an NS home under --muon-tp-mode layer_sharded. ' 'The default of 1 keeps the bit-exact per-matrix path; raise ' '(e.g. to 32) to cut kernel launches on MoE expert homes at ' 'the cost of bitwise parity (baddbmm vs addmm rounding).') group.add_argument('--muon-use-syrk', action='store_true', help='Use the Triton SYRK kernel for the symmetric-output ' 'Newton-Schulz GEMMs in Muon (~1/3 off ' 'NS FLOPs for near-square matrices). Takes effect only with ' '--muon-fp32-matmul-prec medium. Under --muon-tp-mode ' 'layer_sharded, unmet Triton/SM/emerging-optimizers ' 'requirements are rejected at startup.') group.add_argument('--muon-no-concurrent-groups', action='store_false', dest='muon_concurrent_groups', help='Serialize param groups on one CUDA stream under ' '--muon-tp-mode layer_sharded instead of overlapping one group\'s ' 'Newton-Schulz with another\'s all_to_all. Bitwise-neutral; use ' 'when the concurrent transient buffers push peak memory too high.') group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, help='Additional scale factor for the muon update') group.add_argument('--muon-scalar-optimizer', type=str, default='adam', choices=['adam', 'lion'], help='Optimizer for scalar parameters (embeddings, biases, norms) ' 'when using muon. Defaults to adam.') group.add_argument('--lion-beta1', type=float, default=0.95, help='First beta coefficient for Lion optimizer ' '(used in sign update). Default: 0.95.') group.add_argument('--lion-beta2', type=float, default=0.98, help='Second beta coefficient for Lion optimizer ' '(used in momentum EMA update). Default: 0.98.') return parser def _add_rl_args(parser): from megatron.training.config import RLConfig rl_factory = ArgumentGroupFactory(RLConfig) rl_factory.build_group(parser, "rl") return parser def _add_training_args(parser): from megatron.training.config import ProfilingConfig, TrainingConfig prof_factory = ArgumentGroupFactory( ProfilingConfig, exclude=["use_nsys_profiler", "run_workload_inspector_server"] ) prof_group = prof_factory.build_group(parser, "profiling") prof_group.add_argument('--profile', action='store_true', help='Enable nsys profiling. When using this option, nsys ' 'options should be specified in commandline.') train_factory = ArgumentGroupFactory(TrainingConfig) group = train_factory.build_group(parser, "training") # Keep this CLI-only until dataset options have their own config dataclass. group.add_argument( "--train-full-dataset", action="store_true", default=False, help="Train for one complete pass over an externally provided dataset.", ) group.add_argument('--batch-size', type=int, default=None, help='Old batch size parameter, do not use. ' 'Use --micro-batch-size instead') group.add_argument('--recompute-activations', action='store_true', help='recompute activation to allow for training ' 'with larger models, sequences, and batch sizes.') group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', help='Check for NaNs in loss and grad', dest='check_for_nan_in_loss_and_grad') group.add_argument('--check-for-large-grads', action='store_true', help='Check for unexpectedly large grads', dest='check_for_large_grads') group.add_argument('--result-rejected-tracker-filename', type=str, default=None, help='Optional name of file tracking `result_rejected` events.') group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, help='Config file when tp_comm_overlap is enabled.') # deprecated group.add_argument('--checkpoint-activations', action='store_true', help='Checkpoint activation to allow for training ' 'with larger models, sequences, and batch sizes.') group.add_argument('--no-masked-softmax-fusion', action='store_false', help='Disable fusion of query_key_value scaling, ' 'masking, and softmax.', dest='masked_softmax_fusion') group.add_argument('--no-bias-gelu-fusion', action='store_false', help='Disable bias and gelu fusion.', dest='bias_gelu_fusion') group.add_argument('--no-bias-swiglu-fusion', action='store_false', help='Disable bias and swiglu fusion, the fusion is ' 'available only when using megatron-core.', dest='bias_swiglu_fusion') group.add_argument('--no-bias-dropout-fusion', action='store_false', help='Disable bias and dropout fusion.', dest='bias_dropout_fusion') group.add_argument('--no-rope-fusion', action='store_false', help='Disable rope fusion, the fusion is available ' 'only when using megatron-core.', dest='apply_rope_fusion') group.add_argument('--rope-type', type=str, default=None, choices=['rope', 'yarn'], help='Type of rope to use. Note that MLA takes yarn by default, ' 'and common attention takes rope by default.') group.add_argument('--use-flash-attn', action='store_true', help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135') group.add_argument('--optimizer', type=str, default='adam', choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], help='Optimizer function. ' 'Note: dist_muon is deprecated; use --optimizer muon ' 'with --use-distributed-optimizer instead.') group.add_argument('--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU') group.add_argument('--optimizer-cuda-graph', action='store_true', help='Enable CUDA graph for optimizer step') group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, help='Ratio of optimizer state to offload to CPU') group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.") group.add_argument('--overlap-cpu-optimizer-d2h-h2d', action='store_true', default=False, help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.') group.add_argument('--dump-param-to-param-group-map', 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.") group.add_argument('--no-pin-cpu-grads', action='store_false', dest='pin_cpu_grads', help='Disable pinning of CPU memory for gradients.') group.add_argument('--no-pin-cpu-params', action='store_false', dest='pin_cpu_params', help='Disable pinning of CPU memory for parameters.') group.add_argument('--dataloader-type', type=str, default=None, choices=['single', 'cyclic', 'external'], help='Single pass vs multiple pass data loader') group.add_argument('--no-persist-layer-norm', action='store_true', help='Disable using persistent fused layer norm kernel. ' 'This kernel supports only a set of hidden sizes. Please ' 'check persist_ln_hidden_sizes if your hidden ' 'size is supported.') group.add_argument('--no-gradient-accumulation-fusion', action='store_false', help='Disable fusing gradient accumulation to weight ' 'gradient computation of linear layers', dest='gradient_accumulation_fusion') group.add_argument('--use-mcore-models', action='store_true', dest='deprecated_use_mcore_models', help='DEPRECATED. Use the implementation from megatron core.' 'Now ignored and mcore models are the default.') return parser def _add_rerun_machine_args(parser): from megatron.training.config import RerunStateMachineConfig rerun_factory = ArgumentGroupFactory(RerunStateMachineConfig, exclude=["check_for_nan_in_loss"]) group = rerun_factory.build_group(parser, "rerun engine") return parser def _add_initialization_args(parser): from megatron.training.config import RNGConfig rng_factory = ArgumentGroupFactory(RNGConfig) group = rng_factory.build_group(parser, "RNG and initialization") group.add_argument('--init-method-xavier-uniform', action='store_true', help='Enable Xavier uniform parameter initialization') return parser def _add_learning_rate_args(parser): from megatron.training.config import SchedulerConfig sched_factory = ArgumentGroupFactory(SchedulerConfig) group = sched_factory.build_group(parser, title="learning rate and weight decay") group.add_argument('--lr', type=float, default=None, help='Initial learning rate. Depending on decay style ' 'and initial warmup, the learning rate at each ' 'iteration would be different.') group.add_argument('--warmup', type=int, default=None, help='Old lr warmup argument, do not use. Use one of the' '--lr-warmup-* arguments above') group.add_argument('--min-lr', type=float, default=0.0, help='Minimum value for learning rate. The scheduler' 'clip values below this threshold.') group.add_argument('--decoupled-lr', type=float, default=None, help='Separate learning rate for the input and output layer') group.add_argument('--decoupled-min-lr', type=float, default=None, help='Minimum value for learning rate for the input and output layer. The scheduler' 'clip values below this threshold') group.add_argument('--freeze-all-layers', action='store_true', help='Freeze all layers of the model.') return parser def _add_checkpointing_args(parser): from megatron.training.config import CheckpointConfig ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_optim", "save_rng", "load_optim", "load_rng"]) group = ckpt_factory.build_group(parser, "checkpointing") group.add_argument('--no-save-optim', action='store_true', default=None, help='Do not save current optimizer.') group.add_argument('--no-save-rng', action='store_true', default=None, help='Do not save current rng state.') group.add_argument('--no-load-optim', action='store_true', default=None, help='Do not load optimizer when loading checkpoint.') group.add_argument('--no-load-rng', action='store_true', default=None, help='Do not load rng state when loading checkpoint.') group.add_argument('--override-ckpt-iteration', type=int, default=None, help='Override the iteration stored in the loaded checkpoint. ' 'Also resets consumed_train_samples accordingly so the ' 'data loader replays samples from that iteration onward.') group.add_argument('--use-dist-ckpt', action='store_true', dest='use_dist_ckpt_deprecated', help='Deprecated: see --ckpt-format.') group.add_argument('--dist-ckpt-format', dest='dist_ckpt_format_deprecated', help='Deprecated: see --ckpt-format.') group.add_argument('--dist-ckpt-workers', 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.') group.add_argument('--ckpt-fully-parallel-save', action='store_true', dest='ckpt_fully_parallel_save_deprecated', help='Deprecated: see --no-ckpt-fully-parallel-save.') group.add_argument('--ckpt-drop-redundant-extra-state', action='store_true', default=False, help='Drop TE `_extra_state` artifacts that carry no ' 'irreplaceable state (no FP8, or block/current FP8 ' 'scaling) from the distributed checkpoint, keeping them ' 'local instead of writing them. Only delayed-scaling ' '`_extra_state` (amax history + scale) is ever needed and ' 'it is always persisted regardless of this flag. By ' 'default (flag off) all `_extra_state` are persisted, ' 'preserving the previous behavior. Older checkpoints load ' 'unchanged either way.') return parser def _add_mixed_precision_args(parser): group = parser.add_argument_group(title='mixed precision') group.add_argument('--grad-reduce-in-bf16', action='store_true', help='Reduce gradients in bfloat16.') group.add_argument('--loss-scale', type=float, default=None, help='Static loss scaling, positive power of 2 ' 'values can improve fp16 convergence. If None, dynamic' 'loss scaling is used.') group.add_argument('--initial-loss-scale', type=float, default=2**32, help='Initial loss-scale for dynamic loss scaling.') group.add_argument('--min-loss-scale', type=float, default=1.0, help='Minimum loss scale for dynamic loss scaling.') group.add_argument('--loss-scale-window', type=float, default=1000, help='Window over which to raise/lower dynamic scale.') group.add_argument('--hysteresis', type=int, default=2, help='hysteresis for dynamic loss scaling') group.add_argument('--attention-softmax-in-fp32', action='store_true', help='Run attention masking and softmax in fp32.') group.add_argument('--accumulate-allreduce-grads-in-fp32', action='store_true', help='Gradient accumulation and all-reduce in fp32.') group.add_argument('--fp16-lm-cross-entropy', action='store_true', help='Move the cross entropy unreduced loss calculation' 'for lm head to fp16.') group.add_argument('--output-logit-dtype', type=str, choices=['bf16', 'fp32'], default=None, dest='logit_dtype', help='Output dtype for the language-model output-layer GEMM. When the ' 'requested dtype differs from the input dtype, Transformer Engine ' 'general_gemm is used. By default, logits use the output-layer input dtype.') group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') group.add_argument('--mamba-training-ssm-states-dtype', type=str, choices=['fp32', 'bf16'], default=None, help='Dtype of the materialized inter-chunk SSM states in Mamba training') return parser def _add_distributed_args(parser): from megatron.training.config import DistributedInitConfig dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) group = dist_init_factory.build_group(parser, "distributed init") group.add_argument('--use-megatron-fsdp', action='store_true', default=False, help="Use Megatron's Fully Sharded Data Parallel. Cannot be used together with use_torch_fsdp2.") group.add_argument('--decoder-first-pipeline-num-layers', type=int, default=None, help=('The number of transformer layers on the first pipeline stage of the decoder. ' 'Default None is even split of transformer layers across all pipeline stages')) group.add_argument('--decoder-last-pipeline-num-layers', type=int, default=None, help=('The number of transformer layers on the last pipeline stage of the decoder. ' 'Default None is even split of transformer layers across all pipeline stages')) group.add_argument('--pipeline-model-parallel-layout', 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.')) group.add_argument('--model-parallel-size', type=int, default=None, help='Old model parallel argument, do not use. Use ' '--tensor-model-parallel-size instead.') group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, help='Number of layers per virtual pipeline stage') group.add_argument('--num-virtual-stages-per-pipeline-rank', type=int, default=None, help='Number of virtual pipeline stages per pipeline parallelism rank') group.add_argument('--no-overlap-p2p-communication', action='store_false', help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', dest='overlap_p2p_comm') group.add_argument('--overlap-grad-reduce', action='store_true', default=False, help='If set, overlap DDP grad reduce.') group.add_argument('--ddp-num-buckets', type=int, default=None, help='Number of buckets for data-parallel communication') group.add_argument('--ddp-bucket-size', type=int, default=None, help='Bucket size for data-parallel communication') group.add_argument('--ddp-pad-buckets-for-high-nccl-busbw', action='store_true', 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.') group.add_argument('--ddp-reduce-scatter-with-fp32-accumulation', action='store_true', 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.') group.add_argument('--gtp-remat-reduce-scatter-with-fp32-accumulation', action='store_true', default=False, help='Same trade as --ddp-reduce-scatter-with-fp32-accumulation, but for ' 'the wgrad reduce-scatter GTP weight-remat performs over the gtp_remat axis: send ' 'low-precision values over the wire via an all-to-all and accumulate locally in FP32. ' 'Independent of the DDP flag (different collective, different process group). Costs one ' 'extra unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter.') group.add_argument('--ddp-param-name-patterns-for-fp32-local-accumulation', nargs='+', 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.') group.add_argument('--ddp-average-in-collective', action='store_true', default=False, help='If set, average directly in data-parallel communication collective.') group.add_argument('--overlap-param-gather', action='store_true', default=False, help='If set, overlap param all-gather in distributed optimizer.') group.add_argument('--overlap-param-gather-with-optimizer-step', action='store_true', default=False, help='If set, overlap param all-gather of first bucket with optimizer step.') group.add_argument('--no-align-param-gather', action='store_false', help='If not set, all PP stages will launch param all-gathers simultaneously. ' 'Otherwise, each PP stage will independently launch as needed.', dest='align_param_gather') group.add_argument('--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.') group.add_argument('--megatron-fsdp-version', type=int, default=1, choices=[1, 2], help='Megatron-FSDP implementation version. Defaults to 1.') group.add_argument('--overlap-outer-communication', action='store_true', help="Overlap outer-DP and inner-DP gradient reductions " "with Megatron-FSDP v2.") group.add_argument('--no-use-layer-wise-param-layout', action='store_false', dest='use_layer_wise_param_layout', 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.') group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', 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.') group.add_argument('--disable-symmetric-registration', action='store_true', dest='disable_symmetric_registration', 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. ' 'Cannot be combined with --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub, which require symmetric windows.') group.add_argument('--gtp-remat-nccl-ub', action='store_true', dest='gtp_remat_nccl_ub', default=False, help='Register the wgrad reduce-scatter send buffers with NCCL symmetric ' 'memory on the GTP group, independent of --use-nccl-ub (which covers the DP group).') group.add_argument('--gtp-expert-remat-nccl-ub', action='store_true', dest='gtp_expert_remat_nccl_ub', default=False, help='Like --gtp-remat-nccl-ub but for routed-expert (EGTP) groups.') group.add_argument('--fsdp-manual-registration', action='store_true', dest='fsdp_manual_registration', 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.') group.add_argument('--create-all-gather-group', action='store_true', help='Create a separate process group for all-gather operations ' 'to overlap reduce-scatter and all-gather operations.') group.add_argument('--data-parallel-sharding-strategy', type=str, default='optim_grads_params', choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], help='Sharding strategy of data parallelism.') group.add_argument('--expert-data-parallel-sharding-strategy', type=str, default=None, choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], help='Sharding strategy of data parallelism for expert (MoE) parameters. ' 'When set, --data-parallel-sharding-strategy only applies to ' 'non-expert parameters. Expert parameters are sharded over a narrower ' 'DP group than non-expert parameters when expert parallelism is ' 'enabled, so the two classes can warrant different communication / ' 'memory trade-offs. Defaults to None, which applies ' '--data-parallel-sharding-strategy to every parameter.') group.add_argument('--outer-dp-sharding-strategy', type=str, default='no_shard', choices=['no_shard', 'optim'], 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".') group.add_argument('--expert-outer-dp-sharding-strategy', type=str, default=None, choices=['no_shard', 'optim'], help='Sharding strategy for the outer expert data-parallel group in MFSDP v2. ' 'Valid values are "no_shard" (HSDP) and "optim" (HFSDP). ' 'Defaults to --outer-dp-sharding-strategy when omitted.') group.add_argument('--hfsdp-param-gather-overlap', action='store_true', help='Pipeline HFSDP parameter all-gathers across DP-Outer and DP-Inner. ' 'DP-Outer is prefetched one FSDP unit beyond the existing ' 'DP-Inner prefetch frontier. ' 'Only effective with --outer-dp-sharding-strategy=optim.') group.add_argument('--no-gradient-reduce-div-fusion', action='store_false', dest='gradient_reduce_div_fusion', help='If not set, fuse the division in gradient reduce.') group.add_argument('--fsdp-double-buffer', action='store_true', 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.") group.add_argument('--suggested-communication-unit-size', 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.') group.add_argument('--keep-fp8-transpose-cache', action='store_true', help='If set, keep the fp8 transpose cache when using Megatron FSDP.') group.add_argument('--enable-full-sharding-in-hsdp', action='store_true', help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.') group.add_argument('--num-distributed-optimizer-instances', type=int, default=1, help='Number of Distributed Optimizer copies across Data Parallel domain.') group.add_argument('--torch-fsdp2-no-reshard-after-forward', action='store_false', dest='torch_fsdp2_reshard_after_forward', help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' 'Set to enable FSDP ZeRO-2.') group.add_argument('--cp-comm-type', nargs='+', type=str, 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') group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None, help='Maximum sequence length per CP rank. This is used to calculate the ' 'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.') group.add_argument('--hybrid-context-parallel', action='store_true', default=False, help='Deprecated alias for --dynamic-context-parallel.') group.add_argument('--dynamic-context-parallel', action='store_true', default=False, help='Enable variable-sized CP groups selected per packed microbatch.') group.add_argument('--min-dynamic-context-parallel-size', type=int, default=1, help='Minimum runtime CP group size used by dynamic CP.') group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced', 'default_dynamic_cp']) group.add_argument('--pad-packed-seq-alignment', default=None, help="Pad packed THD inputs to a CP-local multiple or 'max'. " "MXFP8 additionally aligns to 32 tokens after SP sharding.") return parser def _add_validation_args(parser): from megatron.training.config import ValidationConfig val_factory = ArgumentGroupFactory(ValidationConfig) group = val_factory.build_group(parser, "validation") return parser def _add_tokenizer_args(parser): from megatron.training.config import TokenizerConfig tokenizer_factory = ArgumentGroupFactory(TokenizerConfig) group = tokenizer_factory.build_group(parser, "tokenizer") return parser def _add_data_args(parser): group = parser.add_argument_group(title='data and dataloader') group.add_argument('--data-path', nargs='*', default=None, help='The weight and prefix list for a set of train, validation, and test' 'datasets which split according to --split. The accepted formats are: ' '(1) a single prefix, ' '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' '(3) a list of prefixes e.g. prefix1 prefix2. ' 'For (3), weights are inferred from the lengths of the contributing datasets. ' 'This argument is exclusive to the other independent --*-data-path arguments.') group.add_argument('--phase-transition-iterations', type=str, default=None, help='Comma-separated list of iterations where phase ' 'transitions occur. Requires fixed global batch size across phases.') group.add_argument('--split', 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.') group.add_argument('--train-data-path', nargs='*', default=None, help='The weight and prefix list for an independent train dataset. ' 'Follows the same pattern rules as --data-path.') group.add_argument('--valid-data-path', nargs='*', default=None, help='The weight and prefix list for an independent validation dataset. ' 'Follows the same pattern rules as --data-path.') group.add_argument('--test-data-path', nargs='*', default=None, help='The weight and prefix list for an independent test dataset. ' 'Follows the same pattern rules as --data-path.') group.add_argument('--data-args-path', 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.') group.add_argument('--per-split-data-args-path', 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') group.add_argument('--per-dataset-sequences-path', 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.') group.add_argument('--dataloader-fast-cache-load', action='store_true', 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.') group.add_argument('--dataloader-defer-npy-index-mmap', action='store_true', 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.') group.add_argument('--data-cache-path', default=None, help='Path to a directory to hold cached index files.') group.add_argument('--no-mmap-bin-files', action='store_false', help='Disable mmap-ing of .bin files.', dest='mmap_bin_files') group.add_argument('--mock-data', action='store_true', help='Skip data loading and validation and opt for artificial ' 'generation of mock data when an implementation is available.') group.add_argument('--seq-length', type=int, default=None, help='Maximum sequence length to process.') group.add_argument('--encoder-seq-length', type=int, default=None, help='Maximum encoder sequence length to process.' 'This should be exclusive of --seq-length') group.add_argument('--decoder-seq-length', type=int, default=None, help="Maximum decoder sequence length to process.") group.add_argument('--sample-rate', type=float, default=1.0, help='sample rate for training data. Supposed to be 0 ' ' < sample_rate < 1') group.add_argument('--mask-prob', type=float, default=0.15, help='Probability of replacing a token with mask.') group.add_argument('--short-seq-prob', type=float, default=0.1, help='Probability of producing a short sequence.') group.add_argument('--num-workers', type=int, default=2, help="Dataloader number of workers.") group.add_argument('--reset-position-ids', action='store_true', help='Reset posistion ids after end-of-document token.') group.add_argument('--reset-attention-mask', action='store_true', help='Reset self attention mask after ' 'end-of-document token.') group.add_argument('--dataloader-inter-document-masking', action='store_true', help='Return cu_seqlens marking document boundaries ' 'within each sample so that attention is restricted ' 'to individual documents.') group.add_argument('--eod-mask-loss', action='store_true', help='Mask loss for the end of document tokens.') group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', help='If set, do not create attention_masks in dataloader.', dest='create_attention_mask_in_dataloader') group.add_argument('--num-dataset-builder-threads', type=int, default=1, help='Number of parallel threads per rank for dataset builder') group.add_argument('--object-storage-cache-path', type=str, default=None, help='Path to cache index files when using s3 or msc dataloader') group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005, help='The sample surplus to build for the mid-level datasets(s)') group.add_argument('--allow-ambiguous-pad-tokens', action='store_true', 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.') group.add_argument('--fim-data', action='store_true', help='Whether to use the FIM dataset.') group.add_argument('--fim-rate', type=float, default=0.5, help='Probability to convert a training sample into a FIM format.') group.add_argument('--fim-spm-rate', type=float, default=0.5, help='Probability that the a FIM sample uses the SPM format over the PSM format.') group.add_argument('--fim-split-sample', type=str, default=None, help='String around which to split the sample for FIM.') group.add_argument('--fim-fragment-rate', type=float, default=None, help='Rate of FIM on each fragment when --fim-split-sample is not None.') group.add_argument('--fim-no-prefix', type=str, default=None, help='Do not apply FIM to fragments that start with this prefix') group.add_argument('--fim-prefix-token', type=str, default='', help='FIM prefix token') group.add_argument('--fim-middle-token', type=str, default='', help='FIM middle token') group.add_argument('--fim-suffix-token', type=str, default='', help='FIM suffix token') group.add_argument('--fim-pad-token', type=str, default='', help='FIM PAD token') group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', help='FIM EOD token') return parser def _add_autoresume_args(parser): group = parser.add_argument_group(title='autoresume') group.add_argument('--adlr-autoresume', action='store_true', help='Enable autoresume on adlr cluster.') group.add_argument('--adlr-autoresume-interval', type=int, default=1000, help='Intervals over which check for autoresume' 'termination signal') return parser def _add_vision_args(parser): group = parser.add_argument_group(title="vision") # general vision arguements group.add_argument('--num-classes', type=int, default=1000, help='num of classes in vision classificaiton task') group.add_argument('--img-h', type=int, default=224, help='Image height for vision classification task') group.add_argument('--img-w', type=int, default=224, help='Image height for vision classification task') group.add_argument('--num-channels', type=int, default=3, help='Number of channels in input image data') group.add_argument('--patch-dim', type=int, default=16, help='patch dimension') group.add_argument('--classes-fraction', type=float, default=1.0, help='training with fraction of classes.') group.add_argument('--data-per-class-fraction', type=float, default=1.0, help='training with fraction of data per class.') group.add_argument('--no-data-sharding', action='store_false', help='Disable data sharding.', dest='data_sharding') group.add_argument('--head-lr-mult', type=float, default=1.0, help='learning rate multiplier for head during finetuning') # pretraining type and backbone selection` group.add_argument('--vision-pretraining', action='store_true', help='flag to indicate vision pretraining') group.add_argument('--vision-pretraining-type', type=str, default='classify', choices=['classify', 'inpaint', 'dino'], help='pretraining objectives') group.add_argument('--vision-backbone-type', type=str, default='vit', choices=['vit', 'mit', 'swin'], help='backbone types types') group.add_argument('--swin-backbone-type', type=str, default='tiny', choices=['tiny', 'base', 'h3'], help='pretraining objectives') # inpainting arguments group.add_argument('--mask-type', type=str, default='random', choices=['random', 'row'], help='mask types') group.add_argument('--mask-factor', type=float, default=1.0, help='mask size scaling parameter') # dino arguments group.add_argument('--iter-per-epoch', type=int, default=1250, help='iterations per epoch') group.add_argument('--dino-local-img-size', type=int, default=96, help='Image size for vision classification task') group.add_argument('--dino-local-crops-number', type=int, default=10, help='Number of local crops') group.add_argument('--dino-head-hidden-size', type=int, default=2048, help='Hidden dimension size in dino head') group.add_argument('--dino-bottleneck-size', type=int, default=256, help='Bottle neck dimension in dino head ') group.add_argument('--dino-freeze-last-layer', type=float, default=1, help='Freezing last layer weights') group.add_argument('--dino-norm-last-layer', action='store_true', help='Disable Norm in last layer.') group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, help='warump teacher temperature') group.add_argument('--dino-teacher-temp', type=float, default=0.07, help='teacher temperature') group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, help='warmup teacher temperaure epochs') return parser def _add_moe_args(parser): group = parser.add_argument_group(title="moe") # General arguments group.add_argument('--num-experts', type=int, default=None, help='Number of Experts in MoE (None means no MoE)') group.add_argument('--moe-layer-freq', type=moe_freq_type, 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 expert layers, ' '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.') group.add_argument('--moe-use-upcycling', action='store_true', help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.') # Router arguments group.add_argument('--moe-router-load-balancing-type', nargs='+', type=str, choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'quantile_balancing', 'none'], default='aux_loss', help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "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; "quantile_balancing" selects either the exact micro-batch estimator or Kimi K3 global-batch histogram estimator via --moe-router-quantile-balancing-estimation-scope; "none" implies no load balancing. The default is "aux_loss".') group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') group.add_argument('--moe-router-quantile-balancing-estimation-scope', type=str, choices=['micro_batch', 'global_batch'], default='micro_batch', help='Population used to estimate quantile-balancing biases. "micro_batch" preserves the existing exact estimator; "global_batch" uses Kimi K3 histogram accumulation.') group.add_argument('--moe-router-qb-num-bins', type=int, default=1000, help='Number of uniform histogram bins per expert for global-batch quantile balancing.') # Token dispatcher arguments # MoE communication overlap arguments group.add_argument('--moe-upcycling-granularity', 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.') return parser def _add_mla_args(parser): group = parser.add_argument_group(title="mla") group.add_argument('--q-lora-rank', type=int, default=None, help="Rank of Query tensor's low rank representation.") group.add_argument('--kv-lora-rank', type=int, default=32, help="Rank of Key and Value tensors' low rank representation.") group.add_argument( '--attention-latent-norm-epsilon', type=float, default=None, help="Epsilon for the primary query and key-value latent norms in attention. " "Defaults to --norm-epsilon when unset.", ) group.add_argument('--qk-head-dim', type=int, default=128, help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim") group.add_argument('--qk-pos-emb-head-dim', type=int, default=64, help="Dimension of the position embedding in the QK projection.") group.add_argument('--v-head-dim', type=int, default=128, help="Dimension of the head in the V projection.") group.add_argument('--rotary-scaling-factor', type=float, default=1.0, help="Rotary scaling factor for the rotary embeddings.") group.add_argument('--original-max-position-embeddings', type=int, default=4096, help="Original maximum position embeddings for the original model, used by YaRN.") group.add_argument('--mscale', type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") group.add_argument('--output-projection-groups', type=int, default=8, help="Number of groups for grouped low-rank output projection (wo_a).") group.add_argument('--output-projection-lora-rank', type=int, default=1024, help="Low-rank dimension per group for grouped output (wo_a). " "Used when --output-projection-groups > 0.") group.add_argument('--cache-mla-latents', action='store_true', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( '--mla-down-proj-fusion', action='store_true', default=False, help="Enable fused q/kv down-projection and fused input layernorm when backend supports. " "Otherwise fall back to the unfused MLA.", ) return parser def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") # Linear attention group.add_argument('--linear-attention-freq', type=la_freq_type, default=None, 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.') group.add_argument( '--csa-compress-ratios', type=compress_ratios_type, default=None, help='Per-layer compress ratios for compressed sparse attention. ' 'Accepts a Python list expression such as "[0,0,4,128,4,128]" or ' '"([0]+[4,128]*2)*3". Valid values are 0, 4, and 128, and the ' 'decoder uses the first num-layers entries. MTP layers use the tail; ' 'HybridModel patterns need one tail entry per inner MTP layer.', ) return parser def _add_heterogeneous_args(parser): """ Heterogeneous models refer to transformer architectures where individual layers can differ in configuration. Specifically: - Attention or MLP layers can be replaced with either a linear layer or a no-op - MLP intermediate dimensions can vary between layers We use the format of the HuggingFace config files in llama nemotron models to define the architecture. For example, https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json Most notably, the "block_config" maps to a list of attention and mlp configurations for each layer. For example, the "block_config" for a 2 layer model is: "block_configs": [ { "attention": { "n_heads_in_group": 8, "no_op": false, "replace_with_linear": false, }, "ffn": { "ffn_mult": 2.625, "no_op": false, "replace_with_linear": false, } }, { "attention": { "n_heads_in_group": null, "no_op": true, "replace_with_linear": false, }, "ffn": { "ffn_mult": 2.625, "no_op": false, "replace_with_linear": false, } } ] """ group = parser.add_argument_group(title="heterogeneous architecture") group.add_argument('--heterogeneous-layers-config-path', 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.') group.add_argument('--heterogeneous-layers-config-encoded-json', 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.') return parser def _add_experimental_args(parser): group = parser.add_argument_group(title='experimental') group.add_argument('--enable-experimental', action='store_true', help='Enable experimental features.') group.add_argument('--spec', type=str, default=None, nargs='*', 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`') group.add_argument('--hybrid-layer-pattern', 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.') group.add_argument('--hybrid-override-pattern', type=str, default=None, help='Deprecated. Use --hybrid-layer-pattern instead. ' 'If specified, its value will be forwarded to --hybrid-layer-pattern.') group.add_argument('--yaml-cfg', type=str, default=None, help = 'Config file to add additional arguments') # Args of precision-aware optimizer. group.add_argument('--use-precision-aware-optimizer', action='store_true', 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.') group.add_argument('--main-grads-dtype', default='fp32', choices=['fp32', 'bf16'], help='Dtype of main grads when enabling precision-aware-optimizer.') group.add_argument('--main-params-dtype', default='fp32', choices=['fp32', 'fp16'], help='Dtype of main params when enabling precision-aware-optimizer.') group.add_argument('--exp-avg-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], 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.') group.add_argument('--exp-avg-sq-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], 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.') # Megatron-FSDP Arguments group.add_argument('--megatron-fsdp-main-params-dtype', default='fp32', choices=['fp32', 'bf16', 'fp16', 'auto'], 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.") group.add_argument('--megatron-fsdp-main-grads-dtype', default='auto', choices=['fp32', 'bf16', 'fp16', '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.") group.add_argument("--megatron-fsdp-grad-comm-dtype", default='auto', choices=['fp32', 'fp16', 'bf16', '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.") group.add_argument("--no-megatron-fsdp-padded-shards", action="store_false", dest="megatron_fsdp_padded_shards", help="Megatron-FSDP v2 only. If not set, pad every parameter group so each " "rank gets an equal shard. If set, pack each group without padding and split " "it into generally unequal per-rank shards at row boundaries (32-row block " "boundaries for MXFP8).") group.add_argument( '--megatron-fsdp-enable-fine-grained-param-gather', action='store_true', default=False, dest='megatron_fsdp_enable_fine_grained_param_gather', 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.' ), ) group.add_argument("--megatron-fsdp-max-pool-double-buffer", action='store_true', help="When using Megatron-FSDP double buffering, use the MaxPoolAllocator instead of " "the FixedPoolAllocator to support asymmetrical FSDP unit configurations. Will " "increase memory overhead to recycle buffers that fit all FSDP units. Enables " "NCCL user buffer registration and CUDA graph replay for mixed-arch models.") group.add_argument("--fsdp-db-use-persist-buf-on-alloc-fail", action='store_true', help="When using Megatron-FSDP double buffering, persist non-unit modules that " "are not included in the symmetric buffer pool. May be necessary for NCCL " "UBR or CUDA Graphs on hybrid architectures.") return parser def _add_msc_args(parser): group = parser.add_argument_group(title="msc") group.add_argument('--enable-msc', default=False, action='store_true', dest='enable_msc', help='Enable the usage of Multi-Storage Client (MSC) in Megatron Core. ' 'Disabled by default; pass this flag to opt in.') group.add_argument('--disable-msc', default=False, action='store_true', dest='disable_msc_deprecated', help='[DEPRECATED] MSC is disabled by default; this flag is a no-op ' 'and will be removed in a future release.') return parser def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): """Add quant-specific arguments to the main parser If kitchen isn't available, nothing to do here, return unchanged parser """ try: from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider except (ImportError, ModuleNotFoundError): HAVE_KITCHEN = False if HAVE_KITCHEN: group = parser.add_argument_group(title="kitchen") recipe_or_config_group = group.add_mutually_exclusive_group(required=False) recipe_or_config_group.add_argument( '--kitchen-config-file', type=str, default=None, help="Use the config .yaml file at the specified location to " "configure kitchen quantization.", ) recipe_or_config_group.add_argument( '--kitchen-recipe-number', type=int, default=None, help="Use a default kitchen recipe for all linear layers as defined by QAT_PARAMS index. " "The argument has no effect on attention layers.", ) return parser def _add_sft_args(parser): group = parser.add_argument_group(title='sft') group.add_argument('--sft', action="store_true", help='Megatron SFT training') group.add_argument( '--sft-loss-log-mode', type=str, default='token-weighted', choices=['token-weighted', 'microbatch'], help=( 'SFT loss logging reduction: average over all trainable tokens or over valid ' 'microbatch losses.' ), ) group.add_argument('--sft-mock-dataset-config-json', type=str, default=None, help='This config provides the necessary information for the mock ' 'dataset. Accepts either an inline JSON literal or a path to a JSON ' 'file containing the same schema. You can either specify a CSV file ' 'that contains sequence lengths, where each line stores the length of ' 'a sequence, for example: {"mode":"file","path":"/path/to/file"}. ' 'Alternatively, you can specify a distribution (currently only ' 'supporting lognormal distribution) along with the required ' 'parameters, for example, {"mode":"distribution","type":"lognormal",' '"min_seq_len":1024,"max_seq_len":2048,"mean_seq_len":1536,' '"lognormal_sigma":1.1}, where sigma controls the variability of the ' 'lognormal distribution. If not specified and --mock-data is set, ' 'defaults to a lognormal distribution with min_seq_len=seq_length//2, ' 'max_seq_len=seq_length, mean_seq_len=seq_length*3//4, ' 'lognormal_sigma=1.1.') return parser def _add_varlen_dataset_args(parser): group = parser.add_argument_group(title='varlen dataset') group.add_argument('--use-varlen-dataset', action="store_true", help='Train with VarlenDataset, a variable-length packed (THD) dataset ' 'that consumes instruction-tuning data from a HuggingFace Hub repo id, ' 'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt ' '/ openai-messages) is auto-detected from the dataset columns. ' 'Mutually exclusive with --sft. Auto-picks a sequence packing ' 'scheduler when none is given: dp_balanced. ' 'Combine with --mock-data for a synthetic lognormal sequence-length ' 'distribution; see --varlen-mock-dataset-config-json.') group.add_argument('--varlen-sbhd-validation', action="store_true", help='Reference SBHD mode for THD numerical verification. When set, ' 'VarlenDataset emits SBHD-style samples right-padded to ' '--seq-length (no cu_seqlens, no packing scheduler), so the run can ' 'be compared against the THD path to validate correctness. ' 'Incompatible with --sequence-packing-scheduler.') group.add_argument('--varlen-mock-dataset-config-json', type=str, default=None, help='Mock-dataset config for --use-varlen-dataset --mock-data. ' 'Accepts either an inline JSON literal or a path to a JSON file ' 'containing the same schema as --sft-mock-dataset-config-json: either ' '{"mode":"file","path":"/path/to/lengths.csv"}, ' '{"mode":"distribution","type":"lognormal","min_seq_len":1024,' '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or ' '{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. ' 'If not specified, defaults to a lognormal distribution with ' 'min_seq_len=seq_length//2, max_seq_len=seq_length, ' 'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.') return parser def _add_logits_distillation_args(parser): group = parser.add_argument_group(title='Logits Distillation') # -- Teacher saving -- group.add_argument('--logits-save-dir', type=str, default=None, help='Directory to save logits.') group.add_argument('--logits-save-top-k', type=int, default=128, help='Number of top logits to save.') group.add_argument('--logits-save-top-p', type=float, default=None, help='Top-P (nucleus) threshold applied after top-K ' 'selection when saving logits. Only the smallest ' 'set of entries whose cumulative probability mass ' 'reaches this threshold is kept. Must be in (0, 1].') group.add_argument('--logits-save-top-p-min-k', type=int, default=1, help='Minimum number of entries kept per token when ' 'top-P masking is active, regardless of ' 'cumulative mass. Default: 1.') group.add_argument('--logits-save-dtype', type=str, default='fp16', choices=['fp16', 'bf16', 'fp32'], help='Dtype for on-disk top-K log-probabilities.') # -- Student loading -- group.add_argument('--logits-load-dir', type=str, default=None, help='Directory to load logits.') group.add_argument('--logits-load-decode-threads', type=int, default=4, help='Number of decode threads for cached-logits zstd ' 'decompression and torch.load processing.') group.add_argument('--logits-load-msc-prefetch-depth', type=int, default=2, help='For MSC/object-storage logits tar shards, number ' 'of whole tar shards to prefetch into the MSC ' 'cache ahead of sequential tar consumption.') group.add_argument('--logits-load-kd-loss-alpha', type=float, default=1.0, help='KD loss alpha for loading logits. Total loss is calculated as ' 'alpha * kd_loss + (1 - alpha) * lm_loss.') group.add_argument('--logits-load-ignore-errors', action='store_true', default=False, help='When set, KD loss errors are logged as warnings and ' 'training falls back to LM-only loss instead of crashing.') group.add_argument('--logits-load-ignore-hash', action='store_true', default=False, help='When set, skip cached-logits tar metadata hash validation. ' 'Use only when intentionally loading logits saved with an ' 'older or otherwise incompatible dataset hash.') return parser def _add_fault_injector_args(parser): from megatron.training.config import FaultInjectorConfig ArgumentGroupFactory(FaultInjectorConfig).build_group(parser, "fault injector") return parser