Add gated GLM-5.2 bit-exact compatibility paths

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -20,6 +20,7 @@
 )
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.torch_norm import AccuracyCompatibleRMSNorm
 from megatron.core.transformer.transformer_block import (
     TransformerBlockSubmodules,
     get_num_layers_to_build,
@@ -107,7 +108,13 @@ def get_dsa_module_spec_for_backend(
     # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm
     # with linear projection and have to use unfused qk layernorm.
     qk_norm = (
-        backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp
+        (
+            AccuracyCompatibleRMSNorm
+            if config.norm_accuracy_compatible
+            else backend.layer_norm(rms_norm=rms_norm, for_qk=True)
+        )
+        if config.qk_layernorm
+        else IdentityOp
     )
 
     attention = ModuleSpec(

megatron/core/models/gpt/gpt_layer_specs.py

@@ -773,7 +773,7 @@ def get_gpt_mtp_block_spec_for_backend(
         raise ValueError(f"Invalid spec: {spec}")
 
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        mtp_model_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec, backend=backend, config=config
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     if config.mtp_use_repeated_layer:

megatron/core/transformer/experimental_attention_variant/dsa.py

@@ -64,6 +64,7 @@ def _unfused_absorbed_dsa_fn(
     varlen_starts: Optional[torch.Tensor] = None,
     varlen_ends: Optional[torch.Tensor] = None,
     key_positions: Optional[torch.Tensor] = None,
+    accuracy_compatible: bool = False,
 ) -> torch.Tensor:
     """Unfused absorbed-MLA attention: output stays [sq, b, np, v_channels]."""
     sq, b, np, hn = query.size()
@@ -99,17 +100,41 @@ def _unfused_absorbed_dsa_fn(
     )
 
     attention_scores = attention_scores + index_mask.unsqueeze(1)
-    valid_index_mask = torch.isfinite(index_mask)
-    attention_scores = dsa_masking.masked_softmax(
-        attention_scores.float(), valid_index_mask.unsqueeze(1).expand(b, np, sq, skv), dim=-1
-    )
+    valid_index_mask = torch.isfinite(index_mask).unsqueeze(1).expand(b, np, sq, skv)
+    if accuracy_compatible:
+        attention_scores = _AccuracyCompatibleSoftmax.apply(
+            attention_scores.float(), valid_index_mask
+        )
+    else:
+        attention_scores = dsa_masking.masked_softmax(
+            attention_scores.float(), valid_index_mask, dim=-1
+        )
 
     # Latent value is the first v_channels slice of absorbed key cache.
     value = key[..., :v_channels].permute(1, 2, 0, 3)  # [b,1,skv,v]
     output = torch.matmul(attention_scores.to(value.dtype), value)  # [b,np,sq,v]
     return output.permute(2, 0, 1, 3).contiguous()
 
 
+class _AccuracyCompatibleSoftmax(torch.autograd.Function):
+    """Masked softmax with an explicit backward formula for DSA alignment."""
+
+    @staticmethod
+    def forward(ctx, logits: torch.Tensor, valid_mask: torch.Tensor) -> torch.Tensor:
+        probabilities = torch.softmax(logits.masked_fill(~valid_mask, float("-inf")), dim=-1)
+        probabilities = probabilities.masked_fill(~valid_mask, 0.0)
+        ctx.save_for_backward(probabilities, valid_mask)
+        return probabilities
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        probabilities, valid_mask = ctx.saved_tensors
+        grad_logits = probabilities * (
+            grad_output - (grad_output * probabilities).sum(dim=-1, keepdim=True)
+        )
+        return grad_logits.masked_fill(~valid_mask, 0.0), None
+
+
 def _run_sparse_attention(
     *,
     absorbed_mla: bool,
@@ -127,6 +152,7 @@ def _run_sparse_attention(
     topk_length: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Run sparse attention for absorbed and non-absorbed MLA paths."""
+    accuracy_compatible = bool(getattr(config, "dsa_accuracy_compatible", False))
     if absorbed_mla:
         latent_v_channels = int(getattr(config, "kv_lora_rank", 0) or 0)
         if latent_v_channels <= 0:
@@ -143,7 +169,7 @@ def _run_sparse_attention(
                 "Received absorbed layout with explicit value tensor."
             )
         output = None
-        if dsa_kernels.use_fused_dsa_kernels(config):
+        if not accuracy_compatible and dsa_kernels.use_fused_dsa_kernels(config):
             output = dsa_kernels.run_fused_absorbed_sparse_attention(
                 config,
                 query,
@@ -166,6 +192,7 @@ def _run_sparse_attention(
                 varlen_starts=varlen_starts,
                 varlen_ends=varlen_ends,
                 key_positions=key_positions,
+                accuracy_compatible=accuracy_compatible,
             )
         assert output is not None
         output = torch.einsum("sbhc,hdc->sbhd", output, up_v_weight).contiguous()
@@ -182,6 +209,7 @@ def _run_sparse_attention(
         varlen_starts=varlen_starts,
         varlen_ends=varlen_ends,
         key_positions=key_positions,
+        accuracy_compatible=accuracy_compatible,
     )
 
 
@@ -1411,6 +1439,7 @@ def unfused_dsa_fn(
     varlen_starts: Optional[torch.Tensor] = None,
     varlen_ends: Optional[torch.Tensor] = None,
     key_positions: Optional[torch.Tensor] = None,
+    accuracy_compatible: bool = False,
 ):
     """
     Unfused sparse attention implementation.
@@ -1457,6 +1486,27 @@ def unfused_dsa_fn(
         device=query.device,
     )
 
+    if accuracy_compatible:
+        index_mask = torch.full((b, sq, skv), float("-inf"), device=query.device)
+        dsa_masking.scatter_topk_into_index_mask(index_mask, topk_indices)
+        index_mask = dsa_masking.apply_sparse_validity_to_index_mask(
+            index_mask,
+            row_mask=row_mask,
+            varlen_starts=varlen_starts,
+            varlen_ends=varlen_ends,
+            key_positions=key_positions,
+        )
+        valid_index_mask = torch.isfinite(index_mask).unsqueeze(1).expand(b, np, sq, skv)
+        attention_scores = (
+            torch.matmul(query_b.float(), key_b.float().transpose(-1, -2)) * softmax_scale
+        )
+        attention_probs = _AccuracyCompatibleSoftmax.apply(
+            attention_scores + index_mask.unsqueeze(1), valid_index_mask
+        )
+        output = torch.matmul(attention_probs.to(value_b.dtype), value_b)
+        output = output.permute(2, 0, 1, 3).contiguous().view(sq, b, np * hnv)
+        return output.squeeze(1) if query_was_thd else output
+
     seq_chunk_size = 512
     head_chunk_size = 16
     topk_chunk_size = 1024

megatron/core/transformer/moe/router.py

@@ -104,6 +104,14 @@ def gating(self, input: torch.Tensor):
             router_dtype = torch.float32
         elif self.config.moe_router_dtype == 'fp64':
             router_dtype = torch.float64
+        if self.config.router_accuracy_compatible:
+            inp_shape = input.shape
+            logits = torch.mm(
+                input.reshape(-1, inp_shape[-1]).float(), self.weight.float().t()
+            )
+            if self.bias is not None:
+                logits = logits + self.bias.float()
+            return logits.view(*inp_shape[:-1], -1)
         logits = router_gating_linear(input, self.weight, self.bias, router_dtype)
         return logits
 

megatron/core/transformer/multi_token_prediction.py

@@ -30,7 +30,7 @@
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-from megatron.core.transformer.torch_norm import LayerNormBuilder
+from megatron.core.transformer.torch_norm import AccuracyCompatibleRMSNorm, LayerNormBuilder
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
@@ -577,7 +577,9 @@ class MultiTokenPredictionLayerSubmodules:
 
 
 def get_mtp_layer_spec(
-    mtp_model_layer_spec: ModuleSpec, use_transformer_engine: bool
+    mtp_model_layer_spec: ModuleSpec,
+    use_transformer_engine: bool,
+    config: Optional[TransformerConfig] = None,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -587,19 +589,26 @@ def get_mtp_layer_spec(
     return get_mtp_layer_spec_for_backend(
         mtp_model_layer_spec,
         backend=TESpecProvider() if use_transformer_engine else LocalSpecProvider(),
+        config=config,
     )
 
 
 def get_mtp_layer_spec_for_backend(
-    mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
+    mtp_model_layer_spec: ModuleSpec,
+    backend: BackendSpecProvider,
+    config: Optional[TransformerConfig] = None,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
     Returns:
         ModuleSpec: Module specification with modules from the backend.
     """
     column_parallel_linear_impl: type = backend.column_parallel_linear()
-    layer_norm_impl = backend.layer_norm()
+    layer_norm_impl = (
+        AccuracyCompatibleRMSNorm
+        if config is not None and config.norm_accuracy_compatible
+        else backend.layer_norm()
+    )
     mtp_layer_spec = ModuleSpec(
         module=MultiTokenPredictionLayer,
         submodules=MultiTokenPredictionLayerSubmodules(

megatron/core/transformer/torch_norm.py

@@ -24,6 +24,54 @@ def __call__(
     ) -> LayerNormInterface: ...
 
 
+class _AccuracyCompatibleRMSNormFunction(torch.autograd.Function):
+    """RMSNorm core with a stable fp32 backward and canonical zero gradients."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, eps: float) -> torch.Tensor:
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        inv_rms = torch.rsqrt(variance + eps)
+        ctx.save_for_backward(x, inv_rms)
+        return x * inv_rms
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        x, inv_rms = ctx.saved_tensors
+        dot = (grad_output * x).sum(dim=-1, keepdim=True)
+        correction_scale = dot * (-0.5) * inv_rms.pow(3) / x.shape[-1]
+        correction = (correction_scale * x) * 2.0
+        grad_input = grad_output * inv_rms + correction
+        grad_input = torch.where(grad_input == 0, torch.zeros_like(grad_input), grad_input)
+        return grad_input, None
+
+
+class AccuracyCompatibleRMSNorm(torch.nn.Module, LayerNormInterface):
+    """RMSNorm with explicit fp32 reduction and one output cast."""
+
+    def __init__(
+        self,
+        normalized_shape: int | None = None,
+        eps: float = 1e-5,
+        *,
+        hidden_size: int | None = None,
+        config: TransformerConfig | None = None,
+        **kwargs,
+    ):
+        super().__init__()
+        normalized_shape = hidden_size if normalized_shape is None else normalized_shape
+        if normalized_shape is None:
+            raise ValueError("normalized_shape or hidden_size is required")
+        self.normalized_shape = (normalized_shape,)
+        self.eps = eps
+        dtype = config.params_dtype if config is not None else None
+        self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_float = x.float()
+        output = _AccuracyCompatibleRMSNormFunction.apply(x_float, self.eps)
+        return (output * self.weight.float()).to(x.dtype)
+
+
 class WrappedTorchNorm:
     """
     A conditional wrapper to initialize an instance of PyTorch's
@@ -56,6 +104,8 @@ def __new__(
         if config.normalization == "LayerNorm":
             norm_cls = torch.nn.LayerNorm
         elif config.normalization == "RMSNorm":
+            if config.norm_accuracy_compatible:
+                return AccuracyCompatibleRMSNorm(normalized_shape=hidden_size, eps=eps)
             assert is_torch_min_version(
                 "2.4.0a0"
             ), 'Torch RMSNorm requires PyTorch version >= 2.4.0'

megatron/core/transformer/transformer_config.py

@@ -186,6 +186,16 @@ class TransformerConfig(ModelParallelConfig):
     )
     """Epsilon value for any LayerNorm/RMSNorm operations."""
 
+    norm_accuracy_compatible: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--norm-accuracy-compatible"]}}
+    )
+    """Use explicit fp32 normalization formulas instead of native norm kernels for alignment."""
+
+    router_accuracy_compatible: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--router-accuracy-compatible"]}}
+    )
+    """Use an explicit fp32 router GEMM instead of the fused Transformer Engine path."""
+
     layernorm_zero_centered_gamma: bool = field(
         default=False, metadata={"argparse_meta": {"arg_names": ["--apply-layernorm-1p"]}}
     )
@@ -318,6 +328,11 @@ class TransformerConfig(ModelParallelConfig):
     ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that
     backend. Unsupported DSA layouts continue to use the PyTorch fallback."""
 
+    dsa_accuracy_compatible: bool = field(
+        default=False, metadata={"argparse_meta": {"arg_names": ["--dsa-accuracy-compatible"]}}
+    )
+    """Use the full-score DSA fallback with explicit softmax backward for alignment."""
+
     dsa_indexer_rope_interleaved: bool = False
     """Whether DSA indexer RoPE should use MLA-style interleaving."""
 

tests/unit_tests/models/test_experimental_attention_variant_module_specs.py

@@ -65,6 +65,7 @@ def _make_config(**overrides):
     defaults = dict(
         num_layers=4,
         normalization="RMSNorm",
+        norm_accuracy_compatible=False,
         qk_layernorm=False,
         multi_latent_attention=False,
         qk_l2_norm=False,
@@ -369,6 +370,23 @@ def test_qk_layernorm_enabled(self, normalization):
         assert spec.submodules.q_layernorm is spec.submodules.kv_layernorm
         backend.layer_norm.assert_any_call(rms_norm=expected_rms, for_qk=True)
 
+    def test_accuracy_compatible_qk_rmsnorm(self):
+        """Verify DSA q/kv norms can use the explicit fp32 RMSNorm path."""
+        from megatron.core.transformer.torch_norm import AccuracyCompatibleRMSNorm
+
+        backend = _make_backend()
+        cfg = _make_config(
+            multi_latent_attention=True,
+            qk_l2_norm=False,
+            qk_layernorm=True,
+            normalization="RMSNorm",
+            norm_accuracy_compatible=True,
+        )
+        spec = self._call(cfg=cfg, backend=backend)
+
+        assert spec.submodules.q_layernorm is AccuracyCompatibleRMSNorm
+        assert spec.submodules.kv_layernorm is AccuracyCompatibleRMSNorm
+
     def test_qk_layernorm_disabled(self):
         """Verify q/kv layernorm becomes IdentityOp, skipping backend.layer_norm for qk."""
         backend = _make_backend()