[Dev] Add configurable attention latent normalization epsilon

megatron/core/extensions/transformer_engine.py

@@ -1009,10 +1009,13 @@ def __init__(
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         stride: int = 1,
         name: str | None = None,
+        eps: float | None = None,
     ):
         """
         Args:
             name (str | None): module instance name passed top-down from its paranet module
+            eps (float | None): Epsilon for the fused layer norm. Defaults to
+                ``config.layernorm_epsilon`` when ``None``.
         """
         if not HAVE_TE:
             raise ImportError(
@@ -1117,7 +1120,7 @@ def __init__(
             super().__init__(
                 in_features=input_size,
                 out_features=output_size,
-                eps=self.config.layernorm_epsilon,
+                eps=self.config.layernorm_epsilon if eps is None else eps,
                 sequence_parallel=self.config.sequence_parallel,
                 fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
                 tp_group=tp_group if torch.distributed.is_initialized() else None,

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}'")

megatron/core/models/hybrid/hybrid_block.py

@@ -5,6 +5,7 @@
 # This source code is licensed under the Apache license found in the
 # LICENSE file in the root directory of this source tree.
 
+import copy
 import warnings
 from contextlib import nullcontext
 from dataclasses import dataclass
@@ -16,7 +17,7 @@
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
 from megatron.core.enums import Fp8Recipe
-from megatron.core.extensions.transformer_engine import TENorm
+from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear, TENorm
 from megatron.core.fp4_utils import get_fp4_context
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.inference.contexts import BaseInferenceContext
@@ -40,6 +41,7 @@
     convert_module_to_dtype_except_fp32_marked,
     mark_keep_in_fp32,
 )
+from megatron.core.transformer.multi_latent_attention import FusedMLASelfAttention
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.transformer.utils import (
@@ -624,6 +626,9 @@ def __init__(
         )
         self.layer_type_list = layer_type_list
 
+        if getattr(self.config, "mla_down_proj_fusion", False):
+            submodules = self._fuse_mla_down_proj(submodules)
+
         # Build layers from the pre-selected segment
         self.layers = nn.ModuleList()
         for i, layer_type in enumerate(self.layer_type_list):
@@ -787,6 +792,26 @@ def __init__(
                 setattr(self.hc_head_base, 'sequence_parallel', True)
                 setattr(self.hc_head_scale, 'sequence_parallel', True)
 
+    def _fuse_mla_down_proj(self, submodules: HybridStackSubmodules) -> HybridStackSubmodules:
+        # Avoid modifying the original object so users don't get surprised about their `submodules`
+        # being modified underneath them.
+        submodules = copy.deepcopy(submodules)
+        mla_spec = submodules.mla_layer
+        # We always fuse the input layernorm because Hybrid always uses TransformerEngine.
+        mla_spec.submodules.input_layernorm = IdentityOp
+        mla_spec.submodules.self_attention.module = FusedMLASelfAttention
+        mla_spec.submodules.self_attention.submodules.linear_qkv_down_proj = (
+            TELayerNormColumnParallelLinear
+        )
+        mla_spec.submodules.self_attention.submodules.linear_q_down_proj = None
+        mla_spec.submodules.self_attention.submodules.linear_kv_down_proj = None
+        mla_spec.submodules.sharded_state_dict_keys_map = {
+            "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.",
+            "self_attention.linear_kv_down_proj.layer_norm_": "input_layernorm.",
+            "self_attention.linear_qkv_down_proj.layer_norm_": "input_layernorm.",
+        }
+        return submodules
+
     @staticmethod
     def _set_mtp_layer_number_for_moe_metrics(
         layer: torch.nn.Module, mtp_layer_number: int

megatron/core/models/hybrid/hybrid_layer_specs.py

@@ -313,6 +313,28 @@
                 self_attn_bda=get_bias_dropout_add,
             ),
         ),
+        mla_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=TEColumnParallelLinear,
+                        linear_q_down_proj=TELinear,
+                        linear_q_up_proj=TEColumnParallelLinear,
+                        linear_kv_down_proj=TELinear,
+                        linear_kv_up_proj=TEColumnParallelLinear,
+                        core_attention=TEDotProductAttention,
+                        linear_proj=InferenceRowParallelLinear,
+                        q_layernorm=IdentityOp,
+                        kv_layernorm=IdentityOp,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         # Started with spec from gpt_layer_specs.py
         # Using the TE spec because we had problems getting the non-TE spec
         # working

megatron/core/tensor_parallel/inference_layers.py

@@ -132,7 +132,14 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
         name: str | None = None,
+        eps: float | None = None,
     ):
+        """Initialize the inference-optimized fused layer norm and linear.
+
+        Args:
+            eps (float | None): Epsilon for the fused layer norm. Defaults to
+                ``config.layernorm_epsilon`` when ``None``.
+        """
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
             input_size,
@@ -148,6 +155,7 @@ def __init__(
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
             name=name,
+            eps=eps,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)
@@ -156,7 +164,7 @@ def __init__(
             output_size % self.tp_size == 0
         ), f"output_size ({output_size}) must be divisible by tp_size ({self.tp_size})"
 
-        self.eps = config.layernorm_epsilon
+        self.eps = config.layernorm_epsilon if eps is None else eps
 
         if self.tp_size > 1:
             assert (

megatron/core/transformer/experimental_attention_variant/absorbed_mla.py

@@ -37,6 +37,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
@@ -164,6 +165,8 @@ def __init__(
             name=name,
         )
 
+        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
@@ -263,7 +266,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,
@@ -309,7 +312,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,
@@ -370,7 +373,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),
@@ -416,17 +419,17 @@ 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,
+                eps=self.config.attention_latent_norm_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,
+            eps=self.config.attention_latent_norm_epsilon,
         )
 
     def get_query_key_value_tensors(

megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py

@@ -33,7 +33,7 @@
 if HAVE_TE:
     from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input
 else:
-    (TEColumnParallelLinear, TELinear, set_save_original_input) = (None, None, None)
+    TEColumnParallelLinear, TELinear, set_save_original_input = (None, None, None)
 
 
 @torch.compile
@@ -571,13 +571,13 @@ def __init__(
         self.kv_layernorm = submodules.kv_layernorm(
             hidden_size=self.config.v_head_dim,
             config=self.config,
-            eps=self.config.layernorm_epsilon,
+            eps=self.config.attention_latent_norm_epsilon,
         )
 
         self.q_layernorm = submodules.q_layernorm(
             hidden_size=self.config.q_lora_rank,
             config=self.config,
-            eps=self.config.layernorm_epsilon,
+            eps=self.config.attention_latent_norm_epsilon,
         )
 
     def get_query_key_value_tensors(