General scaling policy for HP transfer + Depth MuP recipe

docs/api-guide/internal/index.md

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

docs/api-guide/internal/scaling_recipe_contract.md

@@ -0,0 +1,117 @@
+<!---
+   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 Recipe Design Contract
+
+This page documents the internal implementation contract for the current scaling
+recipes. It is intended for maintainers changing model initialization, optimizer
+parameter grouping, checkpoint compatibility, or argument validation.
+
+A scaling recipe is resolved into a single scaling context. Model code,
+optimizer code, checkpoint code, CLI validation, and YAML validation should read
+that resolved context rather than independently interpreting legacy MuP fields.
+
+## Legacy Alias Canonicalization
+
+`--use-mup` is a deprecated alias for exact `--scaling-recipe mup`. The legacy
+fields `mup_base_hidden_size`, `mup_base_head_dim`, and `mup_width_mult` are kept
+for compatibility but should be synchronized from the canonical fields.
+
+`mup_width_mult` is derived state:
+
+```text
+width_mult = hidden_size / scaling_base_hidden_size
+```
+
+A non-default legacy `mup_width_mult` must match the derived value. Conflicting
+legacy and canonical scaling fields are validation errors. CLI and YAML
+validation both warn for legacy aliases and call the same synchronization helper
+so downstream global args have the same canonical shape.
+
+## `depth_mup` Optimizer Contract
+
+`depth_mup` is a v1 Megatron adaptation of the spectral width-depth MuP AdamW
+table. It is supported only for `optimizer='adam'` because the current optimizer
+overrides are defined for Adam/AdamW-style parameter groups.
+
+Because the weight-decay row is derived for decoupled AdamW, nonzero
+`weight_decay` requires `decoupled_weight_decay=True`. Coupled Adam/L2 is allowed
+only when `weight_decay=0.0`. SGD, Muon, and other optimizers are intentionally
+rejected rather than partially mapped.
+
+The default multipliers are:
+
+| Mechanism | Default multiplier |
+| --- | --- |
+| Dense self-attention/MLP residual branch output | `depth_mult^-1` |
+| Hidden matrix-like Adam/AdamW LR | `width_mult^-1` |
+| Hidden matrix-like Adam/AdamW epsilon | `(width_mult * depth_mult)^-1` |
+| Hidden vector-like Adam/AdamW epsilon | `(width_mult * depth_mult)^-1` |
+| Embedding/output-class Adam/AdamW epsilon | `width_mult^-1` |
+| Hidden matrix-like AdamW weight decay | `width_mult` |
+| Dense block output-projection initialization | `depth_mult^+0.5` |
+
+## Parameter-Class Policy
+
+Parameter classification should prefer explicit parameterization metadata
+attached during model construction. Name/shape fallback logic exists only for
+backward compatibility with older unannotated parameters.
+
+| 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 |
+
+This table is deliberate Megatron behavior, not a direct claim that every row is
+spelled out by the paper table. The paper gives base weight decay to hidden
+biases, but Megatron's 1-D tensors also include normalization scale/bias tensors
+and other vectors. Hidden linear/attention/MLP biases therefore keep base weight
+decay, while normalization vectors and otherwise unknown 1-D tensors stay on the
+standard no-weight-decay path unless q/k layernorm is explicitly opted in. The
+hidden-vector epsilon rule applies to hidden vector-like parameters as the
+current v1 policy.
+
+## Initialization and Residual Branches
+
+Megatron already applies layer-count-dependent initialization to residual branch
+output projections. `depth_mup` rebases dense transformer block output projection
+initialization to `scaling_base_num_layers` so the explicit residual-branch
+multiplier carries the intended depth scaling.
+
+The dense residual hook covers:
+
+- self-attention output projection
+- dense MLP output projection
+
+MoE layers do not inherit this hook. Unsupported residual or model-family paths
+should keep failing closed until they have explicit rules and tests.
+
+## Runtime Scope
+
+Training mode is the supported runtime for `depth_mup`. Validation loss can be
+enabled with `allow_depth_mup_eval`, but that switch is validation-only and does
+not make generation or inference a supported path.
+
+YAML configs may not contain newly added argparse fields. YAML validation should
+populate defaults for new global fields that downstream runtime code reads. For
+`allow_depth_mup_eval`, the default is `False`.
+
+## Checkpoint and Optimizer-Group Compatibility
+
+Distributed-optimizer checkpoint preprocessing and optimizer load must identify
+parameter groups through the same tolerant identifier tuple. Optional optimizer
+group fields such as `eps` and per-group `optimizer` may be absent in standard
+Adam/SGD groups.
+
+Missing optional fields should resolve to `None` instead of causing resume-time
+`KeyError`. Sorting code must also be `None`-safe so groups with and without
+optional keys can be preprocessed deterministically.

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,159 @@
+<!---
+   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 Recipes
+
+Megatron-LM supports named scaling recipes through `--scaling-recipe`.
+
+The current named recipes are:
+
+- `none`: standard Megatron parameterization.
+- `mup`: current Megatron width MuP behavior.
+- `depth_mup`: experimental spectral width-depth MuP behavior for dense
+  GPT-style residual Transformer blocks using `--optimizer adam` with AdamW-style
+  semantics. Nonzero weight decay requires `decoupled_weight_decay=True`.
+
+`depth_mup` is intentionally narrow. Megatron rejects unsupported paths instead
+of silently applying unvalidated scaling rules.
+
+## Configuration Surface
+
+New configs should use the canonical scaling fields:
+
+```bash
+--scaling-recipe mup \
+--scaling-base-hidden-size <base-hidden-size> \
+--scaling-base-head-dim <base-head-dim>
+```
+
+`--use-mup` remains a backward-compatible alias for exact
+`--scaling-recipe mup`, but it is deprecated. The legacy aliases
+`--mup-base-hidden-size`, `--mup-base-head-dim`, and `--mup-width-mult` are also
+deprecated where they overlap with the canonical scaling surface.
+
+`--mup-width-mult` is derived from the resolved scaling context:
+
+```text
+width_mult = hidden_size / scaling_base_hidden_size
+```
+
+If a non-default `--mup-width-mult` is supplied, it must match that derived
+value. If legacy MuP fields and canonical scaling fields conflict, Megatron
+raises an error during validation. CLI and YAML configs use the same alias
+warning and canonicalization path.
+
+## `mup`
+
+`mup` preserves the current Megatron width-MuP surface:
+
+- width multiplier from `hidden_size / scaling_base_hidden_size`
+- MuP-family attention softmax scaling through `scaling_base_head_dim`
+- hidden-layer width-scaled initialization
+- MuP-family embedding and logit scaling
+- MuP-family optimizer overrides, including Adam epsilon handling
+
+Example:
+
+```bash
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+    --num-layers 24 \
+    --hidden-size 2048 \
+    --num-attention-heads 16 \
+    --optimizer adam \
+    --scaling-recipe mup \
+    --scaling-base-hidden-size 1024 \
+    --scaling-base-head-dim 128
+```
+
+## `depth_mup`
+
+`depth_mup` extends MuP-family width behavior with depth scaling for dense
+GPT-style residual blocks. It is implemented only for `--optimizer adam`.
+Because its weight-decay rule is AdamW-style, `depth_mup` requires
+`decoupled_weight_decay=True` whenever `weight_decay` is nonzero. Coupled
+Adam/L2 is allowed only with `weight_decay=0.0`.
+
+The main default behaviors are:
+
+- dense self-attention/MLP residual branch output scales as `depth_mult^-1`
+- hidden matrix-like Adam LR scales as `width_mult^-1`
+- hidden matrix-like Adam epsilon scales as `(width_mult * depth_mult)^-1`
+- embedding/output-class Adam epsilon scales as `width_mult^-1`
+- hidden matrix-like AdamW weight decay scales as `width_mult`
+- dense residual output-projection initialization scales as `depth_mult^+0.5`
+
+Megatron's 1-D parameters are not all hidden biases. Under `depth_mup`, hidden
+linear/attention/MLP biases keep base weight decay, while normalization vectors
+and otherwise unknown 1-D tensors keep Megatron's conservative no-weight-decay
+behavior unless q/k layernorm is explicitly opted into weight decay with
+`apply_wd_to_qk_layernorm=True`.
+
+Example:
+
+```bash
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+    --num-layers 24 \
+    --hidden-size 2048 \
+    --num-attention-heads 16 \
+    --optimizer adam \
+    --scaling-recipe depth_mup \
+    --scaling-base-hidden-size 1024 \
+    --scaling-base-num-layers 12 \
+    --scaling-base-head-dim 128
+```
+
+### Current `depth_mup` Scope
+
+`depth_mup` v1 is currently intended for:
+
+- dense GPT-style residual Transformer blocks
+- dense self-attention residual branches
+- dense MLP residual branches
+- `--optimizer adam`, with `decoupled_weight_decay=True` when `weight_decay`
+  is nonzero
+
+`depth_mup` currently rejects unsupported paths, including:
+
+- SGD, Muon, and other non-Adam optimizers
+- residual-branch scaling during inference
+- fused TP inference residual scaling
+- cross-attention
+- `multi_latent_attention`
+- configured experimental attention variants
+- MoE
+- BERT, T5, and Mamba model families
+
+Training mode is the supported runtime. Validation loss can be enabled
+explicitly with `--allow-depth-mup-eval`; this switch is for validation only and
+does not make generation or inference a supported `depth_mup` path. YAML configs
+default `allow_depth_mup_eval` to `False`.
+
+## Manual Overrides
+
+The canonical scaling fields remain overrides on top of the named recipe. For
+example, `--scaling-residual-branch-depth-power 0.0` explicitly disables the
+default `depth_mup` residual multiplier instead of introducing a separate
+recipe.
+
+Overrides only change resolved multipliers inside the supported surface. They do
+not widen the supported-surface contract.
+
+## Non-goals
+
+These recipes do not currently claim:
+
+- HyperP
+- CompleteP
+- MuonH / AdamH
+- the full Muon-Kimi spectral width-depth training setup
+- token-count LR scaling
+- SqrtGate
+- MoE granularity transfer
+- public SGD depth transfer
+- public Muon depth transfer

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2890,6 +2890,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/extensions/transformer_engine.py

