[WEIXILIN0307]

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

What does this PR do?

Add opt-in gdn_core_attn activation offloading for BF16/FLA GDN training, using the existing manager, transfer streams and pinned-memory pool.

Enable with --fine-grained-activation-offloading --offload-modules gdn_core_attn. With TE 2.10 or newer, set NVTE_CPU_OFFLOAD_V1=1 before launch. The scope wraps the FLA recurrence and commits before output normalization. A differentiable recurrence input anchors backward prefetch; shared inputs are not force-released. No-grad and fully frozen recurrence execution bypass the GDN scope.

Selective gdn_norm_out recomputation and existing pre-GDR fusion are supported. GDN2, FP8/FP4, non-BF16 execution, other recurrence backends, full recomputation and CUDA graphs are rejected. GDN2/TE graph support is proposed separately in #8308. Native pre-GDR kernels and autotuning are unchanged.

FP8/NVFP4 support for TE projections and MLPs is proposed separately in #8543, after #8308. It is not included in this PR.

Includes a saved-tensor profiler/benchmark and a local-Qwen training tool with native checkpoint save/load. Megatron Bridge is required only by the Qwen tool. See usage and reproduction commands.

Dependencies and scope

Requires #8243 at 5bc9e72bda70129f088dd78328d4064f1d838c62 and its #7399 prerequisite. The diff against main includes those dependencies; the GDN-only diff isolates this feature.

Projection/convolution/QKV recomputation and convolution-input offloading belong to #7852; combined use is not validated.

Validation

Head: b6fa63ab8c2cb855f0368cd8455245c295cde234, including the #8243 prerequisite and a merge from main at 3ddfe0fe07efc73dbdc40d4a02acbc1c1f8a9f08.

• On one RTX PRO 6000, the scope/configuration/native decoder-DDP-BF16 Adam/tool suite has 104 passes and 84 topology skips. Four targeted pipeline/checkpoint tests also pass.
• On two shared RTX PRO 6000 GPUs, TP2+SP and PP2 native decoder training has 24 passes per rank without skips.
• Training of the complete Qwen3.5-0.8B language model (excluding vision and MTP) runs for three optimizer steps with norm recomputation and four microbatches. Disabled, repeated-disabled and fraction-one arms match every recorded loss, gradient, weight and complete Adam/master-state fingerprint on one GPU and on both ranks with TP2+SP. The enabled arm performs transfers and pinned-buffer ownership is checked before reset.
• At sequence 128 with four microbatches, fraction one and norm recomputation, the same language model runs 203 steps with finite losses and gradient norms. For the 200 post-warmup steps, allocated and reserved GPU peaks are stable at 13937.40 MiB and 14492 MiB; the pinned pool retains 126 buffers (31.92 MiB) and has none in use after each step.
• At the same configuration, six continuous steps match three steps followed by a native checkpoint save, process restart and three resumed steps. Every recorded loss, gradient, weight and full optimizer-state fingerprint matches.

Commands, source import paths, original rank logs and checksums identify the executed head. The current stability and checkpoint evidence includes per-step pool and state checks.

Validation covers fixed-length training initialized by a training step; evaluation-first initialization and tensor-size changes after warmup are not qualified.

The earlier four-rank matrix and TP2+SP+PP2 Qwen runs ran at e96793feecb0fdd54ce00ef09a67e28fce940d06 and were not repeated at this head. Native fused decoder comparisons use rtol=1e-5, atol=1e-8 for atomic gate-gradient variability; unfused comparisons are exact. Upstream GPU CI is pending.

Recorded memory and runtime

At the current head b6fa63ab8c2cb855f0368cd8455245c295cde234, complete Qwen3.5-0.8B language-model training on one RTX PRO 6000 used BF16, sequence 4096, four microbatches, fraction 1, the default threshold and no norm recomputation. Each row uses three fresh-process pairs, three warmup steps and ten measured optimizer steps. Peak values are medians of the per-arm peaks; step changes are paired medians and ranges.

