From d9006c44430967c87e8cde68e60f2960591d6c4d Mon Sep 17 00:00:00 2001 From: jinliangl <975761915@qq.com> Date: Thu, 11 Jun 2026 22:41:08 +0800 Subject: [PATCH 1/3] feat(ssm): add whole-module 'gdn' selective recompute for GatedDeltaNet Add a 'gdn' entry to recompute_modules (selective granularity) that recomputes the entire GatedDeltaNet module (in_proj, CP all-to-all, conv1d, gated delta rule, gated norm, and out_proj) via normal checkpointing, mirroring how 'core_attn' recomputes the attention core. - transformer_config: add 'gdn' to the recompute_modules choices + allowed set, document it as a normal-checkpointing module, and guard that it requires experimental_attention_variant='gated_delta_net'. - GatedDeltaNet: extract the forward compute body into _forward_compute and wrap it in tensor_parallel.checkpoint when recompute_gdn is set and the module is training. - test: add test_selective_recompute_gdn parity check (forward output + input/param gradients match the non-recompute path within bf16 tol). Signed-off-by: jinliangl <975761915@qq.com> --- megatron/core/ssm/gated_delta_net.py | 45 +++++++++++++++++ .../core/transformer/transformer_config.py | 17 ++++++- tests/unit_tests/ssm/test_gated_delta_net.py | 50 +++++++++++++++++++ 3 files changed, 110 insertions(+), 2 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 4dd963eb928..4d213d8d3ef 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -14,6 +14,7 @@ import torch.nn.functional as F from torch import Tensor +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.fp8_utils import get_fp8_align_size @@ -260,6 +261,13 @@ def __init__( name=(name + ".out_proj") if name is not None else None, ) + # Whole-module recompute: when "gdn" is in recompute_modules (selective granularity), + # the entire GatedDeltaNet compute is wrapped in a normal checkpoint and recomputed + # in the backward pass. + self.recompute_gdn = False + if self.config.recompute_granularity == "selective" and self.config.recompute_modules: + self.recompute_gdn = "gdn" in self.config.recompute_modules + self.reset_parameters() def reset_parameters(self): @@ -363,6 +371,43 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + if self.recompute_gdn and self.training: + # Selective full-module recompute ("gdn"): run the whole GDN compute under a + # normal checkpoint so its activations are dropped and recomputed in backward. + # Only the input tensor is passed to checkpoint(); the non-tensor arguments are + # bound via closure because tensor_parallel.checkpoint saves its positional args + # with save_for_backward, which only accepts tensors. + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len, + cp_size, + cp_group, + cu_seqlens_q, + packed_seq_params, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, batch, seq_len, cp_size, cp_group, cu_seqlens_q, packed_seq_params + ) + + return out, out_bias + + def _forward_compute( + self, hidden_states, batch, seq_len, cp_size, cp_group, cu_seqlens_q, packed_seq_params + ): + """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> gated norm -> out_proj). + + Extracted from ``forward`` so the entire module can be wrapped in a recompute + checkpoint when ``recompute_modules`` contains ``"gdn"`` (selective full-module + recompute, normal checkpointing). + + Returns: + Tuple of (output, output_bias). + """ # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 6687d572914..1e5530096a6 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -553,7 +553,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". + "shared_experts", "mhc", "gdn". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -565,8 +565,11 @@ class TransformerConfig(ModelParallelConfig): "mhc": recompute HyperConnection intermediate activations via CheckpointWithoutOutput + CheckpointManager. Requires enable_hyper_connections=True. Cannot be used with "mlp". + "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", and "shared_experts" use normal checkpointing. + "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing. """ #################### @@ -1855,6 +1858,7 @@ def __post_init__(self): "moe", "shared_experts", "mhc", + "gdn", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1873,6 +1877,15 @@ def __post_init__(self): "multi_latent_attention." ) + if ( + "gdn" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index f490e7cfdb8..b55d64d6f5a 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -142,6 +142,56 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_selective_recompute_gdn(self): + """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. + + The same module/input is run twice (recompute off, then on); the forward output and + all parameter / input gradients must agree within bf16 tolerance. + """ + gdn = self.gdn + gdn.train() + + micro_batch_size = 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute): + gdn.recompute_gdn = recompute + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute=False) + out_rc, dinput_rc, pgrad_rc = run(recompute=True) + finally: + gdn.recompute_gdn = False + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-2, atol=1e-2) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-2, + atol=1e-2, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + def test_jit_compiled_helpers(self): import torch._dynamo From 2cd9c12ef5361776b8ffb2e11a8ce04266ec0776 Mon Sep 17 00:00:00 2001 From: jinliangl <975761915@qq.com> Date: Sun, 14 Jun 2026 23:18:17 -0700 Subject: [PATCH 2/3] Remove redundant inline comments in GDN recompute path Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: jinliangl <975761915@qq.com> --- megatron/core/ssm/gated_delta_net.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 4d213d8d3ef..e68269e9af8 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -372,11 +372,7 @@ def forward( cu_seqlens_kv = None if self.recompute_gdn and self.training: - # Selective full-module recompute ("gdn"): run the whole GDN compute under a - # normal checkpoint so its activations are dropped and recomputed in backward. - # Only the input tensor is passed to checkpoint(); the non-tensor arguments are - # bound via closure because tensor_parallel.checkpoint saves its positional args - # with save_for_backward, which only accepts tensors. + def _checkpointed_compute(hidden_states): return self._forward_compute( hidden_states, From 5251bfd13fbbe5af353982842d0efe8e75369f7b Mon Sep 17 00:00:00 2001 From: jinliangl <975761915@qq.com> Date: Mon, 15 Jun 2026 20:20:57 -0700 Subject: [PATCH 3/3] Tighten gdn recompute parity tolerance to 1e-4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The whole-module "gdn" recompute path is run-to-run deterministic on these kernels — empirically bitwise-identical to the non-recompute path (forward output, input grad, and all parameter grads, max_abs_diff 0 across all tp/sp/cp parametrizations on GB200). Tighten rtol/atol from 1e-2 (bf16 floor) to 1e-4 so the parity test actually guards recompute correctness, while keeping a small margin against benign kernel nondeterminism. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: jinliangl <975761915@qq.com> --- tests/unit_tests/ssm/test_gated_delta_net.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index b55d64d6f5a..41535e581a5 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -146,7 +146,9 @@ def test_selective_recompute_gdn(self): """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. The same module/input is run twice (recompute off, then on); the forward output and - all parameter / input gradients must agree within bf16 tolerance. + all parameter / input gradients must agree within a tight tolerance (rtol/atol=1e-4). + The recompute path is run-to-run deterministic on these kernels (empirically bitwise), + so a tolerance well below the bf16 floor is expected to hold. """ gdn = self.gdn gdn.train() @@ -179,16 +181,16 @@ def run(recompute): finally: gdn.recompute_gdn = False - torch.testing.assert_close(out_rc, out_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" assert len(pgrad_ref) > 0, "expected at least one parameter gradient" for name in pgrad_ref: torch.testing.assert_close( pgrad_rc[name], pgrad_ref[name], - rtol=1e-2, - atol=1e-2, + rtol=1e-4, + atol=1e-4, msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", )