Add gdn_qkv whole-block recompute for GatedDeltaNet

[wplf]

What

Adds a new selective-recompute module gdn_qkv that wraps the entire GatedDeltaNet QKV projection + preparation block (in_proj → CP all-to-all → conv1d → _prepare_qkv → g/beta) in a single activation checkpoint, analogous to recompute_modules="moe" wrapping the whole MoE forward.

• gated_delta_net.py: extract the QKV-proj+prep span into _compute_qkv_for_gated_delta_rule(...) returning (query, key, value, g, beta, gate); in forward, run it directly or via tensor_parallel.checkpoint / te_checkpoint (fp8) when "gdn_qkv" is in recompute_modules. seq_len (reassigned to the post-CP-a2a length inside the block) is recovered from value.shape[1] after the call.
• transformer_config.py: register gdn_qkv in the allowed recompute_modules set and guard that it requires experimental_attention_variant="gated_delta_net" (mirrors gdn_norm_out).

Why

The GDN QKV-prep activations previously had no recompute/offload knob and were a large share of per-layer activation memory.

Validation

Qwen3.5-VL 397B proxy, 8×GB200 (EP=8, mbs=4, multimodal entry), 20 iters, stable loss (no NaN/skips):

	max allocated	throughput
baseline stack (no gdn_qkv)	85.9 GB	~510 TFLOP/s
+ gdn_qkv	77.8 GB (−8.1 GB)	493–498 TFLOP/s (−3%)

🤖 Generated with Claude Code

Summary by Sourcery

Add a selective activation recompute option for the GatedDeltaNet QKV projection and preparation block and wire it into the transformer configuration.

Enhancements:

• Introduce a configurable whole-block checkpoint for the GatedDeltaNet QKV projection and preparation path to reduce activation memory usage during training.
• Refactor the GatedDeltaNet QKV projection and preparation logic into a dedicated helper method to enable reuse and checkpoint wrapping.
• Extend transformer configuration validation to recognize the new gdn_qkv recompute module and restrict its use to the gated_delta_net attention variant.

[sourcery-ai]

Reviewer's Guide

Adds a new selective-recompute option gdn_qkv that wraps the entire GatedDeltaNet QKV projection+prep block in a single activation checkpoint, refactoring that span into a helper and wiring it into both standard and TE fp8/fp4 checkpoint paths, plus config validation for the new recompute module.

Sequence diagram for new gdn_qkv selective recompute path

sequenceDiagram
    participant GatedDeltaNet as GatedDeltaNetLayer
    participant Checkpoint as tensor_parallel.checkpoint
    participant TECheckpoint as te_checkpoint
    participant GDRule as gated_delta_rule

    GatedDeltaNet->>GatedDeltaNet: forward(hidden_states)
    alt recompute_qkv and training
        alt fp8 or fp4
            GatedDeltaNet->>TECheckpoint: te_checkpoint(_qkv_proj_and_prepare, False, get_cuda_rng_tracker, tp_group, hidden_states)
            TECheckpoint->>GatedDeltaNet: query,key,value,g,beta,gate
        else non_fp8_fp4
            GatedDeltaNet->>Checkpoint: checkpoint(_qkv_proj_and_prepare, False, hidden_states)
            Checkpoint->>GatedDeltaNet: query,key,value,g,beta,gate
        end
    else no_qkv_recompute
        GatedDeltaNet->>GatedDeltaNet: _compute_qkv_for_gated_delta_rule(hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params)
        GatedDeltaNet-->>GatedDeltaNet: query,key,value,g,beta,gate
    end
    GatedDeltaNet->>GatedDeltaNet: seq_len = value.shape[1]
    GatedDeltaNet->>GDRule: gated_delta_rule(query,key,value,g,g,beta,cu_seqlens_q)
    GDRule-->>GatedDeltaNet: core_attn_out,last_recurrent_state
    GatedDeltaNet->>GatedDeltaNet: _gated_norm_and_a2a(core_attn_out,gate)
    GatedDeltaNet-->>GatedDeltaNet: norm_out
    GatedDeltaNet->>GatedDeltaNet: out_proj(norm_out)
    GatedDeltaNet-->>GatedDeltaNet: out,out_bias

Loading

File-Level Changes

Change	Details	Files
Introduce a reusable QKV projection+preparation helper for GatedDeltaNet and make it optionally recomputed as a single block.	Extract the QKV projection and preparation span (in_proj, CP all-to-all, conv1d, _prepare_qkv, g/beta computation) into a new method _compute_qkv_for_gated_delta_rule that returns (query, key, value, g, beta, gate). Add a local wrapper _qkv_proj_and_prepare in forward that calls the helper with the current hidden_states, batch, seq_len, cu_seqlens_q, and packed_seq_params. Recover seq_len in forward from value.shape[1] after the helper call so downstream reshapes use the post-CP-a2a sequence length. Adjust the end of the original inlined block to return the six QKV-related tensors from the helper instead of continuing the forward path.	megatron/core/ssm/gated_delta_net.py
Wire gdn_qkv into the selective recompute mechanism, including TE checkpoint support for low-precision modes.	Add self.recompute_qkv flag to the GatedDeltaNet module, set when recompute_granularity is "selective" and "gdn_qkv" is present in config.recompute_modules. Wrap the new _qkv_proj_and_prepare helper in tensor_parallel.checkpoint for standard precision when recompute_qkv is enabled and the model is in training mode. Use te_checkpoint with the TP RNG tracker and TP process group instead of tensor_parallel.checkpoint when config.fp8 or config.fp4 is enabled, ensuring compatibility with Transformer Engine checkpointing. Keep the existing gated-norm + CP-all-to-all + out-proj path unchanged, still optionally wrapped by the existing gdn_norm_out checkpoint logic.	megatron/core/ssm/gated_delta_net.py
Register and validate the new gdn_qkv recompute module in transformer configuration.	Extend the allowed set of recompute_modules entries to include "gdn_qkv" when experimental_attention_variant is "gated_delta_net". Add a config-time validation that raises a ValueError if "gdn_qkv" is used with a non-gated_delta_net attention variant, mirroring the existing gdn_norm_out guard.	megatron/core/transformer/transformer_config.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.