[WEIXILIN0307]

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

What does this PR do?

Extend opt-in gdn_core_attn activation offloading to BF16/FLA GDN2 training and Transformer Engine CUDA graphs for GDN and GDN2. The existing recurrence scope, backward prefetch anchor, transfer streams and pinned-memory pool are reused. TE layer-graph replay stays on the caller's autograd stream so shared graph pools remain ordered.

GDN2 keeps sigmoid gate outputs resident when gate backward retains the same storage: the differentiable write gate and the differentiable erase gate without head expansion. This avoids transfers that cannot release the original GPU storage.

Attention and whole-layer TE graphs, gdn_norm_out recomputation and delayed offload are supported. Graph offload requires PyTorch >= 2.9, TE >= 2.14 and positive graph warmup; set NVTE_CPU_OFFLOAD_V1=1 before launch. GDN2 requires fla.ops.gdn2. FP8/FP4, full recomputation and local/full-iteration graphs are outside this PR; low precision is proposed in #8543. See configuration and usage.

Dependencies

Merge #7399, #8243 and #7908 first. This branch also includes the independent mRoPE graph-input fix in #8555, needed for correct Qwen graph training. The original GDN2/graph feature diff excludes that separate fix. Composition with #7852 is not validated.

Validation

At 03b5ef353a1e657a89247f1ae0f67d05d619166f, including #8555 and main at 581248531984920e92f076946ad4629d359442a4:

• On one RTX PRO 6000, the complete 24-layer Qwen3.5-0.8B language model (excluding vision and MTP) matches eager, graph without offload and graph with offload for six BF16 optimizer steps. Sequence length is 128, with four microbatches and an offload threshold of 1024 elements. Input hashes, loss, gradients, weights and full optimizer states match; the offload arm records transfers.
• A single-GPU GDN2 training test with synthetic tokens passes with TE graph capture and replay, comparing offload disabled and enabled. A larger graph test run was interrupted during cold compilation and is not counted as a pass.
• Upstream GPU CI has not run. Fixed-length training initialized by a training step is covered; evaluation-first initialization and tensor-size changes after warmup are not qualified.

Current-head commands, rank logs, source checksums and audit. Earlier Qwen and two-rank GDN/GDN2 tests ran at c5264b66197040f4309edfdb22dc9bbaf62189b1; their original logs remain available. Multi-rank and performance results have not been rerun on the current head.

Memory and runtime

At earlier head c5264b66197040f4309edfdb22dc9bbaf62189b1, BF16 whole-layer graph training on one RTX PRO 6000 used sequence 2048, four microbatches, fraction 1, the default threshold and no norm recomputation. Each fresh-process arm uses three warmup steps and ten measured optimizer steps; capture is timed separately. Peaks are medians of the per-arm peaks, and step changes are paired medians and ranges.

Model	Allocated peak off → on (MiB)	Reserved peak off → on (MiB)	Retained pinned memory on (MiB)	Step-time change, median [range]
Qwen3.5-0.8B language model	16137.82 → 16137.82	22576 → 22122	504	+24.75% [+12.71%, +26.55%]
GDN2, eight blocks	3236.12 → 3237.12	5212 → 4964	432	+97.83% [+96.40%, +98.97%]

The Qwen row uses five pairs because of timing variation; GDN2 uses three. These runs reduce reserved memory but not the measured allocated peak. No other compute process was observed on the measured GPU, but other GPUs were shared; host and PCIe isolation are not established. Commands, rank logs, samples, telemetry and audit.

A separate allocator snapshot at the same head and Qwen configuration locates the measured peak during output projection and cross-entropy: 970 MiB of logits and a 1940 MiB FP32 conversion are live together. In the snapshot after that step, the graph pool reserves 4132 → 3690 MiB, while its allocated count is 46 MiB in both arms. In the GDN2 snapshot, graph-pool reserve also decreases, while the total allocated peak increases by 1 MiB. These observations apply to the measured configurations.

As an exploratory follow-up, smaller offload fractions and the existing delayed-offload option did not reduce allocated peak in this configuration. The delayed-offload arm matches losses, gradients, weights and complete optimizer states for six steps; the timing samples are not sufficient to establish a performance improvement. Snapshot traces, commands, comparison data and audit.

Earlier A6000 measurements and a gate-selection experiment used different revisions and are retained in the historical graph-training evidence and profile. The historical Qwen graph run omitted mRoPE and is not a valid correct-training performance result.

The option remains disabled by default.

Issue tracking

Related to #6757.

Code provenance

No code copied or adapted from external repositories.

Contribution process

Pre-checks

• I have added relevant unit tests.
• I have added relevant functional tests.
• The existing GDN determinism-manifest association covers this dispatch; numerical kernels are unchanged.
• I have added proper typing to the changed code.
• I have added relevant documentation and reproducible commands.
• I have run Black, isort, Ruff and the project Pylint rules on the changed Python files.

[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.