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

docs/envvars.rst

@@ -219,6 +219,12 @@ FP8 Configuration
    :Default: ``0``
    :Description: Use unfused kernel for FP8 amax and scale updates. When set to ``1``, amax and scale updates are computed using separate unfused kernels instead of fused operations.
 
+.. envvar:: NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``0``
+   :Description: Release the columnwise (transposed) copy of frozen quantized weights right after the dgrad GEMM in the backward pass of ``Linear``, ``GroupedLinear``, ``LayerNormLinear``, and ``LayerNormMLP``. Weights that do not require gradients (e.g. the frozen base model in LoRA/PEFT fine-tuning) only need their columnwise copy transiently for dgrad, so releasing it reduces resident weight memory by up to roughly one byte per locally resident frozen parameter, at the cost of rebuilding the copy (a quantization-aware transpose) each time it is needed again. Only applies to weight layouts that can rebuild columnwise data from rowwise data (currently 2D-block-scaled FP8, i.e. ``Float8BlockScaling`` weights); other layouts and trainable weights are unaffected. Note that if the weight's quantizer requests columnwise usage (e.g. primary quantized parameters initialized with gradients enabled), the copy is rebuilt in the next forward pass, so this option then mainly reduces memory residency between training steps; the peak-memory benefit requires initializing frozen parameters without columnwise usage (e.g. under ``torch.no_grad()``). The memory saving materializes after the first backward pass. Not applied during CUDA graph capture.
+
 .. envvar:: NVTE_FP8_DPA_BWD
 
    :Type: ``int`` (0 or 1)

