[2/4] Hyperparameter Transfer: add canonical width-MuP recipe

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,65 @@
+<!---
+   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` and `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 MuP 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-MuP configs, every hook returns the current Megatron default.
+
+## Training Policy
+
+Optimizer code should route per-parameter hyperparameter multipliers through the
+training scaling policy. The policy currently preserves the existing 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.
+
+The public compatibility function `get_mup_config_overrides` remains available
+and delegates to the policy implementation.
+
+## 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 uses a shared helper for the existing stable
+parameter-group identity: `wd_mult`, `lr_mult`, `is_expert_parallel`, and
+`is_decoupled_lr`. The helper tolerates NeMo-style `pre_` field names and missing
+legacy fields without adding newer optional fields such as `eps`, `max_lr`,
+`min_lr`, or per-group `optimizer` to the stable resume identity.

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,76 @@
+<!---
+   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 two recipes:
+
+| Recipe | Behavior |
+| --- | --- |
+| `none` | Standard Megatron parameterization. This is the default. |
+| `mup` | Width MuP for hidden-size transfer. |
+
+## 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.
+
+## 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.
+
+## 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:

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

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
 

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 (