Pre-GDR	Allocated peak off → on (MiB)	Reserved peak off → on (MiB)	Retained pinned memory on (MiB)	Step-time change, median [range]
Unfused	26067.95 → 25115.95	28394 → 27382	1008	+0.67% [−0.08%, +1.39%]
Fused	25770.95 → 24819.58	28914 → 27902	1008	−1.40% [−2.65%, +1.53%]

The fused result does not establish a speedup. No other compute process was observed on the measured GPU, but other GPUs were shared; host and PCIe isolation are not established. The timed runs compare loss and gradient norm; the full-state correctness runs are described above. Commands, rank logs, samples, telemetry and audit.

Earlier performance and checkpoint and stability results used different revisions and configurations; they do not qualify the current head.

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 determinism-manifest association covers GDN dispatch; numerical kernels are unchanged.
• I have added proper typing to the changed code.
• I have added 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.

[lhb8125]

/ok to test b22c9f7

[lhb8125]

Thank you for contributing GDN recurrence activation offloading and for the extensive testing and validation work behind it.

To make this easier to review and merge, I suggest splitting it into three PRs:

1. Shared activation-offload infrastructure fixes: evaluation/warmup handling, fraction-zero behavior, chunk progression, and D2H/H2D ordering and pooled-buffer reuse, including the scheduler changes.
2. Pre-GDR gate backward reproducibility: the fixed tiling and FP32 partial-sum reduction replacing atomic accumulation. This changes the existing fused backward path even when offloading is disabled, so it deserves a separate numerical and performance review.
3. GDN recurrence activation offloading: the new scope, configuration validation, and GDN integration, based on the prerequisite fixes.

Please also consider moving the raw CSV measurements and detailed historical experiment logs into the pinned evidence repository/artifacts already linked here. The in-tree documentation can retain the design and usage guidance, reproducible commands, a concise results summary, and known limitations, with links to the full evidence.

Tests and validation should move with the corresponding changes: each PR should include its relevant regression tests and evidence tied to that PR's code revision, while the final GDN feature PR should validate the combined end-to-end training path. This would preserve the useful validation work while making each change easier to assess independently.

[WEIXILIN0307]

Thank you for contributing GDN recurrence activation offloading and for the extensive testing and validation work behind it.

To make this easier to review and merge, I suggest splitting it into three PRs:

1. Shared activation-offload infrastructure fixes: evaluation/warmup handling, fraction-zero behavior, chunk progression, and D2H/H2D ordering and pooled-buffer reuse, including the scheduler changes.
2. Pre-GDR gate backward reproducibility: the fixed tiling and FP32 partial-sum reduction replacing atomic accumulation. This changes the existing fused backward path even when offloading is disabled, so it deserves a separate numerical and performance review.
3. GDN recurrence activation offloading: the new scope, configuration validation, and GDN integration, based on the prerequisite fixes.

Please also consider moving the raw CSV measurements and detailed historical experiment logs into the pinned evidence repository/artifacts already linked here. The in-tree documentation can retain the design and usage guidance, reproducible commands, a concise results summary, and known limitations, with links to the full evidence.

Tests and validation should move with the corresponding changes: each PR should include its relevant regression tests and evidence tied to that PR's code revision, while the final GDN feature PR should validate the combined end-to-end training path. This would preserve the useful validation work while making each change easier to assess independently.

Thanks @lhb8125, that split makes sense, especially reviewing the pre-GDR backward change separately since it also affects the offload-disabled path.
I can separate the shared offload fixes, pre-GDR backward change, and GDN integration into three PRs, with the relevant tests and revision-specific evidence in each. The backward PR would include numerical and performance checks, and the final GDN PR would cover combined end-to-end training.
I can also move the raw CSVs and historical logs to the pinned evidence repo, keeping the in-tree docs focused on design, usage, reproduction commands, results, and limitations