From a6aefdf0e3a2bc12e0f6ebb0193174c1ce2488b2 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Thu, 7 May 2026 07:58:00 -0700 Subject: [PATCH 1/7] add norm_out recompute in gdn --- megatron/core/ssm/gated_delta_net.py | 21 +++++++++++++++++++ .../core/transformer/transformer_config.py | 16 ++++++++++++-- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 1dbaf161e73..41a6e62772d 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 @@ -238,6 +239,10 @@ def __init__( self.reset_parameters() + self.recompute_norm_out = False + if self.config.recompute_granularity == "selective": + self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + def reset_parameters(self): """Reset the parameters.""" if self.config.perform_initialization: @@ -507,6 +512,9 @@ def forward( 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 @jit_fuser @@ -521,6 +529,19 @@ def _apply_gated_norm(self, x, gate): y = y.to(x_dtype) return y + def _run_gated_norm_and_a2a(self, core_attn_out: Tensor, gate: Tensor) -> Tensor: + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + batch, seq_len = gate.shape[:2] + norm_out = norm_out.reshape(batch, seq_len, -1).transpose(0, 1).contiguous() + + norm_out = tensor_a2a_hp2cp( + norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + return norm_out + @jit_fuser def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): """ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index cab540e6bfd..086116b1e03 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -499,7 +499,8 @@ 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". + choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", + "shared_experts", "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. @@ -508,7 +509,8 @@ class TransformerConfig(ModelParallelConfig): "mlp": recompute the dense MLP submodule. "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. - "moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, + "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. + "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ @@ -1580,6 +1582,7 @@ def __post_init__(self): "mlp", "moe", "shared_experts", + "gdn_norm_out", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1598,6 +1601,15 @@ def __post_init__(self): "multi_latent_attention." ) + 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 "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " From edfe19d00f4747fa9b44f186d0851247fbea6c45 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Tue, 12 May 2026 21:54:31 -0700 Subject: [PATCH 2/7] add gdn test --- tests/unit_tests/ssm/test_gated_delta_net.py | 53 ++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 6bd7e3ea841..b176fe5b05c 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -141,6 +141,59 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_gpu_forward_recompute_norm_out(self): + torch.cuda.manual_seed_all(42) + self.transformer_config.recompute_granularity = "selective" + self.transformer_config.recompute_modules = ["gdn_norm_out"] + + 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) + + gdn_submodules = get_experimental_attention_variant_module_spec( + config=self.transformer_config + ).submodules + gdn = GatedDeltaNet( + self.transformer_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, + ) + gdn = gdn.cuda().bfloat16() + + assert gdn.recompute_norm_out is True + + micro_batch_size = 2 + seq_length = 64 + hidden_states = 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, + requires_grad=True, + ) + attention_mask = None + + output, bias = gdn(hidden_states, attention_mask) + + assert output.shape == ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + gdn.config.hidden_size, + ), f"Output shape mismatch: {output.shape=}" + assert ( + output.dtype == hidden_states.dtype + ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + + # Backward exercises the recompute hook that + # CheckpointWithoutOutput.discard_output_and_register_recompute + # installed on `out`, so the gated norm + a2a runs a second time. + output.sum().backward() + def test_jit_compiled_helpers(self): import torch._dynamo From 8552ddfdea1583720429cdf9476a86c36a019b38 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Mon, 18 May 2026 07:57:04 -0700 Subject: [PATCH 3/7] rebase upstream --- megatron/core/ssm/gated_delta_net.py | 59 ++++++++++++++-------------- 1 file changed, 30 insertions(+), 29 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 41a6e62772d..e4f0e6fbcdf 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -482,30 +482,13 @@ def forward( ) 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") - - # 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() - - # 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, 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, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + if self.recompute_norm_out: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_out = self.norm_out_checkpoint.checkpoint( + self._run_gated_norm_and_a2a, core_attn_out, gate, packed_seq_params ) + else: + norm_out = self._run_gated_norm_and_a2a(core_attn_out, gate, packed_seq_params) # Output projection nvtx_range_push(suffix="out_proj") @@ -529,17 +512,35 @@ def _apply_gated_norm(self, x, gate): y = y.to(x_dtype) return y - def _run_gated_norm_and_a2a(self, core_attn_out: Tensor, gate: Tensor) -> Tensor: + def _run_gated_norm_and_a2a( + self, core_attn_out: Tensor, gate: Tensor, packed_seq_params + ) -> Tensor: nvtx_range_push(suffix="gated_norm") norm_out = self._apply_gated_norm(core_attn_out, gate) nvtx_range_pop(suffix="gated_norm") - batch, seq_len = gate.shape[:2] - norm_out = norm_out.reshape(batch, seq_len, -1).transpose(0, 1).contiguous() + batch, seq_len = norm_out.shape[:2] + + # 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() + + # 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, 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, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) - norm_out = tensor_a2a_hp2cp( - norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) return norm_out @jit_fuser From 2806fecc6900cb437c45faa34820e7632bfe4889 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Mon, 18 May 2026 23:34:32 -0700 Subject: [PATCH 4/7] update test --- tests/unit_tests/ssm/test_gated_delta_net.py | 103 ++++++++++++------- 1 file changed, 65 insertions(+), 38 deletions(-) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index b176fe5b05c..d6b5d7baaf1 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import copy from unittest import mock import pytest @@ -141,58 +142,84 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" - def test_gpu_forward_recompute_norm_out(self): - torch.cuda.manual_seed_all(42) - self.transformer_config.recompute_granularity = "selective" - self.transformer_config.recompute_modules = ["gdn_norm_out"] - + def test_selective_recompute_norm_out(self): 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) - gdn_submodules = get_experimental_attention_variant_module_spec( - config=self.transformer_config - ).submodules - gdn = GatedDeltaNet( - self.transformer_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, - ) - gdn = gdn.cuda().bfloat16() - - assert gdn.recompute_norm_out is True + 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, gdn.config.hidden_size), + ( + 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, ) - attention_mask = None - output, bias = gdn(hidden_states, attention_mask) - - assert output.shape == ( - seq_length // self.sp_size // self.cp_size, - micro_batch_size, - gdn.config.hidden_size, - ), f"Output shape mismatch: {output.shape=}" - assert ( - output.dtype == hidden_states.dtype - ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + # --- 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 + 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) - # Backward exercises the recompute hook that - # CheckpointWithoutOutput.discard_output_and_register_recompute - # installed on `out`, so the gated norm + a2a runs a second time. - output.sum().backward() + 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 From 3955c76fc52594308afa2f1f75b4e7def70df193 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Tue, 19 May 2026 01:03:26 -0700 Subject: [PATCH 5/7] rebase for thd layout --- megatron/core/ssm/gated_delta_net.py | 78 +++++++++++++--------------- 1 file changed, 37 insertions(+), 41 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index e4f0e6fbcdf..b3349a25ee2 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -221,6 +221,9 @@ def __init__( hidden_size=self.value_head_dim, eps=self.config.layernorm_epsilon, ) + self.recompute_norm_out = False + if self.config.recompute_granularity == "selective": + self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules self.out_proj = build_module( submodules.out_proj, @@ -239,10 +242,6 @@ def __init__( self.reset_parameters() - self.recompute_norm_out = False - if self.config.recompute_granularity == "selective": - self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules - def reset_parameters(self): """Reset the parameters.""" if self.config.perform_initialization: @@ -482,22 +481,50 @@ def forward( ) nvtx_range_pop(suffix="gated_delta_rule") + def _gated_norm_and_transform(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() + + return norm_out_hp + if self.recompute_norm_out: self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() - norm_out = self.norm_out_checkpoint.checkpoint( - self._run_gated_norm_and_a2a, core_attn_out, gate, packed_seq_params + norm_out_hp = self.norm_out_checkpoint.checkpoint( + _gated_norm_and_transform, core_attn_out, gate ) else: - norm_out = self._run_gated_norm_and_a2a(core_attn_out, gate, packed_seq_params) + norm_out_hp = _gated_norm_and_transform(core_attn_out, gate) + + # 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 + ) + + if self.recompute_norm_out: + self.norm_out_checkpoint.discard_output_and_register_recompute(norm_out) # 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 @jit_fuser @@ -512,37 +539,6 @@ def _apply_gated_norm(self, x, gate): y = y.to(x_dtype) return y - def _run_gated_norm_and_a2a( - self, core_attn_out: Tensor, gate: Tensor, packed_seq_params - ) -> Tensor: - nvtx_range_push(suffix="gated_norm") - norm_out = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") - - batch, seq_len = norm_out.shape[:2] - - # 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() - - # 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, 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, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - - return norm_out - @jit_fuser def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): """ From 89bffb17c25dc3e77c6ea78e38ea4046f135d123 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Tue, 19 May 2026 02:25:15 -0700 Subject: [PATCH 6/7] include a2a inside the recompute region --- megatron/core/ssm/gated_delta_net.py | 46 +++++++++++++--------------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index b3349a25ee2..39bdd17cbb0 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -481,7 +481,7 @@ def forward( ) nvtx_range_pop(suffix="gated_delta_rule") - def _gated_norm_and_transform(core_attn_out: torch.Tensor, gate: torch.Tensor): + 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) @@ -492,39 +492,37 @@ def _gated_norm_and_transform(core_attn_out: torch.Tensor, gate: torch.Tensor): norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - return norm_out_hp + # 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_hp = self.norm_out_checkpoint.checkpoint( - _gated_norm_and_transform, core_attn_out, gate - ) - else: - norm_out_hp = _gated_norm_and_transform(core_attn_out, gate) - - # 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) + norm_out = self.norm_out_checkpoint.checkpoint(_gated_norm_and_a2a, core_attn_out, gate) else: - norm_out = tensor_a2a_hp2cp( - norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - - if self.recompute_norm_out: - self.norm_out_checkpoint.discard_output_and_register_recompute(norm_out) + 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 @jit_fuser From e3d7d0d5aa1095dde4ae72532206bb95e8cebd0a Mon Sep 17 00:00:00 2001 From: Xuanteng Huang Date: Wed, 27 May 2026 22:26:12 -0700 Subject: [PATCH 7/7] add extra assertion --- megatron/core/ssm/gated_delta_net.py | 1 + tests/unit_tests/ssm/test_gated_delta_net.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 39bdd17cbb0..6c2c3d9fea9 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -222,6 +222,7 @@ def __init__( eps=self.config.layernorm_epsilon, ) self.recompute_norm_out = False + self.norm_out_checkpoint = None if self.config.recompute_granularity == "selective": self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index d6b5d7baaf1..6af2b3ccd42 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -202,6 +202,7 @@ def run(gdn, hidden_states): 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() @@ -211,6 +212,7 @@ def run(gdn, hidden_states): 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=})"