diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 266137b3d5d..a83b8fc4784 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/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( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e1ef8fb344f..01b992846d3 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/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, " diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 5c74825ad34..9876dd505cb 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. +# 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