From 03d0ed3fa0aaf31d73d99b6a98334a2adab701af Mon Sep 17 00:00:00 2001 From: jinliangl <975761915@qq.com> Date: Wed, 3 Jun 2026 17:08:45 +0800 Subject: [PATCH 1/2] Add gdn_qkv discard-output recompute for GatedDeltaNet Recompute the GatedDeltaNet QKV projection + preparation block (in_proj -> CP all-to-all -> conv1d -> _prepare_qkv -> g/beta) as a discard-output checkpoint, selected via recompute_modules="gdn_qkv". The block outputs (query/key/value/ g/beta/gate) are discarded in the forward and regenerated from a grad hook in the backward, freeing the large GDN QKV-prep activations. When gdn_norm_out recompute is also enabled, the two discard-output checkpoints have a forward-order data dependency (the QKV block output `gate` feeds the gated-norm block), so both are registered to a single CheckpointManager that replays their recompute in forward order (qkv -> norm_out) from one unified grad hook on the layer output. Adds "gdn_qkv" to the selective-recompute allowed_modules in TransformerConfig. --- megatron/core/ssm/gated_delta_net.py | 168 ++++++++++++------ .../core/transformer/transformer_config.py | 10 ++ 2 files changed, 121 insertions(+), 57 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 4cf68a92bf6..8f77b6d8b93 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -29,6 +29,7 @@ _undo_attention_load_balancing, ) from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule @@ -221,8 +222,14 @@ def __init__( eps=self.config.layernorm_epsilon, ) self.recompute_norm_out = False + self.recompute_qkv = False if self.config.recompute_granularity == "selective": self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + # gdn_qkv: recompute the whole QKV proj+prep block as a discard-output checkpoint. + self.recompute_qkv = "gdn_qkv" in self.config.recompute_modules + + # Per-forward CheckpointManager for the GDN discard-output recompute (gdn_qkv/gdn_norm_out). + self.gdn_recompute_manager = None self.out_proj = build_module( submodules.out_proj, @@ -336,6 +343,109 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + # gdn_qkv (QKV proj+prep) and gdn_norm_out (gated norm) are discard-output checkpoints; the + # QKV output `gate` feeds the gated-norm block, so when both are on the CheckpointManager + # replays them in forward order (qkv -> norm_out) from one grad hook on `out`. + recompute_qkv = self.recompute_qkv and self.training + recompute_norm_out = self.recompute_norm_out and self.training + self.gdn_recompute_manager = ( + CheckpointManager() if (recompute_qkv or recompute_norm_out) else None + ) + + # QKV projection + prep block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta). + def _qkv_proj_and_prepare(hidden_states): + return self._compute_qkv_for_gated_delta_rule( + hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ) + + if recompute_qkv: + # Discard the QKV outputs now; regenerate them in backward. Synchronous recompute + # (no async reload), so it is safe with the fla/compiled gated_delta_rule backward. + query, key, value, g, beta, gate = tensor_parallel.CheckpointWithoutOutput( + fp8=(self.config.fp8 or self.config.fp4), + ckpt_manager=self.gdn_recompute_manager, + ).checkpoint(_qkv_proj_and_prepare, hidden_states) + else: + query, key, value, g, beta, gate = _qkv_proj_and_prepare(hidden_states) + + # seq_len was reassigned to the post-CP-a2a sequence length inside the block; recover it + # from a produced tensor so the downstream gated-norm reshape uses the correct value. + seq_len = value.shape[1] + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, last_recurrent_state = self.gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): + # RMSNorm + nvtx_range_push(suffix="gated_norm") + norm_out_hp = 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_hp = norm_out_hp.reshape(batch, seq_len, -1) + norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() + + # CP all to all: HP to CP + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_norm_out = _unpack_sequence(norm_out_hp, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( + norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) + else: + norm_out = tensor_a2a_hp2cp( + norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + + return norm_out + + if recompute_norm_out: + norm_out = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self.gdn_recompute_manager + ).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 outputs (now consumed) and register the unified recompute hook on + # `out` — its grad is computed first in backward, before the backwards that need them. + if self.gdn_recompute_manager is not None: + self.gdn_recompute_manager.discard_all_outputs_and_register_unified_recompute(out) + self.gdn_recompute_manager = None + + return out, out_bias + + def _compute_qkv_for_gated_delta_rule( + self, hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ): + """QKV projection + preparation block for the gated delta rule. + + Runs in_proj, CP all-to-all, conv1d, _prepare_qkv and g/beta, producing the tensors consumed + by ``self.gated_delta_rule`` plus the ``gate`` for the gated norm. Extracted so it can be + checkpointed when ``recompute_modules`` contains ``"gdn_qkv"``. + + Returns: + Tuple of (query, key, value, g, beta, gate). + """ # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) @@ -463,63 +573,7 @@ def forward( g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, last_recurrent_state = self.gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="gated_delta_rule") - - def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out_hp = 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_hp = norm_out_hp.reshape(batch, seq_len, -1) - norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - - # CP all to all: HP to CP - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - unpacked_norm_out = _unpack_sequence(norm_out_hp, cu_seqlens_q, dim=0) - outputs = [] - for norm_out_i in unpacked_norm_out: - norm_out_i = tensor_a2a_hp2cp( - norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - outputs.append(norm_out_i) - norm_out = torch.cat(outputs, dim=0) - else: - norm_out = tensor_a2a_hp2cp( - norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - - return norm_out - - if self.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") - - if self.recompute_norm_out: - self.norm_out_checkpoint.discard_output_and_register_recompute(out) - - return out, out_bias + return query, key, value, g, beta, gate @jit_fuser def _apply_gated_norm(self, x, gate): diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c6eca1c5bcf..88d3fa4111e 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1763,6 +1763,7 @@ def __post_init__(self): "shared_experts", "mhc", "gdn_norm_out", + "gdn_qkv", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1790,6 +1791,15 @@ def __post_init__(self): "experimental_attention_variant='gated_delta_net'." ) + if ( + "gdn_qkv" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_qkv 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, " From 2cac1876b5fe09dc7cda94c52e7963d819ccf060 Mon Sep 17 00:00:00 2001 From: jinliangl <975761915@qq.com> Date: Wed, 3 Jun 2026 17:41:23 +0800 Subject: [PATCH 2/2] Add gdn_qkv discard-output recompute unit test Adds test_selective_recompute_gdn_qkv (mirrors test_selective_recompute_norm_out) to TestGatedDeltaNet: builds a no-recompute baseline GatedDeltaNet and one with recompute_modules=["gdn_qkv"], runs forward+backward, and asserts the output, all parameter grads and the input grad match bit-for-bit. Verifies the QKV projection+prep discard-output recompute is numerically exact. --- tests/unit_tests/ssm/test_gated_delta_net.py | 86 ++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index a007a6fe01e..cb0ba4302f5 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -222,6 +222,92 @@ def run(gdn, hidden_states): rec_grads[name], base_grads[name] ), f"Grad not identical for {name} ({rank=})" + def test_selective_recompute_gdn_qkv(self): + """gdn_qkv discard-output recompute must be numerically exact. + + recompute_modules=["gdn_qkv"] recomputes the whole QKV projection + + preparation block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta) + as a discard-output checkpoint. Output, parameter grads and 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_qkv"] + + 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_qkv is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = 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_qkv is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + 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