[WEIXILIN0307]

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

What does this PR do?

Allow gdn_core_attn activation offloading when Transformer Engine projections and MLPs use FP8 or NVFP4. The FLA recurrence and offloaded tensors remain BF16/FP32. Existing TE precision contexts and recipes are reused; the Qwen validation tool adds a precision option and keeps BF16 as its default.

TE graph capture now honors disable_parameter_transpose_cache, matching the existing linear wrappers. NVFP4 graph offloading requires TransformerConfig.disable_parameter_transpose_cache=True through programmatic configuration. With caching enabled, an offload-disabled Qwen graph training baseline fails on the tested SM120 stack; this PR does not fix that failure, and its behavior on supported NVFP4 training GPUs is not yet established. Disabling the cache requantizes weights on each microbatch.

Recurrence kernels, the offload manager and transfer scheduling are unchanged.

Dependencies

Depends on #8308 at 03b5ef353a1e657a89247f1ae0f67d05d619166f, including the mRoPE graph fix in #8555. The incremental diff isolates this PR.

Validation

At current head a04f1e65627ec7a5323bdd56d165a746b59f1c6e, including main at 581248531984920e92f076946ad4629d359442a4, 19 CPU configuration tests and two GPU graph training cases pass. The GPU tests cover GDN with FP8 delayed scaling and GDN2 with FP8 tensorwise scaling, comparing offload disabled and enabled.

On one shared RTX PRO 6000, the complete Qwen3.5-0.8B language model (excluding vision and MTP) with FP8 delayed scaling, sequence 128, four microbatches and threshold 1024 matches eager, graph without offload and graph with offload across six optimizer steps. Input hashes, loss, gradients, weights and full optimizer states match; offload records transfers.

Current-head commands, rank logs and checksum audit.

At the earlier head dfaa6338c65edab118ba415fcc14c73241f503a6, one RTX PRO 6000 passes 111 scope/configuration/tool cases, 54 native low-precision GDN/GDN2 training cases and 23 graph cases. Two shared GPUs pass 18 low-precision GDN2 graph cases per rank, covering TP2+SP and PP2. The selected cases have no skips. They cover FP8 delayed and tensorwise recipes, constrained NVFP4, norm recomputation, fractions 0/0.5/1 and pinned-buffer ownership. TP2 fused-gradient graph cases use the distributed optimizer.

On one GPU, the complete 24-layer Qwen3.5-0.8B language model (excluding vision and MTP) trains for six steps with sequence 128, four microbatches and an offload threshold of 1024 elements. With NVFP4 RHT and stochastic rounding enabled, eager, two repeated graph runs without offload, and graph with offload match input hashes, loss, gradients, weights and full optimizer-state fingerprints. The offload arm performs transfers.

At that earlier head, the same language-model configuration also matches eager, graph without offload and graph with offload for six steps under both delayed and tensorwise FP8, including all recorded training-state fingerprints and actual offload transfers.

Hardware limitation: these GPUs are SM120. TE 2.20.2 documents NVFP4 training support for SM100 and SM103. The SM120 NVFP4 training results are exploratory and must be repeated on a supported training GPU. They do not establish convergence or performance.

Two offload-disabled baseline failures were observed on SM120; whether they reproduce on supported NVFP4 training GPUs remains to be established:

• NVFP4 fused weight-gradient accumulation with the ordinary optimizer can use a misaligned gradient-buffer address. Align NVFP4 fused weight-gradient buffers #8556 proposes a fix; the supported-GPU reproduction is still pending.
• Full Qwen TP2 FP4 backward fails in a separate 4112-wide GEMM with fused weight-gradient accumulation; an unfused baseline also fails in an FP4 GEMM. Both use RHT and stochastic rounding disabled. Explicit NVTE_BACKWARD_OVERRIDE=high_precision passes with and without offload on both ranks; that setting uses high-precision backward and is not default FP4 backward.

Commands, source revisions, original rank logs and checksum audit retain those earlier results and failure controls. Earlier low-precision results ran at 7d52248294fd759dd69b83c68ec6d1c0b36f2836. Upstream GPU CI has not run.

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 and coverage checker pass; 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.