[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. When the base
model is stored as primary FP8 parameters (quantized_model_init) with the
Float8BlockScaling recipe, each weight holds two FP8 copies: rowwise (used by fprop)
and columnwise (used only as the dgrad GEMM operand). For frozen weights
(requires_grad=False) the columnwise copy is needed only transiently, yet today it
stays resident after the first backward pass — up to 1 byte/param of the frozen model
per rank, for the whole run.

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 by the existing
update_usage(columnwise_usage=True) path in the next backward — trading one
quantization-aware transpose each time the columnwise copy is rebuilt (typically once
per frozen layer per training step) for the resident memory.

This is the trade-off suggested by @timmoon10 in #1880 ("reduce training memory usage
to 1 byte/param in exchange for some extra compute"). #1847 addressed initialization
(no columnwise usage when initializing under torch.no_grad()); this PR addresses the
steady state after the first backward. #2168 explores a more general on-demand
columnwise mechanism for all weights; see the "Relationship with #2168" section below.

Known limitation (documented in envvars.rst): for weights whose quantizer requests
columnwise usage (e.g. initialized outside torch.no_grad()), the copy is rebuilt in
the next forward, so this 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, gated on the env
  var; only applies to Float8BlockwiseQTensorStorage with _is_2D_scaled=True and a
  full rowwise representation (data + scale-inv); no-op during CUDA graph capture.
• Call it after the dgrad GEMM in Linear, GroupedLinear (both grouped-tensor and
  legacy paths, including backward_override), LayerNormLinear, and LayerNormMLP
  (FC1 and FC2), gated on the caller's wgrad-requirement flags.
• Document the new env var in docs/envvars.rst.
• Add tests/pytorch/test_frozen_weight_columnwise_release.py (23 tests): release/rebuild
  stability, bitwise-exact numerics vs. flag-off, flag-off regression, trainable weights
  unaffected, non-2D layouts not released, backward_override, columnwise-only guard,
  architecture-independent update_usage spy, workspace/microbatch-cache paths,
  activation recompute, CUDA graphs (make_graphed_callables + direct torch.cuda.graph).
• Add multi-GPU validation scripts tests/pytorch/distributed/run_fsdp2_frozen_release.py
  (both reshard_after_forward modes) and run_tpsp_frozen_release.py (TP+SP), with a
  pytest driver tests/pytorch/distributed/test_frozen_weight_columnwise_release.py
  (three parametrized cases, skipped with a clear reason on <2 GPUs).

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. I am opening this as a draft so
the concrete diff is available; happy to fold this work into #2168 if preferred.

Validation

All targeted validation coverable on H800 (8 GPUs) completed: 23/23 unit tests, FSDP2
both reshard modes (workspace/cache-free path; the columnwise-only guard is covered at
unit level since upstream FSDP2 + primary FP8 params is currently broken), TP+SP smoke,
and module regressions (test_grouped_linear, test_backward_override,
test_cuda_graphs, test_sanity) with zero failures.
Measured on a 16-layer frozen stack (features=4096, tokens=4096, 0.268B params):
steady-state and peak memory drop
by 0.25 GiB under torch.no_grad() initialization, matching the theoretical columnwise
size, at +3.9% step time (CUDA-event timing of fwd+bwd, 50 iterations); with
grad-enabled initialization only between-step residency drops (documented limitation).
No Blackwell hardware was available locally, so Blackwell validation is left for upstream;
a spy-based test provides an architecture-independent guarantee for non-blockwise layouts.

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