From 13334e7422baff6673826765ef6fbfb0a35580e7 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:14 +0530 Subject: [PATCH 01/19] Add scaling parameterization policy primitives --- megatron/core/parameterization/__init__.py | 52 ++++++++ .../core/parameterization/model_policy.py | 116 ++++++++++++++++++ megatron/core/parameterization/roles.py | 74 +++++++++++ megatron/core/parameterization/spec.py | 42 +++++++ .../core/parameterization/training_policy.py | 72 +++++++++++ 5 files changed, 356 insertions(+) create mode 100644 megatron/core/parameterization/__init__.py create mode 100644 megatron/core/parameterization/model_policy.py create mode 100644 megatron/core/parameterization/roles.py create mode 100644 megatron/core/parameterization/spec.py create mode 100644 megatron/core/parameterization/training_policy.py diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py new file mode 100644 index 00000000000..aefab454c11 --- /dev/null +++ b/megatron/core/parameterization/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from .model_policy import ResolvedModelPolicy, build_resolved_model_policy +from .roles import ( + IS_OUTPUT_PARAMETER_ATTR, + PARAMETERIZATION_ROLE_ATTR, + PARAMETERIZATION_SHARED_GROUP_ATTR, + PARAMETERIZATION_TAGS_ATTR, + ROLE_EMBEDDING, + ROLE_HIDDEN_MATRIX, + ROLE_HIDDEN_VECTOR, + ROLE_MUON_MANAGED_MATRIX, + ROLE_OUTPUT, + ROLE_SHARED_EMBEDDING_OUTPUT, + get_parameterization_role, + is_embedding_class_parameter, + is_embedding_or_output_parameter, + is_hidden_matrix_parameter, + is_muon_managed_matrix_parameter, + is_output_parameter, + is_vector_like_parameter, + set_parameterization_metadata, +) +from .spec import ResolvedScalingContext, build_resolved_scaling_context +from .training_policy import ResolvedTrainingPolicy, build_legacy_mup_training_policy + +__all__ = [ + 'IS_OUTPUT_PARAMETER_ATTR', + 'PARAMETERIZATION_ROLE_ATTR', + 'PARAMETERIZATION_SHARED_GROUP_ATTR', + 'PARAMETERIZATION_TAGS_ATTR', + 'ROLE_EMBEDDING', + 'ROLE_HIDDEN_MATRIX', + 'ROLE_HIDDEN_VECTOR', + 'ROLE_MUON_MANAGED_MATRIX', + 'ROLE_OUTPUT', + 'ROLE_SHARED_EMBEDDING_OUTPUT', + 'ResolvedModelPolicy', + 'ResolvedScalingContext', + 'ResolvedTrainingPolicy', + 'build_legacy_mup_training_policy', + 'build_resolved_model_policy', + 'build_resolved_scaling_context', + 'get_parameterization_role', + 'is_embedding_class_parameter', + 'is_embedding_or_output_parameter', + 'is_hidden_matrix_parameter', + 'is_muon_managed_matrix_parameter', + 'is_output_parameter', + 'is_vector_like_parameter', + 'set_parameterization_metadata', +] diff --git a/megatron/core/parameterization/model_policy.py b/megatron/core/parameterization/model_policy.py new file mode 100644 index 00000000000..ed8767cf874 --- /dev/null +++ b/megatron/core/parameterization/model_policy.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +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 ResolvedScalingContext, build_resolved_scaling_context + + +@dataclass(frozen=True) +class ResolvedModelPolicy: + """Model-side policy for existing Megatron scaling behavior.""" + + context: ResolvedScalingContext + + @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 1.0 + + @property + def dense_block_out_proj_init_multiplier(self) -> float: + return 1.0 + + 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, + ): + del init_method_std, num_layers, is_hybrid_model + if output_layer_init_method_is_user_provided: + return default_init_method + return default_init_method + + 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]: + return output_with_bias + + +def build_resolved_model_policy(config) -> ResolvedModelPolicy: + return ResolvedModelPolicy(build_resolved_scaling_context(config)) diff --git a/megatron/core/parameterization/roles.py b/megatron/core/parameterization/roles.py new file mode 100644 index 00000000000..3d5513f9f46 --- /dev/null +++ b/megatron/core/parameterization/roles.py @@ -0,0 +1,74 @@ +# 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_HIDDEN_MATRIX = 'hidden_matrix' +ROLE_HIDDEN_VECTOR = 'hidden_vector' +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)) + + +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 is_hidden_matrix_parameter(param: Any, param_name: Optional[str] = None) -> bool: + if is_embedding_class_parameter(param, param_name): + return False + return param.dim() > 1 + + +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..26dfe3e0edb --- /dev/null +++ b/megatron/core/parameterization/spec.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + + +@dataclass(frozen=True) +class ResolvedScalingContext: + """Resolved internal scaling context for existing Megatron parameterization knobs. + + PR1 intentionally does not add a new public recipe surface. It wraps the existing + ``use_mup`` fields so model and optimizer code can route through a shared policy seam + while preserving current behavior. + """ + + use_mup: bool + width_mult: float = 1.0 + embedding_mult: float = 1.0 + output_mult: float = 1.0 + base_head_dim: Optional[float] = None + attention_scale_power: float = 1.0 + + @property + def enabled(self) -> bool: + return self.use_mup + + @property + def uses_width_mup(self) -> bool: + return self.use_mup + + +def build_resolved_scaling_context(config: Any) -> ResolvedScalingContext: + return ResolvedScalingContext( + use_mup=bool(getattr(config, 'use_mup', False)), + width_mult=getattr(config, 'mup_width_mult', 1.0), + embedding_mult=getattr(config, 'mup_embedding_mult', 1.0), + output_mult=getattr(config, 'mup_output_mult', 1.0), + base_head_dim=getattr(config, 'mup_base_head_dim', None), + attention_scale_power=getattr(config, 'mup_attn_scale_power', 1.0), + ) diff --git a/megatron/core/parameterization/training_policy.py b/megatron/core/parameterization/training_policy.py new file mode 100644 index 00000000000..a979d462b7a --- /dev/null +++ b/megatron/core/parameterization/training_policy.py @@ -0,0 +1,72 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +from dataclasses import dataclass + +from .spec import ResolvedScalingContext + + +@dataclass(frozen=True) +class ResolvedTrainingPolicy: + """Optimizer-side policy for existing Megatron MuP hyperparameter multipliers.""" + + context: ResolvedScalingContext + optimizer_type: str = 'adam' + + @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 'adam' in self.optimizer_type_lower + + @property + def is_muon_optimizer(self) -> bool: + return 'muon' in self.optimizer_type_lower + + @property + def hidden_lr_multiplier(self) -> float: + if not self.enabled or not self.uses_width_mup: + return 1.0 + if self.is_sgd_optimizer: + return 1.0 + return 1.0 / self.context.width_mult + + @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_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 + + @property + def vector_like_lr_multiplier(self) -> float: + return self.hidden_vector_lr_multiplier + + +def build_legacy_mup_training_policy( + *, mup_width_mult: float, optimizer_type: str = 'adam' +) -> ResolvedTrainingPolicy: + return ResolvedTrainingPolicy( + context=ResolvedScalingContext(use_mup=True, width_mult=mup_width_mult), + optimizer_type=optimizer_type, + ) From 1f217de713b0c68747eb3203340033ebb57cd2f8 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:14 +0530 Subject: [PATCH 02/19] Preserve scaling metadata through model paths --- .../megatron_fsdp/param_and_grad_buffer.py | 4 ++ megatron/core/models/T5/t5_model.py | 9 +-- megatron/core/models/bert/bert_model.py | 8 +-- .../embeddings/language_model_embedding.py | 6 +- .../common/language_module/language_module.py | 61 +++++++++++++++---- megatron/core/models/gpt/gpt_model.py | 10 +-- megatron/core/models/hybrid/hybrid_model.py | 10 +-- megatron/core/transformer/attention.py | 10 ++- megatron/core/transformer/mlp.py | 10 ++- .../core/transformer/transformer_config.py | 52 +++++----------- .../core/transformer/transformer_layer.py | 8 +++ 11 files changed, 118 insertions(+), 70 deletions(-) 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 960acc25ef6..5644a0519f9 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 @@ -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): diff --git a/megatron/core/models/T5/t5_model.py b/megatron/core/models/T5/t5_model.py index b2feb974643..13cf5edc41b 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_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 @@ -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_resolved_model_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..989d3152ba7 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_resolved_model_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,9 @@ 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_resolved_model_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 84b0ca2fea3..3d033cd3678 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_resolved_model_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_resolved_model_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_resolved_model_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 b018b805c82..a4863b75fef 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -248,10 +248,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, @@ -668,7 +668,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 bdf5b2684d0..896295e359c 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -288,10 +288,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, @@ -525,7 +525,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/transformer/attention.py b/megatron/core/transformer/attention.py index e920e44d683..bc0b5bd7f4a 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_resolved_model_policy from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) @@ -386,11 +387,18 @@ def __init__( ) # Output. + model_scaling_policy = build_resolved_model_policy(self.config) 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=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=False, + ), 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 46979a8ba8f..e17805f09a4 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_resolved_model_policy from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig @@ -225,13 +226,20 @@ def __init__( else: self.activation_func = self.config.activation_func + model_scaling_policy = build_resolved_model_policy(self.config) 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=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=False, + ), bias=self.config.add_bias_linear, input_is_parallel=True, skip_bias_add=True, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bb044787b9c..e5ddcf4d9b6 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,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend +from megatron.core.parameterization import build_resolved_model_policy from megatron.core.quantization.quant_config import RecipeConfig from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout @@ -18,14 +18,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__) @@ -1788,13 +1781,6 @@ def __post_init__(self): # 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: @@ -1815,6 +1801,11 @@ def __post_init__(self): UserWarning, ) + model_scaling_policy = build_resolved_model_policy(self) + 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), @@ -1838,29 +1829,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 7e69e33318d..4798278fd2f 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -17,6 +17,7 @@ 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_resolved_model_policy from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing from megatron.core.transformer.enums import CudaGraphScope, LayerType @@ -329,6 +330,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_resolved_model_policy(config) # [Module 1: Input Layernorm] Optional Layernorm on the input data # TODO: add pytorch only layernorm @@ -662,6 +664,9 @@ def _forward_attention( hidden_states = attention_output_with_bias[0] else: with self.bias_dropout_add_exec_handler(): + attention_output_with_bias = self.model_scaling_policy.scale_residual_branch_output( + attention_output_with_bias + ) hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)( attention_output_with_bias, residual, self.hidden_dropout ) @@ -911,6 +916,9 @@ def _forward_post_mlp( hidden_states = mlp_output_with_bias[0] else: with self.bias_dropout_add_exec_handler(): + mlp_output_with_bias = self.model_scaling_policy.scale_residual_branch_output( + mlp_output_with_bias + ) hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)( mlp_output_with_bias, residual, self.hidden_dropout ) From 65bb9cecac33af335032ee1be15bcee9acdfefe6 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:14 +0530 Subject: [PATCH 03/19] Route optimizer infrastructure through scaling policies --- megatron/core/optimizer/__init__.py | 121 +++++++------------ megatron/core/optimizer/distrib_optimizer.py | 22 +--- megatron/core/optimizer/optimizer.py | 26 ++-- megatron/training/training.py | 4 +- 4 files changed, 67 insertions(+), 106 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index ebdd42effe2..955d21382c2 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -54,6 +54,15 @@ combine_param_group_overrides, param_group_override_to_tuple, ) +from megatron.core.parameterization import ( + ResolvedTrainingPolicy, + build_legacy_mup_training_policy, + is_embedding_class_parameter, + is_embedding_or_output_parameter, + is_hidden_matrix_parameter, + is_muon_managed_matrix_parameter, + is_vector_like_parameter, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.fsdp_dtensor_checkpoint import get_global_unique_param_name @@ -73,7 +82,6 @@ Float16OptimizerWithFloat16Params, FP32Optimizer, MegatronOptimizer, - param_group_identifier_keys, ) # Subclass aliases kept for backward compatibility; all are OptimizerConfig. @@ -130,41 +138,23 @@ def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, Par 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. +def get_scaling_config_overrides( + config: OptimizerConfig, scaling_policy: ResolvedTrainingPolicy +) -> Dict[ParamKey, ParamGroupOverride]: + """Get optimizer overrides from a resolved internal scaling policy. + + PR1 keeps the public MuP surface unchanged and factors the existing per-parameter + MuP LR/epsilon behavior through this policy seam. """ - 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 + if not scaling_policy.enabled: + return {} decoupled_lr_enabled = config.decoupled_lr is not None if decoupled_lr_enabled: @@ -173,11 +163,11 @@ def get_mup_config_overrides( "absolute LR for embedding+output params, and MuP LR scaling will not " "override those parameters." ) - if is_adam_optimizer: + if scaling_policy.is_adam_optimizer: message += " MuP Adam epsilon scaling remains applied to hidden matrix-like parameters." 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,66 +179,43 @@ def get_mup_config_overrides( "Muon-managed matrices with MuP.", ) - if mup_width_mult == 1.0: - # No scaling needed when width_mult is 1 + if scaling_policy.context.width_mult == 1.0: 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: - return False - return param.dim() == 2 and not getattr(param, 'is_embedding_or_output_parameter', False) - - 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): + 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 - if is_muon_managed_matrix_parameter(param, param_name): + if is_muon_managed_matrix_parameter(param, optimizer_type=scaling_policy.optimizer_type): 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_embedding_class_parameter(param, param_name): + return False if is_vector_like_parameter(param, param_name): return False - if is_muon_managed_matrix_parameter(param, param_name): + 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 mup_overrides: Dict[ParamKey, ParamGroupOverride] = {} - if is_sgd_optimizer: - vector_like_lr_mult = mup_width_mult + if scaling_policy.is_sgd_optimizer: vector_like_lr_override: ParamGroupOverride = {} if base_lr is not None: - vector_like_lr_override["max_lr"] = base_lr * vector_like_lr_mult + vector_like_lr_override["max_lr"] = base_lr * scaling_policy.hidden_vector_lr_multiplier if base_min_lr is not None: - vector_like_lr_override["min_lr"] = base_min_lr * vector_like_lr_mult + vector_like_lr_override["min_lr"] = ( + base_min_lr * scaling_policy.hidden_vector_lr_multiplier + ) if vector_like_lr_override: vector_like_predicate = ParamWithNamePredicate( @@ -263,18 +230,18 @@ def should_scale_eps_with_mup(param: torch.nn.Parameter, param_name: str) -> boo lr_override: ParamGroupOverride = {} if base_lr is not None: - lr_override["max_lr"] = base_lr * hidden_lr_mult + lr_override["max_lr"] = base_lr * scaling_policy.hidden_lr_multiplier if base_min_lr is not None: - lr_override["min_lr"] = base_min_lr * hidden_lr_mult + lr_override["min_lr"] = base_min_lr * scaling_policy.hidden_lr_multiplier 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: + eps_override["eps"] = config.adam_eps * scaling_policy.hidden_eps_multiplier 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 diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 95e8a0c407b..f65457fcba6 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 @@ -884,21 +888,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/optimizer.py b/megatron/core/optimizer/optimizer.py index 9ae23bb4b7f..55a3bc87be1 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -97,6 +97,19 @@ def _multi_tensor_copy_this_to_that( param_group_identifier_keys = ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr') +def get_param_group_identifier_tuple(param_group: Dict) -> tuple: + """Return the stable legacy 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(None) + return tuple(values) + + class MegatronOptimizer(ABC): """ Base class for all Megatron optimizers. @@ -426,22 +439,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/training/training.py b/megatron/training/training.py index ffd2286bcec..5ed27463e2b 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -163,7 +163,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): 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_tuple from megatron.core.optimizer.qk_clip import clip_qk @@ -812,7 +812,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_tuple param_groups.sort(key=key_fn) inner_optimizer["param_groups"] = param_groups From d7b4a96ae58f7521a3caf9518eda15bc03762d07 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:14 +0530 Subject: [PATCH 04/19] Document and test scaling policy infrastructure --- docs/api-guide/internal/index.md | 1 + .../internal/scaling_policy_infrastructure.md | 64 ++++++++++++++++++ tests/unit_tests/test_optimizer.py | 57 ++++++++++++++++ tests/unit_tests/transformer/test_mup.py | 65 ++++++++++++++++++- 4 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 docs/api-guide/internal/scaling_policy_infrastructure.md 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..76fad30c3f9 --- /dev/null +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -0,0 +1,64 @@ + + +# Scaling Policy Infrastructure + +This internal policy layer centralizes Megatron's existing parameterization hooks +without adding a new user-facing recipe surface. + +The current policy resolves the legacy `use_mup` fields into model-side and +optimizer-side decisions. 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 resolved model +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 +resolved training 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 PR1 resume identity. diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 56af8545042..412ce5b414e 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -23,6 +23,10 @@ get_megatron_optimizer, get_standard_config_overrides, ) +from megatron.core.optimizer.optimizer import ( + get_param_group_identifier_sort_key, + 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 +73,59 @@ 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[:4] == (1.0, 1.0, False, False) + assert ident[4:] == (None, None, 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[:4] == (0.0, 0.5, False, True) + assert ident[4:] == (None, None, None, None) + + +def test_param_group_identifier_tuple_tolerates_and_sorts_optional_resume_fields(): + current_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", + } + older_checkpoint_group = { + "wd_mult": 1.0, + "lr_mult": 1.0, + "is_expert_parallel": False, + "is_decoupled_lr": False, + } + mixed_groups = [older_checkpoint_group, current_group] + + current_ident = get_param_group_identifier_tuple(current_group) + older_ident = get_param_group_identifier_tuple(older_checkpoint_group) + + assert current_ident[:4] == older_ident[:4] + assert current_ident[4:] == (1e-4, 1e-6, 1e-8, "adam") + assert older_ident[4:] == (None, None, None, None) + + mixed_groups.sort(key=get_param_group_identifier_sort_key) + assert mixed_groups == [older_checkpoint_group, current_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..99890a033a5 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -16,9 +16,17 @@ 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 ( + build_legacy_mup_training_policy, + build_resolved_model_policy, +) from megatron.core.transformer.multi_token_prediction import process_mtp_loss from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import init_method_normal, mup_scaled_init_method_normal @@ -195,6 +203,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_resolved_training_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 +377,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_resolved_model_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_resolved_model_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_resolved_model_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_resolved_model_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( From 4aecae69a1c3810ea8a23e60babce728fe846682 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Sat, 16 May 2026 10:21:07 +0530 Subject: [PATCH 05/19] Test stable optimizer group resume matching --- tests/unit_tests/test_optimizer.py | 37 +++++++++++++++--------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 412ce5b414e..c2257074c4f 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -24,7 +24,7 @@ get_standard_config_overrides, ) from megatron.core.optimizer.optimizer import ( - get_param_group_identifier_sort_key, + MegatronOptimizer, get_param_group_identifier_tuple, ) from megatron.core.optimizer_param_scheduler import ParamGroupOverride @@ -78,8 +78,7 @@ def test_param_group_identifier_tuple_tolerates_missing_optional_keys(): ident = get_param_group_identifier_tuple(group) - assert ident[:4] == (1.0, 1.0, False, False) - assert ident[4:] == (None, None, None, None) + assert ident == (1.0, 1.0, False, False) def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): @@ -92,38 +91,38 @@ def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): ident = get_param_group_identifier_tuple(group) - assert ident[:4] == (0.0, 0.5, False, True) - assert ident[4:] == (None, None, None, None) + assert ident == (0.0, 0.5, False, True) -def test_param_group_identifier_tuple_tolerates_and_sorts_optional_resume_fields(): +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": 1e-4, - "min_lr": 1e-6, - "eps": 1e-8, - "optimizer": "adam", + "max_lr": 2e-4, + "min_lr": 2e-6, + "params": [0], } - older_checkpoint_group = { + 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, + "params": [7], } - mixed_groups = [older_checkpoint_group, current_group] - current_ident = get_param_group_identifier_tuple(current_group) - older_ident = get_param_group_identifier_tuple(older_checkpoint_group) + assert get_param_group_identifier_tuple(current_group) == get_param_group_identifier_tuple( + checkpoint_group + ) - assert current_ident[:4] == older_ident[:4] - assert current_ident[4:] == (1e-4, 1e-6, 1e-8, "adam") - assert older_ident[4:] == (None, None, None, None) + reordered_groups = MegatronOptimizer._filter_and_reorder_param_groups( + [current_group], [checkpoint_group] + ) - mixed_groups.sort(key=get_param_group_identifier_sort_key) - assert mixed_groups == [older_checkpoint_group, current_group] + assert reordered_groups == [checkpoint_group] @patch('torch.distributed.get_world_size', return_value=1) From 274a4c6addbc410bb59410df990d3cfe1bb81d1a Mon Sep 17 00:00:00 2001 From: plugyawn Date: Sat, 16 May 2026 10:50:29 +0530 Subject: [PATCH 06/19] Default missing optimizer group identifiers --- megatron/core/optimizer/optimizer.py | 8 ++++++- tests/unit_tests/test_optimizer.py | 31 ++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 55a3bc87be1..4aabc6b16fb 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -95,6 +95,12 @@ def _multi_tensor_copy_this_to_that( param_group_identifier_keys = ('wd_mult', 'lr_mult', 'is_expert_parallel', 'is_decoupled_lr') +param_group_identifier_defaults = { + 'wd_mult': 1.0, + 'lr_mult': 1.0, + 'is_expert_parallel': False, + 'is_decoupled_lr': False, +} def get_param_group_identifier_tuple(param_group: Dict) -> tuple: @@ -106,7 +112,7 @@ def get_param_group_identifier_tuple(param_group: Dict) -> tuple: elif f"pre_{key}" in param_group: values.append(param_group[f"pre_{key}"]) else: - values.append(None) + values.append(param_group_identifier_defaults[key]) return tuple(values) diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index c2257074c4f..d2dac021d25 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -94,6 +94,14 @@ def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): assert ident == (0.0, 0.5, False, True) +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) + + def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): current_group = { "wd_mult": 1.0, @@ -125,6 +133,29 @@ def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): 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) From 762c4bfa3b3624f513fadc8b0cceb0e26aecf77b Mon Sep 17 00:00:00 2001 From: plugyawn Date: Wed, 27 May 2026 09:49:37 +0530 Subject: [PATCH 07/19] Remove PR sequencing references from scaling policy docs --- docs/api-guide/internal/scaling_policy_infrastructure.md | 2 +- megatron/core/optimizer/__init__.py | 4 ++-- megatron/core/parameterization/spec.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/api-guide/internal/scaling_policy_infrastructure.md b/docs/api-guide/internal/scaling_policy_infrastructure.md index 76fad30c3f9..c51d6f71ff5 100644 --- a/docs/api-guide/internal/scaling_policy_infrastructure.md +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -61,4 +61,4 @@ 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 PR1 resume identity. +`min_lr`, or per-group `optimizer` to the stable resume identity. diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 955d21382c2..1a7eb142556 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -150,8 +150,8 @@ def get_scaling_config_overrides( ) -> Dict[ParamKey, ParamGroupOverride]: """Get optimizer overrides from a resolved internal scaling policy. - PR1 keeps the public MuP surface unchanged and factors the existing per-parameter - MuP LR/epsilon behavior through this policy seam. + This compatibility path keeps the public MuP surface unchanged and factors the + existing per-parameter MuP LR/epsilon behavior through this policy seam. """ if not scaling_policy.enabled: return {} diff --git a/megatron/core/parameterization/spec.py b/megatron/core/parameterization/spec.py index 26dfe3e0edb..72b929fbf86 100644 --- a/megatron/core/parameterization/spec.py +++ b/megatron/core/parameterization/spec.py @@ -10,9 +10,9 @@ class ResolvedScalingContext: """Resolved internal scaling context for existing Megatron parameterization knobs. - PR1 intentionally does not add a new public recipe surface. It wraps the existing - ``use_mup`` fields so model and optimizer code can route through a shared policy seam - while preserving current behavior. + This infrastructure intentionally does not add a new public recipe surface. It + wraps the existing ``use_mup`` fields so model and optimizer code can route through + a shared policy seam while preserving current behavior. """ use_mup: bool From 568374029e886f2d4b0738b3b73ef2b694fb6071 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Wed, 27 May 2026 10:06:25 +0530 Subject: [PATCH 08/19] Rename scaling policy infrastructure objects --- .../internal/scaling_policy_infrastructure.md | 6 +++--- megatron/core/models/T5/t5_model.py | 4 ++-- .../embeddings/language_model_embedding.py | 4 ++-- .../common/language_module/language_module.py | 6 +++--- megatron/core/optimizer/__init__.py | 6 +++--- megatron/core/parameterization/__init__.py | 16 ++++++++-------- megatron/core/parameterization/model_policy.py | 10 +++++----- megatron/core/parameterization/spec.py | 8 ++++---- .../core/parameterization/training_policy.py | 12 ++++++------ megatron/core/transformer/attention.py | 4 ++-- megatron/core/transformer/mlp.py | 4 ++-- megatron/core/transformer/transformer_config.py | 4 ++-- megatron/core/transformer/transformer_layer.py | 4 ++-- tests/unit_tests/transformer/test_mup.py | 12 ++++++------ 14 files changed, 50 insertions(+), 50 deletions(-) diff --git a/docs/api-guide/internal/scaling_policy_infrastructure.md b/docs/api-guide/internal/scaling_policy_infrastructure.md index c51d6f71ff5..c0b45156c97 100644 --- a/docs/api-guide/internal/scaling_policy_infrastructure.md +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -12,14 +12,14 @@ This internal policy layer centralizes Megatron's existing parameterization hooks without adding a new user-facing recipe surface. -The current policy resolves the legacy `use_mup` fields into model-side and +The current scaling context resolves the legacy `use_mup` fields into model-side and optimizer-side decisions. 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 resolved model +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: @@ -35,7 +35,7 @@ For non-MuP configs, every hook returns the current Megatron default. ## Training Policy Optimizer code should route per-parameter hyperparameter multipliers through the -resolved training policy. The policy currently preserves the existing MuP rules: +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`; diff --git a/megatron/core/models/T5/t5_model.py b/megatron/core/models/T5/t5_model.py index 13cf5edc41b..a85854cbca7 100644 --- a/megatron/core/models/T5/t5_model.py +++ b/megatron/core/models/T5/t5_model.py @@ -15,7 +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_resolved_model_policy +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 @@ -57,7 +57,7 @@ def __init__( config.hidden_size, vocab_size, config=config, - init_method=build_resolved_model_policy(config).output_layer_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, diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py index 989d3152ba7..2e0893631fa 100644 --- a/megatron/core/models/common/embeddings/language_model_embedding.py +++ b/megatron/core/models/common/embeddings/language_model_embedding.py @@ -6,7 +6,7 @@ from torch import Tensor from megatron.core import tensor_parallel -from megatron.core.parameterization import build_resolved_model_policy +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 @@ -129,7 +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. - embeddings = build_resolved_model_policy(self.config).scale_embedding_activations( + embeddings = build_model_scaling_policy(self.config).scale_embedding_activations( embeddings ) diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 19ca66456e2..2650f480e37 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -12,7 +12,7 @@ ROLE_EMBEDDING, ROLE_OUTPUT, ROLE_SHARED_EMBEDDING_OUTPUT, - build_resolved_model_policy, + build_model_scaling_policy, set_parameterization_metadata, ) from megatron.core.transformer.cuda_graphs import CudaGraphManager @@ -53,7 +53,7 @@ def __init__( self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None ) -> None: super().__init__(config=config) - self.model_scaling_policy = build_resolved_model_policy(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() @@ -353,7 +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. """ - return build_resolved_model_policy(self.config).scale_output_logits(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/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index a105119e95d..070b3bd7d33 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -55,7 +55,7 @@ param_group_override_to_tuple, ) from megatron.core.parameterization import ( - ResolvedTrainingPolicy, + TrainingScalingPolicy, build_legacy_mup_training_policy, is_embedding_class_parameter, is_embedding_or_output_parameter, @@ -146,9 +146,9 @@ def get_mup_config_overrides( def get_scaling_config_overrides( - config: OptimizerConfig, scaling_policy: ResolvedTrainingPolicy + config: OptimizerConfig, scaling_policy: TrainingScalingPolicy ) -> Dict[ParamKey, ParamGroupOverride]: - """Get optimizer overrides from a resolved internal scaling policy. + """Get optimizer overrides from an internal training scaling policy. This compatibility path keeps the public MuP surface unchanged and factors the existing per-parameter MuP LR/epsilon behavior through this policy seam. diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py index aefab454c11..fe8a340e194 100644 --- a/megatron/core/parameterization/__init__.py +++ b/megatron/core/parameterization/__init__.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from .model_policy import ResolvedModelPolicy, build_resolved_model_policy +from .model_policy import ModelScalingPolicy, build_model_scaling_policy from .roles import ( IS_OUTPUT_PARAMETER_ATTR, PARAMETERIZATION_ROLE_ATTR, @@ -21,8 +21,8 @@ is_vector_like_parameter, set_parameterization_metadata, ) -from .spec import ResolvedScalingContext, build_resolved_scaling_context -from .training_policy import ResolvedTrainingPolicy, build_legacy_mup_training_policy +from .spec import ScalingContext, build_scaling_context +from .training_policy import TrainingScalingPolicy, build_legacy_mup_training_policy __all__ = [ 'IS_OUTPUT_PARAMETER_ATTR', @@ -35,12 +35,12 @@ 'ROLE_MUON_MANAGED_MATRIX', 'ROLE_OUTPUT', 'ROLE_SHARED_EMBEDDING_OUTPUT', - 'ResolvedModelPolicy', - 'ResolvedScalingContext', - 'ResolvedTrainingPolicy', + 'ModelScalingPolicy', + 'ScalingContext', + 'TrainingScalingPolicy', 'build_legacy_mup_training_policy', - 'build_resolved_model_policy', - 'build_resolved_scaling_context', + 'build_model_scaling_policy', + 'build_scaling_context', 'get_parameterization_role', 'is_embedding_class_parameter', 'is_embedding_or_output_parameter', diff --git a/megatron/core/parameterization/model_policy.py b/megatron/core/parameterization/model_policy.py index ed8767cf874..a3a84dee519 100644 --- a/megatron/core/parameterization/model_policy.py +++ b/megatron/core/parameterization/model_policy.py @@ -15,14 +15,14 @@ scaled_init_method_normal, ) -from .spec import ResolvedScalingContext, build_resolved_scaling_context +from .spec import ScalingContext, build_scaling_context @dataclass(frozen=True) -class ResolvedModelPolicy: +class ModelScalingPolicy: """Model-side policy for existing Megatron scaling behavior.""" - context: ResolvedScalingContext + context: ScalingContext @property def enabled(self) -> bool: @@ -112,5 +112,5 @@ def scale_residual_branch_output( return output_with_bias -def build_resolved_model_policy(config) -> ResolvedModelPolicy: - return ResolvedModelPolicy(build_resolved_scaling_context(config)) +def build_model_scaling_policy(config) -> ModelScalingPolicy: + return ModelScalingPolicy(build_scaling_context(config)) diff --git a/megatron/core/parameterization/spec.py b/megatron/core/parameterization/spec.py index 72b929fbf86..5bbed2051fb 100644 --- a/megatron/core/parameterization/spec.py +++ b/megatron/core/parameterization/spec.py @@ -7,8 +7,8 @@ @dataclass(frozen=True) -class ResolvedScalingContext: - """Resolved internal scaling context for existing Megatron parameterization knobs. +class ScalingContext: + """Internal scaling context for existing Megatron parameterization knobs. This infrastructure intentionally does not add a new public recipe surface. It wraps the existing ``use_mup`` fields so model and optimizer code can route through @@ -31,8 +31,8 @@ def uses_width_mup(self) -> bool: return self.use_mup -def build_resolved_scaling_context(config: Any) -> ResolvedScalingContext: - return ResolvedScalingContext( +def build_scaling_context(config: Any) -> ScalingContext: + return ScalingContext( use_mup=bool(getattr(config, 'use_mup', False)), width_mult=getattr(config, 'mup_width_mult', 1.0), embedding_mult=getattr(config, 'mup_embedding_mult', 1.0), diff --git a/megatron/core/parameterization/training_policy.py b/megatron/core/parameterization/training_policy.py index a979d462b7a..220e46e53fe 100644 --- a/megatron/core/parameterization/training_policy.py +++ b/megatron/core/parameterization/training_policy.py @@ -4,14 +4,14 @@ from dataclasses import dataclass -from .spec import ResolvedScalingContext +from .spec import ScalingContext @dataclass(frozen=True) -class ResolvedTrainingPolicy: +class TrainingScalingPolicy: """Optimizer-side policy for existing Megatron MuP hyperparameter multipliers.""" - context: ResolvedScalingContext + context: ScalingContext optimizer_type: str = 'adam' @property @@ -65,8 +65,8 @@ def vector_like_lr_multiplier(self) -> float: def build_legacy_mup_training_policy( *, mup_width_mult: float, optimizer_type: str = 'adam' -) -> ResolvedTrainingPolicy: - return ResolvedTrainingPolicy( - context=ResolvedScalingContext(use_mup=True, width_mult=mup_width_mult), +) -> TrainingScalingPolicy: + return TrainingScalingPolicy( + context=ScalingContext(use_mup=True, width_mult=mup_width_mult), optimizer_type=optimizer_type, ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index e1e88587434..55da3a1321b 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -28,7 +28,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from megatron.core.parameterization import build_resolved_model_policy +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) @@ -392,7 +392,7 @@ def __init__( ) # Output. - model_scaling_policy = build_resolved_model_policy(self.config) + model_scaling_policy = build_model_scaling_policy(self.config) self.linear_proj = submodules.linear_proj( self.query_projection_size, self.config.hidden_size, diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index 81b5be2c42b..1b7bac4dba7 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -23,7 +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_resolved_model_policy +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 @@ -234,7 +234,7 @@ def __init__( else: self.activation_func = self.config.activation_func - model_scaling_policy = build_resolved_model_policy(self.config) + model_scaling_policy = build_model_scaling_policy(self.config) self.linear_fc2 = submodules.linear_fc2( not_none(self.config.ffn_hidden_size), not_none( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index d3f66773177..dba7fcea1de 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -10,7 +10,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend -from megatron.core.parameterization import build_resolved_model_policy +from megatron.core.parameterization import build_model_scaling_policy from megatron.core.quantization.quant_config import RecipeConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, @@ -1978,7 +1978,7 @@ def __post_init__(self): UserWarning, ) - model_scaling_policy = build_resolved_model_policy(self) + model_scaling_policy = build_model_scaling_policy(self) self.softmax_scale = model_scaling_policy.resolve_attention_softmax_scale( softmax_scale=self.softmax_scale, kv_channels=self.kv_channels ) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index c640a810a7d..4243ec0898d 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -17,7 +17,7 @@ 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_resolved_model_policy +from megatron.core.parameterization import build_model_scaling_policy 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 @@ -336,7 +336,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_resolved_model_policy(config) + self.model_scaling_policy = build_model_scaling_policy(config) # [Module 1: Input Layernorm] Optional Layernorm on the input data # TODO: add pytorch only layernorm diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index 99890a033a5..3eb8f01c941 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -25,7 +25,7 @@ from megatron.core.optimizer_param_scheduler import combine_param_group_overrides from megatron.core.parameterization import ( build_legacy_mup_training_policy, - build_resolved_model_policy, + build_model_scaling_policy, ) from megatron.core.transformer.multi_token_prediction import process_mtp_loss from megatron.core.transformer.transformer_config import TransformerConfig @@ -203,7 +203,7 @@ 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_resolved_training_policy(self): + 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 @@ -377,7 +377,7 @@ def test_mup_with_decoupled_lr_scales_hidden_only_for_lr(self): class TestMuPConfigIntegration: """Integration tests for MuP config with init methods.""" - def test_resolved_model_policy_matches_legacy_config_fields(self): + 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, @@ -387,7 +387,7 @@ def test_resolved_model_policy_matches_legacy_config_fields(self): mup_base_hidden_size=256, mup_embedding_mult=3.0, ) - policy = build_resolved_model_policy(config) + policy = build_model_scaling_policy(config) assert policy.enabled is True assert policy.context.width_mult == pytest.approx(config.mup_width_mult) @@ -399,7 +399,7 @@ def test_resolved_model_policy_matches_legacy_config_fields(self): ) ) - def test_resolved_model_policy_tracks_post_init_multiplier_mutations(self): + 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, @@ -413,7 +413,7 @@ def test_resolved_model_policy_tracks_post_init_multiplier_mutations(self): config.mup_output_mult = 0.25 config.mup_embedding_mult = 3.0 - policy = build_resolved_model_policy(config) + 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) From 9e801dba36717096c6dba6a080bf6f8b13fe6100 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Wed, 27 May 2026 10:50:44 +0530 Subject: [PATCH 09/19] Run autoformatter on scaling policy infra --- .../models/common/embeddings/language_model_embedding.py | 4 +--- tests/unit_tests/test_optimizer.py | 9 ++------- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py index 2e0893631fa..ffdcdb6c651 100644 --- a/megatron/core/models/common/embeddings/language_model_embedding.py +++ b/megatron/core/models/common/embeddings/language_model_embedding.py @@ -129,9 +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. - embeddings = build_model_scaling_policy(self.config).scale_embedding_activations( - embeddings - ) + 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/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index d2dac021d25..ee35b77b701 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -23,10 +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.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 @@ -141,9 +138,7 @@ def test_param_group_matching_normalizes_legacy_missing_identifier_fields(): "is_decoupled_lr": False, "params": [0], } - legacy_checkpoint_group = { - "params": [7], - } + legacy_checkpoint_group = {"params": [7]} assert get_param_group_identifier_tuple(current_group) == get_param_group_identifier_tuple( legacy_checkpoint_group From 10c13dc4456579da1d3e040be5097689b1d6c3ee Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:14 +0530 Subject: [PATCH 10/19] Add canonical width-MuP scaling context --- megatron/core/parameterization/__init__.py | 17 +- megatron/core/parameterization/spec.py | 196 ++++++++++++++++-- .../core/parameterization/training_policy.py | 4 +- .../core/transformer/transformer_config.py | 71 ++++--- megatron/training/arguments.py | 130 ++++++++++++ megatron/training/yaml_arguments.py | 31 ++- 6 files changed, 400 insertions(+), 49 deletions(-) diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py index fe8a340e194..f3a31914eed 100644 --- a/megatron/core/parameterization/__init__.py +++ b/megatron/core/parameterization/__init__.py @@ -21,7 +21,16 @@ is_vector_like_parameter, set_parameterization_metadata, ) -from .spec import ScalingContext, build_scaling_context +from .spec import ( + 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 __all__ = [ @@ -36,11 +45,16 @@ 'ROLE_OUTPUT', 'ROLE_SHARED_EMBEDDING_OUTPUT', 'ModelScalingPolicy', + 'SCALING_RECIPE_MUP', + 'SCALING_RECIPE_NONE', + 'SCALING_RECIPE_VALUES', 'ScalingContext', + 'ScalingUserConfig', 'TrainingScalingPolicy', 'build_legacy_mup_training_policy', 'build_model_scaling_policy', 'build_scaling_context', + 'build_scaling_user_config', 'get_parameterization_role', 'is_embedding_class_parameter', 'is_embedding_or_output_parameter', @@ -49,4 +63,5 @@ 'is_output_parameter', 'is_vector_like_parameter', 'set_parameterization_metadata', + 'sync_legacy_mup_fields', ] diff --git a/megatron/core/parameterization/spec.py b/megatron/core/parameterization/spec.py index 5bbed2051fb..9f730f2cd6e 100644 --- a/megatron/core/parameterization/spec.py +++ b/megatron/core/parameterization/spec.py @@ -2,41 +2,205 @@ from __future__ import annotations +import math from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Literal, Optional + +SCALING_RECIPE_NONE = 'none' +SCALING_RECIPE_MUP = 'mup' +SCALING_RECIPE_VALUES = (SCALING_RECIPE_NONE, SCALING_RECIPE_MUP) @dataclass(frozen=True) -class ScalingContext: - """Internal scaling context for existing Megatron parameterization knobs. +class ScalingUserConfig: + recipe: Optional[Literal['none', 'mup']] = None + base_hidden_size: Optional[int] = None + base_head_dim: 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 - This infrastructure intentionally does not add a new public recipe surface. It - wraps the existing ``use_mup`` fields so model and optimizer code can route through - a shared policy seam while preserving current behavior. - """ - use_mup: bool +@dataclass(frozen=True) +class ScalingContext: + """Internal scaling context for standard and width-MuP parameterization.""" + + recipe: Literal['none', 'mup'] width_mult: float = 1.0 embedding_mult: float = 1.0 output_mult: float = 1.0 + base_hidden_size: Optional[int] = None base_head_dim: Optional[float] = None attention_scale_power: float = 1.0 @property def enabled(self) -> bool: - return self.use_mup + return self.recipe != SCALING_RECIPE_NONE @property def uses_width_mup(self) -> bool: - return self.use_mup + return self.recipe == SCALING_RECIPE_MUP + + @property + def use_mup(self) -> bool: + return self.uses_width_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_head_dim': user_config.base_head_dim, + '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 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_head_dim=getattr(config, 'scaling_base_head_dim', 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. Non-default fields: " + + ", ".join(non_default_fields) + ) + return ScalingContext(recipe=SCALING_RECIPE_NONE) + + 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.') + 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 + return ScalingContext( - use_mup=bool(getattr(config, 'use_mup', False)), - width_mult=getattr(config, 'mup_width_mult', 1.0), - embedding_mult=getattr(config, 'mup_embedding_mult', 1.0), - output_mult=getattr(config, 'mup_output_mult', 1.0), - base_head_dim=getattr(config, 'mup_base_head_dim', None), - attention_scale_power=getattr(config, 'mup_attn_scale_power', 1.0), + recipe=SCALING_RECIPE_MUP, + width_mult=width_mult, + embedding_mult=user_config.mup_embedding_mult, + output_mult=output_mult, + base_hidden_size=base_hidden_size, + base_head_dim=base_head_dim, + attention_scale_power=user_config.mup_attn_scale_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_head_dim = 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_head_dim = context.base_head_dim + 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 index 220e46e53fe..01e98851084 100644 --- a/megatron/core/parameterization/training_policy.py +++ b/megatron/core/parameterization/training_policy.py @@ -4,7 +4,7 @@ from dataclasses import dataclass -from .spec import ScalingContext +from .spec import SCALING_RECIPE_MUP, ScalingContext @dataclass(frozen=True) @@ -67,6 +67,6 @@ def build_legacy_mup_training_policy( *, mup_width_mult: float, optimizer_type: str = 'adam' ) -> TrainingScalingPolicy: return TrainingScalingPolicy( - context=ScalingContext(use_mup=True, width_mult=mup_width_mult), + context=ScalingContext(recipe=SCALING_RECIPE_MUP, width_mult=mup_width_mult), optimizer_type=optimizer_type, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index dba7fcea1de..52bf50b6123 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -10,7 +10,11 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend -from megatron.core.parameterization import build_model_scaling_policy +from megatron.core.parameterization import ( + 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, @@ -358,25 +362,42 @@ class TransformerConfig(ModelParallelConfig): #################### # MuP (Maximal Update Parameterization) #################### + scaling_recipe: Optional[Literal['none', 'mup']] = None + """ + Canonical scaling recipe. ``none`` preserves standard parameterization, and ``mup`` + enables width MuP. If unset, legacy ``use_mup`` selects ``mup`` for backward + compatibility. + """ + + 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_head_dim: Optional[float] = None + """ + Canonical base attention head dimension for MuP attention scaling. This aliases the + deprecated mup_base_head_dim field. + """ + 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 @@ -387,15 +408,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. @@ -405,7 +426,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. """ #################### @@ -1949,20 +1971,11 @@ 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 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) + # MuP (Maximal Update Parameterization) configuration + if scaling_context.uses_width_mup: overridden_init_methods = [] if self.init_method is not None: overridden_init_methods.append("init_method") @@ -1972,7 +1985,7 @@ def __post_init__(self): 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 " + "scaling recipe 'mup' is enabled, but custom " + overridden_init_methods_text + f" {verb} set. This may break MuP initialization assumptions.", UserWarning, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index cd3ce44c3a4..f0ea0b71b73 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -38,6 +38,12 @@ ) from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.parameterization import ( + 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 +65,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 +259,35 @@ 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 _eval_pattern(pattern): """ Validate and evaluate a string containing a Python list expression """ assert isinstance(pattern, str) @@ -1761,6 +1797,9 @@ 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." + warn_deprecated_mup_aliases(args) + sync_legacy_mup_fields(args, build_scaling_context(args)) + # Print arguments. _print_args("arguments", args) @@ -2087,6 +2126,17 @@ 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_head_dim", + "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 +2697,86 @@ 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 " + "or 'mup' for width 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-head-dim', + type=float, + default=None, + help="Base attention head dimension for MuP attention scaling.", + ) + 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/yaml_arguments.py b/megatron/training/yaml_arguments.py index d44f4d31822..3d26f122dfb 100644 --- a/megatron/training/yaml_arguments.py +++ b/megatron/training/yaml_arguments.py @@ -16,8 +16,10 @@ 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 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 +40,20 @@ def env_constructor(loader, node): "bfloat16" : torch.bfloat16 } +DEFAULTABLE_SCALING_FIELDS = { + 'scaling_recipe', + 'scaling_base_hidden_size', + 'scaling_base_head_dim', + '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 +262,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 @@ -380,6 +403,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 +468,3 @@ def load_yaml(yaml_path): getattr(config_namespace, "global_batch_size", None) is not None ) return config_namespace - From 1b7ec227fcad5747df8fb72e81f456743b9979ee Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 11/19] Compare checkpoint scaling by resolved context --- megatron/training/checkpointing.py | 43 ++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 2363b7ae164..9292d4762b2 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_resolved_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,21 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('tensor_model_parallel_size') _compare('pipeline_model_parallel_size') + checkpoint_scaling_context = build_resolved_scaling_context(checkpoint_args) + args_scaling_context = build_resolved_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})." + ) + + +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_resolved_scaling_context(checkpoint_args) + ) + def isfile(filename) -> bool: if MultiStorageClientFeature.is_enabled(): @@ -1502,7 +1521,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 +1532,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 +1564,24 @@ 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_head_dim', 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) From e5e92b11b40755af4d4a1b82f546e2ba60143e38 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 12/19] Document and test canonical width-MuP --- .../internal/scaling_policy_infrastructure.md | 15 +- docs/index.md | 1 + docs/user-guide/index.md | 1 + docs/user-guide/scaling-recipes.md | 76 +++ megatron/training/checkpointing.py | 8 +- tests/unit_tests/transformer/test_mup.py | 475 +++++++++++++++++- 6 files changed, 563 insertions(+), 13 deletions(-) create mode 100644 docs/user-guide/scaling-recipes.md diff --git a/docs/api-guide/internal/scaling_policy_infrastructure.md b/docs/api-guide/internal/scaling_policy_infrastructure.md index c0b45156c97..3a5f06d70a2 100644 --- a/docs/api-guide/internal/scaling_policy_infrastructure.md +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -9,13 +9,14 @@ # Scaling Policy Infrastructure -This internal policy layer centralizes Megatron's existing parameterization hooks -without adding a new user-facing recipe surface. - -The current scaling context resolves the legacy `use_mup` fields into model-side and -optimizer-side decisions. Standard Megatron behavior is represented as the -identity policy, so code paths can call the same hooks whether or not MuP is -active. +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 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..7f1fc045d30 --- /dev/null +++ b/docs/user-guide/scaling-recipes.md @@ -0,0 +1,76 @@ + + +# 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. diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 9292d4762b2..637d6362fc4 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -39,7 +39,7 @@ from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer from megatron.core.parameterization import ( - build_resolved_scaling_context, + build_scaling_context, sync_legacy_mup_fields, ) from megatron.core.rerun_state_machine import get_rerun_state_machine @@ -179,8 +179,8 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('tensor_model_parallel_size') _compare('pipeline_model_parallel_size') - checkpoint_scaling_context = build_resolved_scaling_context(checkpoint_args) - args_scaling_context = build_resolved_scaling_context(args) + 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})." @@ -191,7 +191,7 @@ 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_resolved_scaling_context(checkpoint_args) + checkpoint_args, build_scaling_context(checkpoint_args) ) diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index 3eb8f01c941..3f952efb888 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -9,8 +9,12 @@ 4. LR override computation """ +import argparse +import dataclasses import logging import math +import warnings +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -26,10 +30,13 @@ from megatron.core.parameterization import ( build_legacy_mup_training_policy, build_model_scaling_policy, + build_scaling_context, ) from megatron.core.transformer.multi_token_prediction import process_mtp_loss from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import init_method_normal, mup_scaled_init_method_normal +from megatron.training.arguments import add_megatron_arguments +from megatron.training.yaml_arguments import core_config_from_args class TestMuPConfigValidation: @@ -46,6 +53,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.""" @@ -57,6 +66,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).""" @@ -75,6 +86,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.""" @@ -88,6 +101,464 @@ 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, + ) + + +class TestScalingRecipeSurfaces: + """Tests for public config surfaces that feed the scaling context.""" + + SCALING_FIELD_NAMES = { + 'scaling_recipe', + 'scaling_base_hidden_size', + 'scaling_base_head_dim', + '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 + + 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_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch): + """Checkpoint args must also clear stale CLI explicit-width provenance.""" + 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_head_dim=None, + use_mup=False, + mup_width_mult=1.0, + _mup_width_mult_explicit=True, + 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.mup_width_mult == 1.0 + assert runtime_args._mup_width_mult_explicit is False + 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.""" @@ -177,7 +648,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, @@ -189,7 +660,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, From 8191ade86d4209e04057675b119c6e198f6ada00 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Sat, 16 May 2026 10:23:37 +0530 Subject: [PATCH 13/19] Reset absent checkpoint scaling args --- megatron/training/checkpointing.py | 18 ++++++++++++++++++ tests/unit_tests/transformer/test_mup.py | 21 +++++++++++++++++---- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 637d6362fc4..08146d66839 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -187,12 +187,30 @@ def _compare(arg_name, old_arg_name=None, default=None): ) +_CHECKPOINT_SCALING_ARG_DEFAULTS = { + 'scaling_recipe': 'none', + 'scaling_base_hidden_size': None, + 'scaling_base_head_dim': 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: diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index 3f952efb888..c430324edd4 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -443,7 +443,7 @@ def test_load_checkpoint_args_clears_optional_scaling_fields(self, monkeypatch): assert runtime_args.mup_width_mult == pytest.approx(4.0) def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch): - """Checkpoint args must also clear stale CLI explicit-width provenance.""" + """Old no-scaling checkpoints must clear stale CLI scaling state.""" from megatron.training import checkpointing checkpoint_args = SimpleNamespace(hidden_size=1024) @@ -458,10 +458,15 @@ def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch) iteration=0, scaling_recipe='mup', scaling_base_hidden_size=256, - scaling_base_head_dim=None, - use_mup=False, - mup_width_mult=1.0, + scaling_base_head_dim=64, + 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, ) @@ -470,8 +475,16 @@ def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch) 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_head_dim 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): From 49763fb6687a329483fb5dbcd2411ff9086a2e72 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 14/19] Add depth-MuP config and validation --- megatron/core/parameterization/__init__.py | 36 +++++- megatron/core/parameterization/spec.py | 104 ++++++++++++++++-- .../core/parameterization/training_policy.py | 6 +- .../core/transformer/transformer_config.py | 82 ++++++++++++-- megatron/training/arguments.py | 80 +++++++++++++- megatron/training/yaml_arguments.py | 11 +- 6 files changed, 297 insertions(+), 22 deletions(-) diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py index f3a31914eed..48c9cec4b67 100644 --- a/megatron/core/parameterization/__init__.py +++ b/megatron/core/parameterization/__init__.py @@ -1,27 +1,41 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from .eval_runtime import depth_mup_eval_context, is_depth_mup_eval_enabled 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, @@ -31,20 +45,32 @@ build_scaling_user_config, sync_legacy_mup_fields, ) -from .training_policy import TrainingScalingPolicy, build_legacy_mup_training_policy +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', @@ -55,13 +81,21 @@ 'build_model_scaling_policy', 'build_scaling_context', 'build_scaling_user_config', + 'build_training_scaling_policy', + 'depth_mup_eval_context', 'get_parameterization_role', + 'is_depth_mup_eval_enabled', '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/spec.py b/megatron/core/parameterization/spec.py index 9f730f2cd6e..86d531e6963 100644 --- a/megatron/core/parameterization/spec.py +++ b/megatron/core/parameterization/spec.py @@ -8,14 +8,19 @@ SCALING_RECIPE_NONE = 'none' SCALING_RECIPE_MUP = 'mup' -SCALING_RECIPE_VALUES = (SCALING_RECIPE_NONE, SCALING_RECIPE_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']] = None + 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 @@ -28,15 +33,21 @@ class ScalingUserConfig: @dataclass(frozen=True) class ScalingContext: - """Internal scaling context for standard and width-MuP parameterization.""" + """Internal scaling context for standard, width-MuP, and depth-MuP.""" - recipe: Literal['none', '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: @@ -44,11 +55,15 @@ def enabled(self) -> bool: @property def uses_width_mup(self) -> bool: - return self.recipe == SCALING_RECIPE_MUP + return self.recipe in (SCALING_RECIPE_MUP, SCALING_RECIPE_DEPTH_MUP) @property def use_mup(self) -> bool: - return self.uses_width_mup + return self.recipe == SCALING_RECIPE_MUP + + @property + def is_depth_mup(self) -> bool: + return self.recipe == SCALING_RECIPE_DEPTH_MUP def _resolve_aliased_value( @@ -73,7 +88,11 @@ def _resolve_aliased_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, } @@ -88,6 +107,23 @@ def _non_default_scaling_fields(user_config: ScalingUserConfig) -> list[str]: 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') @@ -103,7 +139,13 @@ def build_scaling_user_config(config: Any) -> ScalingUserConfig: 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, @@ -144,15 +186,30 @@ def build_scaling_context(config: Any) -> ScalingContext: non_default_fields = _non_default_scaling_fields(user_config) if non_default_fields: raise ValueError( - "Scaling overrides require a non-'none' scaling recipe. Non-default fields: " - + ", ".join(non_default_fields) + "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) + 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.') @@ -174,14 +231,32 @@ def build_scaling_context(config: Any) -> ScalingContext: 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=SCALING_RECIPE_MUP, + 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), ) @@ -193,14 +268,23 @@ def sync_legacy_mup_fields(config: Any, context: ScalingContext) -> None: 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 index 01e98851084..d4dc4cec290 100644 --- a/megatron/core/parameterization/training_policy.py +++ b/megatron/core/parameterization/training_policy.py @@ -4,7 +4,7 @@ from dataclasses import dataclass -from .spec import SCALING_RECIPE_MUP, ScalingContext +from .spec import SCALING_RECIPE_MUP, ScalingContext, build_scaling_context @dataclass(frozen=True) @@ -70,3 +70,7 @@ def build_legacy_mup_training_policy( context=ScalingContext(recipe=SCALING_RECIPE_MUP, width_mult=mup_width_mult), optimizer_type=optimizer_type, ) + + +def build_training_scaling_policy(config, optimizer_type: str = 'adam') -> TrainingScalingPolicy: + return TrainingScalingPolicy(context=build_scaling_context(config), optimizer_type=optimizer_type) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 52bf50b6123..87167d0d8ee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -11,6 +11,8 @@ 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, @@ -362,11 +364,11 @@ class TransformerConfig(ModelParallelConfig): #################### # MuP (Maximal Update Parameterization) #################### - scaling_recipe: Optional[Literal['none', 'mup']] = None + scaling_recipe: Optional[Literal['none', 'mup', 'depth_mup']] = None """ - Canonical scaling recipe. ``none`` preserves standard parameterization, and ``mup`` - enables width MuP. If unset, legacy ``use_mup`` selects ``mup`` for backward - compatibility. + 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 @@ -375,12 +377,35 @@ class TransformerConfig(ModelParallelConfig): 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 """ Deprecated alias for scaling_recipe='mup'. Kept for checkpoint and script @@ -1973,25 +1998,66 @@ def __post_init__(self): 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.uses_width_mup: + 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( - "scaling recipe '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, ) - model_scaling_policy = build_model_scaling_policy(self) + 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 ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f0ea0b71b73..9fda3368fcf 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -39,6 +39,7 @@ 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, @@ -288,6 +289,46 @@ def warn_deprecated_mup_aliases(args): ) +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 _eval_pattern(pattern): """ Validate and evaluate a string containing a Python list expression """ assert isinstance(pattern, str) @@ -1797,6 +1838,7 @@ 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) warn_deprecated_mup_aliases(args) sync_legacy_mup_fields(args, build_scaling_context(args)) @@ -2129,7 +2171,11 @@ def _add_network_size_args(parser): # 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", @@ -2711,7 +2757,7 @@ def __call__(self, parser, namespace, values, option_string=None): default=None, help=( "Canonical parameterization recipe. Use 'none' for standard parameterization " - "or 'mup' for width MuP." + "'mup' for width MuP, or 'depth_mup' for dense GPT-style width-depth MuP." ), ) group.add_argument( @@ -2723,12 +2769,44 @@ def __call__(self, parser, namespace, values, option_string=None): "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( + '--allow-depth-mup-eval', + action='store_true', + help=( + "Allow validation/eval loss for depth_mup by entering a validation-only runtime " + "context. This does not enable general inference support." + ), + ) group.add_argument( '--use-mup', action='store_true', diff --git a/megatron/training/yaml_arguments.py b/megatron/training/yaml_arguments.py index 3d26f122dfb..958e0e5ff15 100644 --- a/megatron/training/yaml_arguments.py +++ b/megatron/training/yaml_arguments.py @@ -19,7 +19,7 @@ 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 warn_deprecated_mup_aliases +from megatron.training.arguments import validate_depth_mup_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 @@ -43,7 +43,11 @@ def env_constructor(loader, node): 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', @@ -367,6 +371,11 @@ 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__) + if not hasattr(args, 'allow_depth_mup_eval'): + args.allow_depth_mup_eval = False + validate_depth_mup_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 From 24d04b8240b7b55593612b85fac7ea38644104cf Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 15/19] Apply depth-MuP model and runtime scaling --- megatron/core/parameterization/__init__.py | 6 +- .../core/parameterization/eval_runtime.py | 28 ++++++ .../core/parameterization/model_policy.py | 31 ++++-- megatron/core/transformer/attention.py | 20 ++-- megatron/core/transformer/mlp.py | 23 +++-- .../core/transformer/transformer_layer.py | 96 +++++++++++++++++-- megatron/training/arguments.py | 8 -- megatron/training/training.py | 37 ++++--- megatron/training/yaml_arguments.py | 2 - 9 files changed, 192 insertions(+), 59 deletions(-) create mode 100644 megatron/core/parameterization/eval_runtime.py diff --git a/megatron/core/parameterization/__init__.py b/megatron/core/parameterization/__init__.py index 48c9cec4b67..b5e401ff0eb 100644 --- a/megatron/core/parameterization/__init__.py +++ b/megatron/core/parameterization/__init__.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from .eval_runtime import depth_mup_eval_context, is_depth_mup_eval_enabled +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, @@ -77,14 +77,14 @@ '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', - 'depth_mup_eval_context', 'get_parameterization_role', - 'is_depth_mup_eval_enabled', + 'is_scaling_policy_eval_allowed', 'is_embedding_class_parameter', 'is_embedding_or_output_parameter', 'is_hidden_bias_parameter', 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 index a3a84dee519..a57e0f2e280 100644 --- a/megatron/core/parameterization/model_policy.py +++ b/megatron/core/parameterization/model_policy.py @@ -2,6 +2,7 @@ from __future__ import annotations +import functools import math from dataclasses import dataclass from typing import Iterable, Optional @@ -20,7 +21,7 @@ @dataclass(frozen=True) class ModelScalingPolicy: - """Model-side policy for existing Megatron scaling behavior.""" + """Model-side policy for Megatron scaling recipes.""" context: ScalingContext @@ -34,20 +35,18 @@ def uses_width_mup(self) -> bool: @property def residual_branch_multiplier(self) -> float: - return 1.0 + return self.context.depth_mult**self.context.residual_branch_depth_power @property def dense_block_out_proj_init_multiplier(self) -> float: - return 1.0 + 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 - ) + 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): @@ -73,11 +72,19 @@ def dense_block_output_init_method( num_layers: int, is_hybrid_model: bool, output_layer_init_method_is_user_provided: bool, + apply_depth_hook: bool = True, ): - del init_method_std, num_layers, is_hybrid_model if output_layer_init_method_is_user_provided: return default_init_method - 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, @@ -109,7 +116,13 @@ def scale_output_logits(self, logits: Tensor) -> Tensor: def scale_residual_branch_output( self, output_with_bias: tuple[Tensor, Tensor | None] ) -> tuple[Tensor, Tensor | None]: - return output_with_bias + 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: diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 55da3a1321b..c9bbe6da881 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -393,17 +393,23 @@ def __init__( # Output. model_scaling_policy = build_model_scaling_policy(self.config) - self.linear_proj = submodules.linear_proj( - self.query_projection_size, - self.config.hidden_size, - config=self.config, - init_method=model_scaling_policy.dense_block_output_init_method( + 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=False, - ), + 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=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 1b7bac4dba7..b19c85b0b50 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -171,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, ): @@ -235,19 +236,25 @@ def __init__( 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=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=False, - ), + init_method=fc2_init_method, bias=self.config.add_bias_linear, input_is_parallel=True, skip_bias_add=True, @@ -393,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 @@ -407,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_layer.py b/megatron/core/transformer/transformer_layer.py index 4243ec0898d..be10c3885c3 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -17,7 +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 +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 @@ -44,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 ): @@ -313,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) @@ -405,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, @@ -424,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) @@ -675,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 @@ -682,9 +732,6 @@ def _forward_attention( hidden_states = attention_output_with_bias[0] else: with self.bias_dropout_add_exec_handler(): - attention_output_with_bias = self.model_scaling_policy.scale_residual_branch_output( - attention_output_with_bias - ) hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)( attention_output_with_bias, residual, self.hidden_dropout ) @@ -750,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, @@ -927,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 @@ -934,9 +1015,6 @@ def _forward_post_mlp( hidden_states = mlp_output_with_bias[0] else: with self.bias_dropout_add_exec_handler(): - mlp_output_with_bias = self.model_scaling_policy.scale_residual_branch_output( - mlp_output_with_bias - ) hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)( mlp_output_with_bias, residual, self.hidden_dropout ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 9fda3368fcf..6c869d50af0 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2799,14 +2799,6 @@ def __call__(self, parser, namespace, values, option_string=None): default=None, help="Relative depth exponent for dense block output projection initialization.", ) - group.add_argument( - '--allow-depth-mup-eval', - action='store_true', - help=( - "Allow validation/eval loss for depth_mup by entering a validation-only runtime " - "context. This does not enable general inference support." - ), - ) group.add_argument( '--use-mup', action='store_true', diff --git a/megatron/training/training.py b/megatron/training/training.py index 11c8128792e..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 get_param_group_identifier_tuple +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 = get_param_group_identifier_tuple + 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 958e0e5ff15..a7d1e038a2f 100644 --- a/megatron/training/yaml_arguments.py +++ b/megatron/training/yaml_arguments.py @@ -371,8 +371,6 @@ 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__) - if not hasattr(args, 'allow_depth_mup_eval'): - args.allow_depth_mup_eval = False validate_depth_mup_optimizer_support(args) warn_deprecated_mup_aliases(args) sync_legacy_mup_fields(args, build_scaling_context(args)) From 14cea4981c1962b8c1452b0e531c21918d5816ca Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 16/19] Add depth-MuP optimizer policy --- megatron/core/optimizer/__init__.py | 235 +++++++++++++----- megatron/core/optimizer/optimizer.py | 18 +- megatron/core/parameterization/roles.py | 82 +++++- .../core/parameterization/training_policy.py | 74 +++++- 4 files changed, 327 insertions(+), 82 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 070b3bd7d33..5bfdcd9abdc 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -60,8 +60,10 @@ 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 @@ -97,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} @@ -134,7 +146,6 @@ 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]: @@ -148,23 +159,35 @@ def get_mup_config_overrides( def get_scaling_config_overrides( config: OptimizerConfig, scaling_policy: TrainingScalingPolicy ) -> Dict[ParamKey, ParamGroupOverride]: - """Get optimizer overrides from an internal training scaling policy. - - This compatibility path keeps the public MuP surface unchanged and factors the - existing per-parameter MuP LR/epsilon behavior through this policy seam. - """ + """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 scaling_policy.is_adam_optimizer: - message += " MuP Adam epsilon scaling remains applied to hidden matrix-like parameters." + message += ( + " Adam epsilon scaling remains applied according to the active recipe's " + "parameter-class rules." + ) log_single_rank(logger, logging.WARNING, message) if scaling_policy.is_muon_optimizer: @@ -179,11 +202,16 @@ def get_scaling_config_overrides( "Muon-managed matrices with MuP.", ) - if scaling_policy.context.width_mult == 1.0: - return {} - base_lr = config.lr base_min_lr = config.min_lr + 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): @@ -197,71 +225,144 @@ def should_scale_vector_like_lr_with_mup(param: torch.nn.Parameter, param_name: 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_embedding_class_parameter(param, param_name): - return False - 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 + 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 - return True + return is_hidden_matrix_parameter(param, param_name) - mup_overrides: Dict[ParamKey, ParamGroupOverride] = {} + def should_scale_hidden_vector_eps(param: torch.nn.Parameter, param_name: str) -> bool: + return is_hidden_vector_parameter(param, param_name) + + 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: - vector_like_lr_override: ParamGroupOverride = {} - if base_lr is not None: - vector_like_lr_override["max_lr"] = base_lr * scaling_policy.hidden_vector_lr_multiplier - if base_min_lr is not None: - vector_like_lr_override["min_lr"] = ( - base_min_lr * scaling_policy.hidden_vector_lr_multiplier + 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 + + vector_like_lr_override: ParamGroupOverride = {} + 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: - lr_override["max_lr"] = base_lr * scaling_policy.hidden_lr_multiplier - if base_min_lr is not None: - lr_override["min_lr"] = base_min_lr * scaling_policy.hidden_lr_multiplier + 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 and hidden_lr_mult != 1.0: + lr_override["min_lr"] = base_min_lr * hidden_lr_mult eps_override: ParamGroupOverride = {} - if scaling_policy.is_adam_optimizer and config.adam_eps is not None: - eps_override["eps"] = config.adam_eps * scaling_policy.hidden_eps_multiplier + 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_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], diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 7b98be7b79c..e612ba8cf20 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -94,17 +94,26 @@ 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 the stable legacy identifier for optimizer param-group matching and resume.""" + """Return a stable identifier for optimizer param-group matching and resume.""" values = [] for key in param_group_identifier_keys: if key in param_group: @@ -116,6 +125,11 @@ def get_param_group_identifier_tuple(param_group: Dict) -> tuple: 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): """ Base class for all Megatron optimizers. diff --git a/megatron/core/parameterization/roles.py b/megatron/core/parameterization/roles.py index 3d5513f9f46..28b168c545f 100644 --- a/megatron/core/parameterization/roles.py +++ b/megatron/core/parameterization/roles.py @@ -12,12 +12,31 @@ 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( @@ -62,10 +81,69 @@ def is_vector_like_parameter(param: Any, param_name: Optional[str] = None) -> bo return param.dim() <= 1 -def is_hidden_matrix_parameter(param: Any, param_name: Optional[str] = None) -> bool: +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 - return param.dim() > 1 + 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: diff --git a/megatron/core/parameterization/training_policy.py b/megatron/core/parameterization/training_policy.py index d4dc4cec290..2cfc89dc883 100644 --- a/megatron/core/parameterization/training_policy.py +++ b/megatron/core/parameterization/training_policy.py @@ -6,14 +6,22 @@ from .spec import SCALING_RECIPE_MUP, ScalingContext, build_scaling_context - @dataclass(frozen=True) class TrainingScalingPolicy: - """Optimizer-side policy for existing Megatron MuP hyperparameter multipliers.""" + """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 @@ -32,19 +40,25 @@ def is_sgd_optimizer(self) -> bool: @property def is_adam_optimizer(self) -> bool: - return 'adam' in self.optimizer_type_lower + return self.optimizer_type_lower == 'adam' @property def is_muon_optimizer(self) -> bool: return 'muon' in self.optimizer_type_lower @property - def hidden_lr_multiplier(self) -> float: + 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 - if self.is_sgd_optimizer: - return 1.0 - return 1.0 / self.context.width_mult + 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: @@ -52,17 +66,59 @@ def hidden_vector_lr_multiplier(self) -> float: 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: @@ -70,7 +126,3 @@ def build_legacy_mup_training_policy( context=ScalingContext(recipe=SCALING_RECIPE_MUP, width_mult=mup_width_mult), optimizer_type=optimizer_type, ) - - -def build_training_scaling_policy(config, optimizer_type: str = 'adam') -> TrainingScalingPolicy: - return TrainingScalingPolicy(context=build_scaling_context(config), optimizer_type=optimizer_type) From e0e4884b3c735c4a836bb7a21ccbc242aa08e1b4 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Fri, 15 May 2026 23:10:15 +0530 Subject: [PATCH 17/19] Document and test depth-MuP --- .../internal/scaling_policy_infrastructure.md | 72 ++- docs/user-guide/scaling-recipes.md | 56 ++- megatron/training/checkpointing.py | 4 + tests/unit_tests/transformer/test_mup.py | 443 +++++++++++++++++- 4 files changed, 560 insertions(+), 15 deletions(-) diff --git a/docs/api-guide/internal/scaling_policy_infrastructure.md b/docs/api-guide/internal/scaling_policy_infrastructure.md index 3a5f06d70a2..552ea2d8596 100644 --- a/docs/api-guide/internal/scaling_policy_infrastructure.md +++ b/docs/api-guide/internal/scaling_policy_infrastructure.md @@ -12,11 +12,12 @@ 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. +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 @@ -31,12 +32,24 @@ covers: - output logit scaling; - residual branch output hooks. -For non-MuP configs, every hook returns the current Megatron default. +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. The policy currently preserves the existing MuP rules: +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`; @@ -44,8 +57,24 @@ training scaling policy. The policy currently preserves the existing MuP rules: - Muon-managed matrices stay on Muon scaling rather than Adam-style MuP LR overrides. +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. +and delegates to the policy implementation for the legacy width-MuP surface. ## Parameter Metadata @@ -58,8 +87,25 @@ 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. +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/user-guide/scaling-recipes.md b/docs/user-guide/scaling-recipes.md index 7f1fc045d30..4249314c815 100644 --- a/docs/user-guide/scaling-recipes.md +++ b/docs/user-guide/scaling-recipes.md @@ -12,12 +12,13 @@ Scaling recipes choose the parameterization used to transfer hyperparameters between model sizes. The canonical flag is `--scaling-recipe`. -Megatron currently exposes two recipes: +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 @@ -63,6 +64,59 @@ deprecated as user-facing inputs: 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 diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 08146d66839..fb7fb810c41 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1584,7 +1584,11 @@ def _set_arg(arg_name, old_arg_name=None, force=False, allow_none=False): _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) diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index c430324edd4..4b0c642692a 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -11,6 +11,7 @@ import argparse import dataclasses +import functools import logging import math import warnings @@ -28,17 +29,32 @@ 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 +from megatron.training.arguments import add_megatron_arguments, validate_depth_mup_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: """Tests for MuP config validation and width_mult computation.""" @@ -201,6 +217,119 @@ def test_scaling_override_without_recipe_is_rejected(self): 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, + ) + class TestScalingRecipeSurfaces: """Tests for public config surfaces that feed the scaling context.""" @@ -208,7 +337,11 @@ class TestScalingRecipeSurfaces: 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', @@ -239,6 +372,23 @@ def test_cli_parser_accepts_canonical_scaling_args(self): 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) @@ -442,6 +592,97 @@ def test_load_checkpoint_args_clears_optional_scaling_fields(self, monkeypatch): 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 @@ -923,6 +1164,206 @@ 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.""" From b2c303baa6f7923e5f4e48406e11398267c963c9 Mon Sep 17 00:00:00 2001 From: plugyawn Date: Sat, 16 May 2026 10:27:40 +0530 Subject: [PATCH 18/19] Fix depth-MuP resume and vector weight decay --- megatron/training/checkpointing.py | 4 ++ tests/unit_tests/test_optimizer.py | 8 ++- tests/unit_tests/transformer/test_mup.py | 92 ++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 2 deletions(-) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index fb7fb810c41..ef1c7034b21 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -190,7 +190,11 @@ def _compare(arg_name, old_arg_name=None, default=None): _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, diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index ee35b77b701..0b8fb62fb7b 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -75,7 +75,7 @@ def test_param_group_identifier_tuple_tolerates_missing_optional_keys(): ident = get_param_group_identifier_tuple(group) - assert ident == (1.0, 1.0, False, False) + assert ident == (1.0, 1.0, False, False, None, None) def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): @@ -88,7 +88,7 @@ def test_param_group_identifier_tuple_reads_pre_keys_and_optional_fields(): ident = get_param_group_identifier_tuple(group) - assert ident == (0.0, 0.5, False, True) + assert ident == (0.0, 0.5, False, True, None, None) def test_param_group_identifier_tuple_defaults_missing_legacy_fields(): @@ -107,6 +107,8 @@ def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): "is_decoupled_lr": False, "max_lr": 2e-4, "min_lr": 2e-6, + "eps": 1e-8, + "optimizer": "adam", "params": [0], } checkpoint_group = { @@ -116,6 +118,8 @@ def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): "is_decoupled_lr": False, "max_lr": 1e-4, "min_lr": 1e-6, + "eps": 1e-8, + "optimizer": "adam", "params": [7], } diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index 4b0c642692a..d840ab1aeb8 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -699,7 +699,11 @@ def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch) 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, @@ -717,7 +721,11 @@ def test_load_non_mup_checkpoint_clears_width_mult_provenance(self, monkeypatch) 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 @@ -1368,6 +1376,90 @@ def forward(self, *args, **kwargs): 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_resolved_training_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) From ef280a9308bb27fca397bbf997570071e75b76af Mon Sep 17 00:00:00 2001 From: plugyawn Date: Sat, 16 May 2026 10:51:35 +0530 Subject: [PATCH 19/19] Test depth optimizer identifier defaults --- tests/unit_tests/test_optimizer.py | 2 +- tests/unit_tests/transformer/test_mup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 0b8fb62fb7b..9e16c731b3c 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -96,7 +96,7 @@ def test_param_group_identifier_tuple_defaults_missing_legacy_fields(): ident = get_param_group_identifier_tuple(group) - assert ident == (1.0, 1.0, False, False) + assert ident == (1.0, 1.0, False, False, None, None) def test_param_group_matching_ignores_mutable_scheduler_values_on_resume(): diff --git a/tests/unit_tests/transformer/test_mup.py b/tests/unit_tests/transformer/test_mup.py index d840ab1aeb8..32cb603689a 100644 --- a/tests/unit_tests/transformer/test_mup.py +++ b/tests/unit_tests/transformer/test_mup.py @@ -1396,7 +1396,7 @@ def _depth_mup_adamw_overrides(self, *, apply_wd_to_qk_layernorm=False): mup_base_head_dim=None, mup_attn_scale_power=1.0, ) - scaling_policy = build_resolved_training_policy(config_args, optimizer_type='adam') + scaling_policy = build_training_scaling_policy(config_args, optimizer_type='adam') optimizer_config = OptimizerConfig( optimizer='adam', lr=1e-3,