[Dev] Enable selective recompute for norm_out in GDN layers

megatron/core/ssm/gated_delta_net.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh.
 
 # Some of this code was adopted from https://github.com/huggingface/transformers
@@ -266,8 +266,13 @@ def __init__(
         # the entire GatedDeltaNet compute is wrapped in a normal checkpoint and recomputed
         # in the backward pass.
         self.recompute_gdn = False
+        # gdn_norm_out: recompute only the gated output norm + HP-to-CP all-to-all block as a
+        # discard-output checkpoint.
+        self.recompute_norm_out = False
+        self.norm_out_checkpoint = None
         if self.config.recompute_granularity == "selective" and self.config.recompute_modules:
             self.recompute_gdn = "gdn" in self.config.recompute_modules
+            self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules
 
         self.reset_parameters()
 
@@ -521,25 +526,43 @@ def _forward_compute(
         )
         nvtx_range_pop(suffix="gated_delta_rule")
 
-        # RMSNorm
-        nvtx_range_push(suffix="gated_norm")
-        norm_out = self._apply_gated_norm(core_attn_out, gate)
-        nvtx_range_pop(suffix="gated_norm")
+        def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor):
+            # RMSNorm
+            nvtx_range_push(suffix="gated_norm")
+            norm_out = self._apply_gated_norm(core_attn_out, gate)
+            nvtx_range_pop(suffix="gated_norm")
 
-        # Transpose: b s x --> s b x
-        # From bshd back to sbhd format
-        norm_out = norm_out.reshape(batch, seq_len, -1)
-        norm_out = norm_out.transpose(0, 1).contiguous()
+            # Transpose: b s x --> s b x
+            # From bshd back to sbhd format
+            norm_out = norm_out.reshape(batch, seq_len, -1)
+            norm_out = norm_out.transpose(0, 1).contiguous()
 
-        norm_out = self._a2a_hp_to_cp(
-            norm_out, cp_size, cp_group, packed_seq_params, thd_cp_a2a_inv
-        )
+            norm_out = self._a2a_hp_to_cp(
+                norm_out, cp_size, cp_group, packed_seq_params, thd_cp_a2a_inv
+            )
+
+            return norm_out
+
+        # gdn_norm_out: discard the gated-norm + a2a outputs now and regenerate them from a grad
+        # hook in the backward, freeing the gated-norm activations. Synchronous recompute, so it is
+        # safe with the fla/compiled gated_delta_rule backward.
+        recompute_norm_out = self.recompute_norm_out and self.training
+        if recompute_norm_out:
+            self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            norm_out = self.norm_out_checkpoint.checkpoint(_gated_norm_and_a2a, core_attn_out, gate)
+        else:
+            norm_out = _gated_norm_and_a2a(core_attn_out, gate)
 
         # Output projection
         nvtx_range_push(suffix="out_proj")
         out, out_bias = self.out_proj(norm_out)
         nvtx_range_pop(suffix="out_proj")
 
+        # Discard the checkpointed norm_out (now consumed by out_proj) and register the recompute
+        # hook on `out` — its grad is computed first in backward, before the backward that needs it.
+        if recompute_norm_out:
+            self.norm_out_checkpoint.discard_output_and_register_recompute(out)
+
         return out, out_bias
 
     def _a2a_cp_to_hp(

megatron/core/transformer/transformer_config.py

@@ -554,7 +554,7 @@ class TransformerConfig(ModelParallelConfig):
     recompute_modules: Optional[List[str]] = None
     """The submodules to recompute.
     choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe",
-             "shared_experts", "mhc", "gdn".
+             "shared_experts", "mhc", "gdn", "gdn_norm_out".
     default: ["core_attn"].
     "core_attn": recompute the core attention part of the transformer layer.
     "moe_act": recompute the MoE MLP activation function.
@@ -569,8 +569,11 @@ class TransformerConfig(ModelParallelConfig):
     "gdn": recompute the entire GatedDeltaNet module (in_proj, conv1d, gated delta rule,
             gated norm, CP all-to-all and out_proj). Requires
             experimental_attention_variant="gated_delta_net".
-    "moe_act", "layernorm", "mla_up_proj", and "mhc" use output-discarding checkpointing,
-    "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing.
+    "gdn_norm_out": recompute only the GatedDeltaNet output norm and HP-to-CP all-to-all as a
+            discard-output checkpoint. Requires experimental_attention_variant="gated_delta_net"
+            and cannot be combined with "gdn" (which already recomputes this block).
+    "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use output-discarding
+    checkpointing, "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing.
     """
 
     ####################
@@ -1895,6 +1898,7 @@ def __post_init__(self):
                     "shared_experts",
                     "mhc",
                     "gdn",
+                    "gdn_norm_out",
                 }
                 invalid_modules = set(self.recompute_modules) - allowed_modules
                 assert not invalid_modules, (
@@ -1922,6 +1926,22 @@ def __post_init__(self):
                     "experimental_attention_variant='gated_delta_net'."
                 )
 
