From c5f347820f4fdc5414afa96e819dc12658911e2c Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Wed, 16 Sep 2026 22:56:19 +0000 Subject: [PATCH 1/3] [port] MLA QK-norm resolver and backends.get_backend (upstream #4452) Megatron-Bridge main (radixark bridge-new) imports megatron.core.transformer.mla_qk_norm_config.get_backend and overrides MLASelfAttention._resolve_qk_norm_config in its DeepSeek/GLM-5 attention (models/deepseek/attention.py); on miles-main the import fails and the hook does not exist, so every Bridge-backed path (LoRA, multi-LoRA, bridge-mode weight update) dies at import. Ported from upstream 4464d1c1976b "Port Multi-Latent Attention to HybridModel (#4452)", MLA half only: * new megatron/core/transformer/mla_qk_norm_config.py (identical to upstream main @ f6c33bde4) * MLASelfAttention / FusedMLASelfAttention build q/kv norms and up-projections through _resolve_qk_norm_config -> QKNormConfigResolver * AbsorbedMLA does the same; the combined K/V up-projection goes through the resolver, the miles-main split K/V path keeps its explicit submodules * backends.get_backend(transformer_impl) (Literal import) * the qk_layernorm test refactors that applied cleanly Not ported: the HybridModel MLA support (hybrid_block / hybrid_layer_allocation / hybrid_layer_specs, Symbols.MLA in arguments.py, hybrid tests) -- it conflicts with the miles-main hybrid code and nothing in miles uses HybridModel MLA. (cherry picked from commit 4464d1c1976b2563a70936f9d5d037817e60b2fa, hybrid-model parts dropped) Signed-off-by: Yusheng Su --- megatron/core/models/backends.py | 18 +- .../absorbed_mla.py | 15 +- .../core/transformer/mla_qk_norm_config.py | 293 ++++++++++++++++++ .../transformer/multi_latent_attention.py | 32 +- .../a2a_overlap/test_schedule_layer_1f1b.py | 9 +- .../distributed/test_finalize_model_grads.py | 7 +- .../transformer/test_submodule_callables.py | 4 +- 7 files changed, 358 insertions(+), 20 deletions(-) create mode 100644 megatron/core/transformer/mla_qk_norm_config.py diff --git a/megatron/core/models/backends.py b/megatron/core/models/backends.py index a270161ddd6..e29265c04f4 100644 --- a/megatron/core/models/backends.py +++ b/megatron/core/models/backends.py @@ -4,7 +4,7 @@ import warnings from abc import abstractmethod from functools import partial -from typing import Optional, Protocol, cast +from typing import Literal, Optional, Protocol, cast from megatron.core.extensions.transformer_engine import ( TEColumnParallelGroupedLinear, @@ -200,3 +200,19 @@ def grouped_mlp_modules(self, moe_use_grouped_gemm: bool) -> ExpertsBuilder: activation_func=self.activation_func(), ), ) + + +def get_backend( + transformer_impl: Literal["local", "transformer_engine", "inference_optimized"] +) -> BackendSpecProvider: + """Return the backend that's selected with the given `transformer_impl`.""" + if transformer_impl == "transformer_engine": + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + + return TESpecProvider() + elif transformer_impl == "inference_optimized": + return InferenceSpecProvider() + elif transformer_impl == "local": + return LocalSpecProvider() + else: + raise ValueError(f"unknown transformer_impl='{transformer_impl}'") diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index 6c0f38eaef5..1956dce0ca6 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -36,6 +36,7 @@ ) from megatron.core.transformer.attention import Attention from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import MLATransformerConfig from megatron.core.utils import deprecate_inference_params, get_pg_size, is_te_min_version @@ -167,6 +168,10 @@ def __init__( name=name, ) + # Resolve which classes to use for Q and KV linear up projections and norms, based on + # QK-norm selection. + layer_classes = QKNormConfigResolver(self.config, submodules).resolve() + assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA" assert not ( config.tensor_model_parallel_size > 1 and not config.sequence_parallel @@ -266,7 +271,7 @@ def __init__( if self.config.q_lora_rank is None: # Not projecting query self.linear_q_proj = build_module( - submodules.linear_q_proj, + layer_classes["linear_q_proj"], self.config.hidden_size, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -312,7 +317,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - submodules.linear_q_up_proj, + layer_classes["linear_q_up_proj"], self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -373,7 +378,7 @@ def __init__( if self._uses_combined_kv_up_projection: self.linear_kv_up_proj = build_module( - submodules.linear_kv_up_proj, + layer_classes["linear_kv_up_proj"], self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), @@ -419,14 +424,14 @@ def __init__( if self.config.q_lora_rank is not None: self.q_layernorm = build_module( - submodules.q_layernorm, + layer_classes["q_layernorm"], hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) self.kv_layernorm = build_module( - submodules.kv_layernorm, + layer_classes["kv_layernorm"], hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, diff --git a/megatron/core/transformer/mla_qk_norm_config.py b/megatron/core/transformer/mla_qk_norm_config.py new file mode 100644 index 00000000000..e14066a105d --- /dev/null +++ b/megatron/core/transformer/mla_qk_norm_config.py @@ -0,0 +1,293 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +Resolve MLA and DSA Q/KV norm configuration from a layer specification. +""" + +from typing import NoReturn + +from megatron.core.models.backends import get_backend +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.torch_norm import LayerNormBuilder +from megatron.core.transformer.transformer_config import MLATransformerConfig + +__all__ = [] + +_QKNormResolvedConfig = dict[str, ModuleSpec | type | LayerNormBuilder] + + +class QKNormConfigResolver: + """Validate and resolve Q/KV norm placement for MLA and DSA. + + Q/KV norm can be represented either by a standalone norm module or by + a fused norm+linear projection. MLA can use the fused form; DSA cannot + because it needs the normalized Q/KV values outside the projection. + + Constraints: + - `qk_l2_norm` is unsupported for MLA/DSA. + - A standalone Q norm is only usable when `q_lora_rank` is set. + - Explicit norm modules cannot be paired with fused norm+linear projections. + - Disabled QK norm rejects both explicit norms and fused norm+linear projections. + - DSA with QK norm requires non-fused projections and standalone Q/KV norms. + """ + + def __init__(self, config: MLATransformerConfig, submodules) -> None: + """Capture the configuration, requested modules, and backend implementations.""" + self.config = config + self.submodules = submodules + self.has_q_lora = config.q_lora_rank is not None + self.is_dsa = config.experimental_attention_variant == "dsa" + self.variant_str = "DSA" if self.is_dsa else "MLA" + + backend = get_backend(config.transformer_impl) + self.qk_norm_impl = backend.layer_norm( + rms_norm=config.normalization == "RMSNorm", for_qk=True + ) + self.linear_impl = backend.column_parallel_linear() + self.fused_norm_linear_impl = backend.column_parallel_layer_norm_linear() + + def resolve(self) -> _QKNormResolvedConfig: + """Validate the specification and return the modules to instantiate. + + Returns: + The Q/KV norms and projections after applying the MLA or DSA constraints. + + Raises: + ValueError: If the requested norm placement is unsupported or conflicting. + """ + if self.config.qk_l2_norm: + raise ValueError(f"qk_l2_norm is not supported with {self.variant_str}.") + + self._reject_common_spec_conflicts() + if not self.config.qk_layernorm: + return self._resolve_disabled_qk_layernorm() + if self.is_dsa: + return self._resolve_dsa_qk_layernorm() + return self._resolve_mla_qk_layernorm() + + def _resolve_disabled_qk_layernorm(self) -> _QKNormResolvedConfig: + """Resolve projections when Q/KV normalization is disabled. + + Explicit norm modules and fused norm-linear projections are rejected because + they would still introduce Q/KV normalization. + """ + linear_q_proj_cls = IdentityOp + linear_q_up_proj_cls = IdentityOp + + if self.has_q_lora: + self._reject_disabled_norm( + self.submodules.linear_q_up_proj, + self.submodules.q_layernorm, + "linear_q_up_proj", + "q_layernorm", + ) + linear_q_up_proj_cls = self.submodules.linear_q_up_proj or self.linear_impl + else: + if self._is_fused_norm_linear(self.submodules.linear_q_proj): + raise ValueError( + f"spec sets linear_q_proj={self.submodules.linear_q_proj}, but " + "qk_layernorm/qk_l2_norm are supposed to be disabled" + ) + linear_q_proj_cls = self.submodules.linear_q_proj or self.linear_impl + + self._reject_disabled_norm( + self.submodules.linear_kv_up_proj, + self.submodules.kv_layernorm, + "linear_kv_up_proj", + "kv_layernorm", + ) + return self._result( + linear_q_proj=linear_q_proj_cls, + linear_q_up_proj=linear_q_up_proj_cls, + linear_kv_up_proj=self.submodules.linear_kv_up_proj or self.linear_impl, + q_layernorm=IdentityOp, + kv_layernorm=IdentityOp, + ) + + def _resolve_dsa_qk_layernorm(self) -> _QKNormResolvedConfig: + """Resolve DSA's standalone Q/KV norms and non-fused projections. + + DSA consumes the normalized Q/KV values outside the projection, so it cannot + use fused norm-linear projections. + """ + if not self.has_q_lora: + raise ValueError( + "`qk_layernorm=True` with `q_lora_rank is None` is not supported for DSA " + "because DSA cannot fuse Q norm into `linear_q_proj`." + ) + + return self._result( + linear_q_proj=IdentityOp, + linear_q_up_proj=self._dsa_linear_or_default( + self.submodules.linear_q_up_proj, "linear_q_up_proj" + ), + linear_kv_up_proj=self._dsa_linear_or_default( + self.submodules.linear_kv_up_proj, "linear_kv_up_proj" + ), + q_layernorm=self._default_if_trivial(self.submodules.q_layernorm, self.qk_norm_impl), + kv_layernorm=self._default_if_trivial(self.submodules.kv_layernorm, self.qk_norm_impl), + ) + + def _resolve_mla_qk_layernorm(self) -> _QKNormResolvedConfig: + """Resolve MLA norms, fusing them into projections when no norm is explicit.""" + q_norm_cls = self.submodules.q_layernorm or IdentityOp + linear_q_proj_cls = IdentityOp + linear_q_up_proj_cls = IdentityOp + + if self.has_q_lora: + if self._is_trivial(q_norm_cls): + linear_q_up_proj_cls = self._mla_fused_linear_or_default( + self.submodules.linear_q_up_proj, "linear_q_up_proj" + ) + else: + linear_q_up_proj_cls = self._non_fused_or_default( + self.submodules.linear_q_up_proj, "linear_q_up_proj" + ) + else: + linear_q_proj_cls = self._mla_fused_linear_or_default( + self.submodules.linear_q_proj, "linear_q_proj" + ) + + kv_norm_cls = self.submodules.kv_layernorm or IdentityOp + if self._is_trivial(kv_norm_cls): + linear_kv_up_proj_cls = self._mla_fused_linear_or_default( + self.submodules.linear_kv_up_proj, "linear_kv_up_proj" + ) + else: + linear_kv_up_proj_cls = self._non_fused_or_default( + self.submodules.linear_kv_up_proj, "linear_kv_up_proj" + ) + + return self._result( + linear_q_proj=linear_q_proj_cls, + linear_q_up_proj=linear_q_up_proj_cls, + linear_kv_up_proj=linear_kv_up_proj_cls, + q_layernorm=q_norm_cls, + kv_layernorm=kv_norm_cls, + ) + + def _reject_common_spec_conflicts(self) -> None: + """Reject conflicts that apply regardless of the selected attention variant.""" + if not self.has_q_lora and not self._is_trivial(self.submodules.q_layernorm): + self._raise_unused_q_norm() + if self.has_q_lora: + self._reject_explicit_norm_with_fused_linear( + self.submodules.linear_q_up_proj, + self.submodules.q_layernorm, + "linear_q_up_proj", + "q_layernorm", + ) + self._reject_explicit_norm_with_fused_linear( + self.submodules.linear_kv_up_proj, + self.submodules.kv_layernorm, + "linear_kv_up_proj", + "kv_layernorm", + ) + + def _reject_disabled_norm(self, module_spec, norm_spec, module_name, norm_name) -> None: + """Reject a norm module or fused projection when Q/KV norm is disabled.""" + if self._is_fused_norm_linear(module_spec) or not self._is_trivial(norm_spec): + raise ValueError( + f"spec sets {module_name}={module_spec} and " + f"{norm_name}={norm_spec}, but " + "qk_layernorm/qk_l2_norm are supposed to be disabled" + ) + + def _reject_explicit_norm_with_fused_linear( + self, module_spec, norm_spec, module_name, norm_name + ) -> None: + """Reject specifying the same norm both explicitly and inside a projection.""" + if not self._is_trivial(norm_spec) and self._is_fused_norm_linear(module_spec): + raise ValueError( + f"`{norm_name}={norm_spec}` is non-trivial " + f"and `{module_name}={module_spec}` is a " + f"fused norm+linear; either unset `{norm_name}` or use a " + f"linear layer without norm fusion for `{module_name}`" + ) + + def _non_fused_or_default(self, module_spec, module_name): + """Return a linear implementation, requiring it not to fuse normalization.""" + linear_cls = module_spec or self.linear_impl + self._require_linear(linear_cls, module_name) + if self._is_fused_norm_linear(linear_cls): + raise ValueError( + f"`{module_name}={module_spec}` is fused norm+linear, but a non-fused linear " + f"is required" + ) + return linear_cls + + def _dsa_linear_or_default(self, module_spec, module_name): + """Return DSA's non-fused projection implementation. + + This uses a DSA-specific diagnostic so the rejected constraint is clear. + """ + linear_cls = module_spec or self.linear_impl + self._require_linear(linear_cls, module_name) + if self._is_fused_norm_linear(linear_cls): + raise ValueError( + f"`{module_name}={module_spec}` is fused norm+linear, " + f"which is not supported for DSA." + ) + return linear_cls + + def _mla_fused_linear_or_default(self, module_spec, module_name): + """Return a fused MLA projection, using the backend default when available.""" + if self._is_fused_norm_linear(module_spec): + return module_spec + return self._require_linear(self.fused_norm_linear_impl, module_name) + + def _require_linear(self, module_spec, module_name): + """Return a configured projection or report that no viable implementation exists.""" + if module_spec is None: + raise RuntimeError( + "qk_layernorm requires TransformerEngine or " + "q_layernorm/kv_layernorm to be set in the spec " + f"to build `{module_name}`." + ) + return module_spec + + def _raise_unused_q_norm(self) -> NoReturn: + """Report an explicit Q norm that has no Q-LoRA projection to consume it.""" + help_msg = "" + if not self._is_fused_norm_linear(self.submodules.linear_q_proj): + help_msg = ( + f"Please use a fused norm+linear for " + f"`linear_q_proj={self.submodules.linear_q_proj}` if " + f"you intend to have a Q-norm." + ) + raise ValueError( + f"`q_layernorm={self.submodules.q_layernorm}` is non-trivial, " + f"but `q_lora_rank is None`, meaning it will not be used." + f"{help_msg}" + ) + + def _is_fused_norm_linear(self, module_spec) -> bool: + """Return whether a module specification selects the backend fused projection.""" + module_cls = module_spec.module if isinstance(module_spec, ModuleSpec) else module_spec + return self.fused_norm_linear_impl is not None and module_cls is self.fused_norm_linear_impl + + @staticmethod + def _is_trivial(module_spec) -> bool: + """Return whether a norm slot is unset or explicitly an identity operation.""" + return module_spec in (None, IdentityOp) + + @classmethod + def _default_if_trivial(cls, module_spec, default): + """Replace an unset or identity specification with the supplied default.""" + if cls._is_trivial(module_spec): + return default + return module_spec + + @staticmethod + def _result( + *, linear_q_proj, linear_q_up_proj, linear_kv_up_proj, q_layernorm, kv_layernorm + ) -> _QKNormResolvedConfig: + """Package the resolved Q/KV norms and projections in the caller's schema.""" + return dict( + linear_q_proj=linear_q_proj, + linear_q_up_proj=linear_q_up_proj, + linear_kv_up_proj=linear_kv_up_proj, + q_layernorm=q_layernorm, + kv_layernorm=kv_layernorm, + ) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 4304662db21..d3856418ad6 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -37,6 +37,7 @@ ) from megatron.core.transformer.attention import Attention, LinearProjBuilder from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -531,10 +532,14 @@ def __init__( name=name, ) + # Resolve which classes to use for Q and KV linear up projections and norms, based on + # QK-norm selection. + layer_classes = self._resolve_qk_norm_config(submodules) + if self.config.q_lora_rank is None: # Not projecting query self.linear_q_proj = build_module( - submodules.linear_q_proj, + layer_classes["linear_q_proj"], self.config.hidden_size, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -581,7 +586,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - submodules.linear_q_up_proj, + layer_classes["linear_q_up_proj"], self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -628,7 +633,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - submodules.linear_kv_up_proj, + layer_classes["linear_kv_up_proj"], self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -643,18 +648,24 @@ def __init__( ) if self.config.q_lora_rank is not None: - self.q_layernorm = submodules.q_layernorm( + self.q_layernorm = layer_classes["q_layernorm"]( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = submodules.kv_layernorm( + self.kv_layernorm = layer_classes["kv_layernorm"]( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) + def _resolve_qk_norm_config( + self, submodules + ) -> dict[str, ModuleSpec | type | LayerNormBuilder]: + """Resolve which Q/KV norm and up-projection implementations to build.""" + return QKNormConfigResolver(self.config, submodules).resolve() + def _qkv_down_projection(self, hidden_states): """Unfused q/kv down projection path.""" if self.config.q_lora_rank is not None: @@ -1291,6 +1302,9 @@ def __init__( "FusedMLASelfAttention requires q_lora_rank to be set; " "fallback to MLASelfAttention for q_lora_rank=None." ) + # Resolve which linear class to use for Q and KV up projections, + # based on QK-norm selection. + layer_classes = self._resolve_qk_norm_config(submodules) qkv_down_proj_kwargs = {} if submodules.linear_qkv_down_proj in [TELinear]: @@ -1326,7 +1340,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - submodules.linear_q_up_proj, + layer_classes["linear_q_up_proj"], self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -1341,7 +1355,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - submodules.linear_kv_up_proj, + layer_classes["linear_kv_up_proj"], self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -1355,12 +1369,12 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) - self.q_layernorm = submodules.q_layernorm( + self.q_layernorm = layer_classes["q_layernorm"]( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = submodules.kv_layernorm( + self.kv_layernorm = layer_classes["kv_layernorm"]( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py index d1bb97ca0cd..77996e230ad 100644 --- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py +++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py @@ -450,7 +450,12 @@ def test_mtp_layer_overlap(self, dispatcher_type, flex_backend, fp8_flag): Verifies all-to-all overlap optimization in MTP layer produces the same results as the reference implementation. """ - extra_kwargs = {"mtp_num_layers": 1, "mtp_loss_scaling_factor": 1.1} + qk_layernorm = True + extra_kwargs = { + "mtp_num_layers": 1, + "mtp_loss_scaling_factor": 1.1, + "qk_layernorm": qk_layernorm, + } apply_flex_backend_kwargs(extra_kwargs, dispatcher_type, flex_backend) if fp8_flag is not None: extra_kwargs["fp8_recipe"] = fp8_flag[1] @@ -463,7 +468,7 @@ def test_mtp_layer_overlap(self, dispatcher_type, flex_backend, fp8_flag): transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( num_experts=16, moe_grouped_gemm=True, - qk_layernorm=True, + qk_layernorm=qk_layernorm, multi_latent_attention=True, ) mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec, True) diff --git a/tests/unit_tests/distributed/test_finalize_model_grads.py b/tests/unit_tests/distributed/test_finalize_model_grads.py index ee535c29baf..91a03bf5f41 100644 --- a/tests/unit_tests/distributed/test_finalize_model_grads.py +++ b/tests/unit_tests/distributed/test_finalize_model_grads.py @@ -120,6 +120,7 @@ def test_finalize_model_grads_requires_custom_group_before_grad_sync(self): class TestAllReduceLNGrads: def init_model(self, share_embeddings_and_output_weights: bool = False): + qk_layernorm = True self.transformer_config = TransformerConfig( num_layers=2, hidden_size=12, @@ -127,13 +128,15 @@ def init_model(self, share_embeddings_and_output_weights: bool = False): use_cpu_initialization=True, tensor_model_parallel_size=self.tp_size, pipeline_model_parallel_size=self.pp_size, - qk_layernorm=True, + qk_layernorm=qk_layernorm, pipeline_dtype=torch.float32, ) self.model = GPTModel( config=self.transformer_config, - transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(qk_layernorm=True), + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec( + qk_layernorm=qk_layernorm + ), vocab_size=100, max_sequence_length=4, share_embeddings_and_output_weights=share_embeddings_and_output_weights, diff --git a/tests/unit_tests/transformer/test_submodule_callables.py b/tests/unit_tests/transformer/test_submodule_callables.py index a6a2eb31316..a301a557c54 100644 --- a/tests/unit_tests/transformer/test_submodule_callables.py +++ b/tests/unit_tests/transformer/test_submodule_callables.py @@ -137,9 +137,11 @@ def test_1f1b_overlap(self, dispatcher_type, grouped_gemm, permute_fusion): expert_model_parallel_size=2, virtual_pipeline_model_parallel_size=2, ) + qk_layernorm = True extra_kwargs = { "moe_token_dispatcher_type": dispatcher_type, "moe_permute_fusion": permute_fusion, + "qk_layernorm": qk_layernorm, } if dispatcher_type == "flex": extra_kwargs["moe_flex_dispatcher_backend"] = "deepep" @@ -149,7 +151,7 @@ def test_1f1b_overlap(self, dispatcher_type, grouped_gemm, permute_fusion): transformer_layer_submodules = get_gpt_layer_with_transformer_engine_submodules( num_experts=8, moe_grouped_gemm=grouped_gemm, - qk_layernorm=True, + qk_layernorm=qk_layernorm, multi_latent_attention=True, ) model = TransformerLayer(config, transformer_layer_submodules) From d194d241442ee79989a32942f5a9c1a63ad2b46d Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Thu, 17 Sep 2026 01:04:08 +0000 Subject: [PATCH 2/3] Accept experimental_attention_variant="gdn" as an alias of gated_delta_net Upstream Megatron-LM renamed the Gated DeltaNet variant to "gdn" (#5765, with "gated_delta_net" kept as a deprecated alias). Megatron-Bridge main selects Qwen3.5 linear attention with the new name, so on miles-main every Qwen3.5 model build died in get_experimental_attention_variant_module_spec with "Invalid experimental attention variant: gdn". miles-main keeps "gated_delta_net" as the canonical spelling and now treats "gdn" as an alias: * transformer_config: Literal accepts "gdn"; __post_init__ canonicalizes it; is_gated_delta_net_variant() replaces the literal comparisons so a value assigned after construction (Bridge sets it on the provider and never re-runs __post_init__) is recognized as well. * experimental_attention_variant_module_specs: dispatch and is_linear_attention_variant go through the helper, which is also exported from this module (Megatron-Bridge looks it up there). * training.py FLOP accounting accepts both spellings. "gdn2" (GDN2 kernels, upstream #5765) is intentionally not accepted: miles-main has no implementation for it and should keep failing loudly. Verified with Megatron-Bridge bridge-new: a GPTModelProvider configured like qwen35_bridge (experimental_attention_variant="gdn", linear_attention_freq=4, transformer_layer_spec=get_transformer_block_with_experimental_attention_variant_spec) raised the ValueError before and now resolves to GatedDeltaNet layers. Signed-off-by: Yusheng Su --- ...rimental_attention_variant_module_specs.py | 8 ++-- .../core/transformer/transformer_config.py | 42 ++++++++++++++----- megatron/training/training.py | 2 +- .../transformer/test_gated_delta_net_alias.py | 20 +++++++++ 4 files changed, 57 insertions(+), 15 deletions(-) create mode 100644 tests/unit_tests/transformer/test_gated_delta_net_alias.py diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 3db7c9987b1..d5f8a811213 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -37,7 +37,7 @@ TransformerBlockSubmodules, get_num_layers_to_build, ) -from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_config import TransformerConfig, is_gated_delta_net_variant from megatron.core.transformer.transformer_layer import ( HyperConnectionTransformerLayer, MlpBuilder, @@ -209,7 +209,7 @@ def get_experimental_attention_variant_module_spec( if backend is None: backend = _get_backend_spec_provider(config=config) - if config.experimental_attention_variant == "gated_delta_net": + if is_gated_delta_net_variant(config.experimental_attention_variant): return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) @@ -411,8 +411,8 @@ def get_transformer_block_with_experimental_attention_variant_spec( def is_linear_attention_variant(experimental_attention_variant: Optional[str]) -> bool: """Check if the experimental attention variant is a linear attention variant.""" - linear_attention_variants = ["gated_delta_net"] - return experimental_attention_variant in linear_attention_variants + # gated_delta_net and its upstream alias gdn; see transformer_config.is_gated_delta_net_variant. + return is_gated_delta_net_variant(experimental_attention_variant) def _validate_dsa_index_share_pipeline_split(config: TransformerConfig, local_layer_ids) -> None: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 6071a5694ea..66feda60810 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -53,6 +53,19 @@ HAVE_PACKAGING = False +GATED_DELTA_NET_VARIANTS = ("gated_delta_net", "gdn") + + +def is_gated_delta_net_variant(experimental_attention_variant: Optional[str]) -> bool: + """True for the Gated DeltaNet linear-attention variant under either of its names. + + Upstream Megatron-LM renamed ``gated_delta_net`` to ``gdn`` (NVIDIA/Megatron-LM#5765) and + Megatron-Bridge ``main`` selects Qwen3.5's linear attention with the new name. miles-main + keeps ``gated_delta_net`` as the canonical spelling and treats ``gdn`` as an alias. + """ + return experimental_attention_variant in GATED_DELTA_NET_VARIANTS + + @dataclass @experimental_api class TransformerConfig(ModelParallelConfig): @@ -312,10 +325,13 @@ class TransformerConfig(ModelParallelConfig): # attention variant #################### experimental_attention_variant: Optional[ - Literal['gated_delta_net', 'dsa', 'dsv4_hybrid', 'dsv4'] + Literal['gated_delta_net', 'gdn', 'dsa', 'dsv4_hybrid', 'dsv4'] ] = None """Type of attention variant to use. Currently support gated_delta_net, dsa, dsv4_hybrid, and - dsv4 (miles' DeepSeek-V4 sparse-attention path).""" + dsv4 (miles' DeepSeek-V4 sparse-attention path). 'gdn' is upstream Megatron-LM's name for + gated_delta_net (NVIDIA/Megatron-LM#5765) and is accepted as an alias: it is normalized to + 'gated_delta_net' in __post_init__, and every check goes through is_gated_delta_net_variant so a + value assigned after construction (Megatron-Bridge sets it on the provider) is recognized too.""" cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" """How THD sequence rows are partitioned across context-parallel ranks. @@ -1490,6 +1506,10 @@ def __post_init__(self): """ super().__post_init__() + if self.experimental_attention_variant == "gdn": + # Upstream spelling (NVIDIA/Megatron-LM#5765); canonicalize so every check below matches. + self.experimental_attention_variant = "gated_delta_net" + # When fp32 residual connections are enabled, pipeline parallel communication must # use fp32 to match the dtype of the residual stream between pipeline stages. if self.fp32_residual_connection and self.pipeline_dtype is not None: @@ -1619,14 +1639,16 @@ def __post_init__(self): "cp_partition_mode='contiguous' is not supported with " "multi_latent_attention outside dsv4_hybrid." ) - if self.experimental_attention_variant not in ("dsv4_hybrid", "gated_delta_net"): + if self.experimental_attention_variant != "dsv4_hybrid" and not is_gated_delta_net_variant( + self.experimental_attention_variant + ): raise ValueError( "cp_partition_mode='contiguous' with context parallelism currently " "requires experimental_attention_variant to be either 'dsv4_hybrid' " "or 'gated_delta_net'." ) if ( - self.experimental_attention_variant == "gated_delta_net" + is_gated_delta_net_variant(self.experimental_attention_variant) and self.linear_cp_mode == "headwise" ): raise ValueError( @@ -1670,12 +1692,12 @@ def __post_init__(self): ) self.dsa_kernel_backend = legacy_backend - if self.experimental_attention_variant in ["gated_delta_net"]: + if is_gated_delta_net_variant(self.experimental_attention_variant): assert ( self.linear_attention_freq is not None ), f"linear_attention_freq must be set for linear attention." - if self.experimental_attention_variant == "gated_delta_net": + if is_gated_delta_net_variant(self.experimental_attention_variant): if self.pad_packed_seq_alignment is not None: tail_policy = self.thd_tail_padding_policy or 'append_dummy_seq' assert tail_policy == 'append_dummy_seq', ( @@ -1846,7 +1868,7 @@ def __post_init__(self): if ( self.gdn_pre_gated_delta_rule_fusion - and self.experimental_attention_variant != "gated_delta_net" + and not is_gated_delta_net_variant(self.experimental_attention_variant) ): raise ValueError( "gdn_pre_gated_delta_rule_fusion is only supported with " @@ -2252,7 +2274,7 @@ def __post_init__(self): if ( "gdn_norm_out" in self.recompute_modules - and self.experimental_attention_variant != "gated_delta_net" + and not is_gated_delta_net_variant(self.experimental_attention_variant) ): raise ValueError( "gdn_norm_out in recompute_modules is only supported with " @@ -2261,7 +2283,7 @@ def __post_init__(self): if ( "gdn" in self.recompute_modules - and self.experimental_attention_variant != "gated_delta_net" + and not is_gated_delta_net_variant(self.experimental_attention_variant) ): raise ValueError( "gdn in recompute_modules is only supported with " @@ -3201,7 +3223,7 @@ def _scope_to_str(s): and (not self.cuda_graph_modules or CudaGraphModule.attn in self.cuda_graph_modules) ) - cp_layout_conversion_required = self.experimental_attention_variant == "gated_delta_net" + cp_layout_conversion_required = is_gated_delta_net_variant(self.experimental_attention_variant) # TODO: Extend this predicate as GDN2/KDA are introduced, and for DSv4 when # dsa_cp_balance_indexer is introduced; those paths will also require module-local THD CP # layout conversion. diff --git a/megatron/training/training.py b/megatron/training/training.py index 7594a921ca0..29d49e65d3c 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1027,7 +1027,7 @@ def transformer_flops(): num_linear_attention_layers = sum(linear_attention_pattern) num_standard_attention_layers = num_layers - num_linear_attention_layers - if args.experimental_attention_variant == "gated_delta_net": + if args.experimental_attention_variant in ("gated_delta_net", "gdn"): # Calculate the FLOPs for the gated delta net attention. qk_head_dim = args.linear_key_head_dim v_head_dim = args.linear_value_head_dim diff --git a/tests/unit_tests/transformer/test_gated_delta_net_alias.py b/tests/unit_tests/transformer/test_gated_delta_net_alias.py new file mode 100644 index 00000000000..9007b8e5005 --- /dev/null +++ b/tests/unit_tests/transformer/test_gated_delta_net_alias.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + is_gated_delta_net_variant as module_specs_is_gdn, + is_linear_attention_variant, +) +from megatron.core.transformer.transformer_config import is_gated_delta_net_variant + + +def test_gdn_is_an_alias_of_gated_delta_net(): + for name in ("gated_delta_net", "gdn"): + assert is_gated_delta_net_variant(name) + assert module_specs_is_gdn(name) + assert is_linear_attention_variant(name) + + +def test_other_variants_are_not_gated_delta_net(): + for name in (None, "dsa", "dsv4_hybrid", "dsv4", "gdn2"): + assert not is_gated_delta_net_variant(name) + assert not is_linear_attention_variant(name) From 00fac22f39bed0b9d08bf7e632bc9440919e846f Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Fri, 18 Sep 2026 11:47:12 -0700 Subject: [PATCH 3/3] [port] FullyShardedDataParallelV1/V2 names from upstream #5865 for Megatron-Bridge Upstream Megatron-Core split Megatron-FSDP into FullyShardedDataParallelV1 (the implementation miles-main ships as FullyShardedDataParallel) and FullyShardedDataParallelV2, and turned FullyShardedDataParallel into a factory (NVIDIA/Megatron-LM#5865). Megatron-Bridge main imports both names from mcore_fsdp_adapter since NVIDIA-NeMo/Megatron-Bridge#5933 (conversion/utils.py::unwrap_model, training/setup.py, training/train.py), so on miles-main every bridge-mode weight load / export raised ImportError with bridge-new @ 2e09c234a. Add FullyShardedDataParallelV1 as an alias of the existing class and a FullyShardedDataParallelV2 placeholder whose constructor raises NotImplementedError. The v2 implementation is not ported and FullyShardedDataParallel stays the v1 class, so the existing isinstance checks in megatron/core/utils.py and megatron/training/training.py are unchanged. Signed-off-by: Yusheng Su --- .../distributed/fsdp/mcore_fsdp_adapter.py | 23 ++++++++++++++ .../mfsdp_v1/test_mcore_fsdp_adapter_names.py | 31 +++++++++++++++++++ 2 files changed, 54 insertions(+) create mode 100644 tests/unit_tests/distributed/mfsdp_v1/test_mcore_fsdp_adapter_names.py diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 703dac7db46..36f745db827 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -484,6 +484,29 @@ def sync_rng_states_across_tp_group(self): _load_rng_state_dict(broadcast_list[0]) +# Name compatibility with upstream Megatron-Core (NVIDIA/Megatron-LM#5865), which split +# Megatron-FSDP into ``FullyShardedDataParallelV1`` (this implementation) and +# ``FullyShardedDataParallelV2`` and turned ``FullyShardedDataParallel`` into a factory. +# Callers written against upstream, such as Megatron-Bridge's ``unwrap_model``, import both +# names from this module. ``miles-main`` keeps ``FullyShardedDataParallel`` as the v1 class +# (nothing here selects v2) and does not carry the v2 implementation. +FullyShardedDataParallelV1 = FullyShardedDataParallel + + +class FullyShardedDataParallelV2(_BaseDataParallel): + """Placeholder for Megatron-FSDP v2 (NVIDIA/Megatron-LM#5865), which is not ported. + + The class exists so that upstream-style ``isinstance`` checks import and evaluate to + ``False``; constructing it raises. + """ + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "Megatron-FSDP v2 (NVIDIA/Megatron-LM#5865) is not ported to miles-main; " + "use FullyShardedDataParallel (FullyShardedDataParallelV1)." + ) + + def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group, ep_size=1): assert HAVE_EINOPS, "einops is not installed. Please install it with `pip install einops`." world_size = dist.get_world_size() diff --git a/tests/unit_tests/distributed/mfsdp_v1/test_mcore_fsdp_adapter_names.py b/tests/unit_tests/distributed/mfsdp_v1/test_mcore_fsdp_adapter_names.py new file mode 100644 index 00000000000..35f070b4af3 --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v1/test_mcore_fsdp_adapter_names.py @@ -0,0 +1,31 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +import pytest + +from megatron.core.distributed.data_parallel_base import _BaseDataParallel +from megatron.core.distributed.fsdp import mcore_fsdp_adapter +from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( + FullyShardedDataParallel, + FullyShardedDataParallelV1, + FullyShardedDataParallelV2, +) + + +def test_v1_is_the_shipped_implementation(): + assert FullyShardedDataParallelV1 is FullyShardedDataParallel + + +def test_v2_name_is_importable_but_not_constructible(): + assert issubclass(FullyShardedDataParallelV2, _BaseDataParallel) + assert not issubclass(FullyShardedDataParallel, FullyShardedDataParallelV2) + with pytest.raises(NotImplementedError, match="5865"): + FullyShardedDataParallelV2(None, None, None) + + +def test_upstream_style_isinstance_tuple_imports(): + # The tuple Megatron-Bridge's unwrap_model builds after NVIDIA-NeMo/Megatron-Bridge#5933. + module_instances = ( + mcore_fsdp_adapter.FullyShardedDataParallelV1, + mcore_fsdp_adapter.FullyShardedDataParallelV2, + ) + assert not isinstance(object(), module_instances)