@@ -2457,6 +2457,7 @@ def as_mlp_submodule(
                 is_expert=is_expert,
                 input_size=input_size,
                 ffn_hidden_size=ffn_hidden_size,
+                apply_block_output_init_scaling=True,
             )
 
 else:

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 SCALING_RECIPE_DEPTH_MUP, build_resolved_model_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
@@ -51,15 +52,16 @@ def __init__(
             log_config_to_disk(config, locals(), prefix=type(self).__name__)
 
         self.parallel_output = parallel_output
+        self.model_scaling_policy = build_resolved_model_policy(config)
 
         self.output_layer = tensor_parallel.ColumnParallelLinear(
             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=self.model_scaling_policy.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,
@@ -158,6 +160,12 @@ def __init__(
         pg_collection: ProcessGroupCollection = None,
     ):
 
+        if config.scaling_recipe == SCALING_RECIPE_DEPTH_MUP:
+            raise NotImplementedError(
+                "scaling_recipe='depth_mup' currently supports dense GPT-style residual "
+                "Transformer blocks only. T5Model is out of scope for v1."
+            )
+
         super(T5Model, self).__init__(config=config)
 
         self.config: TransformerConfig = config
@@ -427,6 +435,7 @@ def forward(
             if self.share_embeddings_and_output_weights:
                 output_weight = self.shared_embedding_or_output_weight()
             lm_logits = self.lm_head(decoder_hidden_states, word_embeddings_weight=output_weight)
+            lm_logits = self._scale_logits(lm_logits)
 
             if lm_labels is None:
                 # [s b h] => [b s h]