Add gdn_in_proj_conv wide selective recompute for GatedDeltaNet

[maxyzhu]

• [✅] I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds a fine-grained selective activation-recompute option, gdn_in_proj_conv, to the GatedDeltaNet (GDN) module. It recomputes the front-end region of the GDN layer — input projection, CP all-to-all, split, and the causal conv1d — as a single activation checkpoint, instead of stashing its (large) intermediate activations for the backward pass.

On a 20-layer GDN stack this saves ~20.3% peak memory and ~27.9% activation memory at ~17.6% iteration-time overhead, versus the no-recompute baseline.

Enable with recompute_granularity="selective" and recompute_modules=["gdn_in_proj_conv"].

How we arrived at this boundary (measured on a 4×GPU environment)

Activation recompute trades compute for memory; the memory it saves is whatever the checkpointed region would otherwise retain across the forward→backward boundary. The interesting question was which region to recompute. All numbers below are per single GDN layer unless stated otherwise (bf16, hidden=2048, seq=8192, batch=2, TP=CP=1); act = activation memory retained after forward, peak = full fwd+bwd step peak.

1. Recomputing any single sub-module alone does not pay off

segment (recomputed alone)	act saved	peak saved	note
in_proj linear only	3.98%	2.02%	only frees the linear's small internals
conv1d only	~0%	~0%	conv output is not a peak driver
gated_delta_rule (deterministic)	78.8%	−4.8%	recompute of the memory-heavy torch path spikes the backward peak; deterministic mode is not used in production
gated_delta_rule (non-det, FLA)	3.7%	−3.9%	redundant — the FLA kernel already recomputes its state internally

The large tensor in this region is qkvzba (the in_proj output, in_proj_dim wide). A narrow checkpoint cannot free it, because qkvzba is the checkpoint's output and is stashed by the downstream ops regardless.

2. Widening the region from in_proj toward the delta-rule input

We swept the checkpoint end-point (region always starts at the layer input):

boundary	act saved	peak saved	time ovhd
W1: in_proj linear only	3.98%	2.02%	~17%
W2: + CP all-to-all + split	3.98%	2.02%	~16%
W3: + conv1d	27.95%	14.20%	~17%
W4: + l2norm + repeat_interleave	36.03%	18.30%	~20%

The step change at W3 is because once conv1d is inside the region, qkvzba becomes an internal activation and is freed; the region now retains only the small layer input (hidden_states) plus the narrower post-conv output. W3 gives the cleanest saving-per-overhead ratio, so this PR ships W3.

3. End-to-end check on a 20-layer GDN stack

Per-layer savings compound across depth (each layer frees its front-end activations early, lowering the whole-forward peak). On a 20-layer GDN stack (seq=2048, batch=1, TP=CP=1):

config	time (ms)	peak (MB)	act (MB)
baseline	198.6	5531.9	4033.2
gdn_in_proj_conv	233.5	4406.7	2908.0
saved / overhead	+17.57%	−20.34%	−27.90%

Correctness

gdn_in_proj_conv recomputes a region that contains a CP collective (all-to-all); the recompute re-issues it symmetrically across ranks (same pattern as the existing gdn_norm_out recompute). Verified by unit tests over TP/CP ≤ 2:

• Deterministic mode — output, input grad and all parameter grads are bit-exact vs the no-recompute baseline (torch.equal). Hard gate on all architectures.
• Non-deterministic (production) mode — recompute error must stay within the kernels' own run-to-run noise floor, measured per-run (self-calibrating, no fixed tolerance), so the check remains valid across architectures. The printed deviations also serve as a per-hardware kernel-determinism record.

What's included

• gdn_in_proj_conv selective-recompute option (config validation guard: requires experimental_attention_variant="gated_delta_net").
• Unit tests: test_selective_recompute_in_proj_conv_deterministic and test_selective_recompute_in_proj_conv_within_kernel_noise (parametrized over TP/CP ≤ 2).
• Measurement methodology and raw records (see Reproducibility below).

What's deferred (follow-ups)

• W4 (+ l2norm + repeat_interleave) — an extra ~8% activation / ~4% peak at higher overhead; left as a follow-up boundary option.
• gdn_gated_delta_rule recompute — deterministic mode is counterproductive (peak rises) and non-deterministic mode is redundant with the FLA kernel's internal recompute; a genuinely useful variant needs a memory-efficient non-deterministic recompute path (e.g. a bwd-friendly GDN kernel). Tracked separately.
• Whole-model / mixed-layer numbers — the stack above is GDN-only (an upper bound); in a model with ~3:1 GDN:attention the whole-model peak saving is diluted. A full GPTModel run with per-layer profiling is a follow-up.
• Packed (thd) sequences — the non-packed path is covered and tested; the thd path is structurally supported but not yet validated, so it is guarded.

Reproducibility

Measurement scripts and the boundary-sweep data are kept out of this PR (they are measurement-only and decoupled from the feature), preserved on branch archive/local-boundary-bench / tag evidence/gdn-boundary-sweep.

Issue tracking

Linked issue: New Feature: fine-grained selective recompute for in_proj, conv1d, gated delta rule in [2026 Q3] Megatron Core Roadmap #5676

Contribution process

Pre-checks

• [✅] I have added relevant unit tests
• I have added relevant functional tests
• [✅] I have added proper typing to my code Typing guidelines
• [✅] I have added relevant documentation
• [✅] I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxyzhu]

@NVIDIA/mcore-oncall PR is ready for review whenever you get a chance, thanks!

[maxyzhu]

hi @NVIDIA/mcore-oncall PR ok to add label and test/merge. Thanks!