tests/pytorch/distributed/run_fsdp2_frozen_release.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Multi-GPU FSDP2 validation for NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE.
+
+Launch: torchrun --nproc_per_node=2 run_fsdp2_frozen_release.py --reshard-after-forward {0,1}
+
+Model: frozen TE Linear stack (bf16 Parameters, Float8BlockScaling autocast)
++ trainable bf16 head, wrapped with FSDP2 fully_shard. TE intentionally
+disables FP8 weight-workspace caching under FSDP2 (linear.py:
+``cache_name = None`` when ``is_fsdp2``), so there is no persistent cache to
+inspect after the step; transient backward workspaces may still be handled by
+the release helper and the existing FSDP2 cleanup, without forming resident
+state. This path validates safety and numerics. For both
+reshard_after_forward settings this script verifies:
+  1. 3 training steps run without error with the flag on;
+  2. input grads are bitwise identical to a flag-off run;
+  3. no cached FP8 weight workspaces exist (cache-free expectation).
+
+Note: quantized_model_init (primary FP8 params) + FSDP2 + Float8BlockScaling is
+currently broken upstream (scale-inv padding in all-gather slice ops, see
+fsdp2_tests/run_fsdp2_model.py xfail), so that combination cannot be run here;
+the columnwise-only all-gather guard is covered at unit level in
+tests/pytorch/test_frozen_weight_columnwise_release.py.
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import fully_shard
+
+import transformer_engine.pytorch as te
+from transformer_engine.common import recipe
+from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage
+
+SEED = 1234
+FEATURES = 256
+LAYERS = 4
+TOKENS = 256
+
+
+def _columnwise_present(ws) -> bool:
+    return ws.get_usages()["columnwise"]
+
+
+def build_model(device):
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    layers = [
+        te.Linear(FEATURES, FEATURES, bias=False, params_dtype=torch.bfloat16)
+        for _ in range(LAYERS)
+    ]
+    for layer in layers:
+        for p in layer.parameters():
+            p.requires_grad_(False)
+    head = torch.nn.Linear(FEATURES, FEATURES, bias=False, dtype=torch.bfloat16, device=device)
+    return torch.nn.Sequential(*layers, head)
+
+
+def run(flag: str, reshard: bool, device):
+    os.environ["NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE"] = flag
+    fp8_recipe = recipe.Float8BlockScaling(fp8_format=recipe.Format.E4M3)
+    model = build_model(device)
+    for layer in model:
+        fully_shard(layer, reshard_after_forward=reshard)
+    fully_shard(model, reshard_after_forward=reshard)
+
+    grads = []
+    for step in range(3):
+        torch.manual_seed(SEED + step + dist.get_rank())
+        inp = torch.randn(TOKENS, FEATURES, device=device, dtype=torch.bfloat16, requires_grad=True)
+        with te.autocast(enabled=True, recipe=fp8_recipe):
+            # is_first_microbatch is passed, but under FSDP2 TE disables the
+            # weight-workspace cache (cache_name=None when is_fsdp2), so no
+            # workspace is retained across the step; transient backward
+            # workspaces may still be processed by the release helper.
+            out = inp
+            for layer in model[:-1]:
+                out = layer(out, is_first_microbatch=(step == 0))
+            out = model[-1](out)
+        out.float().pow(2).mean().backward()
+        grads.append(inp.grad.detach().clone())
+
+    cached_workspaces = []
+    for module in model.modules():
+        workspaces = getattr(module, "_fp8_workspaces", None)
+        if not workspaces:
+            continue
+        for ws in workspaces.values():
+            if isinstance(ws, QuantizedTensorStorage):
+                cached_workspaces.append(ws)
+    return grads, cached_workspaces
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--reshard-after-forward", type=int, required=True, choices=(0, 1))
+    args = parser.parse_args()
+    reshard = bool(args.reshard_after_forward)
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device("cuda", local_rank)
+    dist.init_process_group(backend="nccl", device_id=device)
+
+    grads_off, workspaces_off = run("0", reshard, device)
+    grads_on, workspaces_on = run("1", reshard, device)
+
+    # TE intentionally disables weight-workspace caching under FSDP2
+    # (linear.py: cache_name=None when is_fsdp2). Assert that cache-free
+    # expectation explicitly (no empty-`all()` vacuous pass): there is no
+    # persistent cache to inspect; this path validates safety and numerics.
+    assert len(workspaces_off) == 0, f"unexpected cached workspaces: {len(workspaces_off)}"
+    assert len(workspaces_on) == 0, f"unexpected cached workspaces: {len(workspaces_on)}"
+    for ref, rel in zip(grads_off, grads_on):
+        assert torch.equal(ref, rel), "dgrad mismatch between flag off/on under FSDP2"
+
+    dist.barrier()
+    if dist.get_rank() == 0:
+        print(
+            f"FSDP2 OK reshard_after_forward={reshard} world={dist.get_world_size()}"
+            " steps=3 cached_workspaces=0 (explicit cache-free expectation)"
+            " dgrad=bitwise-equal",
+            flush=True,
+        )
+    dist.destroy_process_group()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/pytorch/distributed/run_tpsp_frozen_release.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Multi-GPU TP/SP smoke for NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE.
+
+Launch: torchrun --nproc_per_node=2 run_tpsp_frozen_release.py
+
+Frozen column-parallel + row-parallel te.Linear pair (Float8BlockScaling,
+quantized_model_init) with sequence_parallel=True. One fwd+bwd per step,
+3 steps: must not raise, frozen weights must have columnwise released,
+and dgrad must be bitwise identical to a flag-off run.
+"""
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+import transformer_engine.pytorch as te
+from transformer_engine.common import recipe
+from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage
+
+SEED = 1234
+FEATURES = 256
+TOKENS_PER_RANK = 128
+
+
+def run(flag: str, device, tp_group, world):
+    os.environ["NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE"] = flag
+    fp8_recipe = recipe.Float8BlockScaling(fp8_format=recipe.Format.E4M3)
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    with torch.no_grad(), te.quantized_model_init(enabled=True, recipe=fp8_recipe):
+        col = te.Linear(
+            FEATURES,
+            FEATURES,
+            bias=False,
+            params_dtype=torch.bfloat16,
+            parallel_mode="column",
+            tp_group=tp_group,
+            tp_size=world,
+            sequence_parallel=True,
+        )
+        row = te.Linear(
+            FEATURES,
+            FEATURES,
+            bias=False,
+            params_dtype=torch.bfloat16,
+            parallel_mode="row",
+            tp_group=tp_group,
+            tp_size=world,
+            sequence_parallel=True,
+        )
+    weights = []
+    for module in (col, row):
+        for p in module.parameters():
+            p.requires_grad_(False)
+            if isinstance(p, QuantizedTensorStorage):
+                weights.append(p)
+
+    grads = []
+    for step in range(3):
+        torch.manual_seed(SEED + step)  # same across ranks: SP gathers along sequence
+        inp = torch.randn(
+            TOKENS_PER_RANK, FEATURES, device=device, dtype=torch.bfloat16, requires_grad=True
+        )
+        with te.autocast(enabled=True, recipe=fp8_recipe):
+            out = row(col(inp))
+        out.float().pow(2).mean().backward()
+        grads.append(inp.grad.detach().clone())
+    released = [not w.get_usages()["columnwise"] for w in weights]
+    return grads, released
+
+
+def main():
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device("cuda", local_rank)
+    dist.init_process_group(backend="nccl", device_id=device)
+    world = dist.get_world_size()
+    tp_group = dist.new_group(ranks=list(range(world)))
+
+    grads_off, _ = run("0", device, tp_group, world)
+    grads_on, released_on = run("1", device, tp_group, world)
+
+    assert released_on, "expected quantized frozen weights"
+    assert all(released_on), f"frozen weights not released: {released_on}"
+    for ref, rel in zip(grads_off, grads_on):
+        assert torch.equal(ref, rel), "dgrad mismatch between flag off/on under TP/SP"
+
+    dist.barrier()
+    if dist.get_rank() == 0:
+        print(
+            f"TP/SP OK world={world} sequence_parallel=True steps=3"
+            f" released={sum(released_on)}/{len(released_on)} dgrad=bitwise-equal",
+            flush=True,
+        )
+    dist.destroy_process_group()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/pytorch/distributed/test_frozen_weight_columnwise_release.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Distributed tests for NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE.
+
+Each case launches a torchrun subprocess (see run_fsdp2_frozen_release.py
+and run_tpsp_frozen_release.py for the executed checks).
+"""
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+import transformer_engine.pytorch as te
+
+if torch.cuda.device_count() < 2:
+    pytest.skip(
+        "NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE distributed tests need at least 2 GPUs.",
+        allow_module_level=True,
+    )
+
+fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available(
+    return_reason=True
+)
+
+TEST_ROOT = Path(__file__).parent.resolve()
+NUM_PROCS: int = 2
+LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"]
+
+_CASES = (
+    ("fsdp2_reshard_true", "run_fsdp2_frozen_release.py", ["--reshard-after-forward", "1"]),
+    ("fsdp2_reshard_false", "run_fsdp2_frozen_release.py", ["--reshard-after-forward", "0"]),
+    ("tp_sp_smoke", "run_tpsp_frozen_release.py", []),
+)
+
+
+@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling)
+@pytest.mark.parametrize("case", _CASES, ids=lambda case: case[0])
+def test_distributed_frozen_columnwise_release(case):
+    _, script, extra_args = case
+    test_cmd = LAUNCH_CMD + [str(TEST_ROOT / script)] + extra_args
+    result = subprocess.run(test_cmd, env=os.environ, check=False)
+    assert result.returncode == 0