+            if (
+                "gdn_norm_out" in self.recompute_modules
+                and self.experimental_attention_variant != "gated_delta_net"
+            ):
+                raise ValueError(
+                    "gdn_norm_out in recompute_modules is only supported with "
+                    "experimental_attention_variant='gated_delta_net'."
+                )
+
+            if "gdn" in self.recompute_modules and "gdn_norm_out" in self.recompute_modules:
+                raise ValueError(
+                    "gdn and gdn_norm_out in recompute_modules cannot be used together: 'gdn' "
+                    "recomputes the entire GatedDeltaNet module, which already includes the "
+                    "gated output norm + all-to-all block that 'gdn_norm_out' targets."
+                )
+
             if "core_attn" in self.recompute_modules:
                 warnings.warn(
                     "If you are using transformer_engine as the transformer implementation, "

tests/unit_tests/ssm/test_gated_delta_net.py

@@ -1,5 +1,6 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import copy
 import os
 from functools import partial
 from unittest import mock
@@ -218,6 +219,95 @@ def run(recompute):
                 msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}",
             )
 
+    def test_selective_recompute_norm_out(self):
+        """gdn_norm_out discard-output recompute must be numerically exact.
+
+        recompute_modules=["gdn_norm_out"] recomputes only the gated output norm +
+        HP-to-CP all-to-all block as a discard-output checkpoint. The block is re-run
+        with the same RNG state and shares storage with the discarded output, so the
+        output, all parameter grads and the input grad must match the no-recompute
+        baseline bit-for-bit.
+        """
+        tp_group = parallel_state.get_tensor_model_parallel_group()
+        cp_group = parallel_state.get_context_parallel_group()
+        pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group)
+
+        def build_gdn(config):
+            gdn_submodules = get_experimental_attention_variant_module_spec(
+                config=config
+            ).submodules
+            gdn = GatedDeltaNet(
+                config,
+                submodules=gdn_submodules,
+                layer_number=1,
+                bias=False,
+                conv_bias=False,
+                conv_init=1.0,
+                use_qk_l2norm=True,
+                A_init_range=(1, 16),
+                pg_collection=pg_collection,
+            )
+            return gdn.cuda().bfloat16()
+
+        def run(gdn, hidden_states):
+            output, _ = gdn(hidden_states, None)
+            output.float().sum().backward()
+            grads = {
+                name: param.grad.detach()
+                for name, param in gdn.named_parameters()
+                if param.grad is not None
+            }
+            input_grad = hidden_states.grad.detach().clone()
+            return output.detach(), grads, input_grad
+
+        micro_batch_size = 2
+        seq_length = 64
+        base_config = copy.deepcopy(self.transformer_config)
+        rec_config = copy.deepcopy(self.transformer_config)
+        rec_config.recompute_granularity = "selective"
+        rec_config.recompute_modules = ["gdn_norm_out"]
+
+        model_parallel_cuda_manual_seed(42)
+        torch.manual_seed(42)
+        hidden_states = torch.randn(
+            (
+                seq_length // self.sp_size // self.cp_size,
+                micro_batch_size,
+                self.gdn.config.hidden_size,
+            ),
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
+            requires_grad=True,
+        )
+
+        # --- Baseline (no recompute) ---
+        model_parallel_cuda_manual_seed(42)
+        torch.manual_seed(42)
+        base_gdn = build_gdn(base_config)
+        assert base_gdn.recompute_norm_out is False
+        base_output, base_grads, base_input_grad = run(base_gdn, hidden_states)
+        hidden_states.grad = None
+        assert base_gdn.norm_out_checkpoint is None
+        del base_gdn
+        torch.cuda.empty_cache()
+
+        # --- Recompute ---
+        model_parallel_cuda_manual_seed(42)
+        torch.manual_seed(42)
+        rec_gdn = build_gdn(rec_config)
+        assert rec_gdn.recompute_norm_out is True
+        rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states)
+        assert rec_gdn.norm_out_checkpoint is not None
+
+        rank = torch.distributed.get_rank()
+        assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})"
+        assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})"
+        assert set(rec_grads.keys()) == set(base_grads.keys())
+        for name in base_grads:
+            assert torch.equal(
+                rec_grads[name], base_grads[name]
+            ), f"Grad not identical for {name} ({rank=})"
+
     def test_jit_compiled_helpers(self):
         import torch._dynamo