[PyTorch] Optionally release columnwise copy of frozen FP8 block-scaled weights after dgrad

[1tex]

Description

PEFT/LoRA-style fine-tuning keeps a large frozen base model in memory. With primary FP8
parameters (quantized_model_init) and the Float8BlockScaling recipe, each weight holds
two FP8 copies: rowwise (used by fprop) and columnwise (used only as the dgrad GEMM
operand). For frozen (requires_grad=False) weights the columnwise copy is needed only
transiently, yet it stays resident after the first backward pass — up to 1 byte per
locally resident frozen parameter. This is the compute-for-memory trade-off suggested by
@timmoon10 in #1880; #1847 addressed initialization, this PR addresses the steady state
during training.

This PR adds an opt-in environment variable NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE
(default off). When enabled, the columnwise copy of frozen 2D-block-scaled weights is
released right after the dgrad GEMM and rebuilt on demand via the existing update_usage
path — one quantization-aware transpose per frozen layer per step in exchange for the
resident memory.

Known limitation (documented in envvars.rst): if the weight's quantizer requests
columnwise usage (e.g. parameters initialized outside torch.no_grad()), the copy is
rebuilt in the next forward, so the option mainly reduces between-step residency; the
peak-memory benefit requires initializing frozen parameters without columnwise usage.

Type of change

• Documentation change
• Bug fix
• New feature (non-breaking change which adds functionality)
• Breaking change
• Infra/Build change
• Code refactoring

Changes

Core library change is ~80 lines; the rest is tests and validation scripts.

• Add release_frozen_weight_columnwise helper in module/base.py (env-var gated; only
  2D-block-scaled Float8BlockwiseQTensorStorage with a full rowwise representation;
  no-op during CUDA graph capture), called after the dgrad GEMM in Linear,
  GroupedLinear, LayerNormLinear, and LayerNormMLP (FC1/FC2), gated on the caller's
  wgrad-requirement flags.
• Document the new env var in docs/envvars.rst.
• Add a 23-test unit suite (bitwise-exact numerics vs. flag-off, trainable weights,
  non-2D layouts, backward_override, workspace caching, activation recompute, CUDA
  graphs) plus FSDP2/TP+SP multi-GPU validation scripts with a pytest driver.

Relationship with #2168

#2168 proposes a more general on-demand columnwise mechanism for blockwise FP8 weights.
This PR is intentionally narrower: it applies only to frozen weights, reuses the existing
update_usage path, releases immediately after dgrad, and additionally covers
LayerNormMLP.

I asked about coordination in #2168 on Jul 18. This PR provides a concrete frozen-only
implementation based on current main; happy to fold this work into #2168 if preferred.

Validation

Everything coverable on H800 (8 GPUs) passes: the 23-test suite (on both the default and
the fused grouped-GEMM paths), the distributed driver
(FSDP2 with both reshard_after_forward modes, TP+SP), and module regressions, all with
zero failures. No Blackwell hardware was available locally; a spy-based test provides an
architecture-independent guarantee that non-blockwise layouts are never touched.

On a 16-layer frozen stack (features=4096, tokens=4096, 0.268B params), steady-state and
peak memory drop by 0.25 GiB — matching the theoretical columnwise size — at +3.9% step
time under torch.no_grad() initialization; with grad-enabled initialization only
between-step residency drops (the documented limitation).

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

Adds an opt-in mechanism to release and lazily rebuild columnwise copies of frozen 2D block-scaled FP8 weights.

• Introduces NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE and documents its scope and CUDA-graph limitation.
• Applies post-dgrad release handling to Linear, GroupedLinear, LayerNormLinear, and both LayerNormMLP weights.
• Adds single-GPU and distributed coverage to the curated L0 and L1 QA jobs, resolving the previous CI-coverage finding.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds the environment-gated helper that restricts release to rebuildable 2D block-scaled FP8 storage outside CUDA graph capture.
transformer_engine/pytorch/module/linear.py	Releases frozen Linear weight columnwise storage immediately after the dgrad GEMM.
transformer_engine/pytorch/module/grouped_linear.py	Integrates release into grouped dgrad paths, including backward-override handling.
transformer_engine/pytorch/module/layernorm_linear.py	Adds post-dgrad release for frozen LayerNormLinear weights.
transformer_engine/pytorch/module/layernorm_mlp.py	Independently releases frozen FC1 and FC2 columnwise representations after their dgrad GEMMs.
tests/pytorch/test_frozen_weight_columnwise_release.py	Covers release eligibility, retained layouts, numerical equivalence, workspace caching, recomputation, and CUDA graph behavior.
tests/pytorch/distributed/test_frozen_weight_columnwise_release.py	Drives the FSDP2 and tensor/sequence-parallel validation scripts through two-GPU subprocesses.
qa/L0_pytorch_unittest/test.sh	Adds the new single-GPU suite to the curated L0 PyTorch CI job.
qa/L1_pytorch_distributed_unittest/test.sh	Adds the new multi-GPU suite to the curated L1 distributed CI job.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Backward requires dgrad] --> B[Run dgrad GEMM]
    B --> C{Weight gradient required?}
    C -- Yes --> D[Keep columnwise representation]
    C -- No --> E{Release option enabled?}
    E -- No --> D
    E -- Yes --> F{2D block-scaled weight with complete rowwise data?}
    F -- No --> D
    F -- Yes --> G{CUDA graph capture active?}
    G -- Yes --> D
    G -- No --> H[Release columnwise representation]
    H --> I[Rebuild on demand through update_usage]

Loading

_{Reviews (2): Last reviewed commit: "[PyTorch] Add frozen FP8 weight release ..." | Re-trigger Greptile}