[4/4] Hyperparameter Transfer across Muon and preconditioned optimizers.

docs/api-guide/internal/index.md

@@ -16,4 +16,5 @@ Internal utility APIs.
 
 num_microbatches_calculator
 optimizer_param_scheduler
+scaling_policy_infrastructure
 ```

docs/api-guide/internal/scaling_policy_infrastructure.md

@@ -0,0 +1,113 @@
+<!---
+   Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.
+   NVIDIA CORPORATION and its licensors retain all intellectual property
+   and proprietary rights in and to this software, related documentation
+   and any modifications thereto. Any use, reproduction, disclosure or
+   distribution of this software and related documentation without an express
+   license agreement from NVIDIA CORPORATION is strictly prohibited.
+-->
+
+# 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.

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
 ```
 

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
 ```

docs/user-guide/scaling-recipes.md

@@ -0,0 +1,135 @@
+<!---
+   Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+   NVIDIA CORPORATION and its licensors retain all intellectual property
+   and proprietary rights in and to this software, related documentation
+   and any modifications thereto. Any use, reproduction, disclosure or
+   distribution of this software and related documentation without an express
+   license agreement from NVIDIA CORPORATION is strictly prohibited.
+-->
+
+# 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.

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):

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,

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,

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: