# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import fnmatch import logging import math from dataclasses import dataclass, field from typing import Callable, Optional, Tuple, Union import torch from ..utils import is_te_min_version logger = logging.getLogger(__name__) @dataclass(frozen=True) class ParamPredicate: """Wraps a matching function to make it hashable for ParamKey. Example: >>> shape_1_param = ParamPredicate(name="s1", fn=lambda param: len(param.shape) == 1) >>> shape_1_param(torch.empty(10)) True >>> shape_1_param_copy = ParamPredicate(name="s1", fn=lambda param: len(param.shape) == 1) >>> shape_1_param == shape_1_param_copy # name is used to match True >>> {shape_1_param, shape_1_param_copy} == {shape_1_param} # set hashing works properly NOTE: __hash__ and __eq__ are automatically generated by @dataclass(frozen=True) based solely on 'name' because we set compare=False/hash=False on 'fn'. """ name: str fn: Callable[[torch.nn.Parameter], bool] = field(compare=False, hash=False) def __call__(self, param: torch.nn.Parameter) -> bool: return self.fn(param) @dataclass(frozen=True) class ParamWithNamePredicate: """Wraps a matching function to make it hashable for ParamKey. Example: >>> shape_1_not_qkln_param = ParamWithNamePredicate( name="s1_not_qkln", fn=lambda param, name: ( len(param.shape) == 1 or name.endswith(".bias") and not ("q_layernorm." in name or "k_layernorm." in name) ) ) >>> shape_1_not_qkln_param(torch.empty(10), "interesting.bias") True >>> shape_1_not_qkln_param(torch.empty(10), "interesting.q_layernorm.bias") False NOTE: __hash__ and __eq__ are automatically generated by @dataclass(frozen=True) based solely on 'name' because we set compare=False/hash=False on 'fn'. """ name: str fn: Callable[[torch.nn.Parameter, str], bool] = field(compare=False, hash=False) def __call__(self, param: torch.nn.Parameter, name: str) -> bool: return self.fn(param, name) @dataclass(frozen=True, slots=True) class ParamKey: """Key to group parameters by. All such grouped parameters can share an optimizer config specification.""" # TODO: Can add layer_id here later. name: Union[str, Tuple[str]] = field(default_factory=tuple) """Parameter name(s), will use unix filesystem path syntax for matching.""" attr: Union[str, Tuple[str]] = field(default_factory=tuple) """Parameter attribute(s).""" predicate: Union[ParamPredicate, Tuple[ParamPredicate]] = field(default_factory=tuple) """Predicate(s) to match parameters by. If multiple predicates are provided, any must match.""" with_name_predicate: Union[ParamWithNamePredicate, Tuple[ParamWithNamePredicate]] = field( default_factory=tuple ) """ Predicate(s) to match parameters with their name. If multiple predicates are provided, any must match. This is useful if you need to filter out some parameters from an otherwise positive match by their name. """ def matches(self, param: torch.nn.Parameter, param_name: str) -> bool: """Returns true if passed-in parameter (with name) matches `param_key`. Args: param (torch.nn.Parameter): Handle to parameter object. param_name (str): Name of parameter in underlying PyTorch module. Returns: bool: True if parameter matches passed-in param_key. """ # Check if name matches. if isinstance(self.name, str): target_names = [self.name] else: target_names = list(self.name) for target_name in target_names: if fnmatch.fnmatch(param_name, target_name): return True # Check if attribute matches. if isinstance(self.attr, str): target_attrs = [self.attr] else: target_attrs = list(self.attr) for target_attr in target_attrs: if getattr(param, target_attr, False): return True # Check if predicate matches. if isinstance(self.predicate, ParamPredicate): if self.predicate(param): return True else: for predicate in self.predicate: if predicate(param): return True # Check if with_name_predicate matches. if isinstance(self.with_name_predicate, ParamWithNamePredicate): if self.with_name_predicate(param, param_name): return True else: for predicate in self.with_name_predicate: if predicate(param, param_name): return True return False @dataclass class OptimizerConfig: """Configuration object for Megatron optimizers.""" ############## # General ############## lr: Optional[float] = None """Initial learning rate. Depending on decay style and initial warmup, the learning rate at each iteration would be different. """ min_lr: Optional[float] = None """Minumum value for learning rate. The scheduler clip values below this threshold.""" decoupled_lr: Optional[float] = None """Separate learning rate for the input and output layer.""" decoupled_min_lr: Optional[float] = None """Minimum value for learning rate for the input and output layer. The scheduler clip values below this threshold. """ weight_decay: float = 0.01 """Weight decay coefficient for L2 regularization.""" rmsnorm_weight_decay_mult: Optional[float] = None """Multiplier on scheduled global weight decay for RMSNorm gamma parameters. ``None`` preserves the default weight-decay policy. Set to 1.0 for full decay, 0.5 for half decay, or 0.0 to disable decay on RMSNorm gamma. Must be finite and nonnegative. When set, takes precedence over ``apply_wd_to_qk_layernorm`` for RMSNorm gamma only; other parameters retain their existing policy. """ embedding_min_norm: Optional[float] = None """Trigger threshold for input word-embedding row L2 norms after an optimizer update. Rows with nonzero norms below this value are rescaled to twice this value. Disabled when ``None`` and requires the distributed optimizer. Exact-zero rows are left unchanged because they have no direction along which to rescale. """ apply_wd_to_qk_layernorm: bool = False """If true, apply weight decay to qk layernorm as a special case.""" ############## # Precision ############## fp8_recipe: Optional[str] = None """The type of fp8 recipe will affect the processing logic inside distributed optimizer.""" fp16: bool = False """If true, train with fp16 mixed precision training. Defaults to False.""" bf16: bool = False """If true, train with bf16 mixed precision training. Defaults to False.""" reuse_grad_buf_for_mxfp8_param_ag: bool = False """If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True.""" params_dtype: torch.dtype = torch.float32 """dtype used when intializing the weights. Defaults to torch.float32.""" use_precision_aware_optimizer: bool = False """If true, allows optimizer-related tensors (master_param, gradients and optimizer states) to be set to lower precision. Defaults to False. """ store_param_remainders: bool = True """If true, store the 16-bit FP32 parameter remainders in the optimizer state, excluding the 16 bits shared with the BF16 parameters. This lowers GPU memory usage. Defaults to True. """ main_grads_dtype: torch.dtype = torch.float32 """dtype of main grads when enabling precision-aware-optimizer""" main_params_dtype: torch.dtype = torch.float32 """dtype of main params when enabling precision-aware-optimizer""" exp_avg_dtype: torch.dtype = torch.float32 """dtype of exp_avg when enabling precision-aware-optimizer""" exp_avg_sq_dtype: torch.dtype = torch.float32 """dtype of exp_avg_sq when enabling precision-aware-optimizer""" optimizer: str = 'adam' """Optimizer name (e.g., 'adam', 'sgd', 'muon'). Can be overridden per-parameter group via config_overrides to use different optimizers for different parameters.""" ############### # Loss scaling ############### loss_scale: Optional[float] = None """Static loss scaling, positive power of 2 values can improve fp16 convergence. If None, dynamic loss scaling is used. """ initial_loss_scale: float = 2**32 """Initial loss-scale for dynamic loss scaling.""" min_loss_scale: float = 1.0 """Minimum loss scale for dynamic loss scaling.""" loss_scale_window: float = 1000 """Window over which to raise/lower dynamic scale.""" hysteresis: int = 2 """Hysteresis for dynamic loss scaling.""" ################################################################################### # Optimizer-specific parameters. ################################################################################### # Adam. adam_beta1: float = 0.9 """First coefficient for computing running averages of gradient and its square in Adam optimizer. """ adam_beta2: float = 0.999 """Second coefficient for computing running averages of gradient and its square in Adam optimizer. """ adam_eps: float = 1e-08 """Term added to the denominator to improve numerical stability in Adam optimizer.""" decoupled_weight_decay: bool = True """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false, original Adam update rule will be used. Defaults to True. """ # SGD. sgd_momentum: float = 0.9 """Momentum factor for SGD optimizer.""" # emerging optimizers. muon_momentum: float = 0.95 """The momentum used by the internal SGD in Muon optimizer.""" muon_split_qkv: bool = True """Whether to split QKV parameters for Muon optimizer.""" muon_nesterov: bool = False """Whether to use Nesterov-style momentum in the internal SGD.""" muon_scale_mode: str = "spectral" """The mode to use for the scale factor. Defaults to "spectral".""" muon_fp32_matmul_prec: str = "medium" """The precision to use for the fp32 matmul. Defaults to "medium".""" muon_coefficient_type: str = "quintic" """Newton-Schulz coefficient type for the Muon optimizer. Valid types are discovered dynamically from the installed ``emerging_optimizers`` package. Defaults to "quintic".""" muon_num_ns_steps: int = 5 """The number of iteration steps to use in the Newton-Schulz iteration.""" muon_tp_mode: str = "duplicated" """How to perform NS calculation for tensor parallel weights. "blockwise" orthogonalizes each shard independently, which makes the update rule depend on the parallelism config; "duplicated" and "distributed" both orthogonalize the whole matrix, so results do not change as TP changes. "auto" select between duplicated and distributed mode per-weight for dense weights. "layer_sharded" uses LayerShardedMuon instead of TensorParallelMuon: each 2D weight is assigned one NS home rank in the (GTP_remat x TP) domain — i.e. the GTP domain — all_to_all stages over the gtp_remat and tp groups assemble the complete (P, Q) momentum on the home, the exact same full-matrix Newton-Schulz as duplicated mode runs there with zero communication and zero redundancy, and reverse all_to_all stages scatter the result back to the original shards. layer_sharded requires optimizer='muon', the layer-wise distributed optimizer path and muon_split_qkv=False (split-QKV Newton-Schulz is not implemented on the layer-sharded path); validate_run checks all three. Defaults to "duplicated".""" muon_use_syrk: bool = False """Use the Triton SYRK kernel for the Gram matrix in Newton-Schulz iteration.""" muon_extra_scale_factor: float = 1.0 """Additional scale factor for the muon update.""" muon_ns_batch_size: int = 1 """Max number of same-shape matrices fused into one batched Newton-Schulz under muon_tp_mode='layer_sharded', on each NS home or in a single-rank (GTP_remat x TP) domain. MoE gives such a rank hundreds of identically shaped expert weights, where the per-matrix loop is kernel-launch bound; batching trades transient memory for fewer launches. Local NS inputs and updates are consumed one batch at a time. The value is an upper bound per same-shape bucket, so a value above the matrix count adds no memory. Batches of more than one use baddbmm instead of addmm and differ from the per-matrix path by kernel-level rounding, so bitwise parity with duplicated mode is lost; the default of 1 keeps the bit-exact path. Values > 1 require emerging-optimizers >= 0.3.0.""" muon_concurrent_groups: bool = True """Run each param group's layer-sharded pipeline (exchange + Newton-Schulz + update) on its own CUDA stream so one group's compute fills another group's all_to_all stall. Bitwise-neutral (op order within a group is unchanged), but the transient buffers of all groups are live at once — disable (or lower muon_ns_batch_size) if peak memory is tight. Only used when muon_tp_mode='layer_sharded'. Defaults to True.""" muon_scalar_optimizer: str = 'adam' """Optimizer for nonlinear parameters (embeddings, biases, norms) when using muon. One of 'adam' or 'lion'. Defaults to 'adam'.""" # Lion. lion_beta1: float = 0.95 """First beta coefficient for Lion optimizer (used in sign update). Defaults to 0.95.""" lion_beta2: float = 0.98 """Second beta coefficient for Lion optimizer (used in momentum EMA update). Defaults to 0.98.""" soap_shampoo_beta: float = 0.95 """The beta parameter for the Shampoo preconditioner.""" soap_use_kl_shampoo: bool = True """Whether to use the KL-Shampoo preconditioner.""" adaptive_muon_moment2_method: str = "adamuon" """The method to use for the moment2 update in Adaptive Muon optimizer.""" adaptive_muon_beta2: float = 0.95 """The beta2 parameter for the Adaptive Muon optimizer.""" adaptive_muon_eps: float = 1e-8 """The eps parameter for the Adaptive Muon optimizer.""" ####################### # Distributed optimizer ####################### use_distributed_optimizer: bool = False """Distribute optimizer state over data-parallel replicas.""" use_layer_wise_distributed_optimizer: bool = False """Use :class:`LayerWiseDistributedOptimizer` for emerging optimizers (e.g. Muon). When set via ``--use-distributed-optimizer`` with an emerging optimizer, the training arguments layer sets this flag and resets ``use_distributed_optimizer`` to False so that the standard distributed-optimizer path is not triggered.""" overlap_param_gather: bool = False """If true, overlap param all-gather with forward compute. This argument is intended to have the same value as the "overlap_param_gather" argument in the "distributed_data_parallel_config.py" file. In the optimizer, this argument is only used when "reuse_grad_buf_for_mxfp8_param_ag=True & fp8_param_gather=True". """ overlap_param_gather_with_optimizer_step: bool = False """If true, overlap param all-gather of first bucket with optimizer step.""" ####################### # Optimizer Offload ####################### optimizer_cpu_offload: bool = False """If True, offload optimizer states tensor and compute to CPU.""" optimizer_offload_fraction: float = 0.0 """Specifies the fraction of optimizer states to offload from GPU memory to CPU.""" use_torch_optimizer_for_cpu_offload: bool = False """If True, use torch.optim.Optimizer for CPU offload.""" overlap_cpu_optimizer_d2h_h2d: bool = False """ When set to `True`, this flag enables overlapping of the CPU optimizer update process with the data transfer operations. This can help improve overall training efficiency by reducing idle time during data movement, allowing the optimizer to perform updates while gradients and parameters are being transferred between devices. """ pin_cpu_grads: bool = True """If True, pin the optimizer gradients to CPU memory.""" pin_cpu_params: bool = True """If True, pin the optimizer parameters to CPU memory.""" ################ # Miscellaneous ################ clip_grad: float = 1.0 """Gradient clipping based on global L2 norm.""" grad_norm_skip_threshold: float = float('inf') """Skip gradient update if the gradient norm exceeds this threshold. Disabled by default. Set a finite value to enable skip-on-large-grad behavior. """ dsa_indexer_clip_grad: Optional[float] = None """Threshold for clipping DSA indexer gradients with an L2 norm taken separately from the non-indexer gradients, which keep the ``clip_grad`` threshold. ``None`` clips every parameter together under ``clip_grad``.""" log_num_zeros_in_grad: bool = False """If true, calculate and log the number of zeros in gradient.""" barrier_with_L1_time: bool = False """If true, use barrier with level 1 time measurements.""" timers: Optional[Callable] = None """Function to get timers.""" config_logger_dir: str = "" """When non-empty, dumps entry-point configs to config_logger_dir""" optimizer_cuda_graph: bool = False """If true, enables CUDA graph for optimizer step.""" def validate(self) -> None: """Check optimizer-local settings without modifying the config.""" if self.embedding_min_norm is not None and ( not math.isfinite(self.embedding_min_norm) or self.embedding_min_norm <= 0.0 ): raise ValueError("embedding_min_norm must be a finite positive value") if self.embedding_min_norm is not None and self.optimizer_cuda_graph: raise ValueError("embedding_min_norm is not supported with optimizer_cuda_graph") # HybridDeviceOptimizer keeps the authoritative master copy of offloaded parameters # on the CPU and overwrites the GPU copy from it, which would undo the projection. # Offload also requires the precision-aware optimizer, which has no FP32 master here. if self.embedding_min_norm is not None and self.optimizer_cpu_offload: raise ValueError("embedding_min_norm is not supported with optimizer CPU offload") if self.rmsnorm_weight_decay_mult is not None and ( not math.isfinite(self.rmsnorm_weight_decay_mult) or self.rmsnorm_weight_decay_mult < 0.0 ): raise ValueError("rmsnorm_weight_decay_mult must be a finite nonnegative value") if self.optimizer == 'dist_muon': logger.warning( "optimizer='dist_muon' is deprecated. " "Use --optimizer muon --use-distributed-optimizer instead." ) if self.lr is not None and self.min_lr is not None: assert self.min_lr <= self.lr if self.fp8_recipe == "mxfp8": if not self.reuse_grad_buf_for_mxfp8_param_ag: import warnings warnings.warn( "mxfp8 without using reuse_grad_buf_for_mxfp8_param_ag and fp8_param_gather" "will use significant amount additional GPU memory." "Setting --reuse-grad-buf-for-mxfp8-param-ag and --fp8-param-gather is " "recommended for mxfp8 training." ) if self.reuse_grad_buf_for_mxfp8_param_ag and self.overlap_param_gather_with_optimizer_step: raise ValueError( "overlap_param_gather_with_optimizer_step is not supported with " "reuse_grad_buf_for_mxfp8_param_ag." ) if self.optimizer_cuda_graph: assert not math.isfinite( self.grad_norm_skip_threshold ), 'Setting grad_norm_skip_threshold not supported with optimizer CUDA graph' # finalize_optimizer_config() enables the layer-wise path for non-sgd/adam optimizers # that use the distributed optimizer (or the deprecated dist_muon alias), so check the # value it will produce, not only the pre-finalize flag. will_be_layer_wise = self.use_layer_wise_distributed_optimizer or ( self.optimizer not in ('sgd', 'adam') and (self.use_distributed_optimizer or self.optimizer == 'dist_muon') ) if self.muon_tp_mode == 'layer_sharded': # 'dist_muon' is the deprecated alias for muon + the layer-wise path. if self.optimizer not in ('muon', 'dist_muon'): raise ValueError( f"muon_tp_mode='layer_sharded' requires optimizer='muon' (got " f"{self.optimizer!r}); other optimizers, including adaptive_muon, do not " "implement layer sharding." ) if not will_be_layer_wise: raise ValueError( "muon_tp_mode='layer_sharded' requires the layer-wise distributed " "optimizer path (use_layer_wise_distributed_optimizer=True)." ) if self.muon_split_qkv: raise ValueError( "muon_tp_mode='layer_sharded' does not implement split-QKV Newton-Schulz; " "set muon_split_qkv=False." ) # Dtypes may still be CLI strings here; finalize_optimizer_config maps them to torch.dtype. fp32 = (torch.float32, 'fp32') if self.use_precision_aware_optimizer: assert ( self.optimizer == 'adam' ), '--use-precision-aware-optimizer only supported with adam' # Only the FusedAdam in TE and HybridDeviceOptimizer supports # --use-precision-aware-optimizer. # TODO: Remove this check when apex's FusedAdam is no longer used. if not self.optimizer_cpu_offload: try: import inspect # TODO: Move this below? from transformer_engine.pytorch.optimizers import FusedAdam as Adam adam_args = inspect.signature(Adam).parameters arg_names = [ 'master_weight_dtype', 'exp_avg_dtype', 'exp_avg_sq_dtype', 'use_decoupled_grad', ] for name in arg_names: assert name in adam_args, ( "Current FusedAdam of TE doesn't support " "--use-precision-aware-optimizer, please update TE version." ) except ImportError: raise RuntimeError( '--use-precision-aware-optimizer requires FusedAdam from ' 'TransformerEngine, but not found.' ) else: assert ( self.main_grads_dtype in fp32 ), "main_grads_dtype can only be fp32 when not using precision-aware optimizer" assert ( self.main_params_dtype in fp32 ), "main_params_dtype can only be fp32 when not using precision-aware optimizer" assert ( self.exp_avg_dtype in fp32 ), "exp_avg_dtype can only be fp32 when not using precision-aware optimizer" assert ( self.exp_avg_sq_dtype in fp32 ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer" # Optimizer CPU offload check if self.optimizer_cpu_offload: assert self.use_precision_aware_optimizer, ( "The optimizer cpu offload must be used in conjunction with " "`--use-precision-aware-optimizer`, as the hybrid device optimizer reuses the " "code path of this flag." ) if self.overlap_param_gather_with_optimizer_step: assert self.use_distributed_optimizer, ( '--overlap-param-gather-with-optimizer-step only supported with ' 'distributed optimizer' ) assert ( self.overlap_param_gather ), 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' def validate_run( self, *, ddp_config: "DistributedDataParallelConfig", use_torch_fsdp2: bool, ckpt_format: str | None, ) -> None: """Check settings shared with DDP and checkpointing without modifying them.""" if self.optimizer_cpu_offload: assert not ddp_config.fp8_param_gather or self.fp8_recipe == "delayed", ( "When `--fp8-param-gather` is enabled, the optimizer cpu offload " "must be used in conjunction with `--fp8-recipe delayed`." ) if self.overlap_param_gather: # finalize_optimizer_config() may already have turned use_distributed_optimizer into # use_layer_wise_distributed_optimizer (muon / dist_muon), which shares the same # param-gather overlap, so accept either form. assert ( self.use_distributed_optimizer or self.use_layer_wise_distributed_optimizer or ddp_config.use_megatron_fsdp or self.optimizer == 'dist_muon' ), ( '--overlap-param-gather only supported with distributed optimizer, ' 'megatron fsdp, or dist_muon' ) assert ( ddp_config.overlap_grad_reduce ), 'Must use --overlap-param-gather with --overlap-grad-reduce' if self.use_distributed_optimizer: assert ( not use_torch_fsdp2 ), "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" if self.overlap_param_gather_with_optimizer_step: assert ckpt_format == "torch", ( '--overlap-param-gather-with-optimizer-step not supported with ' 'distributed checkpointing yet' ) def validate_model( self, *, moe_single_grouped_weight: bool, virtual_pipeline_model_parallel_size: int | None ) -> None: """Check model-dependent constraints at the native or legacy model boundary.""" will_be_layer_wise = self.use_layer_wise_distributed_optimizer or ( self.optimizer not in ('sgd', 'adam') and (self.use_distributed_optimizer or self.optimizer == 'dist_muon') ) assert not (will_be_layer_wise and moe_single_grouped_weight), ( "The LayerWise distributed optimizer does not support --moe-single-grouped-weight: " "Muon semantics for a single grouped [E, N, K] expert weight are not defined. " "Disable --moe-single-grouped-weight or use Adam/DistributedOptimizer." ) if self.overlap_param_gather_with_optimizer_step: assert virtual_pipeline_model_parallel_size is not None, ( '--overlap-param-gather-with-optimizer-step only supported with ' 'interleaved pipeline parallelism' ) @property def use_precision_aware_optimizer_no_fp8_or_ds_fp8(self) -> bool: """Whether decoupled grads / precision-aware handling applies outside FP8 (except delayed). The following condition is used to avoid repetition in distrib_optimizer.py. This is because in distrib_optimizer.py, the process to handle parameters are different for different training precision settings. FP8 cases require different handling while FP8 delayed scaling is an exception because the Adam optimizer in TransformerEngine supports it in the kernel computation. This is also the flag to determine the usage of param.grad or param.decoupled_grad Derived on access so configs built directly (without finalize_optimizer_config) work. """ return bool( self.use_precision_aware_optimizer and ( self.main_params_dtype not in (torch.float32, 'fp32') or (self.fp8_recipe is None or self.fp8_recipe == "delayed") or self.optimizer_cpu_offload ) ) def finalize_optimizer_config(self) -> None: """Resolve dtypes, derived flags and optimizer aliases once the config is validated.""" # 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] self.main_grads_dtype = map_dtype(self.main_grads_dtype) self.main_params_dtype = map_dtype(self.main_params_dtype) self.exp_avg_dtype = map_dtype(self.exp_avg_dtype) self.exp_avg_sq_dtype = map_dtype(self.exp_avg_sq_dtype) if self.use_precision_aware_optimizer and not is_te_min_version("2.1.0"): self.store_param_remainders = False # emerging optimizer check if self.optimizer not in ('sgd', 'adam'): if self.optimizer == 'dist_muon': self.optimizer = 'muon' self.use_layer_wise_distributed_optimizer = True if self.use_distributed_optimizer: self.use_layer_wise_distributed_optimizer = True self.use_distributed_optimizer = False # Backward-compatible aliases (deprecated; use OptimizerConfig directly). AdamOptimizerConfig = OptimizerConfig SGDOptimizerConfig = OptimizerConfig