diff --git a/docs/api-guide/internal/index.md b/docs/api-guide/internal/index.md index 312081ce70b..0f71bac266c 100644 --- a/docs/api-guide/internal/index.md +++ b/docs/api-guide/internal/index.md @@ -16,4 +16,5 @@ Internal utility APIs. num_microbatches_calculator optimizer_param_scheduler +scaling_policy_infrastructure ``` diff --git a/docs/api-guide/internal/scaling_policy_infrastructure.md b/docs/api-guide/internal/scaling_policy_infrastructure.md new file mode 100644 index 00000000000..00b2dfb18a3 --- /dev/null +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -0,0 +1,113 @@ + + +# Scaling Policy Infrastructure + +This internal policy layer centralizes Megatron's parameterization hooks behind a +scaling context. + +The current public recipes are `none`, `mup`, and `depth_mup`. The policy +resolver also accepts legacy MuP aliases, then syncs them to the canonical recipe +fields so model, optimizer, YAML, and checkpoint paths see the same effective +scaling context. Standard Megatron behavior is represented as the identity +policy, so code paths can call the same hooks whether or not a scaling recipe is +active. + +## Model Policy + +Model code should route scaling-sensitive decisions through the model scaling +policy instead of reading `use_mup` at each call site. The policy currently +covers: + +- hidden-weight initialization; +- output-projection initialization; +- attention softmax scale; +- embedding activation scaling; +- output logit scaling; +- residual branch output hooks. + +For non-scaling configs, every hook returns the current Megatron default. + +`depth_mup` adds depth-aware model hooks for dense GPT-style residual blocks: + +- dense self-attention residual branch output scaling; +- dense MLP residual branch output scaling; +- dense block output-projection initialization rebased to the base depth. + +Unsupported residual paths must fail closed. Cross-attention, MoE, fused TP +inference residual scaling, hybrid/Mamba layer patterns, MTP, and TE fused MLPs +without an explicit depth-init implementation should not silently inherit +dense-block hooks. + +## Training Policy + +Optimizer code should route per-parameter hyperparameter multipliers through the +training scaling policy. For `mup`, the policy preserves the existing width-MuP +rules: + +- Adam-family hidden matrix parameters use `lr / mup_width_mult`; +- Adam-family hidden matrix parameters use `eps / mup_width_mult`; +- SGD vector-like parameters use `lr * mup_width_mult`; +- Muon-managed matrices stay on Muon scaling rather than Adam-style MuP LR + overrides. +- Muon-family nonlinear and embedding-class scalar parameters are routed through + the configured scalar optimizer, currently `adam` or `lion`. + +For `depth_mup`, the policy is Adam/AdamW-only. Nonzero weight decay requires +`decoupled_weight_decay=True`; coupled Adam/L2 is allowed only with +`weight_decay=0.0`. The default multipliers are: + +| Parameter class | LR policy | Epsilon policy | Weight-decay policy | +| --- | --- | --- | --- | +| Embedding/output class | Preserve embedding/output LR policy, including `decoupled_lr` precedence | `width_mult^-1` | Base Megatron policy | +| Hidden matrix-like weights | `width_mult^-1` | `(width_mult * depth_mult)^-1` | `width_mult` | +| Hidden linear/attention/MLP biases | Base LR | `(width_mult * depth_mult)^-1` | Base weight decay | +| Norm scale/bias and unknown 1-D tensors | Base LR | `(width_mult * depth_mult)^-1` as current v1 policy | No weight decay | +| q/k layernorm vectors with `apply_wd_to_qk_layernorm=True` | Base LR | `(width_mult * depth_mult)^-1` | Base weight decay | + +The 1-D parameter policy is deliberate. Tensor rank alone is not semantic: +hidden biases, norm scales, q/k layernorm vectors, and unknown vectors are all +1-D tensors but do not share the same weight-decay rule. + +The public compatibility function `get_mup_config_overrides` remains available +and delegates to the policy implementation for the legacy width-MuP surface. + +## Parameter Metadata + +Model construction may attach explicit parameterization metadata to parameters. +Optimizer grouping should prefer this metadata and keep existing name/shape +fallbacks only for compatibility with unannotated parameters. + +FSDP and other parameter-rewriting paths must preserve the metadata attributes so +optimizer grouping remains stable after wrapping or sharding. + +## Checkpoint Resume + +Distributed optimizer resume and optimizer load must use the same tolerant +parameter-group identifier helper. The identifier includes optimizer-group fields +that can distinguish scaling-policy groups, while treating optional absent fields +as `None`. + +Call sites must not sort groups with direct indexing over the identifier key +list. Standard Adam groups may not carry per-group `optimizer`, and SGD groups +may not carry `eps`; direct indexing turns those valid checkpoints into +resume-time `KeyError`s. Sorting must also use the None-safe sort key rather +than the raw identifier tuple, because Python cannot order `None` against floats +or strings when optional fields are present in only some groups. + +## CLI, YAML, and Checkpoints + +CLI validation, YAML validation, and checkpoint argument restore should all +resolve the same scaling context before downstream code reads global args. + +- Legacy MuP aliases should warn and synchronize to canonical scaling fields. +- `mup_width_mult` is derived from `hidden_size / scaling_base_hidden_size`, not + an independent user input. +- Checkpoint compatibility should compare effective scaling contexts + rather than raw legacy spelling. diff --git a/docs/index.md b/docs/index.md index 11337315588..4f1e5526ec6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -49,6 +49,7 @@ get-started/quickstart user-guide/data-preparation user-guide/training-examples +user-guide/scaling-recipes user-guide/parallelism-guide ``` diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index 2a7ee2eeab9..6dc60f54a6f 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -21,6 +21,7 @@ Guides for using Megatron Core and Megatron-LM. msc_integration data-preparation training-examples +scaling-recipes parallelism-guide features/index ``` diff --git a/docs/user-guide/scaling-recipes.md b/docs/user-guide/scaling-recipes.md new file mode 100644 index 00000000000..128298be960 --- /dev/null +++ b/docs/user-guide/scaling-recipes.md @@ -0,0 +1,135 @@ + + +# Scaling Recipes + +Scaling recipes choose the parameterization used to transfer hyperparameters +between model sizes. The canonical flag is `--scaling-recipe`. + +Megatron currently exposes three recipes: + +| Recipe | Behavior | +| --- | --- | +| `none` | Standard Megatron parameterization. This is the default. | +| `mup` | Width MuP for hidden-size transfer. | +| `depth_mup` | Experimental dense-transformer width-depth MuP for AdamW-style training. | + +## Standard Parameterization + +Use `--scaling-recipe none`, or omit `--scaling-recipe`, to keep the standard +parameterization. Scaling-specific fields such as `--scaling-base-hidden-size` +are rejected unless a scaling recipe is selected. + +## Width MuP + +Use `--scaling-recipe mup` when transferring hyperparameters from a base width to +a target width. + +```bash +--scaling-recipe mup \ +--scaling-base-hidden-size 1024 \ +--scaling-base-head-dim 64 +``` + +For MuP, Megatron derives the width multiplier internally: + +```text +width_mult = hidden_size / scaling_base_hidden_size +``` + +This derived value controls MuP initialization, attention scale, output-logit +scale, and optimizer multipliers. `--mup-width-mult` is no longer an independent +input. If it is provided on the CLI for compatibility, it must match the derived +value. + +When MuP is combined with Muon-family optimizers, Muon-managed matrix parameters +keep Muon's spectral scaling. Nonlinear and embedding-class scalar parameters +are routed through `--muon-scalar-optimizer`, which currently accepts `adam` or +`lion`. + +## Legacy MuP Flags + +The following flags are accepted for checkpoint and script compatibility, but are +deprecated as user-facing inputs: + +| Deprecated flag | Canonical replacement | +| --- | --- | +| `--use-mup` | `--scaling-recipe mup` | +| `--mup-base-hidden-size` | `--scaling-base-hidden-size` | +| `--mup-base-head-dim` | `--scaling-base-head-dim` | +| `--mup-width-mult` | derived from `hidden_size / scaling_base_hidden_size` | + +`--mup-embedding-mult`, `--mup-output-mult`, and `--mup-attn-scale-power` remain +MuP-specific tuning knobs. When `--mup-output-mult` is left at `1.0`, Megatron +sets it to `1 / width_mult` for non-base widths. + +## Depth MuP + +Use `--scaling-recipe depth_mup` when transferring from a base width and depth to +a target dense GPT-style transformer width and depth. + +```bash +--scaling-recipe depth_mup \ +--scaling-base-hidden-size 1024 \ +--scaling-base-num-layers 12 \ +--scaling-base-head-dim 64 +``` + +Megatron derives both multipliers internally: + +```text +width_mult = hidden_size / scaling_base_hidden_size +depth_mult = num_layers / scaling_base_num_layers +``` + +`depth_mup` includes the width-MuP model-side behavior, plus depth-aware residual +branch scaling, dense block output-projection initialization, and Adam/AdamW +optimizer multipliers. The default depth behavior is: + +| Mechanism | Default multiplier | +| --- | --- | +| Dense self-attention and dense MLP residual branch output | `depth_mult^-1` | +| Hidden matrix Adam LR | `width_mult^-1` | +| Hidden matrix Adam epsilon | `(width_mult * depth_mult)^-1` | +| Hidden vector Adam epsilon | `(width_mult * depth_mult)^-1` | +| Embedding/output-class Adam epsilon | `width_mult^-1` | +| Hidden matrix AdamW weight decay | `width_mult` | +| Dense block output-projection initialization | `depth_mult^+0.5` | + +`depth_mup` is intentionally narrow. It currently supports `--optimizer adam`. +If `weight_decay` is nonzero, the optimizer must use AdamW-style decoupled +weight decay (`decoupled_weight_decay=True`). Coupled Adam/L2 is allowed only +with `weight_decay=0.0`. + +Megatron also keeps the standard distinction between hidden biases and +normalization vectors. Under `depth_mup`, hidden linear/attention/MLP biases keep +base weight decay, while normalization vectors and otherwise unknown 1-D tensors +stay on the conservative no-weight-decay path. q/k layernorm vectors use weight +decay only when `apply_wd_to_qk_layernorm=True`. + +The supported runtime path is training. Megatron's validation-loss path enables +the required internal scaling-policy eval context automatically. This does not +make generation, inference, or fused TP inference residual scaling supported. + +The current implementation fails closed for unsupported surfaces, including +cross-attention, hybrid/Mamba layer patterns, MTP, multi-latent attention, +experimental attention variants, MoE, non-Adam optimizers, and TE fused MLPs +when nontrivial dense block output-init depth scaling would be required. + +## Checkpoints and YAML + +Megatron stores and compares the resolved scaling recipe, not just the raw flag +spelling. A checkpoint created with legacy MuP aliases is compatible with the +canonical spelling when both resolve to the same effective recipe and base size. + +YAML configs use the same effective resolution rules as CLI configs. Existing +YAML files that omit the new canonical scaling fields default to +`--scaling-recipe none`. For compatibility with full legacy YAML files that +materialized old defaults, `mup_width_mult: 1.0` is treated as an omitted default; +non-`1.0` YAML values are still validated against the derived width multiplier. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 266da6b74c4..777eb188f41 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -2891,6 +2891,10 @@ def set_param_attribute(): "partition_stride", "is_embedding_or_output_parameter", "is_embedding_parameter", + "is_output_parameter", + "parameterization_role", + "parameterization_shared_group", + "parameterization_tags", "_tensor_parallel_mode", ]: if hasattr(orig_param, attr_name): diff --git a/megatron/core/models/T5/t5_model.py b/megatron/core/models/T5/t5_model.py index b2feb974643..a85854cbca7 100644 --- a/megatron/core/models/T5/t5_model.py +++ b/megatron/core/models/T5/t5_model.py @@ -15,6 +15,7 @@ from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding from megatron.core.models.common.language_module.language_module import LanguageModule from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region from megatron.core.transformer.module import MegatronModule @@ -56,10 +57,10 @@ def __init__( config.hidden_size, vocab_size, config=config, - init_method=( - config.embedding_init_method - if config.use_mup and not share_embeddings_and_output_weights - else config.init_method + init_method=build_model_scaling_policy(config).output_layer_init_method( + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + default_init_method=config.init_method, + embedding_init_method=config.embedding_init_method, ), bias=share_embeddings_and_output_weights, skip_bias_add=not share_embeddings_and_output_weights, diff --git a/megatron/core/models/bert/bert_model.py b/megatron/core/models/bert/bert_model.py index 3fd1e01f4a1..75875f9f7d3 100644 --- a/megatron/core/models/bert/bert_model.py +++ b/megatron/core/models/bert/bert_model.py @@ -135,10 +135,10 @@ def __init__( config.hidden_size, self.vocab_size, config=config, - init_method=( - config.embedding_init_method - if config.use_mup and not self.share_embeddings_and_output_weights - else config.init_method + init_method=self.model_scaling_policy.output_layer_init_method( + share_embeddings_and_output_weights=self.share_embeddings_and_output_weights, + default_init_method=config.init_method, + embedding_init_method=config.embedding_init_method, ), bias=True, skip_bias_add=False, diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py index 7e49ec6c02d..ffdcdb6c651 100644 --- a/megatron/core/models/common/embeddings/language_model_embedding.py +++ b/megatron/core/models/common/embeddings/language_model_embedding.py @@ -6,6 +6,7 @@ from torch import Tensor from megatron.core import tensor_parallel +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none, nvtx_decorator @@ -128,8 +129,7 @@ def forward(self, input_ids: Tensor, position_ids: Tensor, tokentype_ids: int = assert self.tokentype_embeddings is None # MuP: scale embeddings by alpha_input. - if self.config.use_mup and self.config.mup_embedding_mult != 1.0: - embeddings = embeddings * self.config.mup_embedding_mult + embeddings = build_model_scaling_policy(self.config).scale_embedding_activations(embeddings) # If the input flag for fp32 residual connection is set, convert for float. if self.config.fp32_residual_connection: diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 34e3f6b1ba4..2650f480e37 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -8,6 +8,13 @@ from megatron.core import parallel_state, tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.parameterization import ( + ROLE_EMBEDDING, + ROLE_OUTPUT, + ROLE_SHARED_EMBEDDING_OUTPUT, + build_model_scaling_policy, + set_parameterization_metadata, +) from megatron.core.transformer.cuda_graphs import CudaGraphManager try: @@ -46,6 +53,7 @@ def __init__( self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None ) -> None: super().__init__(config=config) + self.model_scaling_policy = build_model_scaling_policy(config) self._set_attention_backend() if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -204,27 +212,55 @@ def setup_embeddings_and_output_layer(self) -> None: # This is the original Megatron attribute used by decoupled_lr, Muon, FSDP, etc. if self.pre_process and hasattr(self, 'embedding'): self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True + if self.share_embeddings_and_output_weights: + self.embedding.word_embeddings.weight.is_output_parameter = True + set_parameterization_metadata( + self.embedding.word_embeddings.weight, + role=( + ROLE_SHARED_EMBEDDING_OUTPUT + if self.share_embeddings_and_output_weights + else ROLE_EMBEDDING + ), + shared_group=( + 'lm_embedding_output' if self.share_embeddings_and_output_weights else None + ), + ) if ( self.post_process and hasattr(self, 'output_layer') and self.output_layer.weight is not None ): self.output_layer.weight.is_embedding_or_output_parameter = True + self.output_layer.weight.is_output_parameter = True + set_parameterization_metadata( + self.output_layer.weight, + role=( + ROLE_SHARED_EMBEDDING_OUTPUT + if self.share_embeddings_and_output_weights + else ROLE_OUTPUT + ), + shared_group=( + 'lm_embedding_output' if self.share_embeddings_and_output_weights else None + ), + ) # Mark embedding-class parameters for MuP optimizer grouping. # Under MuP table-8-style grouping, embeddings/output use base LR/eps while # hidden matrix-like params use width-scaled LR/eps. mtp_process = getattr(self, 'mtp_process', False) - if self.config.use_mup and (self.pre_process or mtp_process) and hasattr(self, 'embedding'): - for param in self.embedding.parameters(): - param.is_embedding_parameter = True if ( - self.config.use_mup + self.model_scaling_policy.enabled + and (self.pre_process or mtp_process) + and hasattr(self, 'embedding') + ): + self.model_scaling_policy.mark_embedding_class_parameters(self.embedding.parameters()) + if ( + self.model_scaling_policy.enabled and self.post_process and hasattr(self, 'output_layer') and self.output_layer.weight is not None ): - self.output_layer.weight.is_embedding_parameter = True + self.model_scaling_policy.mark_embedding_class_parameters([self.output_layer.weight]) # If share_embeddings_and_output_weights is True, we need to maintain duplicated # embedding weights in post processing stage. If use Multi-Token Prediction (MTP), @@ -264,9 +300,14 @@ def setup_embeddings_and_output_layer(self) -> None: weight.data.fill_(0) weight.shared = True weight.shared_embedding = True + weight.is_embedding_or_output_parameter = True + weight.is_output_parameter = True + set_parameterization_metadata( + weight, role=ROLE_SHARED_EMBEDDING_OUTPUT, shared_group='lm_embedding_output' + ) # Keep optimizer grouping consistent for tied embedding/output copies. - if self.config.use_mup: - weight.is_embedding_parameter = True + if self.model_scaling_policy.enabled: + self.model_scaling_policy.mark_embedding_class_parameters([weight]) # Parameters are shared between the word embeddings layers, and the # heads at the end of the model. In a pipelined setup with more than @@ -312,11 +353,7 @@ def _scale_logits(self, logits: Tensor) -> Tensor: Tensor: Scaled logits if MuP is enabled and mup_output_mult != 1.0, otherwise unchanged logits. """ - if not self.config.use_mup: - return logits - if self.config.mup_output_mult != 1.0: - return logits * self.config.mup_output_mult - return logits + return build_model_scaling_policy(self.config).scale_output_logits(logits) def shared_embedding_or_output_weight(self) -> Tensor: """Gets the embedding weight or output logit weights when share embedding and output weights set to True diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 9f8d9da4a10..48d58cf2fe5 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -252,10 +252,10 @@ def __init__( config.hidden_size, self.vocab_size, config=config, - init_method=( - config.embedding_init_method - if config.use_mup and not self.share_embeddings_and_output_weights - else config.init_method + init_method=self.model_scaling_policy.output_layer_init_method( + share_embeddings_and_output_weights=self.share_embeddings_and_output_weights, + default_init_method=config.init_method, + embedding_init_method=config.embedding_init_method, ), bias=False, skip_bias_add=False, @@ -676,7 +676,7 @@ def _postprocess( config=self.config, cp_group=self.pg_collection.cp, packed_seq_params=packed_seq_params, - scale_logits_fn=self._scale_logits if self.config.use_mup else None, + scale_logits_fn=(self._scale_logits if self.config.use_mup else None), ) sequence_parallel_override = False diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 511b24673b0..f0df0f60a6e 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -296,10 +296,10 @@ def __init__( config.hidden_size, self.vocab_size, config=config, - init_method=( - config.embedding_init_method - if config.use_mup and not self.share_embeddings_and_output_weights - else config.init_method + init_method=self.model_scaling_policy.output_layer_init_method( + share_embeddings_and_output_weights=self.share_embeddings_and_output_weights, + default_init_method=config.init_method, + embedding_init_method=config.embedding_init_method, ), bias=False, skip_bias_add=False, @@ -543,7 +543,7 @@ def forward( config=self.config, cp_group=self.pg_collection.cp, packed_seq_params=packed_seq_params, - scale_logits_fn=self._scale_logits if self.config.use_mup else None, + scale_logits_fn=(self._scale_logits if self.config.use_mup else None), ) sequence_parallel_override = False if ( diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 1598f6ed95c..b248610dd92 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -54,6 +54,17 @@ combine_param_group_overrides, param_group_override_to_tuple, ) +from megatron.core.parameterization import ( + TrainingScalingPolicy, + build_legacy_mup_training_policy, + is_embedding_class_parameter, + is_embedding_or_output_parameter, + is_hidden_matrix_parameter, + is_hidden_vector_parameter, + is_muon_managed_matrix_parameter, + is_vector_like_parameter, + should_skip_depth_mup_vector_weight_decay, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.fsdp_dtensor_checkpoint import get_global_unique_param_name @@ -73,7 +84,6 @@ Float16OptimizerWithFloat16Params, FP32Optimizer, MegatronOptimizer, - param_group_identifier_keys, ) # Subclass aliases kept for backward compatibility; all are OptimizerConfig. @@ -89,33 +99,43 @@ logger = logging.getLogger(__name__) -def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, ParamGroupOverride]: - """Get standard config overrides for the optimizer, handling decoupled LR and common wd skips. - - Args: - config (OptimizerConfig): optimizer configuration object. - - Returns: - Dict[ParamKey, ParamGroupOverride]: standard config overrides. - """ +def get_standard_config_overrides( + config: OptimizerConfig, scaling_policy: Optional[TrainingScalingPolicy] = None +) -> Dict[ParamKey, ParamGroupOverride]: + """Get standard config overrides for optimizer LR and weight-decay skips.""" config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = {} - # First, figure out how we are going to do wd skipping. The two main approaches are: - # 1. The classic megatron approach of skipping all len 1 and bias parameters. - # 2. The Qwen3-Next approach of doing 1, other than qk layernorm parameters. - if config.apply_wd_to_qk_layernorm: - 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), + + use_depth_mup_adamw_table = bool( + scaling_policy and scaling_policy.context.is_depth_mup and scaling_policy.is_adam_optimizer + ) + if use_depth_mup_adamw_table: + depth_mup_vector_wd_skip = ParamWithNamePredicate( + name="depth_mup_norm_and_unknown_vector_wd_skip", + fn=lambda param, name: should_skip_depth_mup_vector_weight_decay( + param, name, apply_wd_to_qk_layernorm=config.apply_wd_to_qk_layernorm + ), ) - param_wd_mult_key = ParamKey(with_name_predicate=shape_1_not_qkln_param) - else: - param_length_1_match = ParamPredicate( - name="param_len_1", fn=lambda param: len(param.shape) == 1 + config_overrides[ParamKey(with_name_predicate=depth_mup_vector_wd_skip)] = ( + ParamGroupOverride(wd_mult=0.0) ) - param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match) + else: + # First, figure out how we are going to do wd skipping. The two main approaches are: + # 1. The classic megatron approach of skipping all len 1 and bias parameters. + # 2. The Qwen3-Next approach of doing 1, other than qk layernorm parameters. + if config.apply_wd_to_qk_layernorm: + 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), + ) + param_wd_mult_key = ParamKey(with_name_predicate=shape_1_not_qkln_param) + else: + param_length_1_match = ParamPredicate( + name="param_len_1", fn=lambda param: len(param.shape) == 1 + ) + param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match) - config_overrides[param_wd_mult_key] = ParamGroupOverride(wd_mult=0.0) + config_overrides[param_wd_mult_key] = ParamGroupOverride(wd_mult=0.0) if config.decoupled_lr is not None: decoupled_lr_config: ParamGroupOverride = {"max_lr": config.decoupled_lr} @@ -126,58 +146,51 @@ def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, Par return config_overrides - def get_mup_config_overrides( config: OptimizerConfig, mup_width_mult: float, optimizer_type: str = 'adam' ) -> Dict[ParamKey, ParamGroupOverride]: - """Get MuP config overrides for per-layer LR and Adam epsilon scaling. - - In MuP, optimizer learning rates are adjusted by parameter class to ensure - stable update scales across model widths and enable hyperparameter transfer. - - MuP optimizer scaling rules (as implemented here): - - Adam/AdamW: - - hidden (matrix-like) lr = base_lr / width_mult - - hidden (matrix-like) eps = base_eps / width_mult - - vector-like params keep base lr and eps - - SGD: - - vector-like lr = base_lr * width_mult - - hidden (matrix-like) lr keeps base_lr in the current uniform-width setup - - no eps override is applied - - Non-Adam optimizers: - - hidden (matrix-like) lr = base_lr / width_mult - - no eps override is applied. - - for Muon optimizers, matrix-like params managed by Muon itself are - excluded from these Adam-style MuP overrides. - - With decoupled_lr enabled, embedding/output params continue using decoupled LR - and MuP will not override those explicit decoupled values. + """Compatibility wrapper for the existing MuP optimizer override surface.""" + scaling_policy = build_legacy_mup_training_policy( + mup_width_mult=mup_width_mult, optimizer_type=optimizer_type + ) + return get_scaling_config_overrides(config=config, scaling_policy=scaling_policy) - Args: - config (OptimizerConfig): optimizer configuration object. - mup_width_mult (float): Width multiplier (hidden_size / base_hidden_size). - optimizer_type (str): Optimizer type string from config.optimizer. - Returns: - Dict[ParamKey, ParamGroupOverride]: MuP optimizer overrides. - """ - optimizer_type_lower = optimizer_type.lower() - is_sgd_optimizer = optimizer_type_lower == 'sgd' - is_adam_optimizer = 'adam' in optimizer_type_lower - is_muon_optimizer = 'muon' in optimizer_type_lower +def get_scaling_config_overrides( + config: OptimizerConfig, scaling_policy: TrainingScalingPolicy +) -> Dict[ParamKey, ParamGroupOverride]: + """Get scaling-policy overrides for per-parameter optimizer settings.""" + if not scaling_policy.enabled: + return {} + + if ( + scaling_policy.context.is_depth_mup + and scaling_policy.is_adam_optimizer + and config.weight_decay != 0.0 + and not config.decoupled_weight_decay + ): + raise ValueError( + "scaling_recipe='depth_mup' with nonzero weight_decay requires " + "decoupled_weight_decay=True because the width-depth weight-decay scaling " + "is derived for AdamW. Use weight_decay=0.0 for coupled Adam, or enable " + "decoupled_weight_decay." + ) decoupled_lr_enabled = config.decoupled_lr is not None if decoupled_lr_enabled: message = ( - "Both decoupled_lr and MuP LR scaling are enabled. decoupled_lr sets an " - "absolute LR for embedding+output params, and MuP LR scaling will not " - "override those parameters." + "Both decoupled_lr and scaling-recipe LR scaling are enabled. decoupled_lr " + "sets an absolute LR for embedding+output params, and the active scaling " + "recipe will not override those parameters." ) - if is_adam_optimizer: - message += " MuP Adam epsilon scaling remains applied to hidden matrix-like parameters." + if scaling_policy.is_adam_optimizer: + message += ( + " Adam epsilon scaling remains applied according to the active recipe's " + "parameter-class rules." + ) log_single_rank(logger, logging.WARNING, message) - if is_muon_optimizer: + if scaling_policy.is_muon_optimizer: muon_scale_mode = getattr(config, 'muon_scale_mode', 'spectral') if muon_scale_mode == 'spectral': log_single_rank( @@ -189,112 +202,167 @@ def get_mup_config_overrides( "Muon-managed matrices with MuP.", ) - if mup_width_mult == 1.0: - # No scaling needed when width_mult is 1 - return {} - - hidden_lr_mult = 1.0 / mup_width_mult base_lr = config.lr base_min_lr = config.min_lr - - # Hidden matrix-like layers get scaled LR/eps; vector-like params keep base values. - # Prefer the explicit parameter attribute set by LanguageModule. Fall back to - # a conservative name check for older or non-language modules. - def is_embedding_parameter(param: torch.nn.Parameter, param_name: str) -> bool: - if getattr(param, 'shared_embedding', False): - return True - if hasattr(param, 'is_embedding_parameter'): - return bool(param.is_embedding_parameter) - return 'embedding' in param_name.lower() - - def is_vector_like_parameter(param: torch.nn.Parameter, param_name: str) -> bool: - if is_embedding_parameter(param, param_name): - return True - if param.dim() <= 1: - return True - return False - - def is_muon_managed_matrix_parameter(param: torch.nn.Parameter, _: str) -> bool: - if not is_muon_optimizer: + hidden_lr_mult = scaling_policy.hidden_lr_multiplier + hidden_vector_lr_mult = scaling_policy.hidden_vector_lr_multiplier + hidden_eps_mult = scaling_policy.hidden_eps_multiplier + hidden_vector_eps_mult = scaling_policy.hidden_vector_eps_multiplier + embedding_class_eps_mult = scaling_policy.embedding_class_eps_multiplier + hidden_matrix_wd_mult = scaling_policy.hidden_matrix_wd_multiplier + hidden_vector_wd_mult = scaling_policy.hidden_vector_wd_multiplier + embedding_class_wd_mult = scaling_policy.embedding_class_wd_multiplier + + def should_scale_hidden_matrix(param: torch.nn.Parameter, param_name: str) -> bool: + if decoupled_lr_enabled and is_embedding_or_output_parameter(param): return False - return is_managed_by_layer_wise_optimizer(param) - - def should_scale_lr_with_mup(param: torch.nn.Parameter, param_name: str) -> bool: - if decoupled_lr_enabled and getattr(param, 'is_embedding_or_output_parameter', False): + if is_muon_managed_matrix_parameter(param, optimizer_type=scaling_policy.optimizer_type): return False - if is_muon_managed_matrix_parameter(param, param_name): - return False - return not is_vector_like_parameter(param, param_name) + return is_hidden_matrix_parameter(param, param_name) def should_scale_vector_like_lr_with_mup(param: torch.nn.Parameter, param_name: str) -> bool: - if decoupled_lr_enabled and getattr(param, 'is_embedding_or_output_parameter', False): + if decoupled_lr_enabled and is_embedding_or_output_parameter(param): return False return is_vector_like_parameter(param, param_name) - def should_scale_eps_with_mup(param: torch.nn.Parameter, param_name: str) -> bool: - if is_vector_like_parameter(param, param_name): + def should_scale_hidden_vector(param: torch.nn.Parameter, param_name: str) -> bool: + if decoupled_lr_enabled and is_embedding_or_output_parameter(param): return False - if is_muon_managed_matrix_parameter(param, param_name): + return is_hidden_vector_parameter(param, param_name) + + def should_scale_hidden_matrix_eps(param: torch.nn.Parameter, param_name: str) -> bool: + if is_muon_managed_matrix_parameter(param, optimizer_type=scaling_policy.optimizer_type): return False - # MuP Appendix B.3: eps scales with fan_in when non-negligible. - # This implementation follows the common denominator form: sqrt(v) + eps. - return True + return is_hidden_matrix_parameter(param, param_name) + + def should_scale_hidden_vector_eps(param: torch.nn.Parameter, param_name: str) -> bool: + return is_hidden_vector_parameter(param, param_name) - mup_overrides: Dict[ParamKey, ParamGroupOverride] = {} + def should_scale_embedding_class_eps(param: torch.nn.Parameter, param_name: str) -> bool: + return is_embedding_class_parameter(param, param_name) + + scaling_overrides: Dict[ParamKey, ParamGroupOverride] = {} + + if scaling_policy.is_sgd_optimizer: + hidden_lr_override: ParamGroupOverride = {} + if base_lr is not None and hidden_lr_mult != 1.0: + hidden_lr_override["max_lr"] = base_lr * hidden_lr_mult + if base_min_lr is not None and hidden_lr_mult != 1.0: + hidden_lr_override["min_lr"] = base_min_lr * hidden_lr_mult + if hidden_lr_override: + hidden_predicate = ParamWithNamePredicate( + name="scaling_hidden_only_excluding_embedding_output", fn=should_scale_hidden_matrix + ) + scaling_overrides[ParamKey(with_name_predicate=hidden_predicate)] = hidden_lr_override - if is_sgd_optimizer: - vector_like_lr_mult = mup_width_mult vector_like_lr_override: ParamGroupOverride = {} - if base_lr is not None: - vector_like_lr_override["max_lr"] = base_lr * vector_like_lr_mult - if base_min_lr is not None: - vector_like_lr_override["min_lr"] = base_min_lr * vector_like_lr_mult + if base_lr is not None and hidden_vector_lr_mult != 1.0: + vector_like_lr_override["max_lr"] = base_lr * hidden_vector_lr_mult + if base_min_lr is not None and hidden_vector_lr_mult != 1.0: + vector_like_lr_override["min_lr"] = base_min_lr * hidden_vector_lr_mult if vector_like_lr_override: vector_like_predicate = ParamWithNamePredicate( name="mup_sgd_vector_like_excluding_embedding_output", fn=should_scale_vector_like_lr_with_mup, ) - mup_overrides[ParamKey(with_name_predicate=vector_like_predicate)] = ( + scaling_overrides[ParamKey(with_name_predicate=vector_like_predicate)] = ( vector_like_lr_override ) - return mup_overrides + return scaling_overrides + + if scaling_policy.context.is_depth_mup and scaling_policy.is_adam_optimizer: + hidden_matrix_override: ParamGroupOverride = {} + if base_lr is not None and hidden_lr_mult != 1.0: + hidden_matrix_override["max_lr"] = base_lr * hidden_lr_mult + if base_min_lr is not None and hidden_lr_mult != 1.0: + hidden_matrix_override["min_lr"] = base_min_lr * hidden_lr_mult + if config.adam_eps is not None and hidden_eps_mult != 1.0: + hidden_matrix_override["eps"] = config.adam_eps * hidden_eps_mult + if hidden_matrix_wd_mult != 1.0: + hidden_matrix_override["wd_mult"] = hidden_matrix_wd_mult + if hidden_matrix_override: + hidden_matrix_predicate = ParamWithNamePredicate( + name="depth_mup_hidden_matrix_adamw", fn=should_scale_hidden_matrix_eps + ) + scaling_overrides[ParamKey(with_name_predicate=hidden_matrix_predicate)] = ( + hidden_matrix_override + ) + + hidden_vector_override: ParamGroupOverride = {} + if config.adam_eps is not None and hidden_vector_eps_mult != 1.0: + hidden_vector_override["eps"] = config.adam_eps * hidden_vector_eps_mult + if hidden_vector_wd_mult != 1.0: + hidden_vector_override["wd_mult"] = hidden_vector_wd_mult + if hidden_vector_override: + hidden_vector_predicate = ParamWithNamePredicate( + name="depth_mup_hidden_vector_adamw", fn=should_scale_hidden_vector_eps + ) + scaling_overrides[ParamKey(with_name_predicate=hidden_vector_predicate)] = ( + hidden_vector_override + ) + + embedding_class_override: ParamGroupOverride = {} + if config.adam_eps is not None and embedding_class_eps_mult != 1.0: + embedding_class_override["eps"] = config.adam_eps * embedding_class_eps_mult + if embedding_class_wd_mult != 1.0: + embedding_class_override["wd_mult"] = embedding_class_wd_mult + if embedding_class_override: + embedding_class_predicate = ParamWithNamePredicate( + name="depth_mup_embedding_output_adamw", fn=should_scale_embedding_class_eps + ) + scaling_overrides[ParamKey(with_name_predicate=embedding_class_predicate)] = ( + embedding_class_override + ) + + return scaling_overrides lr_override: ParamGroupOverride = {} - if base_lr is not None: + if base_lr is not None and hidden_lr_mult != 1.0: lr_override["max_lr"] = base_lr * hidden_lr_mult - if base_min_lr is not None: + if base_min_lr is not None and hidden_lr_mult != 1.0: lr_override["min_lr"] = base_min_lr * hidden_lr_mult eps_override: ParamGroupOverride = {} - if is_adam_optimizer and config.adam_eps is not None: - eps_override["eps"] = config.adam_eps * hidden_lr_mult + if scaling_policy.is_adam_optimizer and config.adam_eps is not None and hidden_eps_mult != 1.0: + eps_override["eps"] = config.adam_eps * hidden_eps_mult if decoupled_lr_enabled: if lr_override: hidden_predicate = ParamWithNamePredicate( - name="mup_hidden_only_excluding_embedding_output", fn=should_scale_lr_with_mup + name="mup_hidden_only_excluding_embedding_output", fn=should_scale_hidden_matrix ) - mup_overrides[ParamKey(with_name_predicate=hidden_predicate)] = lr_override + scaling_overrides[ParamKey(with_name_predicate=hidden_predicate)] = lr_override if eps_override: hidden_output_predicate = ParamWithNamePredicate( - name="mup_hidden_only_for_adam_eps", fn=should_scale_eps_with_mup + name="mup_hidden_only_for_adam_eps", fn=should_scale_hidden_matrix_eps ) - mup_overrides[ParamKey(with_name_predicate=hidden_output_predicate)] = eps_override + scaling_overrides[ParamKey(with_name_predicate=hidden_output_predicate)] = eps_override else: - combined_override: ParamGroupOverride = {} - combined_override.update(lr_override) - combined_override.update(eps_override) - if combined_override: + if lr_override and eps_override: + combined_override: ParamGroupOverride = {} + combined_override.update(lr_override) + combined_override.update(eps_override) hidden_output_predicate = ParamWithNamePredicate( - name="mup_hidden_and_output", fn=should_scale_eps_with_mup + name="mup_hidden_and_output", fn=should_scale_hidden_matrix_eps ) - mup_overrides[ParamKey(with_name_predicate=hidden_output_predicate)] = combined_override - - return mup_overrides + scaling_overrides[ParamKey(with_name_predicate=hidden_output_predicate)] = ( + combined_override + ) + elif lr_override: + hidden_predicate = ParamWithNamePredicate( + name="scaling_hidden_and_output_lr", fn=should_scale_hidden_matrix + ) + scaling_overrides[ParamKey(with_name_predicate=hidden_predicate)] = lr_override + elif eps_override: + hidden_output_predicate = ParamWithNamePredicate( + name="mup_hidden_and_output_eps", fn=should_scale_hidden_matrix_eps + ) + scaling_overrides[ParamKey(with_name_predicate=hidden_output_predicate)] = eps_override + return scaling_overrides def _get_param_groups( model_chunks: List[MegatronModule], @@ -784,8 +852,12 @@ def _get_megatron_emerging_optimizer( if 'linear_qkv.weight' in name and len(param.shape) == 2: param.is_qkv = True - # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam). - config_overrides.update(_EMERGING_OPTIMIZERS[eopt_name].default_param_overrides) + # Apply optimizer-specific param overrides (e.g. muon: non-linear -> scalar optimizer). + entry = _EMERGING_OPTIMIZERS[eopt_name] + if entry.config_to_param_overrides is not None: + config_overrides.update(entry.config_to_param_overrides(config)) + else: + config_overrides.update(entry.default_param_overrides) # Build param groups and bucket by (optimizer_name, is_expert_parallel). # Layer-wise distributed optimizer handles expert params internally so we skip that split. @@ -859,7 +931,7 @@ def _get_megatron_emerging_optimizer( if opt_name in _EMERGING_OPTIMIZERS: optimizer, init_state_fn = _create_emerging_optimizer( - config, groups, eopt_name, model_chunks, pg_collection + config, groups, opt_name, model_chunks, pg_collection ) if use_layer_wise: layer_wise_base_results.append((optimizer, init_state_fn)) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index b388161a610..622603d0151 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -57,7 +57,11 @@ from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict from ..transformer.module import MegatronModule from .grad_scaler import MegatronGradScaler -from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper, param_group_identifier_keys +from .optimizer import ( + MixedPrecisionOptimizer, + _zero_grad_group_helper, + get_param_group_identifier_tuple, +) from .optimizer_config import OptimizerConfig from .param_layout import FullParamLayout, PerBufferParamLayout, pad_bucket_end, pad_param_start @@ -888,21 +892,7 @@ def load_state_dict(self, state_dict): # the ordering of parameters within its flattened parameter state # list. def make_needed_groups(param_group): - needed_groups = [] - for key in param_group_identifier_keys: - # NeMo changes these variable names from `lr_mult` and `wd_mult` - # to `pre_lr_mult` and `pre_wd_mult`, so we need to check both. - if key in param_group: - pass - elif f"pre_{key}" in param_group: - key = f"pre_{key}" - else: - raise ValueError( - f"Key {key} (or pre_{key}) not found in param_group {param_group}." - ) - needed_groups.append(param_group[key]) - needed_groups = tuple(needed_groups) - return needed_groups + return get_param_group_identifier_tuple(param_group) param_groups_map = {} for param_group in state_dict["optimizer"]["param_groups"]: diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index cc218d6ba40..6212f2b4624 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -80,13 +80,23 @@ def _eopt_init_state_fn(opt, config=None): def _default_param_overrides_factory() -> Dict[ParamKey, Dict[str, Any]]: """Default param overrides: route non-linear/embedding params to Adam.""" + return _nonlinear_or_embedding_param_overrides('adam') + + +def _nonlinear_or_embedding_param_overrides(optimizer_name: str) -> Dict[ParamKey, Dict[str, Any]]: + """Route non-matrix and embedding-class parameters to the requested scalar optimizer.""" return { ParamKey( predicate=ParamPredicate(name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding) - ): {'optimizer': 'adam'} + ): {'optimizer': optimizer_name} } +def _muon_default_param_overrides(config) -> Dict[ParamKey, Dict[str, Any]]: + """Respect the configured scalar optimizer for Muon-family nonlinear parameters.""" + return _nonlinear_or_embedding_param_overrides(config.muon_scalar_optimizer) + + @dataclass class EmergingOptimizerEntry: """Everything needed to create and configure an emerging optimizer. @@ -95,6 +105,8 @@ class EmergingOptimizerEntry: optimizer_cls: The torch optimizer class. init_state_fn: Lazily initialises optimizer state (needed for checkpoint formats). config_to_kwargs: ``(config, model_chunks, pg_collection) -> dict`` of constructor kwargs. + config_to_param_overrides: ``config -> dict`` of per-parameter overrides derived from the + resolved optimizer config. default_param_overrides: Per-parameter config overrides applied automatically (e.g. route non-linear params to Adam). """ @@ -102,6 +114,7 @@ class EmergingOptimizerEntry: optimizer_cls: type init_state_fn: Callable = _eopt_init_state_fn config_to_kwargs: Callable | None = None + config_to_param_overrides: Callable | None = None default_param_overrides: Dict[ParamKey, Dict[str, Any]] = field( default_factory=_default_param_overrides_factory ) @@ -402,10 +415,17 @@ def _default_adam_based_eopt_config_to_kwargs( ) -> Dict[str, Any]: """Convert OptimizerConfig to default emerging optimizer constructor kwargs.""" kwargs = _kwargs_from_config(registry.get_optimizer_cls(eopt_name), eopt_name, config) - kwargs["betas"] = (config.adam_beta1, config.adam_beta2) + kwargs["betas"] = _default_betas_for_eopt(eopt_name, config) return kwargs +def _default_betas_for_eopt(eopt_name, config) -> tuple[float, float]: + """Return the default beta pair for an emerging optimizer.""" + if eopt_name == "lion": + return (config.lion_beta1, config.lion_beta2) + return (config.adam_beta1, config.adam_beta2) + + # ----------------------------------------------------------------------- # Register emerging optimizers # ----------------------------------------------------------------------- @@ -415,25 +435,13 @@ def _default_adam_based_eopt_config_to_kwargs( optimizer_cls=TensorParallelMuon, init_state_fn=_eopt_init_state_fn, config_to_kwargs=_muon_config_to_kwargs, - default_param_overrides={ - ParamKey( - predicate=ParamPredicate( - name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding - ) - ): {'optimizer': 'adam'} - }, + config_to_param_overrides=_muon_default_param_overrides, ), "adaptive_muon": EmergingOptimizerEntry( optimizer_cls=TensorParallelAdaptiveMuon, init_state_fn=_eopt_init_state_fn, config_to_kwargs=_adaptive_muon_config_to_kwargs, - default_param_overrides={ - ParamKey( - predicate=ParamPredicate( - name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding - ) - ): {'optimizer': 'adam'} - }, + config_to_param_overrides=_muon_default_param_overrides, ), } ) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index ddc3dd8620e..e612ba8cf20 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -94,7 +94,40 @@ def _multi_tensor_copy_this_to_that( that_.copy_(this_) -param_group_identifier_keys = ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr') +param_group_identifier_keys = ( + 'wd_mult', + 'lr_mult', + 'is_expert_parallel', + 'is_decoupled_lr', + 'eps', + 'optimizer', +) +param_group_identifier_defaults = { + 'wd_mult': 1.0, + 'lr_mult': 1.0, + 'is_expert_parallel': False, + 'is_decoupled_lr': False, + 'eps': None, + 'optimizer': None, +} + + +def get_param_group_identifier_tuple(param_group: Dict) -> tuple: + """Return a stable identifier for optimizer param-group matching and resume.""" + values = [] + for key in param_group_identifier_keys: + if key in param_group: + values.append(param_group[key]) + elif f"pre_{key}" in param_group: + values.append(param_group[f"pre_{key}"]) + else: + values.append(param_group_identifier_defaults[key]) + return tuple(values) + + +def get_param_group_identifier_sort_key(param_group: Dict) -> tuple: + """Return a None-safe ordering key for optimizer param-group identifiers.""" + return tuple((value is not None, value) for value in get_param_group_identifier_tuple(param_group)) class MegatronOptimizer(ABC): @@ -426,22 +459,13 @@ def _filter_and_reorder_param_groups( ValueError: If parameter groups in state dict don't match current optimizer. """ # Define groups order that is needed in the current optimizer (coming from runtime) - needed_groups = [ - # NeMo may have different key for required fields, e.g., "wd_mult" to "pre_wd_mult" - tuple(g[key] if key in g else g[f"pre_{key}"] for key in param_group_identifier_keys) - for g in current_groups - ] + needed_groups = [get_param_group_identifier_tuple(g) for g in current_groups] # Keep state_dict param group order since groups are LocalNonpersistentObject # and their order is determined at runtime, not from the checkpoint. params_in_state_dict_order = [g['params'] for g in state_dict_groups] loaded_groups_map = { - tuple( - # NeMo may have different key for required fields, e.g., "wd_mult" to "pre_wd_mult" - group[key] if key in group else group[f"pre_{key}"] - for key in param_group_identifier_keys - ): group - for group in state_dict_groups + get_param_group_identifier_tuple(group): group for group in state_dict_groups } final_groups = [] diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py new file mode 100644 index 00000000000..b5e401ff0eb --- /dev/null +++ b/megatron/core/parameterization/__init__.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from .eval_runtime import allow_scaling_policy_eval, is_scaling_policy_eval_allowed +from .model_policy import ModelScalingPolicy, build_model_scaling_policy +from .roles import ( + IS_OUTPUT_PARAMETER_ATTR, + PARAMETERIZATION_ROLE_ATTR, + PARAMETERIZATION_SHARED_GROUP_ATTR, + PARAMETERIZATION_TAGS_ATTR, + ROLE_BLOCK_OUT_PROJ, + ROLE_EMBEDDING, + ROLE_HIDDEN_BIAS, + ROLE_HIDDEN_MATRIX, + ROLE_HIDDEN_VECTOR, + ROLE_HIDDEN_VECTOR_OTHER, + ROLE_MUON_MANAGED_MATRIX, + ROLE_NORM_BIAS, + ROLE_NORM_SCALE, + ROLE_OUTPUT, + ROLE_QK_NORM_SCALE, + ROLE_SHARED_EMBEDDING_OUTPUT, + ROLE_VECTOR_LIKE, + get_parameterization_role, + is_embedding_class_parameter, + is_embedding_or_output_parameter, + is_hidden_bias_parameter, + is_hidden_matrix_parameter, + is_hidden_vector_parameter, + is_muon_managed_matrix_parameter, + is_norm_parameter, + is_output_parameter, + is_qk_norm_parameter, + is_vector_like_parameter, + set_parameterization_metadata, + should_skip_depth_mup_vector_weight_decay, +) +from .spec import ( + SCALING_RECIPE_DEPTH_MUP, + SCALING_RECIPE_MUP, + SCALING_RECIPE_NONE, + SCALING_RECIPE_VALUES, + ScalingContext, + ScalingUserConfig, + build_scaling_context, + build_scaling_user_config, + sync_legacy_mup_fields, +) +from .training_policy import ( + TrainingScalingPolicy, + build_legacy_mup_training_policy, + build_training_scaling_policy, +) + +__all__ = [ + 'IS_OUTPUT_PARAMETER_ATTR', + 'PARAMETERIZATION_ROLE_ATTR', + 'PARAMETERIZATION_SHARED_GROUP_ATTR', + 'PARAMETERIZATION_TAGS_ATTR', + 'ROLE_BLOCK_OUT_PROJ', + 'ROLE_EMBEDDING', + 'ROLE_HIDDEN_BIAS', + 'ROLE_HIDDEN_MATRIX', + 'ROLE_HIDDEN_VECTOR', + 'ROLE_HIDDEN_VECTOR_OTHER', + 'ROLE_MUON_MANAGED_MATRIX', + 'ROLE_NORM_BIAS', + 'ROLE_NORM_SCALE', + 'ROLE_OUTPUT', + 'ROLE_QK_NORM_SCALE', + 'ROLE_SHARED_EMBEDDING_OUTPUT', + 'ModelScalingPolicy', + 'ROLE_VECTOR_LIKE', + 'SCALING_RECIPE_DEPTH_MUP', + 'SCALING_RECIPE_MUP', + 'SCALING_RECIPE_NONE', + 'SCALING_RECIPE_VALUES', + 'ScalingContext', + 'ScalingUserConfig', + 'TrainingScalingPolicy', + 'allow_scaling_policy_eval', + 'build_legacy_mup_training_policy', + 'build_model_scaling_policy', + 'build_scaling_context', + 'build_scaling_user_config', + 'build_training_scaling_policy', + 'get_parameterization_role', + 'is_scaling_policy_eval_allowed', + 'is_embedding_class_parameter', + 'is_embedding_or_output_parameter', + 'is_hidden_bias_parameter', + 'is_hidden_matrix_parameter', + 'is_hidden_vector_parameter', + 'is_muon_managed_matrix_parameter', + 'is_norm_parameter', + 'is_output_parameter', + 'is_qk_norm_parameter', + 'is_vector_like_parameter', + 'set_parameterization_metadata', + 'should_skip_depth_mup_vector_weight_decay', + 'sync_legacy_mup_fields', +] diff --git a/megatron/core/parameterization/eval_runtime.py b/megatron/core/parameterization/eval_runtime.py new file mode 100644 index 00000000000..67abec81da0 --- /dev/null +++ b/megatron/core/parameterization/eval_runtime.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Iterator + +_SCALING_POLICY_EVAL_DEPTH: ContextVar[int] = ContextVar( + 'scaling_policy_eval_depth', default=0 +) + + +def is_scaling_policy_eval_allowed() -> bool: + return _SCALING_POLICY_EVAL_DEPTH.get() > 0 + + +@contextmanager +def allow_scaling_policy_eval(enabled: bool) -> Iterator[None]: + if not enabled: + yield + return + + token = _SCALING_POLICY_EVAL_DEPTH.set(_SCALING_POLICY_EVAL_DEPTH.get() + 1) + try: + yield + finally: + _SCALING_POLICY_EVAL_DEPTH.reset(token) diff --git a/megatron/core/parameterization/model_policy.py b/megatron/core/parameterization/model_policy.py new file mode 100644 index 00000000000..a57e0f2e280 --- /dev/null +++ b/megatron/core/parameterization/model_policy.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import functools +import math +from dataclasses import dataclass +from typing import Iterable, Optional + +import torch +from torch import Tensor + +from megatron.core.utils import ( + init_method_normal, + mup_scaled_init_method_normal, + scaled_init_method_normal, +) + +from .spec import ScalingContext, build_scaling_context + + +@dataclass(frozen=True) +class ModelScalingPolicy: + """Model-side policy for Megatron scaling recipes.""" + + context: ScalingContext + + @property + def enabled(self) -> bool: + return self.context.enabled + + @property + def uses_width_mup(self) -> bool: + return self.context.uses_width_mup + + @property + def residual_branch_multiplier(self) -> float: + return self.context.depth_mult**self.context.residual_branch_depth_power + + @property + def dense_block_out_proj_init_multiplier(self) -> float: + return self.context.depth_mult**self.context.block_out_proj_init_depth_power + + def resolve_attention_softmax_scale( + self, *, softmax_scale: Optional[float], kv_channels: int + ) -> Optional[float]: + if softmax_scale is not None or not self.uses_width_mup: + return softmax_scale + base_head_scale = 1.0 if self.context.base_head_dim is None else self.context.base_head_dim**0.5 + return base_head_scale / (kv_channels**self.context.attention_scale_power) + + def build_hidden_init_method(self, *, init_method_std: float): + if not self.uses_width_mup: + return init_method_normal(init_method_std) + return init_method_normal(init_method_std / math.sqrt(self.context.width_mult)) + + def build_default_output_layer_init_method( + self, *, init_method_std: float, num_layers: int, is_hybrid_model: bool + ): + multiplier = 2.0 if not is_hybrid_model else 1.0 + if self.uses_width_mup: + return mup_scaled_init_method_normal( + init_method_std, num_layers, self.context.width_mult, multiplier=multiplier + ) + return scaled_init_method_normal(init_method_std, num_layers, multiplier=multiplier) + + def dense_block_output_init_method( + self, + *, + default_init_method, + init_method_std: float, + num_layers: int, + is_hybrid_model: bool, + output_layer_init_method_is_user_provided: bool, + apply_depth_hook: bool = True, + ): + if output_layer_init_method_is_user_provided: + return default_init_method + if not apply_depth_hook or self.dense_block_out_proj_init_multiplier == 1.0: + return default_init_method + + multiplier = 2.0 if not is_hybrid_model else 1.0 + std = init_method_std / math.sqrt(multiplier * num_layers) + if self.uses_width_mup: + std = std / math.sqrt(self.context.width_mult) + std = std * self.dense_block_out_proj_init_multiplier + return functools.partial(torch.nn.init.normal_, mean=0.0, std=std) + + def output_layer_init_method( + self, + *, + share_embeddings_and_output_weights: bool, + default_init_method, + embedding_init_method, + ): + if self.uses_width_mup and not share_embeddings_and_output_weights: + return embedding_init_method + return default_init_method + + def mark_embedding_class_parameters(self, parameters: Iterable[torch.nn.Parameter]) -> None: + if not self.uses_width_mup: + return + for param in parameters: + param.is_embedding_parameter = True + + def scale_embedding_activations(self, embeddings: Tensor) -> Tensor: + if not self.uses_width_mup or self.context.embedding_mult == 1.0: + return embeddings + return embeddings * self.context.embedding_mult + + def scale_output_logits(self, logits: Tensor) -> Tensor: + if not self.uses_width_mup or self.context.output_mult == 1.0: + return logits + return logits * self.context.output_mult + + def scale_residual_branch_output( + self, output_with_bias: tuple[Tensor, Tensor | None] + ) -> tuple[Tensor, Tensor | None]: + if self.residual_branch_multiplier == 1.0: + return output_with_bias + + output, bias = output_with_bias + scaled_output = output * self.residual_branch_multiplier + scaled_bias = None if bias is None else bias * self.residual_branch_multiplier + return scaled_output, scaled_bias + + +def build_model_scaling_policy(config) -> ModelScalingPolicy: + return ModelScalingPolicy(build_scaling_context(config)) diff --git a/megatron/core/parameterization/roles.py b/megatron/core/parameterization/roles.py new file mode 100644 index 00000000000..28b168c545f --- /dev/null +++ b/megatron/core/parameterization/roles.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +PARAMETERIZATION_ROLE_ATTR = 'parameterization_role' +PARAMETERIZATION_SHARED_GROUP_ATTR = 'parameterization_shared_group' +PARAMETERIZATION_TAGS_ATTR = 'parameterization_tags' +IS_OUTPUT_PARAMETER_ATTR = 'is_output_parameter' + +ROLE_EMBEDDING = 'embedding' +ROLE_OUTPUT = 'output' +ROLE_SHARED_EMBEDDING_OUTPUT = 'shared_embedding_output' +ROLE_BLOCK_OUT_PROJ = 'block_out_proj' +ROLE_HIDDEN_MATRIX = 'hidden_matrix' +ROLE_HIDDEN_VECTOR = 'hidden_vector' +ROLE_HIDDEN_BIAS = 'hidden_bias' +ROLE_NORM_SCALE = 'norm_scale' +ROLE_NORM_BIAS = 'norm_bias' +ROLE_QK_NORM_SCALE = 'qk_norm_scale' +ROLE_HIDDEN_VECTOR_OTHER = 'hidden_vector_other' +ROLE_VECTOR_LIKE = 'vector_like' +ROLE_MUON_MANAGED_MATRIX = 'muon_managed_matrix' + +_EMBEDDING_CLASS_ROLES = frozenset((ROLE_EMBEDDING, ROLE_OUTPUT, ROLE_SHARED_EMBEDDING_OUTPUT)) +_OUTPUT_ROLES = frozenset((ROLE_OUTPUT, ROLE_SHARED_EMBEDDING_OUTPUT)) +_HIDDEN_VECTOR_ROLES = frozenset( + ( + ROLE_HIDDEN_VECTOR, + ROLE_HIDDEN_BIAS, + ROLE_NORM_SCALE, + ROLE_NORM_BIAS, + ROLE_QK_NORM_SCALE, + ROLE_HIDDEN_VECTOR_OTHER, + ROLE_VECTOR_LIKE, + ) +) +_NORM_ROLES = frozenset((ROLE_NORM_SCALE, ROLE_NORM_BIAS, ROLE_QK_NORM_SCALE)) + + +def set_parameterization_metadata( + param: Any, *, role: str, shared_group: Optional[str] = None, tags: Iterable[str] = () +) -> None: + setattr(param, PARAMETERIZATION_ROLE_ATTR, role) + if shared_group is not None: + setattr(param, PARAMETERIZATION_SHARED_GROUP_ATTR, shared_group) + if tags: + setattr(param, PARAMETERIZATION_TAGS_ATTR, tuple(tags)) + + +def get_parameterization_role(param: Any) -> Optional[str]: + return getattr(param, PARAMETERIZATION_ROLE_ATTR, None) + + +def is_output_parameter(param: Any) -> bool: + if hasattr(param, IS_OUTPUT_PARAMETER_ATTR): + return bool(getattr(param, IS_OUTPUT_PARAMETER_ATTR)) + return get_parameterization_role(param) in _OUTPUT_ROLES + + +def is_embedding_or_output_parameter(param: Any) -> bool: + if hasattr(param, 'is_embedding_or_output_parameter'): + return bool(param.is_embedding_or_output_parameter) + return get_parameterization_role(param) in _EMBEDDING_CLASS_ROLES + + +def is_embedding_class_parameter(param: Any, param_name: Optional[str] = None) -> bool: + if getattr(param, 'shared_embedding', False): + return True + if hasattr(param, 'is_embedding_parameter'): + return bool(param.is_embedding_parameter) + if get_parameterization_role(param) in _EMBEDDING_CLASS_ROLES: + return True + return bool(param_name and 'embedding' in param_name.lower()) + + +def is_vector_like_parameter(param: Any, param_name: Optional[str] = None) -> bool: + if is_embedding_class_parameter(param, param_name): + return True + return param.dim() <= 1 + + +def _lower_name(param_name: Optional[str]) -> str: + return param_name.lower() if param_name else '' + + +def is_qk_norm_parameter(param: Any, param_name: Optional[str] = None) -> bool: + role = get_parameterization_role(param) + if role == ROLE_QK_NORM_SCALE: + return True + name = _lower_name(param_name) + return param.dim() <= 1 and ('q_layernorm.' in name or 'k_layernorm.' in name) + + +def is_norm_parameter(param: Any, param_name: Optional[str] = None) -> bool: + role = get_parameterization_role(param) + if role in _NORM_ROLES: + return True + if param.dim() > 1: + return False + name = _lower_name(param_name) + return 'layernorm' in name or 'layer_norm' in name or 'rmsnorm' in name or '.norm.' in name + + +def is_hidden_bias_parameter(param: Any, param_name: Optional[str] = None) -> bool: + role = get_parameterization_role(param) + if role == ROLE_HIDDEN_BIAS: + return True + if role in _EMBEDDING_CLASS_ROLES or role in _NORM_ROLES: + return False + name = _lower_name(param_name) + return param.dim() <= 1 and name.endswith('.bias') and not is_norm_parameter(param, param_name) + + +def should_skip_depth_mup_vector_weight_decay( + param: Any, param_name: Optional[str] = None, *, apply_wd_to_qk_layernorm: bool = False +) -> bool: + """Return true for vector-like params outside the depth-MuP hidden-bias table row.""" + if is_embedding_class_parameter(param, param_name): + return False + if is_hidden_bias_parameter(param, param_name): + return False + if apply_wd_to_qk_layernorm and is_qk_norm_parameter(param, param_name): + return False + if is_norm_parameter(param, param_name): + return True + return param.dim() <= 1 + + +def is_hidden_vector_parameter(param: Any, param_name: Optional[str] = None) -> bool: + role = get_parameterization_role(param) + if role in _HIDDEN_VECTOR_ROLES: + return True + if role in _EMBEDDING_CLASS_ROLES: + return False + return param.dim() <= 1 and not is_embedding_class_parameter(param, param_name) + + +def is_hidden_matrix_parameter(param: Any, param_name: Optional[str] = None) -> bool: + role = get_parameterization_role(param) + if role in (ROLE_HIDDEN_MATRIX, ROLE_BLOCK_OUT_PROJ): + return True + if role in _HIDDEN_VECTOR_ROLES or role in _EMBEDDING_CLASS_ROLES: + return False + return param.dim() > 1 and not is_embedding_class_parameter(param, param_name) + + +def is_muon_managed_matrix_parameter(param: Any, *, optimizer_type: str) -> bool: + if 'muon' not in optimizer_type.lower(): + return False + return param.dim() == 2 and not is_embedding_or_output_parameter(param) diff --git a/megatron/core/parameterization/spec.py b/megatron/core/parameterization/spec.py new file mode 100644 index 00000000000..86d531e6963 --- /dev/null +++ b/megatron/core/parameterization/spec.py @@ -0,0 +1,290 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal, Optional + +SCALING_RECIPE_NONE = 'none' +SCALING_RECIPE_MUP = 'mup' +SCALING_RECIPE_DEPTH_MUP = 'depth_mup' +SCALING_RECIPE_VALUES = (SCALING_RECIPE_NONE, SCALING_RECIPE_MUP, SCALING_RECIPE_DEPTH_MUP) + + +@dataclass(frozen=True) +class ScalingUserConfig: + recipe: Optional[Literal['none', 'mup', 'depth_mup']] = None + base_hidden_size: Optional[int] = None + base_num_layers: Optional[int] = None + base_head_dim: Optional[float] = None + residual_branch_depth_power: Optional[float] = None + hidden_lr_depth_power: Optional[float] = None + block_out_proj_init_depth_power: Optional[float] = None + use_mup_alias: bool = False + mup_width_mult: Optional[float] = None + mup_width_mult_explicit: bool = False + mup_base_hidden_size: Optional[int] = None + mup_embedding_mult: float = 1.0 + mup_output_mult: float = 1.0 + mup_base_head_dim: Optional[float] = None + mup_attn_scale_power: float = 1.0 + + +@dataclass(frozen=True) +class ScalingContext: + """Internal scaling context for standard, width-MuP, and depth-MuP.""" + + recipe: Literal['none', 'mup', 'depth_mup'] + width_mult: float = 1.0 + depth_mult: float = 1.0 + embedding_mult: float = 1.0 + output_mult: float = 1.0 + base_hidden_size: Optional[int] = None + base_num_layers: Optional[int] = None + base_head_dim: Optional[float] = None + current_head_dim: Optional[int] = None + attention_scale_power: float = 1.0 + residual_branch_depth_power: float = 0.0 + hidden_lr_depth_power: float = 0.0 + block_out_proj_init_depth_power: float = 0.0 + + @property + def enabled(self) -> bool: + return self.recipe != SCALING_RECIPE_NONE + + @property + def uses_width_mup(self) -> bool: + return self.recipe in (SCALING_RECIPE_MUP, SCALING_RECIPE_DEPTH_MUP) + + @property + def use_mup(self) -> bool: + return self.recipe == SCALING_RECIPE_MUP + + @property + def is_depth_mup(self) -> bool: + return self.recipe == SCALING_RECIPE_DEPTH_MUP + + +def _resolve_aliased_value( + explicit_value: Optional[float | int], + legacy_value: Optional[float | int], + *, + explicit_name: str, + legacy_name: str, +) -> Optional[float | int]: + if explicit_value is None: + return legacy_value + if legacy_value is None: + return explicit_value + if explicit_value != legacy_value: + raise ValueError( + f"{explicit_name} ({explicit_value}) conflicts with {legacy_name} ({legacy_value}). " + f"Specify only one or set them to the same value." + ) + return explicit_value + + +def _non_default_scaling_fields(user_config: ScalingUserConfig) -> list[str]: + candidates: dict[str, object] = { + 'scaling_base_hidden_size': user_config.base_hidden_size, + 'scaling_base_num_layers': user_config.base_num_layers, + 'scaling_base_head_dim': user_config.base_head_dim, + 'scaling_residual_branch_depth_power': user_config.residual_branch_depth_power, + 'scaling_hidden_lr_depth_power': user_config.hidden_lr_depth_power, + 'scaling_block_out_proj_init_depth_power': user_config.block_out_proj_init_depth_power, + 'mup_base_hidden_size': user_config.mup_base_hidden_size, + 'mup_base_head_dim': user_config.mup_base_head_dim, + } + if user_config.mup_embedding_mult != 1.0: + candidates['mup_embedding_mult'] = user_config.mup_embedding_mult + if user_config.mup_output_mult != 1.0: + candidates['mup_output_mult'] = user_config.mup_output_mult + if user_config.mup_attn_scale_power != 1.0: + candidates['mup_attn_scale_power'] = user_config.mup_attn_scale_power + if user_config.mup_width_mult_explicit: + candidates['mup_width_mult'] = user_config.mup_width_mult + return [name for name, value in candidates.items() if value is not None] + + +def _infer_current_head_dim(config: Any) -> Optional[int]: + kv_channels = getattr(config, 'kv_channels', None) + if kv_channels is not None: + return kv_channels + + hidden_size = getattr(config, 'hidden_size', None) + num_attention_heads = getattr(config, 'num_attention_heads', None) + if hidden_size is None or num_attention_heads is None: + return None + if num_attention_heads is None or num_attention_heads <= 0: + raise AttributeError( + "Cannot resolve current head dimension without kv_channels or a positive " + "num_attention_heads value." + ) + return hidden_size // num_attention_heads + + +def build_scaling_user_config(config: Any) -> ScalingUserConfig: + raw_mup_width_mult = getattr(config, 'mup_width_mult', None) + marker_present = hasattr(config, '_mup_width_mult_explicit') + mup_width_mult_explicit = bool(getattr(config, '_mup_width_mult_explicit', False)) + if ( + not marker_present + and not mup_width_mult_explicit + and raw_mup_width_mult not in (None, 1.0) + ): + # Direct TransformerConfig construction has no argparse provenance marker. + mup_width_mult_explicit = True + + return ScalingUserConfig( + recipe=getattr(config, 'scaling_recipe', None), + base_hidden_size=getattr(config, 'scaling_base_hidden_size', None), + base_num_layers=getattr(config, 'scaling_base_num_layers', None), + base_head_dim=getattr(config, 'scaling_base_head_dim', None), + residual_branch_depth_power=getattr(config, 'scaling_residual_branch_depth_power', None), + hidden_lr_depth_power=getattr(config, 'scaling_hidden_lr_depth_power', None), + block_out_proj_init_depth_power=getattr( + config, 'scaling_block_out_proj_init_depth_power', None + ), + use_mup_alias=bool(getattr(config, 'use_mup', False)), + mup_width_mult=raw_mup_width_mult if mup_width_mult_explicit else None, + mup_width_mult_explicit=mup_width_mult_explicit, + mup_base_hidden_size=getattr(config, 'mup_base_hidden_size', None), + mup_embedding_mult=getattr(config, 'mup_embedding_mult', 1.0), + mup_output_mult=getattr(config, 'mup_output_mult', 1.0), + mup_base_head_dim=getattr(config, 'mup_base_head_dim', None), + mup_attn_scale_power=getattr(config, 'mup_attn_scale_power', 1.0), + ) + +def build_scaling_context(config: Any) -> ScalingContext: + user_config = build_scaling_user_config(config) + recipe = user_config.recipe + if recipe is None: + recipe = SCALING_RECIPE_MUP if user_config.use_mup_alias else SCALING_RECIPE_NONE + elif user_config.use_mup_alias and recipe != SCALING_RECIPE_MUP: + raise ValueError( + f"--scaling-recipe {recipe} conflicts with --use-mup. " + "Use either the canonical MuP recipe or the legacy MuP alias, not both." + ) + if recipe not in SCALING_RECIPE_VALUES: + raise ValueError(f"Unsupported scaling recipe: {recipe}") + + base_hidden_size = _resolve_aliased_value( + user_config.base_hidden_size, + user_config.mup_base_hidden_size, + explicit_name='--scaling-base-hidden-size', + legacy_name='--mup-base-hidden-size', + ) + base_head_dim = _resolve_aliased_value( + user_config.base_head_dim, + user_config.mup_base_head_dim, + explicit_name='--scaling-base-head-dim', + legacy_name='--mup-base-head-dim', + ) + + if recipe == SCALING_RECIPE_NONE: + non_default_fields = _non_default_scaling_fields(user_config) + if non_default_fields: + raise ValueError( + "Scaling overrides require a non-'none' scaling recipe (for example `mup` " + "or `depth_mup`). Non-default fields: " + ", ".join(non_default_fields) + ) + return ScalingContext( + recipe=SCALING_RECIPE_NONE, + current_head_dim=_infer_current_head_dim(config), + ) + + if base_hidden_size is None: + base_hidden_size = config.hidden_size + if base_hidden_size <= 0: + raise AssertionError('--scaling-base-hidden-size must be positive.') + + current_num_layers = getattr(config, 'num_layers', None) + if current_num_layers is None: + if recipe == SCALING_RECIPE_DEPTH_MUP: + raise AttributeError("Cannot resolve depth_mup without num_layers.") + current_num_layers = 1 + + base_num_layers = user_config.base_num_layers + if base_num_layers is None: + base_num_layers = current_num_layers + if base_num_layers <= 0: + raise AssertionError('--scaling-base-num-layers must be positive.') + if base_head_dim is not None and base_head_dim <= 0: + raise AssertionError('--scaling-base-head-dim must be positive.') + + width_mult = config.hidden_size / base_hidden_size + if ( + user_config.mup_width_mult_explicit + and user_config.mup_width_mult is not None + and not math.isclose( + user_config.mup_width_mult, width_mult, rel_tol=1e-12, abs_tol=1e-12 + ) + ): + raise ValueError( + "--mup-width-mult is deprecated as an input and must match the derived " + f"hidden_size / scaling_base_hidden_size value ({width_mult}). " + f"Got --mup-width-mult={user_config.mup_width_mult}." + ) + + output_mult = user_config.mup_output_mult + if output_mult == 1.0 and width_mult != 1.0: + output_mult = 1.0 / width_mult + + residual_branch_depth_power = user_config.residual_branch_depth_power + if residual_branch_depth_power is None: + residual_branch_depth_power = -1.0 if recipe == SCALING_RECIPE_DEPTH_MUP else 0.0 + + hidden_lr_depth_power = user_config.hidden_lr_depth_power + if hidden_lr_depth_power is None: + hidden_lr_depth_power = 0.0 + + block_out_proj_init_depth_power = user_config.block_out_proj_init_depth_power + if block_out_proj_init_depth_power is None: + block_out_proj_init_depth_power = 0.5 if recipe == SCALING_RECIPE_DEPTH_MUP else 0.0 + + return ScalingContext( + recipe=recipe, + width_mult=width_mult, + depth_mult=current_num_layers / base_num_layers, + embedding_mult=user_config.mup_embedding_mult, + output_mult=output_mult, + base_hidden_size=base_hidden_size, + base_num_layers=base_num_layers, + base_head_dim=base_head_dim, + current_head_dim=_infer_current_head_dim(config), + attention_scale_power=user_config.mup_attn_scale_power, + residual_branch_depth_power=float(residual_branch_depth_power), + hidden_lr_depth_power=float(hidden_lr_depth_power), + block_out_proj_init_depth_power=float(block_out_proj_init_depth_power), + ) + + +def sync_legacy_mup_fields(config: Any, context: ScalingContext) -> None: + config.scaling_recipe = context.recipe + config.use_mup = context.recipe == SCALING_RECIPE_MUP + config.mup_width_mult = context.width_mult + config._mup_width_mult_explicit = False + config.mup_embedding_mult = context.embedding_mult + config.mup_output_mult = context.output_mult + config.mup_attn_scale_power = context.attention_scale_power + + if context.recipe == SCALING_RECIPE_NONE: + config.scaling_base_hidden_size = None + config.scaling_base_num_layers = None + config.scaling_base_head_dim = None + config.scaling_residual_branch_depth_power = None + config.scaling_hidden_lr_depth_power = None + config.scaling_block_out_proj_init_depth_power = None + config.mup_base_hidden_size = None + config.mup_base_head_dim = None + return + + config.scaling_base_hidden_size = context.base_hidden_size + config.scaling_base_num_layers = context.base_num_layers + config.scaling_base_head_dim = context.base_head_dim + config.scaling_residual_branch_depth_power = context.residual_branch_depth_power + config.scaling_hidden_lr_depth_power = context.hidden_lr_depth_power + config.scaling_block_out_proj_init_depth_power = context.block_out_proj_init_depth_power + config.mup_base_hidden_size = context.base_hidden_size + config.mup_base_head_dim = context.base_head_dim diff --git a/megatron/core/parameterization/training_policy.py b/megatron/core/parameterization/training_policy.py new file mode 100644 index 00000000000..2cfc89dc883 --- /dev/null +++ b/megatron/core/parameterization/training_policy.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from dataclasses import dataclass + +from .spec import SCALING_RECIPE_MUP, ScalingContext, build_scaling_context + +@dataclass(frozen=True) +class TrainingScalingPolicy: + """Optimizer-side policy for Megatron scaling recipes.""" + + context: ScalingContext + optimizer_type: str = 'adam' + + def __post_init__(self) -> None: + if self.context.is_depth_mup and not self.is_adam_optimizer: + raise ValueError( + "scaling_recipe='depth_mup' currently supports optimizer='adam' only. " + "AdamW semantics should continue to use decoupled_weight_decay. " + "SGD depth-mup requires explicit hidden-weight, hidden-bias, norm/vector, " + "and input/output-bias rules and is intentionally out of scope for v1." + ) + + @property + def enabled(self) -> bool: + return self.context.enabled + + @property + def optimizer_type_lower(self) -> str: + return self.optimizer_type.lower() + + @property + def uses_width_mup(self) -> bool: + return self.context.uses_width_mup + + @property + def is_sgd_optimizer(self) -> bool: + return self.optimizer_type_lower == 'sgd' + + @property + def is_adam_optimizer(self) -> bool: + return self.optimizer_type_lower == 'adam' + + @property + def is_muon_optimizer(self) -> bool: + return 'muon' in self.optimizer_type_lower + + @property + def hidden_lr_width_power(self) -> float: + if not self.enabled or not self.uses_width_mup: + return 0.0 + return 0.0 if self.is_sgd_optimizer else -1.0 + + @property + def hidden_lr_multiplier(self) -> float: + if not self.enabled: + return 1.0 + return (self.context.width_mult**self.hidden_lr_width_power) * ( + self.context.depth_mult**self.context.hidden_lr_depth_power + ) + + @property + def hidden_vector_lr_multiplier(self) -> float: + if not (self.enabled and self.uses_width_mup and self.is_sgd_optimizer): + return 1.0 + return self.context.width_mult + + @property + def hidden_eps_depth_power(self) -> float: + if not (self.enabled and self.uses_width_mup and self.is_adam_optimizer): + return 0.0 + return -1.0 if self.context.is_depth_mup else 0.0 + + @property + def hidden_eps_multiplier(self) -> float: + if not (self.enabled and self.uses_width_mup and self.is_adam_optimizer): + return 1.0 + return (1.0 / self.context.width_mult) * ( + self.context.depth_mult**self.hidden_eps_depth_power + ) + + @property + def hidden_vector_eps_multiplier(self) -> float: + if not (self.enabled and self.uses_width_mup and self.is_adam_optimizer): + return 1.0 + if not self.context.is_depth_mup: + return 1.0 + return self.hidden_eps_multiplier + + @property + def embedding_class_eps_multiplier(self) -> float: + if not (self.enabled and self.uses_width_mup and self.is_adam_optimizer): + return 1.0 + if not self.context.is_depth_mup: + return 1.0 + return 1.0 / self.context.width_mult + + @property + def hidden_matrix_wd_multiplier(self) -> float: + if not (self.enabled and self.context.is_depth_mup and self.is_adam_optimizer): + return 1.0 + return self.context.width_mult + + @property + def hidden_vector_wd_multiplier(self) -> float: + return 1.0 + + @property + def embedding_class_wd_multiplier(self) -> float: + return 1.0 + + @property + def vector_like_lr_multiplier(self) -> float: + return self.hidden_vector_lr_multiplier + + +def build_training_scaling_policy(config, optimizer_type: str = 'adam') -> TrainingScalingPolicy: + return TrainingScalingPolicy(context=build_scaling_context(config), optimizer_type=optimizer_type) + + +def build_legacy_mup_training_policy( + *, mup_width_mult: float, optimizer_type: str = 'adam' +) -> TrainingScalingPolicy: + return TrainingScalingPolicy( + context=ScalingContext(recipe=SCALING_RECIPE_MUP, width_mult=mup_width_mult), + optimizer_type=optimizer_type, + ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index b27f90c53d0..c9bbe6da881 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -28,6 +28,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) @@ -391,11 +392,24 @@ def __init__( ) # Output. + model_scaling_policy = build_model_scaling_policy(self.config) + linear_proj_init_method = self.config.output_layer_init_method + if self.attention_type != "cross": + linear_proj_init_method = model_scaling_policy.dense_block_output_init_method( + default_init_method=not_none(self.config.output_layer_init_method), + init_method_std=self.config.init_method_std, + num_layers=self.config.num_layers, + is_hybrid_model=self.config.is_hybrid_model, + output_layer_init_method_is_user_provided=getattr( + self.config, '_parameterization_output_layer_init_method_user_provided', False + ), + apply_depth_hook=True, + ) self.linear_proj = submodules.linear_proj( self.query_projection_size, self.config.hidden_size, config=self.config, - init_method=not_none(self.config.output_layer_init_method), + init_method=linear_proj_init_method, bias=self.config.add_bias_linear, input_is_parallel=True, skip_bias_add=True, diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 1a578151f1e..b19c85b0b50 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -23,6 +23,7 @@ ) from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl, weighted_bias_swiglu_impl +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig @@ -170,6 +171,7 @@ def __init__( is_expert: bool = False, input_size: Optional[int] = None, ffn_hidden_size: Optional[int] = None, + apply_block_output_init_scaling: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, ): @@ -233,13 +235,26 @@ def __init__( else: self.activation_func = self.config.activation_func + model_scaling_policy = build_model_scaling_policy(self.config) + fc2_init_method = not_none(self.config.output_layer_init_method) + if apply_block_output_init_scaling and not is_expert: + fc2_init_method = model_scaling_policy.dense_block_output_init_method( + default_init_method=not_none(self.config.output_layer_init_method), + init_method_std=self.config.init_method_std, + num_layers=self.config.num_layers, + is_hybrid_model=self.config.is_hybrid_model, + output_layer_init_method_is_user_provided=getattr( + self.config, '_parameterization_output_layer_init_method_user_provided', False + ), + apply_depth_hook=True, + ) self.linear_fc2 = submodules.linear_fc2( not_none(self.config.ffn_hidden_size), not_none( self.config.hidden_size if not use_latent_size else self.config.moe_latent_size ), config=self.config, - init_method=not_none(self.config.output_layer_init_method), + init_method=fc2_init_method, bias=self.config.add_bias_linear, input_is_parallel=True, skip_bias_add=True, @@ -385,6 +400,7 @@ def as_mlp_submodule( input_size: int | None = None, ffn_hidden_size: int | None = None, name: str | None = None, + apply_block_output_init_scaling: bool = False, ) -> MLP: """Helper function to build an MLP as a TransformerLayer's mlp submodule.""" del is_mtp_layer @@ -399,6 +415,7 @@ def as_mlp_submodule( input_size=input_size, ffn_hidden_size=ffn_hidden_size, name=name, + apply_block_output_init_scaling=apply_block_output_init_scaling, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 3e91a2b8042..87167d0d8ee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,7 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging -import math import warnings from dataclasses import dataclass, field from typing import Callable, List, Literal, Optional, Tuple, Union @@ -11,6 +10,13 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend +from megatron.core.parameterization import ( + SCALING_RECIPE_DEPTH_MUP, + SCALING_RECIPE_MUP, + build_model_scaling_policy, + build_scaling_context, + sync_legacy_mup_fields, +) from megatron.core.quantization.quant_config import RecipeConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, @@ -30,14 +36,7 @@ from .._rank_utils import log_single_rank from ..fusions.fused_bias_geglu import quick_gelu from ..model_parallel_config import ModelParallelConfig -from ..utils import ( - get_te_version, - init_method_normal, - is_te_min_version, - is_torch_min_version, - mup_scaled_init_method_normal, - scaled_init_method_normal, -) +from ..utils import get_te_version, init_method_normal, is_te_min_version, is_torch_min_version logger = logging.getLogger(__name__) @@ -365,25 +364,65 @@ class TransformerConfig(ModelParallelConfig): #################### # MuP (Maximal Update Parameterization) #################### + scaling_recipe: Optional[Literal['none', 'mup', 'depth_mup']] = None + """ + Canonical scaling recipe. ``none`` preserves standard parameterization, ``mup`` + enables width MuP, and ``depth_mup`` enables the narrow spectral width-depth + MuP recipe for dense GPT-style residual Transformer blocks. + """ + + scaling_base_hidden_size: Optional[int] = None + """ + Canonical base hidden size for width scaling. For MuP, the width multiplier is + derived as hidden_size / scaling_base_hidden_size. + """ + + scaling_base_num_layers: Optional[int] = None + """ + Canonical base transformer-layer count for depth-based scaling recipes. + Defaults to ``num_layers`` when omitted. + """ + + scaling_base_head_dim: Optional[float] = None + """ + Canonical base attention head dimension for MuP attention scaling. This aliases the + deprecated mup_base_head_dim field. + """ + + scaling_residual_branch_depth_power: Optional[float] = None + """ + Relative depth exponent for dense self-attention/MLP residual-branch outputs. + Under ``depth_mup``, the default is ``-1.0``. + """ + + scaling_hidden_lr_depth_power: Optional[float] = None + """ + Relative depth exponent for hidden matrix-like LR overrides. + """ + + scaling_block_out_proj_init_depth_power: Optional[float] = None + """ + Relative depth exponent for dense transformer block output projection + initialization. Under ``depth_mup``, the default is ``+0.5``. + """ + use_mup: bool = False """ - Enable Maximal Update Parameterization (MuP) for hyperparameter transfer across - model widths. When enabled, learning rates and initialization are scaled according - to the width multiplier to ensure consistent training dynamics. + Deprecated alias for scaling_recipe='mup'. Kept for checkpoint and script + compatibility. """ mup_width_mult: float = 1.0 """ - Width multiplier for MuP scaling, computed as hidden_size / mup_base_hidden_size. - This value is automatically computed in __post_init__ when use_mup is enabled. + Deprecated derived MuP width multiplier. The canonical value is computed as + hidden_size / scaling_base_hidden_size. If this legacy input is non-default, it + must match the derived value. """ mup_base_hidden_size: Optional[int] = None """ - Base hidden size for MuP width scaling. This is the reference width from which - scaling factors are computed. Defaults to hidden_size if not specified (base model - case where width_mult=1.0). Set this to your base/proxy model's hidden size when - scaling up. + Deprecated alias for scaling_base_hidden_size. Set scaling_recipe='mup' and + scaling_base_hidden_size for new configs. """ mup_embedding_mult: float = 1.0 @@ -394,15 +433,15 @@ class TransformerConfig(ModelParallelConfig): mup_output_mult: float = 1.0 """ - Multiplier for output logits before softmax. When MuP is enabled and this is left - at 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable across - widths. Override to customize output scaling. + Multiplier for output logits before softmax. When scaling_recipe='mup' and this is + left at 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable + across widths. Override to customize output scaling. Default: 1.0. """ mup_base_head_dim: Optional[float] = None """ - Base head dimension for MuP attention scaling. When set, + Deprecated alias for scaling_base_head_dim. When set, softmax_scale = sqrt(mup_base_head_dim) / (kv_channels ** mup_attn_scale_power). Set to base model's d_head (e.g., 64) to match standard 1/sqrt(d_head) scaling at the base width, ensuring non-MuP compatibility for that specific value. @@ -412,7 +451,8 @@ class TransformerConfig(ModelParallelConfig): """ Power for attention scaling: softmax_scale = 1 / (kv_channels ** mup_attn_scale_power). 0.5 = standard attention (1/sqrt(d_head)), 1.0 = MuP attention (1/d_head). - Default: 1.0 (MuP scaling when use_mup is True). Set to 0.5 for standard scaling. + Default: 1.0 (MuP scaling when scaling_recipe='mup'). Set to 0.5 for standard + scaling. """ #################### @@ -1956,42 +1996,72 @@ def __post_init__(self): if self.multi_latent_attention and self.rotary_interleaved: raise ValueError("rotary_interleaved does not work with multi_latent_attention.") - # MuP (Maximal Update Parameterization) configuration - if self.use_mup: - # Default base_hidden_size to hidden_size (base model case, width_mult=1.0) - if self.mup_base_hidden_size is None: - self.mup_base_hidden_size = self.hidden_size - assert self.mup_base_hidden_size > 0, "--mup-base-hidden-size must be positive." - # Compute width multiplier - self.mup_width_mult = self.hidden_size / self.mup_base_hidden_size - - # MuP attention scaling: 1/d_head instead of 1/sqrt(d_head). - if self.softmax_scale is None: - base_head_scale = ( - 1.0 if self.mup_base_head_dim is None else self.mup_base_head_dim**0.5 - ) - self.softmax_scale = base_head_scale / (self.kv_channels**self.mup_attn_scale_power) - - # MuP output scaling: scale logits by 1/width_mult to keep outputs O(1). - # Only auto-set if user hasn't explicitly configured it. - if self.mup_output_mult == 1.0 and self.mup_width_mult != 1.0: - self.mup_output_mult = 1.0 / self.mup_width_mult + scaling_context = build_scaling_context(self) + sync_legacy_mup_fields(self, scaling_context) + model_scaling_policy = build_model_scaling_policy(self) + if scaling_context.is_depth_mup: + if self.is_hybrid_model: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports dense GPT-style residual " + "Transformer blocks only. Hybrid/Mamba layer patterns are out of scope for v1." + ) + if self.mtp_num_layers is not None and self.mtp_num_layers > 0: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports standard next-token " + "training only. MTP depth transfer is out of scope for v1." + ) + if self.multi_latent_attention: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports dense GPT-style residual " + "self-attention only. multi_latent_attention is out of scope for v1." + ) + if self.experimental_attention_variant is not None: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports dense GPT-style residual " + "self-attention only. experimental attention variants are out of scope for v1." + ) + if self.num_moe_experts is not None: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports dense GPT-style residual " + "Transformer blocks only. MoE depth transfer is out of scope for v1." + ) + + # MuP (Maximal Update Parameterization) configuration + if scaling_context.recipe in (SCALING_RECIPE_MUP, SCALING_RECIPE_DEPTH_MUP): overridden_init_methods = [] if self.init_method is not None: overridden_init_methods.append("init_method") + if self.embedding_init_method is not None: + overridden_init_methods.append("embedding_init_method") if self.output_layer_init_method is not None: overridden_init_methods.append("output_layer_init_method") if overridden_init_methods: overridden_init_methods_text = " and ".join(overridden_init_methods) verb = "is" if len(overridden_init_methods) == 1 else "are" warnings.warn( - "use_mup is enabled, but custom " + f"scaling recipe {scaling_context.recipe!r} is enabled, but custom " + overridden_init_methods_text - + f" {verb} set. This may break MuP initialization assumptions.", + + f" {verb} set. This may break scaling initialization assumptions.", UserWarning, ) + self._parameterization_output_layer_init_method_user_provided = ( + self.output_layer_init_method is not None + ) + if ( + self._parameterization_output_layer_init_method_user_provided + and model_scaling_policy.dense_block_out_proj_init_multiplier != 1.0 + ): + warnings.warn( + "Custom output_layer_init_method is set, so dense block output projection " + "depth scaling will be ignored.", + UserWarning, + ) + self.softmax_scale = model_scaling_policy.resolve_attention_softmax_scale( + softmax_scale=self.softmax_scale, kv_channels=self.kv_channels + ) + # Set the embedding init method. # NOTE: This block must run AFTER the MuP block above but BEFORE the init_method # block below. When MuP is enabled and init_method is None (the common case), @@ -2015,29 +2085,18 @@ def __post_init__(self): self.embedding_init_method = self.init_method if self.init_method is None: - if self.use_mup: - # MuP: scale std by 1/sqrt(width_mult). - self.init_method = init_method_normal( - self.init_method_std / math.sqrt(self.mup_width_mult) - ) - else: - self.init_method = init_method_normal(self.init_method_std) + self.init_method = model_scaling_policy.build_hidden_init_method( + init_method_std=self.init_method_std + ) if self.output_layer_init_method is None: - if self.use_mup: - # MuP: depth and width scaling for output layers. - self.output_layer_init_method = mup_scaled_init_method_normal( - self.init_method_std, - self.num_layers, - self.mup_width_mult, - multiplier=2.0 if not self.is_hybrid_model else 1.0, - ) - else: - self.output_layer_init_method = scaled_init_method_normal( - self.init_method_std, - self.num_layers, - multiplier=2.0 if not self.is_hybrid_model else 1.0, + self.output_layer_init_method = ( + model_scaling_policy.build_default_output_layer_init_method( + init_method_std=self.init_method_std, + num_layers=self.num_layers, + is_hybrid_model=self.is_hybrid_model, ) + ) if self.num_moe_experts is not None and self.add_bias_linear: assert ( diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..be10c3885c3 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -17,6 +17,10 @@ from megatron.core.dist_checkpointing.utils import apply_prefix_mapping from megatron.core.inference.utils import InferenceMode from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parameterization import ( + build_model_scaling_policy, + is_scaling_policy_eval_allowed, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType @@ -43,6 +47,19 @@ logger = logging.getLogger(__name__) +def _get_mlp_builder_module(mlp_builder: Any, te_fused_mlp_cls: Optional[type] = None): + """Return the MLP module class hidden behind ModuleSpec or classmethod partial builders.""" + if isinstance(mlp_builder, ModuleSpec): + return mlp_builder.module + if isinstance(mlp_builder, functools.partial): + owner = getattr(mlp_builder.func, "__self__", None) + if owner is MLP: + return MLP + if te_fused_mlp_cls is not None and owner is te_fused_mlp_cls: + return te_fused_mlp_cls + return None + + def get_transformer_layer_offset( config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None ): @@ -312,6 +329,20 @@ def __init__( Args: name (str | None): module instance name passed top-down from its paranet module """ + cross_attention_spec = submodules.cross_attention + uses_cross_attention = not ( + cross_attention_spec is IdentityOp + or ( + isinstance(cross_attention_spec, ModuleSpec) + and cross_attention_spec.module is IdentityOp + ) + ) + if config.scaling_recipe == 'depth_mup' and uses_cross_attention: + raise NotImplementedError( + "scaling_recipe='depth_mup' currently supports dense GPT-style residual " + "self-attention-only Transformer blocks. Cross-attention is out of scope for v1." + ) + self.submodules_config = submodules super().__init__(config=config, vp_stage=vp_stage) @@ -335,6 +366,7 @@ def __init__( ) self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout self.is_mtp_layer = is_mtp_layer + self.model_scaling_policy = build_model_scaling_policy(config) # [Module 1: Input Layernorm] Optional Layernorm on the input data # TODO: add pytorch only layernorm @@ -403,8 +435,22 @@ def __init__( # MLP expects tp_group but MoELayer expects pg_collection to be passed in. # We can change MLP to accept pg_collection but it makes the logic implicit # The conditional below is to make the logic explicit - # if submodules.mlp is not a ModuleSpec,we dont have to handle passing additional kwargs - if isinstance(submodules.mlp, ModuleSpec) and submodules.mlp.module in (MLP, TEFusedMLP): + # Dense GPT specs usually pass partial(MLP.as_mlp_submodule, ...), while some tests + # and extensions still use ModuleSpec. Both forms need the same depth-MuP handling. + additional_mlp_kwargs = {} + mlp_module = _get_mlp_builder_module(submodules.mlp, TEFusedMLP) + if mlp_module is MLP: + additional_mlp_kwargs["apply_block_output_init_scaling"] = True + elif ( + TEFusedMLP is not None + and mlp_module is TEFusedMLP + and self.model_scaling_policy.dense_block_out_proj_init_multiplier != 1.0 + ): + raise NotImplementedError( + "Dense block output-projection init scaling is not implemented for " + "TEFusedMLP. Use unfused MLP or disable the depth init hook." + ) + if isinstance(submodules.mlp, ModuleSpec) and mlp_module in (MLP, TEFusedMLP): submodules.mlp = functools.partial( submodules.mlp.module.as_mlp_submodule, submodules=submodules.mlp.submodules, @@ -422,6 +468,7 @@ def __init__( pg_collection=pg_collection, is_mtp_layer=self.is_mtp_layer, name=(name + ".mlp") if name is not None else None, + **additional_mlp_kwargs, ) if hasattr(self.mlp, 'set_layer_number'): self.mlp.set_layer_number(self.layer_number) @@ -673,6 +720,11 @@ def _forward_attention( # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="self_attn_bda") + attention_output_with_bias = self._scale_dense_residual_branch_output( + attention_output_with_bias, + branch_name="self attention", + using_fused_tp_inference_kernel=using_fused_tp_inference_kernel, + ) if using_fused_tp_inference_kernel: # In inference optimized transformer layer, there is no bias and dropout # The remaining residual add is already handled inside the @@ -745,6 +797,34 @@ def forward(self, *args, **kwargs): ) return output, context + def _scale_dense_residual_branch_output( + self, + output_with_bias: tuple[Tensor, Tensor | None], + *, + branch_name: str, + using_fused_tp_inference_kernel: bool, + apply_depth_hook: bool = True, + ) -> tuple[Tensor, Tensor | None]: + if not apply_depth_hook: + return output_with_bias + if self.model_scaling_policy.context.is_depth_mup and not getattr(self, 'training', True): + if using_fused_tp_inference_kernel: + raise NotImplementedError( + f"Residual-branch scaling is not supported with fused TP inference for {branch_name}." + ) + if not is_scaling_policy_eval_allowed(): + raise NotImplementedError( + f"Residual-branch scaling is not supported during inference for {branch_name}. " + "Validation loss must run inside Megatron's scaling-policy eval context." + ) + if self.model_scaling_policy.residual_branch_multiplier == 1.0: + return output_with_bias + if using_fused_tp_inference_kernel: + raise NotImplementedError( + f"Residual-branch scaling is not supported with fused TP inference for {branch_name}." + ) + return self.model_scaling_policy.scale_residual_branch_output(output_with_bias) + def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, @@ -922,6 +1002,12 @@ def _forward_post_mlp( # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="mlp_bda") + mlp_output_with_bias = self._scale_dense_residual_branch_output( + mlp_output_with_bias, + branch_name="mlp", + using_fused_tp_inference_kernel=using_fused_tp_inference_kernel, + apply_depth_hook=not self.is_moe_layer, + ) if using_fused_tp_inference_kernel: # In inference optimized transformer layer, there is no bias and dropout # The remaining residual add is already handled inside the diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index cd3ce44c3a4..a9ded30541a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -38,6 +38,13 @@ ) from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.parameterization import ( + SCALING_RECIPE_DEPTH_MUP, + SCALING_RECIPE_MUP, + SCALING_RECIPE_VALUES, + build_scaling_context, + sync_legacy_mup_fields, +) from megatron.training.global_vars import set_global_variables from megatron.training.utils import ( get_device_arch_version, @@ -59,6 +66,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): # Standard arguments. parser = _add_network_size_args(parser) + parser = _add_scaling_args(parser) parser = _add_regularization_args(parser) parser = _add_training_args(parser) parser = _add_rl_args(parser) @@ -252,6 +260,87 @@ def validate_model_config_args_from_heterogeneous_config(args): f"Arguments differ from heterogeneous config: {incompatible_args_str}" ) + +def warn_deprecated_mup_aliases(args): + """Warn when users select the legacy MuP flag surface instead of scaling recipes.""" + + deprecated_aliases = [] + if getattr(args, '_use_mup_explicit', getattr(args, 'use_mup', False)): + deprecated_aliases.append('--use-mup') + if getattr( + args, '_mup_base_hidden_size_explicit', getattr(args, 'mup_base_hidden_size', None) is not None + ): + deprecated_aliases.append('--mup-base-hidden-size') + if getattr( + args, '_mup_base_head_dim_explicit', getattr(args, 'mup_base_head_dim', None) is not None + ): + deprecated_aliases.append('--mup-base-head-dim') + if getattr(args, '_mup_width_mult_explicit', False): + deprecated_aliases.append('--mup-width-mult') + + if deprecated_aliases: + warn_rank_0( + "Deprecated MuP argument(s) " + + ", ".join(deprecated_aliases) + + " were provided. Use --scaling-recipe mup with " + "--scaling-base-hidden-size and --scaling-base-head-dim instead. " + "--mup-width-mult is derived from hidden_size / scaling_base_hidden_size.", + getattr(args, 'rank', 0), + ) + + +def _resolve_validation_attr(args, attr_name): + """Resolve validation fields from either flat argparse args or nested YAML namespaces.""" + if hasattr(args, attr_name): + value = getattr(args, attr_name) + if value is not None: + return value + language_model = getattr(args, 'language_model', None) + if language_model is not None and hasattr(language_model, attr_name): + return getattr(language_model, attr_name) + return None + + +def validate_depth_mup_optimizer_support(args) -> None: + """Enforce the public optimizer support surface for depth_mup.""" + if _resolve_validation_attr(args, 'scaling_recipe') != SCALING_RECIPE_DEPTH_MUP: + return + + if _resolve_validation_attr(args, 'optimizer') not in ('adam',): + raise ValueError( + "scaling_recipe='depth_mup' currently supports optimizer='adam' only. " + "AdamW semantics should continue to use decoupled_weight_decay. " + "SGD depth-mup requires explicit hidden-weight, hidden-bias, norm/vector, " + "and input/output-bias rules and is intentionally out of scope for v1." + ) + + weight_decay = _resolve_validation_attr(args, 'weight_decay') + decoupled_weight_decay = _resolve_validation_attr(args, 'decoupled_weight_decay') + if decoupled_weight_decay is None: + # CLI argparse does not expose this field directly; OptimizerConfig defaults + # to AdamW semantics. + decoupled_weight_decay = True + if weight_decay is not None and weight_decay != 0.0 and decoupled_weight_decay is not True: + raise ValueError( + "scaling_recipe='depth_mup' with nonzero weight_decay requires " + "decoupled_weight_decay=True because the width-depth weight-decay scaling " + "is derived for AdamW. Use weight_decay=0.0 for coupled Adam, or enable " + "decoupled_weight_decay." + ) + + +def validate_muon_scalar_optimizer_support(args) -> None: + """Keep CLI and YAML validation aligned for Muon scalar optimizer selection.""" + muon_scalar_optimizer = _resolve_validation_attr(args, 'muon_scalar_optimizer') + if muon_scalar_optimizer is None: + return + if muon_scalar_optimizer not in ('adam', 'lion'): + raise ValueError( + "muon_scalar_optimizer must be one of ('adam', 'lion'). " + f"Got {muon_scalar_optimizer!r}." + ) + + def _eval_pattern(pattern): """ Validate and evaluate a string containing a Python list expression """ assert isinstance(pattern, str) @@ -1761,6 +1850,11 @@ def validate_args(args, defaults={}): 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." + validate_depth_mup_optimizer_support(args) + validate_muon_scalar_optimizer_support(args) + warn_deprecated_mup_aliases(args) + sync_legacy_mup_fields(args, build_scaling_context(args)) + # Print arguments. _print_args("arguments", args) @@ -2087,6 +2181,21 @@ def _add_network_size_args(parser): "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", + # generated by the explicit scaling argument group + "scaling_recipe", + "scaling_base_hidden_size", + "scaling_base_num_layers", + "scaling_base_head_dim", + "scaling_residual_branch_depth_power", + "scaling_hidden_lr_depth_power", + "scaling_block_out_proj_init_depth_power", + "use_mup", + "mup_width_mult", + "mup_base_hidden_size", + "mup_embedding_mult", + "mup_output_mult", + "mup_base_head_dim", + "mup_attn_scale_power", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -2647,6 +2756,110 @@ def _add_learning_rate_args(parser): return parser +def _add_scaling_args(parser): + group = parser.add_argument_group(title='scaling') + + class _StoreMupWidthMult(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, values) + setattr(namespace, '_mup_width_mult_explicit', True) + + group.add_argument( + '--scaling-recipe', + choices=SCALING_RECIPE_VALUES, + default=None, + help=( + "Canonical parameterization recipe. Use 'none' for standard parameterization " + "'mup' for width MuP, or 'depth_mup' for dense GPT-style width-depth MuP." + ), + ) + group.add_argument( + '--scaling-base-hidden-size', + type=int, + default=None, + help=( + "Base hidden size for scaling recipes. For MuP, width multiplier is derived as " + "hidden_size / scaling_base_hidden_size." + ), + ) + group.add_argument( + '--scaling-base-num-layers', + type=int, + default=None, + help="Base transformer-layer count for depth-based scaling recipes.", + ) + group.add_argument( + '--scaling-base-head-dim', + type=float, + default=None, + help="Base attention head dimension for MuP attention scaling.", + ) + group.add_argument( + '--scaling-residual-branch-depth-power', + type=float, + default=None, + help="Relative depth exponent for dense residual branch outputs.", + ) + group.add_argument( + '--scaling-hidden-lr-depth-power', + type=float, + default=None, + help="Relative depth exponent for hidden matrix-like LR scaling.", + ) + group.add_argument( + '--scaling-block-out-proj-init-depth-power', + type=float, + default=None, + help="Relative depth exponent for dense block output projection initialization.", + ) + group.add_argument( + '--use-mup', + action='store_true', + help=f"Deprecated alias for --scaling-recipe {SCALING_RECIPE_MUP}.", + ) + group.add_argument( + '--mup-width-mult', + type=float, + default=1.0, + action=_StoreMupWidthMult, + help=( + "Deprecated derived MuP width multiplier. If supplied, it must match " + "hidden_size / scaling_base_hidden_size." + ), + ) + group.add_argument( + '--mup-base-hidden-size', + type=int, + default=None, + help="Deprecated alias for --scaling-base-hidden-size.", + ) + group.add_argument( + '--mup-embedding-mult', + type=float, + default=1.0, + help="MuP embedding activation multiplier.", + ) + group.add_argument( + '--mup-output-mult', + type=float, + default=1.0, + help="MuP output logit multiplier. Defaults to 1 / width_mult when left at 1.0.", + ) + group.add_argument( + '--mup-base-head-dim', + type=float, + default=None, + help="Deprecated alias for --scaling-base-head-dim.", + ) + group.add_argument( + '--mup-attn-scale-power', + type=float, + default=1.0, + help="MuP attention scale power. The default uses 1 / d_head.", + ) + return parser + + def _add_checkpointing_args(parser): from megatron.training.config import CheckpointConfig diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 2363b7ae164..ef1c7034b21 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -38,6 +38,10 @@ from megatron.core.msc_utils import MultiStorageClientFeature, open_file from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer +from megatron.core.parameterization import ( + build_scaling_context, + sync_legacy_mup_fields, +) from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model @@ -175,6 +179,43 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('tensor_model_parallel_size') _compare('pipeline_model_parallel_size') + checkpoint_scaling_context = build_scaling_context(checkpoint_args) + args_scaling_context = build_scaling_context(args) + assert checkpoint_scaling_context == args_scaling_context, ( + f"Scaling recipe from checkpoint ({checkpoint_scaling_context}) is not equal to " + f"the input argument value ({args_scaling_context})." + ) + + +_CHECKPOINT_SCALING_ARG_DEFAULTS = { + 'scaling_recipe': 'none', + 'scaling_base_hidden_size': None, + 'scaling_base_num_layers': None, + 'scaling_base_head_dim': None, + 'scaling_residual_branch_depth_power': None, + 'scaling_hidden_lr_depth_power': None, + 'scaling_block_out_proj_init_depth_power': None, + 'use_mup': False, + 'mup_width_mult': 1.0, + '_mup_width_mult_explicit': False, + 'mup_base_hidden_size': None, + 'mup_embedding_mult': 1.0, + 'mup_output_mult': 1.0, + 'mup_base_head_dim': None, + 'mup_attn_scale_power': 1.0, +} + + +def _sync_checkpoint_scaling_args(checkpoint_args): + """Populate canonical scaling fields on checkpoint args before force-copying them.""" + + sync_legacy_mup_fields( + checkpoint_args, build_scaling_context(checkpoint_args) + ) + for arg_name, default_value in _CHECKPOINT_SCALING_ARG_DEFAULTS.items(): + if not hasattr(checkpoint_args, arg_name): + setattr(checkpoint_args, arg_name, default_value) + def isfile(filename) -> bool: if MultiStorageClientFeature.is_enabled(): @@ -1502,7 +1543,9 @@ def load_args_from_checkpoint( if hasattr(checkpoint_args, 'num_layers'): setattr(checkpoint_args, 'num_layers', None) - def _set_arg(arg_name, old_arg_name=None, force=False): + _sync_checkpoint_scaling_args(checkpoint_args) + + def _set_arg(arg_name, old_arg_name=None, force=False, allow_none=False): if not force and getattr(args, arg_name, None) is not None: return @@ -1511,7 +1554,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False): else: checkpoint_value = getattr(checkpoint_args, arg_name, None) - if checkpoint_value is not None: + if checkpoint_value is not None or allow_none: print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint") setattr(args, arg_name, checkpoint_value) else: @@ -1543,6 +1586,28 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('apply_query_key_layer_scaling', force=True) _set_arg('attention_dropout', force=True) _set_arg('hidden_dropout', force=True) + _set_arg('scaling_recipe', force=True) + _set_arg('scaling_base_hidden_size', force=True, allow_none=True) + _set_arg('scaling_base_num_layers', force=True, allow_none=True) + _set_arg('scaling_base_head_dim', force=True, allow_none=True) + _set_arg('scaling_residual_branch_depth_power', force=True, allow_none=True) + _set_arg('scaling_hidden_lr_depth_power', force=True, allow_none=True) + _set_arg('scaling_block_out_proj_init_depth_power', force=True, allow_none=True) + _set_arg('use_mup', force=True) + setattr(args, '_use_mup_explicit', False) + _set_arg('mup_width_mult', force=True) + setattr( + args, + '_mup_width_mult_explicit', + getattr(checkpoint_args, '_mup_width_mult_explicit', False), + ) + _set_arg('mup_base_hidden_size', force=True, allow_none=True) + setattr(args, '_mup_base_hidden_size_explicit', False) + _set_arg('mup_embedding_mult', force=True) + _set_arg('mup_output_mult', force=True) + _set_arg('mup_base_head_dim', force=True, allow_none=True) + setattr(args, '_mup_base_head_dim_explicit', False) + _set_arg('mup_attn_scale_power', force=True) # Legacy MTP pattern for old checkpoints _set_arg('mtp_hybrid_override_pattern', force=True) diff --git a/megatron/training/training.py b/megatron/training/training.py index f69e4f30f6a..c699508eee8 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -146,7 +146,11 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( is_linear_attention_variant, ) -from megatron.core.optimizer import get_mup_config_overrides, get_standard_config_overrides +from megatron.core.optimizer import ( + get_mup_config_overrides, + get_scaling_config_overrides, + get_standard_config_overrides, +) from megatron.core.optimizer.optimizer import param_group_identifier_keys from megatron.core.optimizer.optimizer_cuda_graph import OptimizerCudaGraphWrapper from megatron.core.optimizer.qk_clip import clip_qk @@ -157,13 +161,14 @@ def set_startup_timestamps(program_start=None, main_entry=None): is_vp_last_stage, ) from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.parameterization import allow_scaling_policy_eval, build_training_scaling_policy from megatron.core.transformer.cuda_graphs import TECudaGraphHelper from megatron.core.transformer.module import Float16Module from megatron.core.transformer.moe.paged_stash import PagedStashRunner from megatron.core.distributed import DistributedDataParallelConfig, TorchFullyShardedDataParallelConfig from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel as megatron_FSDP -from megatron.core.optimizer.optimizer import param_group_identifier_keys +from megatron.core.optimizer.optimizer import get_param_group_identifier_sort_key from megatron.core.optimizer.qk_clip import clip_qk from megatron.core.utils import ( @@ -1009,7 +1014,7 @@ def reorder_inner_param_groups(optimizer_state_dict): if "param_groups" not in inner_optimizer: return param_groups = inner_optimizer["param_groups"] - key_fn = lambda pg: [pg[key] for key in param_group_identifier_keys] + key_fn = get_param_group_identifier_sort_key param_groups.sort(key=key_fn) inner_optimizer["param_groups"] = param_groups @@ -1953,18 +1958,18 @@ def setup_model_and_optimizer( else: config, config_overrides = get_megatron_optimizer_config(args) config.timers = timers - if getattr(args, "use_mup", False): - model_config_source = ( - unwrapped_model[0] if isinstance(unwrapped_model, list) else unwrapped_model + model_config_source = unwrapped_model[0] if isinstance(unwrapped_model, list) else unwrapped_model + model_config = get_model_config(model_config_source) + scaling_policy = build_training_scaling_policy(model_config, optimizer_type=config.optimizer) + if scaling_policy.enabled: + config_overrides = get_standard_config_overrides( + config=config, scaling_policy=scaling_policy ) - model_config = get_model_config(model_config_source) - mup_overrides = get_mup_config_overrides( - config=config, - mup_width_mult=model_config.mup_width_mult, - optimizer_type=config.optimizer, + scaling_overrides = get_scaling_config_overrides( + config=config, scaling_policy=scaling_policy ) - if mup_overrides: - config_overrides = {**(config_overrides or {}), **mup_overrides} + if scaling_overrides: + config_overrides = {**(config_overrides or {}), **scaling_overrides} optimizer = get_megatron_optimizer( config, @@ -3811,6 +3816,10 @@ def trace_handler(p): return iteration, num_floating_point_operations_so_far +def _should_allow_scaling_policy_eval(args): + return getattr(args, 'scaling_recipe', None) == 'depth_mup' + + def evaluate( forward_step_func, data_iterator, @@ -3874,7 +3883,7 @@ def evaluate( if eval_iters is None: eval_iters = args.eval_iters - with torch.no_grad(): + with allow_scaling_policy_eval(_should_allow_scaling_policy_eval(args)), torch.no_grad(): iteration = 0 if verbose: print_rank_0(f'Evaluating on {eval_iters * eval_batch_size} samples') diff --git a/megatron/training/yaml_arguments.py b/megatron/training/yaml_arguments.py index d44f4d31822..5da76fe2f31 100644 --- a/megatron/training/yaml_arguments.py +++ b/megatron/training/yaml_arguments.py @@ -16,8 +16,14 @@ import torch.nn.functional as F +from megatron.core.parameterization import build_scaling_context, sync_legacy_mup_fields from megatron.core.transformer import TransformerConfig, MLATransformerConfig from megatron.core.utils import get_torch_version, is_torch_min_version +from megatron.training.arguments import ( + validate_depth_mup_optimizer_support, + validate_muon_scalar_optimizer_support, + warn_deprecated_mup_aliases, +) # Taken from https://stackoverflow.com/questions/65414773/parse-environment-variable-from-yaml-with-pyyaml # Allows for yaml to use environment variables @@ -38,6 +44,24 @@ def env_constructor(loader, node): "bfloat16" : torch.bfloat16 } +DEFAULTABLE_SCALING_FIELDS = { + 'scaling_recipe', + 'scaling_base_hidden_size', + 'scaling_base_num_layers', + 'scaling_base_head_dim', + 'scaling_residual_branch_depth_power', + 'scaling_hidden_lr_depth_power', + 'scaling_block_out_proj_init_depth_power', + 'use_mup', + 'mup_width_mult', + 'mup_base_hidden_size', + 'mup_embedding_mult', + 'mup_output_mult', + 'mup_base_head_dim', + 'mup_attn_scale_power', +} + + def validate_yaml(args, defaults={}): # This is for legacy script env var setting @@ -246,6 +270,13 @@ def validate_yaml(args, defaults={}): assert args.language_model.hidden_size % args.language_model.num_attention_heads == 0 args.language_model.kv_channels = args.language_model.hidden_size // args.language_model.num_attention_heads + if getattr(args.language_model, 'mup_width_mult', 1.0) != 1.0: + args.language_model._mup_width_mult_explicit = True + warn_deprecated_mup_aliases(args.language_model) + sync_legacy_mup_fields( + args.language_model, build_scaling_context(args.language_model) + ) + #TODO: Implement arguments for encoder-decoder if args.seq_length is not None: assert args.encoder_seq_length is None @@ -344,6 +375,10 @@ def validate_yaml(args, defaults={}): #TODO: Added as much of the global initialization requires the model parallel arguments args = SimpleNamespace(**args.__dict__, **args.model_parallel.__dict__) args = SimpleNamespace(**args.__dict__, **args.language_model.__dict__) + validate_depth_mup_optimizer_support(args) + validate_muon_scalar_optimizer_support(args) + warn_deprecated_mup_aliases(args) + sync_legacy_mup_fields(args, build_scaling_context(args)) # For GPT Layer spec in pretrain_gpt args.num_experts = args.language_model.num_moe_experts @@ -380,6 +415,13 @@ def core_config_from_args(args, dataclass=TransformerConfig): for f in dataclasses.fields(dataclass): if hasattr(args, f.name): kw_args[f.name] = getattr(args, f.name) + elif f.name in DEFAULTABLE_SCALING_FIELDS: + if f.default is not dataclasses.MISSING: + kw_args[f.name] = f.default + elif f.default_factory is not dataclasses.MISSING: + kw_args[f.name] = f.default_factory() + else: + raise Exception(f"Missing argument {f.name} for {str(dataclass)} config") else: raise Exception(f"Missing argument {f.name} for {str(dataclass)} config") return kw_args @@ -438,4 +480,3 @@ def load_yaml(yaml_path): getattr(config_namespace, "global_batch_size", None) is not None ) return config_namespace - diff --git a/tests/unit_tests/test_lion_optimizer.py b/tests/unit_tests/test_lion_optimizer.py index b0df91073ed..c942ee5bdfd 100644 --- a/tests/unit_tests/test_lion_optimizer.py +++ b/tests/unit_tests/test_lion_optimizer.py @@ -13,12 +13,18 @@ import torch import torch.nn as nn +import megatron.core.optimizer as opt_module from megatron.core.optimizer import ( HAVE_EMERGING_OPTIMIZERS, OptimizerConfig, _get_megatron_optimizer_based_on_param_groups, _get_param_groups, ) +from megatron.core.optimizer.emerging_optimizers import ( + _EMERGING_OPTIMIZERS, + _default_betas_for_eopt, + _muon_default_param_overrides, +) from megatron.core.optimizer.optimizer import FP32Optimizer requires_emerging_optimizers = pytest.mark.skipif( @@ -79,6 +85,97 @@ def test_lion_config_defaults(self): assert config.lion_beta2 == 0.98 assert config.muon_scalar_optimizer == "adam" + def test_muon_scalar_optimizer_controls_nonlinear_param_override(self): + """Muon scalar optimizer selection should flow into nonlinear param overrides.""" + entry = _EMERGING_OPTIMIZERS["muon"] + config = OptimizerConfig(muon_scalar_optimizer="lion") + + overrides = entry.config_to_param_overrides(config) + + assert len(overrides) == 1 + ((_, override),) = overrides.items() + assert override["optimizer"] == "lion" + + def test_muon_scalar_optimizer_routes_lion_groups_to_lion_entry(self): + """Muon scalar-optimizer overrides must create a real Lion optimizer bucket.""" + model = SimpleModel() + config = OptimizerConfig( + optimizer="muon", + lr=1e-4, + muon_scalar_optimizer="lion", + adam_beta1=0.81, + adam_beta2=0.88, + lion_beta1=0.91, + lion_beta2=0.97, + ) + recorded = [] + + def fake_create(_config, _groups, eopt_name, _model_chunks, _pg_collection): + if eopt_name == "lion": + recorded.append((eopt_name, _default_betas_for_eopt(eopt_name, _config))) + else: + recorded.append((eopt_name, None)) + return SimpleNamespace(param_groups=[]), (lambda *_args, **_kwargs: None) + + fake_pg_collection = SimpleNamespace(mp=None, tp=None, tp_ep_pp=None) + fake_muon_entry = SimpleNamespace( + config_to_param_overrides=_muon_default_param_overrides, + default_param_overrides={}, + optimizer_cls=object, + init_state_fn=lambda *_args, **_kwargs: None, + config_to_kwargs=None, + ) + fake_lion_entry = SimpleNamespace( + config_to_param_overrides=None, + default_param_overrides={}, + optimizer_cls=object, + init_state_fn=lambda *_args, **_kwargs: None, + config_to_kwargs=None, + ) + + with ( + patch("torch.distributed.get_world_size", return_value=1), + patch( + "torch.distributed.all_gather_object", + lambda output_list, obj: output_list.__setitem__(0, obj), + ), + patch.object(opt_module, "HAVE_EMERGING_OPTIMIZERS", True), + patch.dict( + opt_module._EMERGING_OPTIMIZERS, + {"muon": fake_muon_entry, "lion": fake_lion_entry}, + clear=False, + ), + patch.object(opt_module, "_create_emerging_optimizer", side_effect=fake_create), + patch.object( + opt_module, + "FP32Optimizer", + side_effect=lambda optimizer, *_args, **_kwargs: optimizer, + ), + patch.object(opt_module, "ChainedOptimizer", side_effect=lambda optimizers: optimizers), + ): + results = opt_module._get_megatron_emerging_optimizer( + config=config, + model_chunks=[model], + config_overrides={}, + pg_collection=fake_pg_collection, + ) + + assert set(recorded) == {("muon", None), ("lion", (0.91, 0.97))} + assert len(results) == 2 + + def test_default_emerging_lion_betas_use_lion_betas(self): + """Shared beta selection must keep Lion on lion_beta{1,2}.""" + config = OptimizerConfig( + optimizer="muon", + lr=1e-4, + adam_beta1=0.81, + adam_beta2=0.88, + lion_beta1=0.91, + lion_beta2=0.97, + ) + + assert _default_betas_for_eopt("lion", config) == (0.91, 0.97) + @patch("torch.distributed.get_world_size", return_value=1) @patch( "torch.distributed.all_gather_object", diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 56af8545042..9e16c731b3c 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -23,6 +23,7 @@ get_megatron_optimizer, get_standard_config_overrides, ) +from megatron.core.optimizer.optimizer import MegatronOptimizer, get_param_group_identifier_tuple from megatron.core.optimizer_param_scheduler import ParamGroupOverride from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer import TransformerConfig @@ -69,6 +70,91 @@ def forward(self, x): return x +def test_param_group_identifier_tuple_tolerates_missing_optional_keys(): + group = {"wd_mult": 1.0, "lr_mult": 1.0, "is_expert_parallel": False, "is_decoupled_lr": False} + + ident = get_param_group_identifier_tuple(group) + + assert ident == (1.0, 1.0, False, False, None, None) + + +def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): + group = { + "pre_wd_mult": 0.0, + "pre_lr_mult": 0.5, + "pre_is_expert_parallel": False, + "pre_is_decoupled_lr": True, + } + + ident = get_param_group_identifier_tuple(group) + + assert ident == (0.0, 0.5, False, True, None, None) + + +def test_param_group_identifier_tuple_defaults_missing_legacy_fields(): + group = {} + + ident = get_param_group_identifier_tuple(group) + + assert ident == (1.0, 1.0, False, False, None, None) + + +def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): + current_group = { + "wd_mult": 1.0, + "lr_mult": 1.0, + "is_expert_parallel": False, + "is_decoupled_lr": False, + "max_lr": 2e-4, + "min_lr": 2e-6, + "eps": 1e-8, + "optimizer": "adam", + "params": [0], + } + checkpoint_group = { + "wd_mult": 1.0, + "lr_mult": 1.0, + "is_expert_parallel": False, + "is_decoupled_lr": False, + "max_lr": 1e-4, + "min_lr": 1e-6, + "eps": 1e-8, + "optimizer": "adam", + "params": [7], + } + + assert get_param_group_identifier_tuple(current_group) == get_param_group_identifier_tuple( + checkpoint_group + ) + + reordered_groups = MegatronOptimizer._filter_and_reorder_param_groups( + [current_group], [checkpoint_group] + ) + + assert reordered_groups == [checkpoint_group] + + +def test_param_group_matching_normalizes_legacy_missing_identifier_fields(): + current_group = { + "wd_mult": 1.0, + "lr_mult": 1.0, + "is_expert_parallel": False, + "is_decoupled_lr": False, + "params": [0], + } + legacy_checkpoint_group = {"params": [7]} + + assert get_param_group_identifier_tuple(current_group) == get_param_group_identifier_tuple( + legacy_checkpoint_group + ) + + reordered_groups = MegatronOptimizer._filter_and_reorder_param_groups( + [current_group], [legacy_checkpoint_group] + ) + + assert reordered_groups == [legacy_checkpoint_group] + + @patch('torch.distributed.get_world_size', return_value=1) @patch( 'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj) diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index f1d99cad1e6..ac303677ea0 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -9,19 +9,54 @@ 4. LR override computation """ +import argparse +import dataclasses +import functools import logging import math +import warnings +from types import SimpleNamespace from unittest.mock import patch import pytest import torch -from megatron.core.optimizer import get_mup_config_overrides, get_standard_config_overrides +from megatron.core.optimizer import ( + get_mup_config_overrides, + get_scaling_config_overrides, + get_standard_config_overrides, +) from megatron.core.optimizer.optimizer_config import OptimizerConfig from megatron.core.optimizer_param_scheduler import combine_param_group_overrides +from megatron.core.parameterization import ( + allow_scaling_policy_eval, + build_legacy_mup_training_policy, + build_model_scaling_policy, + build_scaling_context, + build_training_scaling_policy, +) +from megatron.core.transformer.mlp import MLP, MLPSubmodules from megatron.core.transformer.multi_token_prediction import process_mtp_loss from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, + _get_mlp_builder_module, +) from megatron.core.utils import init_method_normal, mup_scaled_init_method_normal +from megatron.training.arguments import ( + add_megatron_arguments, + validate_depth_mup_optimizer_support, + validate_muon_scalar_optimizer_support, +) +from megatron.training.yaml_arguments import core_config_from_args + + +def _combined_override_for_param(overrides, param, param_name): + matches = [ + override for param_key, override in overrides.items() if param_key.matches(param, param_name) + ] + return combine_param_group_overrides(matches) class TestMuPConfigValidation: @@ -38,6 +73,8 @@ def test_mup_defaults_base_hidden_size(self): ) assert config.mup_base_hidden_size == 512 assert config.mup_width_mult == 1.0 + assert config.scaling_recipe == 'mup' + assert config.scaling_base_hidden_size == 512 def test_mup_width_mult_calculation(self): """width_mult = hidden_size / base_hidden_size.""" @@ -49,6 +86,8 @@ def test_mup_width_mult_calculation(self): mup_base_hidden_size=256, ) assert config.mup_width_mult == 4.0 + assert config.scaling_recipe == 'mup' + assert config.scaling_base_hidden_size == 256 def test_mup_width_mult_fractional(self): """width_mult can be fractional (smaller than base).""" @@ -67,6 +106,8 @@ def test_mup_backward_compatible(self): assert config.use_mup is False assert config.mup_width_mult == 1.0 assert config.mup_base_hidden_size is None + assert config.scaling_recipe == 'none' + assert config.scaling_base_hidden_size is None def test_mup_base_hidden_size_must_be_positive(self): """mup_base_hidden_size must be positive.""" @@ -80,6 +121,721 @@ def test_mup_base_hidden_size_must_be_positive(self): ) assert "positive" in str(exc_info.value).lower() + def test_scaling_recipe_mup_sets_legacy_fields(self): + """Canonical MuP fields populate the legacy fields used by existing call sites.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_head_dim=64, + ) + + assert config.use_mup is True + assert config.mup_base_hidden_size == 256 + assert config.mup_base_head_dim == 64 + assert config.mup_width_mult == pytest.approx(4.0) + assert config.scaling_base_hidden_size == 256 + assert config.scaling_base_head_dim == 64 + + def test_legacy_mup_fields_resolve_to_canonical_recipe(self): + """Legacy MuP flags remain compatible but are not separate state.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + use_mup=True, + mup_base_hidden_size=256, + mup_base_head_dim=64, + ) + + assert config.scaling_recipe == 'mup' + assert config.scaling_base_hidden_size == 256 + assert config.scaling_base_head_dim == 64 + assert build_scaling_context(config).width_mult == pytest.approx(4.0) + + def test_scaling_recipe_none_rejects_scaling_overrides(self): + """Scaling fields cannot silently affect standard parameterization.""" + with pytest.raises(ValueError, match="Scaling overrides"): + TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='none', + scaling_base_hidden_size=256, + ) + + def test_use_mup_conflicts_with_scaling_recipe_none(self): + """The deprecated MuP boolean cannot override an explicit canonical recipe.""" + with pytest.raises(ValueError, match="conflicts"): + TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='none', + use_mup=True, + ) + + def test_canonical_and_legacy_base_hidden_must_match(self): + """Canonical and deprecated base hidden-size fields are aliases.""" + with pytest.raises(ValueError, match="conflicts"): + TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + mup_base_hidden_size=512, + ) + + def test_deprecated_width_mult_must_match_derived_value(self): + """mup_width_mult is accepted only when it matches the derived width.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + mup_width_mult=4.0, + ) + assert config.mup_width_mult == pytest.approx(4.0) + + with pytest.raises(ValueError, match="must match the derived"): + TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + mup_width_mult=2.0, + ) + + def test_scaling_override_without_recipe_is_rejected(self): + """Base scaling fields do not implicitly enable MuP.""" + with pytest.raises(ValueError, match="Scaling overrides"): + TransformerConfig( + hidden_size=1024, + num_layers=4, + num_attention_heads=16, + scaling_base_hidden_size=256, + ) + + def test_depth_mup_resolves_distinct_recipe_defaults(self): + """Depth-MuP is width-MuP-family behavior without setting the legacy use_mup bit.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_base_head_dim=64, + ) + context = build_scaling_context(config) + + assert config.scaling_recipe == 'depth_mup' + assert config.use_mup is False + assert context.uses_width_mup is True + assert context.is_depth_mup is True + assert context.width_mult == pytest.approx(4.0) + assert context.depth_mult == pytest.approx(2.0) + assert context.base_hidden_size == 256 + assert context.base_num_layers == 6 + assert context.base_head_dim == 64 + assert context.residual_branch_depth_power == pytest.approx(-1.0) + assert context.hidden_lr_depth_power == pytest.approx(0.0) + assert context.block_out_proj_init_depth_power == pytest.approx(0.5) + assert context.output_mult == pytest.approx(0.25) + + def test_depth_mup_manual_overrides_can_zero_recipe_defaults(self): + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_residual_branch_depth_power=0.0, + scaling_block_out_proj_init_depth_power=0.0, + ) + context = build_scaling_context(config) + + assert context.residual_branch_depth_power == pytest.approx(0.0) + assert context.block_out_proj_init_depth_power == pytest.approx(0.0) + + def test_mup_does_not_inherit_depth_mup_defaults(self): + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + ) + context = build_scaling_context(config) + + assert context.recipe == 'mup' + assert context.depth_mult == pytest.approx(2.0) + assert context.residual_branch_depth_power == pytest.approx(0.0) + assert context.hidden_lr_depth_power == pytest.approx(0.0) + assert context.block_out_proj_init_depth_power == pytest.approx(0.0) + + def test_depth_mup_conflicts_with_legacy_use_mup_alias(self): + with pytest.raises(ValueError, match='conflicts with --use-mup'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + use_mup=True, + ) + + def test_depth_mup_rejects_unsupported_attention_and_moe_surfaces(self): + with pytest.raises(NotImplementedError, match='multi_latent_attention'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + multi_latent_attention=True, + ) + with pytest.raises(NotImplementedError, match='experimental attention variants'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + experimental_attention_variant='gated_delta_net', + linear_attention_freq=1, + ) + with pytest.raises(NotImplementedError, match='MoE depth transfer'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + num_moe_experts=4, + ) + with pytest.raises(NotImplementedError, match='Hybrid/Mamba'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + is_hybrid_model=True, + ) + with pytest.raises(NotImplementedError, match='MTP depth transfer'): + TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + mtp_num_layers=1, + ) + + def test_muon_scalar_optimizer_gate_rejects_invalid_yaml_value(self): + with pytest.raises(ValueError, match='muon_scalar_optimizer'): + validate_muon_scalar_optimizer_support( + SimpleNamespace(muon_scalar_optimizer='sgd') + ) + + def test_muon_scalar_optimizer_gate_accepts_supported_values(self): + validate_muon_scalar_optimizer_support(SimpleNamespace(muon_scalar_optimizer='adam')) + validate_muon_scalar_optimizer_support(SimpleNamespace(muon_scalar_optimizer='lion')) + validate_muon_scalar_optimizer_support(SimpleNamespace()) + + +class TestScalingRecipeSurfaces: + """Tests for public config surfaces that feed the scaling context.""" + + SCALING_FIELD_NAMES = { + 'scaling_recipe', + 'scaling_base_hidden_size', + 'scaling_base_num_layers', + 'scaling_base_head_dim', + 'scaling_residual_branch_depth_power', + 'scaling_hidden_lr_depth_power', + 'scaling_block_out_proj_init_depth_power', + 'use_mup', + 'mup_width_mult', + 'mup_base_hidden_size', + 'mup_embedding_mult', + 'mup_output_mult', + 'mup_base_head_dim', + 'mup_attn_scale_power', + } + + def test_cli_parser_accepts_canonical_scaling_args(self): + """The explicit scaling arg group owns canonical and legacy MuP flags.""" + parser = argparse.ArgumentParser(allow_abbrev=False) + add_megatron_arguments(parser) + + args, _ = parser.parse_known_args( + [ + '--scaling-recipe', + 'mup', + '--scaling-base-hidden-size', + '256', + '--mup-base-head-dim', + '64', + ] + ) + + assert args.scaling_recipe == 'mup' + assert args.scaling_base_hidden_size == 256 + assert args.mup_base_head_dim == 64 + assert args.mup_width_mult == 1.0 + + depth_args, _ = parser.parse_known_args( + [ + '--scaling-recipe', + 'depth_mup', + '--scaling-base-hidden-size', + '256', + '--scaling-base-num-layers', + '6', + '--scaling-residual-branch-depth-power', + '-1.0', + ] + ) + + assert depth_args.scaling_recipe == 'depth_mup' + assert depth_args.scaling_base_num_layers == 6 + assert depth_args.scaling_residual_branch_depth_power == pytest.approx(-1.0) + + def test_cli_explicit_mup_width_mult_one_is_validated(self): + """Explicit legacy width multiplier must match the derived value, even at 1.0.""" + parser = argparse.ArgumentParser(allow_abbrev=False) + add_megatron_arguments(parser) + + args, _ = parser.parse_known_args( + [ + '--scaling-recipe', + 'mup', + '--scaling-base-hidden-size', + '256', + '--mup-width-mult', + '1.0', + ] + ) + args.hidden_size = 1024 + + with pytest.raises(ValueError, match="must match the derived"): + build_scaling_context(args) + + def test_yaml_core_config_defaults_missing_scaling_fields(self): + """Existing YAML files may omit the new scaling fields.""" + values = {} + for field in dataclasses.fields(TransformerConfig): + if field.name in self.SCALING_FIELD_NAMES: + continue + if field.default is not dataclasses.MISSING: + values[field.name] = field.default + elif field.default_factory is not dataclasses.MISSING: + values[field.name] = field.default_factory() + elif field.type is int: + values[field.name] = 1 + else: + values[field.name] = None + values['hidden_size'] = 512 + values['num_layers'] = 2 + values['num_attention_heads'] = 8 + + kwargs = core_config_from_args(SimpleNamespace(**values), TransformerConfig) + + assert kwargs['scaling_recipe'] is None + assert kwargs['scaling_base_hidden_size'] is None + assert kwargs['mup_width_mult'] == 1.0 + + def test_yaml_default_width_mult_is_not_treated_as_explicit(self): + """Full legacy YAML files may materialize the old default width multiplier.""" + yaml_args = SimpleNamespace( + hidden_size=1024, + scaling_recipe=None, + scaling_base_hidden_size=None, + scaling_base_head_dim=None, + use_mup=True, + mup_width_mult=1.0, + mup_base_hidden_size=256, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=None, + mup_attn_scale_power=1.0, + ) + + context = build_scaling_context(yaml_args) + + assert context.width_mult == pytest.approx(4.0) + + def test_scaling_context_matches_legacy_checkpoint_and_canonical_args(self): + """Checkpoint compatibility compares effective scaling, not flag spelling.""" + legacy_checkpoint_args = SimpleNamespace( + hidden_size=1024, + use_mup=True, + mup_base_hidden_size=256, + mup_width_mult=1.0, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=64, + mup_attn_scale_power=1.0, + ) + canonical_args = SimpleNamespace( + hidden_size=1024, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_head_dim=64, + use_mup=False, + mup_width_mult=1.0, + mup_base_hidden_size=None, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=None, + mup_attn_scale_power=1.0, + ) + + assert build_scaling_context( + legacy_checkpoint_args + ) == build_scaling_context(canonical_args) + + def test_checkpoint_scaling_sync_populates_canonical_fields(self): + """Old checkpoints with only legacy MuP fields become canonical before copy.""" + from megatron.training.checkpointing import _sync_checkpoint_scaling_args + + legacy_checkpoint_args = SimpleNamespace( + hidden_size=1024, + use_mup=True, + mup_base_hidden_size=256, + mup_width_mult=1.0, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=64, + mup_attn_scale_power=1.0, + ) + + _sync_checkpoint_scaling_args(legacy_checkpoint_args) + + assert legacy_checkpoint_args.scaling_recipe == 'mup' + assert legacy_checkpoint_args.scaling_base_hidden_size == 256 + assert legacy_checkpoint_args.scaling_base_head_dim == 64 + assert legacy_checkpoint_args.mup_width_mult == pytest.approx(4.0) + + def test_check_checkpoint_args_compares_effective_scaling_context(self, monkeypatch): + """The real checkpoint check accepts legacy and canonical spellings if equivalent.""" + from megatron.training import checkpointing + + runtime_args = SimpleNamespace( + num_layers=2, + hidden_size=1024, + num_attention_heads=16, + add_position_embedding=True, + vocab_file=None, + data_parallel_random_init=False, + phase_transition_iterations=None, + use_dist_ckpt=False, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_head_dim=64, + use_mup=False, + mup_width_mult=1.0, + mup_base_hidden_size=None, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=None, + mup_attn_scale_power=1.0, + ) + checkpoint_args = SimpleNamespace( + num_layers=2, + hidden_size=1024, + num_attention_heads=16, + add_position_embedding=True, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + use_mup=True, + mup_base_hidden_size=256, + mup_width_mult=1.0, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=64, + mup_attn_scale_power=1.0, + ) + monkeypatch.setattr(checkpointing, 'get_args', lambda: runtime_args) + monkeypatch.setattr(checkpointing, 'get_checkpoint_version', lambda: 3.0) + + checkpointing.check_checkpoint_args(checkpoint_args) + + def test_load_checkpoint_args_clears_optional_scaling_fields(self, monkeypatch): + """use-checkpoint-args must clear stale optional canonical fields.""" + from megatron.training import checkpointing + + checkpoint_args = SimpleNamespace( + num_layers=2, + hidden_size=1024, + num_attention_heads=16, + use_mup=True, + mup_base_hidden_size=256, + mup_width_mult=1.0, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_attn_scale_power=1.0, + ) + state_dict = {'args': checkpoint_args, 'iteration': 7} + monkeypatch.setattr( + checkpointing, + '_load_base_checkpoint', + lambda *args, **kwargs: (state_dict, 'model_optim_rng.pt', False, None), + ) + runtime_args = SimpleNamespace( + load='dummy-checkpoint', + iteration=0, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_head_dim=64, + mup_base_head_dim=64, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + + checkpointing.load_args_from_checkpoint(runtime_args) + + assert runtime_args.iteration == 7 + assert runtime_args.scaling_recipe == 'mup' + assert runtime_args.scaling_base_hidden_size == 256 + assert runtime_args.scaling_base_head_dim is None + assert runtime_args.mup_base_head_dim is None + assert runtime_args.mup_width_mult == pytest.approx(4.0) + + def test_load_checkpoint_args_restores_depth_mup_scaling_fields(self, monkeypatch): + """use-checkpoint-args must preserve the canonical depth-MuP surface.""" + from megatron.training import checkpointing + + checkpoint_args = SimpleNamespace( + num_layers=12, + hidden_size=1024, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_base_head_dim=64, + scaling_residual_branch_depth_power=-1.0, + scaling_hidden_lr_depth_power=0.0, + scaling_block_out_proj_init_depth_power=0.5, + use_mup=False, + mup_width_mult=1.0, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_attn_scale_power=1.0, + ) + state_dict = {'args': checkpoint_args, 'iteration': 11} + monkeypatch.setattr( + checkpointing, + '_load_base_checkpoint', + lambda *args, **kwargs: (state_dict, 'model_optim_rng.pt', False, None), + ) + runtime_args = SimpleNamespace( + load='dummy-checkpoint', + iteration=0, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + + checkpointing.load_args_from_checkpoint(runtime_args) + + assert runtime_args.iteration == 11 + assert runtime_args.use_mup is False + assert runtime_args.scaling_recipe == 'depth_mup' + assert runtime_args.scaling_base_hidden_size == 256 + assert runtime_args.scaling_base_num_layers == 6 + assert runtime_args.scaling_base_head_dim == 64 + assert runtime_args.scaling_residual_branch_depth_power == pytest.approx(-1.0) + assert runtime_args.scaling_hidden_lr_depth_power == pytest.approx(0.0) + assert runtime_args.scaling_block_out_proj_init_depth_power == pytest.approx(0.5) + + def test_distributed_resume_preprocessing_tolerates_missing_optional_group_keys(self): + """Distributed resume sorting must tolerate groups without eps/optimizer.""" + from megatron.training.training import preprocess_common_state_dict + + common_state_dict = { + 'args': SimpleNamespace( + use_distributed_optimizer=True, + rank=3, + local_rank=1, + ), + 'optimizer': { + 'optimizer': { + 'param_groups': [ + { + 'wd_mult': 1.0, + 'lr_mult': 1.0, + 'is_expert_parallel': False, + 'is_decoupled_lr': False, + 'max_lr': 1.0e-3, + 'min_lr': 1.0e-5, + 'eps': 1.0e-8, + 'optimizer': 'adam', + 'params': [1], + }, + { + 'wd_mult': 1.0, + 'lr_mult': 1.0, + 'is_expert_parallel': False, + 'is_decoupled_lr': False, + 'max_lr': 1.0e-3, + 'min_lr': 1.0e-5, + 'params': [0], + }, + ] + } + }, + } + + preprocessed = preprocess_common_state_dict(common_state_dict) + + param_groups = preprocessed['optimizer']['optimizer']['param_groups'] + assert [group['params'] for group in param_groups] == [[0], [1]] + assert 'rank' not in preprocessed['args'] + assert 'local_rank' not in preprocessed['args'] + + def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch): + """Old no-scaling checkpoints must clear stale CLI scaling state.""" + from megatron.training import checkpointing + + checkpoint_args = SimpleNamespace(hidden_size=1024) + state_dict = {'args': checkpoint_args, 'iteration': 3} + monkeypatch.setattr( + checkpointing, + '_load_base_checkpoint', + lambda *args, **kwargs: (state_dict, 'model_optim_rng.pt', False, None), + ) + runtime_args = SimpleNamespace( + load='dummy-checkpoint', + iteration=0, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_base_head_dim=64, + scaling_residual_branch_depth_power=-1.0, + scaling_hidden_lr_depth_power=0.25, + scaling_block_out_proj_init_depth_power=0.5, + use_mup=True, + mup_width_mult=4.0, + _mup_width_mult_explicit=True, + mup_base_hidden_size=256, + mup_embedding_mult=2.0, + mup_output_mult=0.25, + mup_base_head_dim=64, + mup_attn_scale_power=-0.5, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + + checkpointing.load_args_from_checkpoint(runtime_args) + + assert runtime_args.iteration == 3 + assert runtime_args.scaling_recipe == 'none' + assert runtime_args.scaling_base_hidden_size is None + assert runtime_args.scaling_base_num_layers is None + assert runtime_args.scaling_base_head_dim is None + assert runtime_args.scaling_residual_branch_depth_power is None + assert runtime_args.scaling_hidden_lr_depth_power is None + assert runtime_args.scaling_block_out_proj_init_depth_power is None + assert runtime_args.use_mup is False + assert runtime_args.mup_width_mult == 1.0 + assert runtime_args._mup_width_mult_explicit is False + assert runtime_args.mup_base_hidden_size is None + assert runtime_args.mup_embedding_mult == 1.0 + assert runtime_args.mup_output_mult == 1.0 + assert runtime_args.mup_base_head_dim is None + assert runtime_args.mup_attn_scale_power == 1.0 + assert build_scaling_context(runtime_args).recipe == 'none' + + def test_checkpoint_derived_width_mult_does_not_warn_as_deprecated_cli(self): + """Normalized checkpoint state should not look like user-provided --mup-width-mult.""" + from megatron.training.arguments import warn_deprecated_mup_aliases + + checkpoint_derived_args = SimpleNamespace( + rank=0, + use_mup=False, + mup_base_hidden_size=None, + mup_base_head_dim=None, + mup_width_mult=4.0, + _mup_width_mult_explicit=False, + ) + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + warn_deprecated_mup_aliases(checkpoint_derived_args) + + assert len(caught_warnings) == 0 + + def test_false_width_mult_marker_overrides_stale_non_default_value(self): + """Marker-present False means checkpoint/internal provenance, not explicit user input.""" + checkpoint_derived_args = SimpleNamespace( + hidden_size=1024, + scaling_recipe='none', + scaling_base_hidden_size=None, + scaling_base_head_dim=None, + use_mup=False, + mup_width_mult=4.0, + _mup_width_mult_explicit=False, + mup_base_hidden_size=None, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=None, + mup_attn_scale_power=1.0, + ) + + assert build_scaling_context(checkpoint_derived_args).recipe == 'none' + + def test_checkpoint_derived_legacy_aliases_do_not_warn_as_deprecated_cli(self): + """Checkpoint-synced legacy fields should not masquerade as user CLI aliases.""" + from megatron.training.arguments import warn_deprecated_mup_aliases + + checkpoint_derived_args = SimpleNamespace( + rank=0, + use_mup=True, + _use_mup_explicit=False, + mup_base_hidden_size=256, + _mup_base_hidden_size_explicit=False, + mup_base_head_dim=64, + _mup_base_head_dim_explicit=False, + mup_width_mult=4.0, + _mup_width_mult_explicit=False, + ) + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + warn_deprecated_mup_aliases(checkpoint_derived_args) + + assert len(caught_warnings) == 0 + + def test_user_provided_legacy_aliases_warn_as_deprecated_cli(self): + """Real user-provided legacy aliases should still produce a deprecation warning.""" + from megatron.training.arguments import warn_deprecated_mup_aliases + + user_args = SimpleNamespace( + rank=0, + use_mup=True, + _use_mup_explicit=True, + mup_base_hidden_size=256, + _mup_base_hidden_size_explicit=True, + mup_base_head_dim=64, + _mup_base_head_dim_explicit=True, + mup_width_mult=1.0, + _mup_width_mult_explicit=True, + ) + + with pytest.warns(UserWarning) as caught_warnings: + warn_deprecated_mup_aliases(user_args) + + warning_text = str(caught_warnings[0].message) + assert '--use-mup' in warning_text + assert '--mup-base-hidden-size' in warning_text + assert '--mup-base-head-dim' in warning_text + assert '--mup-width-mult' in warning_text + class TestMuPInitMethods: """Tests for MuP initialization methods.""" @@ -169,7 +925,7 @@ class TestMuPWarnings: def test_mup_warns_with_custom_init_method(self): """Warn when MuP is enabled and init_method is user-provided.""" - with pytest.warns(UserWarning, match="use_mup is enabled"): + with pytest.warns(UserWarning, match="scaling recipe 'mup' is enabled"): TransformerConfig( hidden_size=512, num_layers=4, @@ -181,7 +937,7 @@ def test_mup_warns_with_custom_init_method(self): def test_mup_warns_with_custom_output_layer_init_method(self): """Warn when MuP is enabled and output_layer_init_method is user-provided.""" - with pytest.warns(UserWarning, match="use_mup is enabled"): + with pytest.warns(UserWarning, match="scaling recipe 'mup' is enabled"): TransformerConfig( hidden_size=512, num_layers=4, @@ -195,6 +951,20 @@ def test_mup_warns_with_custom_output_layer_init_method(self): class TestMuPLRScaling: """Tests for MuP learning rate and Adam epsilon scaling.""" + def test_mup_overrides_route_through_training_scaling_policy(self): + """The new policy seam preserves the legacy MuP override surface.""" + optimizer_config = OptimizerConfig(lr=1e-3, min_lr=1e-5) + width_mult = 4.0 + + legacy_overrides = get_mup_config_overrides(optimizer_config, width_mult) + policy_overrides = get_scaling_config_overrides( + optimizer_config, + build_legacy_mup_training_policy(mup_width_mult=width_mult, optimizer_type='adam'), + ) + + assert legacy_overrides.keys() == policy_overrides.keys() + assert list(legacy_overrides.values()) == list(policy_overrides.values()) + def test_mup_lr_override_computation(self): """Hidden LR and Adam eps scale as 1/width_mult.""" optimizer_config = OptimizerConfig(lr=1e-3, min_lr=1e-5) @@ -355,6 +1125,47 @@ def test_mup_with_decoupled_lr_scales_hidden_only_for_lr(self): class TestMuPConfigIntegration: """Integration tests for MuP config with init methods.""" + def test_model_scaling_policy_matches_legacy_config_fields(self): + """Model policy exposes the same effective MuP values as TransformerConfig.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=8, + num_attention_heads=16, + use_mup=True, + mup_base_hidden_size=256, + mup_embedding_mult=3.0, + ) + policy = build_model_scaling_policy(config) + + assert policy.enabled is True + assert policy.context.width_mult == pytest.approx(config.mup_width_mult) + assert policy.context.output_mult == pytest.approx(config.mup_output_mult) + assert policy.context.embedding_mult == pytest.approx(config.mup_embedding_mult) + assert config.softmax_scale == pytest.approx( + policy.resolve_attention_softmax_scale( + softmax_scale=None, kv_channels=config.kv_channels + ) + ) + + def test_model_scaling_policy_tracks_post_init_multiplier_mutations(self): + """Policy resolution should preserve legacy live reads of mutable config fields.""" + config = TransformerConfig( + hidden_size=1024, + num_layers=8, + num_attention_heads=16, + use_mup=True, + mup_base_hidden_size=256, + ) + logits = torch.ones(2, 4) + embeddings = torch.ones(2, 4) + + config.mup_output_mult = 0.25 + config.mup_embedding_mult = 3.0 + policy = build_model_scaling_policy(config) + + assert torch.equal(policy.scale_output_logits(logits), logits * 0.25) + assert torch.equal(policy.scale_embedding_activations(embeddings), embeddings * 3.0) + def test_mup_output_layer_init(self): """Output layer init should also scale with MuP.""" config = TransformerConfig( @@ -376,10 +1187,294 @@ def test_mup_output_layer_init(self): assert abs(actual_std - expected_std) < expected_std * 0.15 + def test_depth_mup_residual_multiplier_exact_depth_factors(self): + base_depth_config = TransformerConfig( + hidden_size=512, + num_layers=12, + num_attention_heads=8, + scaling_recipe='depth_mup', + scaling_base_hidden_size=512, + scaling_base_num_layers=12, + ) + double_depth_config = TransformerConfig( + hidden_size=512, + num_layers=24, + num_attention_heads=8, + scaling_recipe='depth_mup', + scaling_base_hidden_size=512, + scaling_base_num_layers=12, + ) + + assert build_model_scaling_policy(base_depth_config).residual_branch_multiplier == pytest.approx(1.0) + assert build_model_scaling_policy(double_depth_config).residual_branch_multiplier == pytest.approx(0.5) + + def test_depth_mup_default_block_output_init_rebases_to_base_depth(self): + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + ) + policy = build_model_scaling_policy(config) + init_fn = policy.dense_block_output_init_method( + default_init_method=config.output_layer_init_method, + init_method_std=config.init_method_std, + num_layers=config.num_layers, + is_hybrid_model=config.is_hybrid_model, + output_layer_init_method_is_user_provided=False, + ) + weights = torch.empty(200_000) + init_fn(weights) + + expected_std = config.init_method_std / ( + math.sqrt(2 * config.scaling_base_num_layers) * math.sqrt(policy.context.width_mult) + ) + assert abs(weights.std().item() - expected_std) < expected_std * 0.05 + + def test_plain_mlp_requires_explicit_block_output_init_scaling_opt_in(self): + class DummyLinear(torch.nn.Module): + def __init__(self, init_method): + super().__init__() + self.init_method = init_method + + def forward(self, hidden_states): + return hidden_states, None + + def backward_dw(self): + return None + + def fc1_builder(input_size, output_size, *, init_method, **kwargs): + return DummyLinear(init_method) + + def fc2_builder(input_size, output_size, *, init_method, **kwargs): + return DummyLinear(init_method) + + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_block_out_proj_init_depth_power=-0.5, + ) + submodules = MLPSubmodules(linear_fc1=fc1_builder, linear_fc2=fc2_builder) + + plain_mlp = MLP(config, submodules, apply_block_output_init_scaling=False) + scaled_mlp = MLP(config, submodules, apply_block_output_init_scaling=True) + + assert plain_mlp.linear_fc2.init_method is config.output_layer_init_method + assert scaled_mlp.linear_fc2.init_method is not config.output_layer_init_method + + def test_partial_mlp_builder_receives_block_output_init_scaling_opt_in(self): + class DummyLinear(torch.nn.Module): + def __init__(self, init_method): + super().__init__() + self.init_method = init_method + + def forward(self, hidden_states): + return hidden_states, None + + def backward_dw(self): + return None + + def fc1_builder(input_size, output_size, *, init_method, **kwargs): + return DummyLinear(init_method) + + def fc2_builder(input_size, output_size, *, init_method, **kwargs): + return DummyLinear(init_method) + + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + ) + submodules = MLPSubmodules(linear_fc1=fc1_builder, linear_fc2=fc2_builder) + mlp_builder = functools.partial(MLP.as_mlp_submodule, submodules=submodules) + + assert _get_mlp_builder_module(mlp_builder) is MLP + + additional_mlp_kwargs = {"apply_block_output_init_scaling": True} + mlp = mlp_builder( + config=config, + pg_collection=SimpleNamespace(tp=None), + is_mtp_layer=False, + **additional_mlp_kwargs, + ) + + assert mlp.linear_fc2.init_method is not config.output_layer_init_method + + def test_transformer_layer_residual_branch_scaling_helper(self): + config = TransformerConfig( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_residual_branch_depth_power=-0.5, + ) + layer = object.__new__(TransformerLayer) + layer.model_scaling_policy = build_model_scaling_policy(config) + + output = torch.ones(4, 8) + bias = torch.ones(8) + scaled_output, scaled_bias = layer._scale_dense_residual_branch_output( + (output, bias), branch_name='self attention', using_fused_tp_inference_kernel=False + ) + expected_mult = ( + config.num_layers / config.scaling_base_num_layers + ) ** config.scaling_residual_branch_depth_power + assert torch.equal(scaled_output, output * expected_mult) + assert torch.equal(scaled_bias, bias * expected_mult) + + with pytest.raises(NotImplementedError, match='Residual-branch scaling'): + layer._scale_dense_residual_branch_output( + (output, bias), branch_name='self attention', using_fused_tp_inference_kernel=True + ) + + def test_allow_scaling_policy_eval_allows_unfused_validation_scaling(self): + config = TransformerConfig( + hidden_size=1024, + num_layers=24, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=12, + ) + layer = object.__new__(TransformerLayer) + layer.model_scaling_policy = build_model_scaling_policy(config) + layer.training = False + + output = torch.ones(2, 2) + bias = torch.full((2, 2), 3.0) + with pytest.raises(NotImplementedError, match='during inference'): + layer._scale_dense_residual_branch_output( + (output, bias), branch_name='self attention', using_fused_tp_inference_kernel=False + ) + + with allow_scaling_policy_eval(True): + scaled_output, scaled_bias = layer._scale_dense_residual_branch_output( + (output, bias), branch_name='self attention', using_fused_tp_inference_kernel=False + ) + + assert torch.equal(scaled_output, output * 0.5) + assert torch.equal(scaled_bias, bias * 0.5) + + def test_transformer_layer_rejects_cross_attention_for_depth_mup(self): + class DummyCrossAttention(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, *args, **kwargs): + return torch.ones(1, 1, 1), None + + config = TransformerConfig( + hidden_size=16, + num_layers=12, + num_attention_heads=4, + scaling_recipe='depth_mup', + scaling_base_hidden_size=8, + scaling_base_num_layers=6, + ) + submodules = TransformerLayerSubmodules(cross_attention=DummyCrossAttention) + + with pytest.raises(NotImplementedError, match='Cross-attention is out of scope for v1'): + TransformerLayer(config=config, submodules=submodules) + class TestMuPOptimizerTypeHandling: """Tests for MuP optimizer-specific override behavior.""" + def _depth_mup_adamw_overrides(self, *, apply_wd_to_qk_layernorm=False): + config_args = SimpleNamespace( + hidden_size=1024, + num_layers=12, + num_attention_heads=16, + scaling_recipe='depth_mup', + scaling_base_hidden_size=256, + scaling_base_num_layers=6, + scaling_base_head_dim=None, + scaling_residual_branch_depth_power=None, + scaling_hidden_lr_depth_power=None, + scaling_block_out_proj_init_depth_power=None, + use_mup=False, + mup_width_mult=1.0, + mup_base_hidden_size=None, + mup_embedding_mult=1.0, + mup_output_mult=1.0, + mup_base_head_dim=None, + mup_attn_scale_power=1.0, + ) + scaling_policy = build_training_scaling_policy(config_args, optimizer_type='adam') + optimizer_config = OptimizerConfig( + optimizer='adam', + lr=1e-3, + min_lr=1e-5, + weight_decay=0.1, + decoupled_weight_decay=True, + apply_wd_to_qk_layernorm=apply_wd_to_qk_layernorm, + ) + standard_overrides = get_standard_config_overrides( + optimizer_config, scaling_policy=scaling_policy + ) + scaling_overrides = get_scaling_config_overrides(optimizer_config, scaling_policy) + return {**standard_overrides, **scaling_overrides} + + def test_depth_mup_adamw_keeps_norm_and_unknown_vectors_no_decay(self): + """Depth-MuP AdamW must not re-enable WD for norms or unclassified vectors.""" + overrides = self._depth_mup_adamw_overrides() + + hidden_matrix = torch.nn.Parameter(torch.zeros(4, 4)) + hidden_bias = torch.nn.Parameter(torch.zeros(4)) + norm_weight = torch.nn.Parameter(torch.zeros(4)) + qk_norm_weight = torch.nn.Parameter(torch.zeros(4)) + unknown_vector = torch.nn.Parameter(torch.zeros(4)) + + hidden_matrix_override = _combined_override_for_param( + overrides, hidden_matrix, 'decoder.layers.0.mlp.linear_fc2.weight' + ) + hidden_bias_override = _combined_override_for_param( + overrides, hidden_bias, 'decoder.layers.0.mlp.linear_fc2.bias' + ) + norm_override = _combined_override_for_param( + overrides, norm_weight, 'decoder.layers.0.input_layernorm.weight' + ) + qk_norm_override = _combined_override_for_param( + overrides, qk_norm_weight, 'decoder.layers.0.self_attention.q_layernorm.weight' + ) + unknown_vector_override = _combined_override_for_param( + overrides, unknown_vector, 'decoder.layers.0.unclassified_vector' + ) + + assert hidden_matrix_override['wd_mult'] == pytest.approx(4.0) + assert 'wd_mult' not in hidden_bias_override + assert norm_override['wd_mult'] == 0.0 + assert qk_norm_override['wd_mult'] == 0.0 + assert unknown_vector_override['wd_mult'] == 0.0 + + def test_depth_mup_adamw_qk_layernorm_wd_opt_in_is_narrow(self): + """The q/k layernorm WD opt-in must not affect ordinary norm no-decay skips.""" + overrides = self._depth_mup_adamw_overrides(apply_wd_to_qk_layernorm=True) + + qk_norm_weight = torch.nn.Parameter(torch.zeros(4)) + norm_weight = torch.nn.Parameter(torch.zeros(4)) + + qk_norm_override = _combined_override_for_param( + overrides, qk_norm_weight, 'decoder.layers.0.self_attention.k_layernorm.weight' + ) + norm_override = _combined_override_for_param( + overrides, norm_weight, 'decoder.layers.0.post_attention_layernorm.weight' + ) + + assert qk_norm_override.get('wd_mult') is None + assert norm_override['wd_mult'] == 0.0 + def test_sgd_scales_vector_like_lr_only(self): """SGD scales vector-like params by width_mult; hidden params keep base LR.""" optimizer_config = OptimizerConfig(lr=1e-3, min_lr=1e-5)