From 16f4f7391fc0186a2be6853330e85d4fb2e62a63 Mon Sep 17 00:00:00 2001 From: Tianyu Zhao Date: Wed, 22 Jul 2026 09:43:37 +0800 Subject: [PATCH] [PyTorch] Optionally release columnwise copy of frozen FP8 block-scaled weights after dgrad Add an opt-in path for frozen 2D-block-scaled FP8 weights to release their columnwise representation after dgrad and rebuild it on demand. The behavior is disabled by default and covers Linear, GroupedLinear, LayerNormLinear, and LayerNormMLP. Preserve trainable-weight, CUDA graph, and FSDP2 behavior, and add unit and distributed coverage for rebuild numerics, workspace caching, activation recompute, CUDA graphs, FSDP2, and TP/SP. Signed-off-by: Tianyu Zhao --- docs/envvars.rst | 6 + .../distributed/run_fsdp2_frozen_release.py | 135 ++++++ .../distributed/run_tpsp_frozen_release.py | 105 +++++ .../test_frozen_weight_columnwise_release.py | 46 ++ .../test_frozen_weight_columnwise_release.py | 437 ++++++++++++++++++ transformer_engine/pytorch/module/base.py | 52 ++- .../pytorch/module/grouped_linear.py | 8 + .../pytorch/module/layernorm_linear.py | 3 + .../pytorch/module/layernorm_mlp.py | 7 + transformer_engine/pytorch/module/linear.py | 3 + 10 files changed, 801 insertions(+), 1 deletion(-) create mode 100644 tests/pytorch/distributed/run_fsdp2_frozen_release.py create mode 100644 tests/pytorch/distributed/run_tpsp_frozen_release.py create mode 100644 tests/pytorch/distributed/test_frozen_weight_columnwise_release.py create mode 100644 tests/pytorch/test_frozen_weight_columnwise_release.py diff --git a/docs/envvars.rst b/docs/envvars.rst index b3765a06bd..0e7e169e87 100644 --- a/docs/envvars.rst +++ b/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) diff --git a/tests/pytorch/distributed/run_fsdp2_frozen_release.py b/tests/pytorch/distributed/run_fsdp2_frozen_release.py new file mode 100644 index 0000000000..755abc2a45 --- /dev/null +++ b/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()) diff --git a/tests/pytorch/distributed/run_tpsp_frozen_release.py b/tests/pytorch/distributed/run_tpsp_frozen_release.py new file mode 100644 index 0000000000..29c8ebada9 --- /dev/null +++ b/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()) diff --git a/tests/pytorch/distributed/test_frozen_weight_columnwise_release.py b/tests/pytorch/distributed/test_frozen_weight_columnwise_release.py new file mode 100644 index 0000000000..7845f4021f --- /dev/null +++ b/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 diff --git a/tests/pytorch/test_frozen_weight_columnwise_release.py b/tests/pytorch/test_frozen_weight_columnwise_release.py new file mode 100644 index 0000000000..cb3c330fc7 --- /dev/null +++ b/tests/pytorch/test_frozen_weight_columnwise_release.py @@ -0,0 +1,437 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE. + +Frozen (``requires_grad=False``) quantized weights only need their +columnwise (transposed) copy transiently, as the dgrad GEMM operand. +With ``NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE=1`` the copy is released +right after the dgrad GEMM and rebuilt on demand from rowwise data, +saving up to 1 byte/param of resident GPU memory for PEFT-style +fine-tuning of large frozen base models. +""" + +from __future__ import annotations + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + GroupedLinear, + LayerNormLinear, + LayerNormMLP, + Linear, + autocast, + quantized_model_init, +) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage + +from utils import assert_close, reset_rng_states + +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) + +SEED = 1234 +IN_FEATURES = 256 +OUT_FEATURES = 256 +NUM_GEMMS = 4 +TOKENS_PER_GEMM = 64 + +_MODULE_IDS = ("linear", "grouped_linear", "layernorm_linear", "layernorm_mlp") +_MODULE_CLASSES = (Linear, GroupedLinear, LayerNormLinear, LayerNormMLP) + + +@pytest.fixture(autouse=True) +def reset_global_state(): + reset_rng_states() + yield + + +def _make_block_scaling_recipe(backward_override=None) -> recipe.Float8BlockScaling: + kwargs = {"fp8_format": recipe.Format.E4M3} + if backward_override is not None: + kwargs["backward_override"] = backward_override + return recipe.Float8BlockScaling(**kwargs) + + +def _make_module(module_cls, fp8_recipe, frozen: bool): + """Construct a quantized module, optionally freezing all its weights.""" + ctx = torch.no_grad() if frozen else torch.enable_grad() + with ctx, quantized_model_init(enabled=True, recipe=fp8_recipe): + if module_cls is GroupedLinear: + module = GroupedLinear( + NUM_GEMMS, IN_FEATURES, OUT_FEATURES, bias=False, params_dtype=torch.bfloat16 + ) + elif module_cls is LayerNormMLP: + module = LayerNormMLP( + IN_FEATURES, 4 * IN_FEATURES, bias=False, params_dtype=torch.bfloat16 + ) + else: + module = module_cls(IN_FEATURES, OUT_FEATURES, bias=False, params_dtype=torch.bfloat16) + if frozen: + for param in module.parameters(): + param.requires_grad_(False) + return module + + +def _quantized_weights(module): + return [p for p in module.parameters() if isinstance(p, QuantizedTensorStorage)] + + +def _run_fwd_bwd(module, fp8_recipe): + """One fwd+bwd step where dgrad is required.""" + inp = torch.randn( + NUM_GEMMS * TOKENS_PER_GEMM, + IN_FEATURES, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + with autocast(enabled=True, recipe=fp8_recipe): + if isinstance(module, GroupedLinear): + out = module(inp, [TOKENS_PER_GEMM] * NUM_GEMMS) + else: + out = module(inp) + out.float().pow(2).mean().backward() + return inp.grad.detach().clone() + + +def _columnwise_present(weight) -> bool: + """Whether the tensor currently satisfies columnwise usage.""" + return weight.get_usages()["columnwise"] + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +@pytest.mark.parametrize("module_cls", _MODULE_CLASSES, ids=_MODULE_IDS) +def test_frozen_columnwise_released_after_dgrad(module_cls, monkeypatch): + """Columnwise data of frozen 2D-block-scaled weights is released after backward.""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe() + module = _make_module(module_cls, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + assert weights, "expected quantized weights" + + for _ in range(3): # release -> rebuild -> release must be stable + _run_fwd_bwd(module, fp8_recipe) + for weight in weights: + assert weight._columnwise_data is None, "columnwise copy should be released" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +@pytest.mark.parametrize("module_cls", _MODULE_CLASSES, ids=_MODULE_IDS) +def test_frozen_columnwise_kept_when_disabled(module_cls, monkeypatch): + """Default behavior (env unset/0) keeps the columnwise copy resident.""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "0") + fp8_recipe = _make_block_scaling_recipe() + module = _make_module(module_cls, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + + _run_fwd_bwd(module, fp8_recipe) + for weight in weights: + assert _columnwise_present(weight), "columnwise copy should be kept by default" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +@pytest.mark.parametrize("module_cls", _MODULE_CLASSES, ids=_MODULE_IDS) +def test_frozen_columnwise_release_preserves_numerics(module_cls, monkeypatch): + """dgrad is identical with and without the release flag.""" + fp8_recipe = _make_block_scaling_recipe() + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "0") + reset_rng_states() + module_ref = _make_module(module_cls, fp8_recipe, frozen=True) + dgrad_ref = [_run_fwd_bwd(module_ref, fp8_recipe) for _ in range(2)] + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + reset_rng_states() + module_rel = _make_module(module_cls, fp8_recipe, frozen=True) + dgrad_rel = [_run_fwd_bwd(module_rel, fp8_recipe) for _ in range(2)] + + for ref, rel in zip(dgrad_ref, dgrad_rel): + assert_close(rel, ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +def test_grouped_linear_release_with_backward_override(monkeypatch): + """Release also works when dgrad runs on dequantized weight copies.""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe(backward_override="dequantized") + module = _make_module(GroupedLinear, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + + _run_fwd_bwd(module, fp8_recipe) + for weight in weights: + assert not _columnwise_present(weight), ( + "columnwise copy of the original quantized weight should be released" + " even when dgrad uses dequantized copies" + ) + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +@pytest.mark.parametrize("module_cls", (Linear, GroupedLinear), ids=("linear", "grouped_linear")) +def test_trainable_weights_unaffected(module_cls, monkeypatch): + """Trainable weights keep their columnwise copy and receive wgrad.""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe() + module = _make_module(module_cls, fp8_recipe, frozen=False) + assert all(p.requires_grad for p in module.parameters()) + + _run_fwd_bwd(module, fp8_recipe) + for param in module.parameters(): + assert param.grad is not None or getattr(param, "main_grad", None) is not None + for weight in _quantized_weights(module): + assert _columnwise_present(weight), "trainable columnwise copy must not be released" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +def test_columnwise_only_tensor_skipped(monkeypatch): + """Columnwise-only tensors (FSDP2 backward all-gather shape) are skipped, not crashed. + + Upstream FSDP2 + Float8BlockScaling + quantized_model_init is currently + broken (scale-inv padding in all-gather slice ops), so the guard is + exercised here at unit level with a hand-built columnwise-only tensor. + """ + from transformer_engine.pytorch.module.base import release_frozen_weight_columnwise + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe() + module = _make_module(Linear, fp8_recipe, frozen=True) + (weight,) = _quantized_weights(module) + + # Mimic the FSDP2 backward all-gather product: columnwise-only. + # (Materialize columnwise from rowwise first, then drop rowwise — + # the columnwise-only update_usage path requires the copy to exist.) + weight.update_usage(rowwise_usage=True, columnwise_usage=True) + weight.update_usage(rowwise_usage=False, columnwise_usage=True) + assert weight._rowwise_data is None + assert _columnwise_present(weight) + + release_frozen_weight_columnwise((weight,)) # must be a safe no-op + assert _columnwise_present(weight), "columnwise-only tensor must not be touched" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_non_blockwise_never_receives_release_call(monkeypatch): + """Spy on Float8TensorStorage.update_usage: no columnwise release is ever requested. + + Architecture-independent proof (works even where releasing would be a + physical no-op, e.g. non-TN-capable GEMM hardware). + """ + from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( + Float8TensorStorage, + ) + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + calls = [] + original = Float8TensorStorage.update_usage + + def spy(self, *args, **kwargs): + calls.append((args, kwargs)) + return original(self, *args, **kwargs) + + monkeypatch.setattr(Float8TensorStorage, "update_usage", spy) + + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.E4M3) + module = _make_module(Linear, fp8_recipe, frozen=True) + _run_fwd_bwd(module, fp8_recipe) + + for args, kwargs in calls: + columnwise = kwargs.get("columnwise_usage", args[1] if len(args) > 1 else None) + assert columnwise is not False, "non-blockwise tensor received a columnwise release" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +@pytest.mark.parametrize("frozen", (True, False), ids=("frozen", "trainable")) +def test_workspace_path_with_microbatch_cache(frozen, monkeypatch): + """bf16 Parameters + FP8 workspaces (no quantized_model_init), with microbatch cache. + + Trainable workspaces must keep their columnwise copy; frozen workspaces + are released and rebuilt with numerics identical to flag-off. + """ + fp8_recipe = _make_block_scaling_recipe() + + def build_and_run(flag: str): + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", flag) + reset_rng_states() + module = Linear(IN_FEATURES, OUT_FEATURES, bias=False, params_dtype=torch.bfloat16) + if frozen: + for param in module.parameters(): + param.requires_grad_(False) + grads = [] + for step in range(3): + inp = torch.randn( + NUM_GEMMS * TOKENS_PER_GEMM, + IN_FEATURES, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + with autocast(enabled=True, recipe=fp8_recipe): + out = module(inp, is_first_microbatch=(step == 0)) + out.float().pow(2).mean().backward() + grads.append(inp.grad.detach().clone()) + workspaces = [ + ws for ws in module._fp8_workspaces.values() if isinstance(ws, QuantizedTensorStorage) + ] + return grads, workspaces + + grads_on, workspaces_on = build_and_run("1") + assert workspaces_on, "expected FP8 weight workspaces" + for workspace in workspaces_on: + if frozen: + assert not _columnwise_present(workspace), "frozen workspace should be released" + else: + assert _columnwise_present(workspace), "trainable workspace must be kept" + + grads_off, _ = build_and_run("0") + for ref, rel in zip(grads_off, grads_on): + assert_close(rel, ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +def test_fp8_activation_recompute(monkeypatch): + """te.checkpoint (activation recompute) + release flag: stable and numerics match.""" + from transformer_engine.pytorch import checkpoint + + fp8_recipe = _make_block_scaling_recipe() + + def run(flag: str): + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", flag) + reset_rng_states() + module = _make_module(Linear, fp8_recipe, frozen=True) + grads = [] + for _ in range(2): + inp = torch.randn( + NUM_GEMMS * TOKENS_PER_GEMM, + IN_FEATURES, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + with autocast(enabled=True, recipe=fp8_recipe): + out = checkpoint(module, inp) + out.float().pow(2).mean().backward() + grads.append(inp.grad.detach().clone()) + return grads, _quantized_weights(module) + + grads_on, weights_on = run("1") + for weight in weights_on: + assert not _columnwise_present(weight), "frozen weight should be released after recompute" + grads_off, _ = run("0") + for ref, rel in zip(grads_off, grads_on): + assert_close(rel, ref, rtol=0, atol=0) + + +def _eager_reference(inp_data, flag: str, monkeypatch): + """Same-seed frozen module, one eager fwd+bwd; returns (output, dgrad).""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", flag) + fp8_recipe = _make_block_scaling_recipe() + reset_rng_states() + module = _make_module(Linear, fp8_recipe, frozen=True) + inp = inp_data.clone().requires_grad_(True) + with autocast(enabled=True, recipe=fp8_recipe): + out = module(inp) + out.float().pow(2).mean().backward() + return out.detach().clone(), inp.grad.detach().clone() + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +def test_cuda_graph_make_graphed_callables(monkeypatch): + """make_graphed_callables: guard keeps columnwise, numerics match eager flag-off.""" + from transformer_engine.pytorch import make_graphed_callables + + torch.manual_seed(SEED) + inp_data = torch.randn( + NUM_GEMMS * TOKENS_PER_GEMM, IN_FEATURES, device="cuda", dtype=torch.bfloat16 + ) + ref_out, ref_dgrad = _eager_reference(inp_data, "0", monkeypatch) + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe() + reset_rng_states() + module = _make_module(Linear, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + + sample = inp_data.clone().requires_grad_(True) + graphed = make_graphed_callables( + module, (sample,), num_warmup_iters=3, enabled=True, recipe=fp8_recipe + ) + for _ in range(3): # capture + replays must be stable + inp = inp_data.clone().requires_grad_(True) + out = graphed(inp) + out.float().pow(2).mean().backward() + + # Guard is the only thing standing between capture-time dgrad and the + # release: columnwise must still be present after capture + replays + # (replay executes no Python, so the state must be stable). + for weight in weights: + assert _columnwise_present(weight), "columnwise must survive graph capture/replay" + assert_close(out.detach(), ref_out, rtol=0, atol=0) + assert_close(inp.grad, ref_dgrad, rtol=0, atol=0) + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +def test_cuda_graph_direct_capture(monkeypatch): + """Direct torch.cuda.graph capture/replay: fixed grad buffer updated, numerics match.""" + torch.manual_seed(SEED) + inp_data = torch.randn( + NUM_GEMMS * TOKENS_PER_GEMM, IN_FEATURES, device="cuda", dtype=torch.bfloat16 + ) + ref_out, ref_dgrad = _eager_reference(inp_data, "0", monkeypatch) + + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = _make_block_scaling_recipe() + reset_rng_states() + module = _make_module(Linear, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + + static_inp = inp_data.clone().requires_grad_(True) + # Warmup on a side stream (per torch.cuda.graph docs). Do NOT reset + # .grad to None afterwards: replay writes into the captured buffers, + # so the grad buffer address must stay fixed to observe the update. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + with autocast(enabled=True, recipe=fp8_recipe): + out = module(static_inp) + out.float().pow(2).mean().backward() + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + with autocast(enabled=True, recipe=fp8_recipe): + static_out = module(static_inp) + static_loss = static_out.float().pow(2).mean() + static_loss.backward() + + for _ in range(3): # replays must be stable and re-write the fixed buffer + static_inp.grad.zero_() # in place: keep the captured address + graph.replay() + torch.cuda.synchronize() + assert not torch.all(static_inp.grad == 0), "replay must update the fixed grad buffer" + + for weight in weights: + assert _columnwise_present(weight), "columnwise must survive graph capture/replay" + assert_close(static_out.detach(), ref_out, rtol=0, atol=0) + assert_close(static_inp.grad.clone(), ref_dgrad, rtol=0, atol=0) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_non_2d_scaled_weights_not_released(monkeypatch): + """Delayed-scaling weights (not rebuildable from rowwise) keep columnwise usage.""" + monkeypatch.setenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "1") + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.E4M3) + module = _make_module(Linear, fp8_recipe, frozen=True) + weights = _quantized_weights(module) + + for _ in range(2): # would fail on step 2 if the transpose had been dropped + _run_fwd_bwd(module, fp8_recipe) + for weight in weights: + assert not getattr(weight, "_is_2D_scaled", False) + assert _columnwise_present(weight), "columnwise usage must not be released" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5eab45b7ac..7b94e27c89 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -10,7 +10,7 @@ import warnings from enum import Enum from abc import ABC, abstractmethod -from typing import Any, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union from contextlib import contextmanager from types import MethodType @@ -756,6 +756,56 @@ def _is_weight_workspace_valid( return True +def release_frozen_weight_columnwise( + weights: Iterable[Union[torch.Tensor, QuantizedTensorStorage]], +) -> None: + """Release the columnwise copy of frozen quantized weights after dgrad. + + Frozen weights (e.g. the base model in LoRA/PEFT fine-tuning) never + need a wgrad, so their columnwise (transposed) copy is only used + transiently as the dgrad GEMM operand. Keeping it resident costs up + to 1 byte/param of GPU memory for the whole run. When + ``NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE=1``, release the columnwise + copy right after the dgrad GEMM; it is rebuilt on demand from the + rowwise data by the next ``update_usage`` call. + + Callers must only pass weights that do not require a wgrad in this + backward pass (e.g. gated on ``ctx.weights_requires_grad``); the + weight objects seen in backward may be workspaces or restored saved + tensors whose ``requires_grad`` attribute is not meaningful. + + Note: if the tensor's quantizer requests columnwise usage (e.g. a + primary quantized parameter initialized with gradients enabled), the + next forward's ``quantize_weight`` rebuilds the copy, so the release + then mainly reduces between-step residency rather than the full-step + peak. Initializing frozen parameters under ``torch.no_grad()`` + avoids the rebuild and yields the peak-memory benefit. + + Only weights that can rebuild the columnwise copy from a complete + rowwise representation are released: 2D-block-scaled + ``Float8BlockwiseQTensorStorage`` with both rowwise data and scale-inv + present. Columnwise-only tensors (e.g. FSDP2 all-gathered weights in + backward) and other layouts are skipped. No-op during CUDA graph + capture. + """ + if os.getenv("NVTE_RELEASE_FROZEN_WEIGHT_COLUMNWISE", "0") != "1": + return + if FP8GlobalStateManager.fp8_graph_capturing(): + return + for weight in weights: + if not isinstance(weight, Float8BlockwiseQTensorStorage): + continue + if not weight._is_2D_scaled: + continue + # Only weights holding a complete rowwise representation can be + # released and rebuilt later. FSDP2 all-gathered weights in + # backward are columnwise-only and must be skipped (they are + # cleaned up by ``clear_columnwise_cache`` right after this). + if weight._rowwise_data is None or weight._rowwise_scale_inv is None: + continue + weight.update_usage(rowwise_usage=True, columnwise_usage=False) + + def quantize_weight( *, tensor: Optional[torch.Tensor] = None, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a65ee3b5c3..895aef8941 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -23,6 +23,7 @@ from .base import ( get_dummy_wgrad, quantize_weight, + release_frozen_weight_columnwise, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -882,6 +883,8 @@ def _backward_grouped_tensor( layout="NN", use_split_accumulator=dgrad_gemm_use_split_accumulator, ) + if not ctx.weights_requires_grad: + release_frozen_weight_columnwise(weights) if ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1113,6 +1116,11 @@ def backward( grad=True, use_split_accumulator=dgrad_gemm_use_split_accumulator, ) + # Release the original quantized weights, not + # ``weights_for_dgrad``, which may hold dequantized + # copies under ``backward_override``. + if not ctx.weights_requires_grad: + release_frozen_weight_columnwise(weights) if ctx.weights_requires_grad: wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..ceaf2de6b8 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -25,6 +25,7 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, + release_frozen_weight_columnwise, TransformerEngineBaseModule, get_dummy_wgrad, _2X_ACC_FPROP, @@ -851,6 +852,8 @@ def backward( bulk_overlap=ctx.ub_bulk_dgrad, ) nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + if not ctx.requires_wgrad: + release_frozen_weight_columnwise((weight,)) # FSDP2 only handles deallocation all-gathered weights that it allocates. # Columnwise data is derived from rowwise data after allgather for fp8 diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ad629c3f11..31820607c3 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -27,6 +27,7 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, + release_frozen_weight_columnwise, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -1258,6 +1259,9 @@ def backward( ub_type=tex.CommOverlapType.AG if ctx.ub_overlap_ag else None, ) + if not ctx.fc2_weight_requires_grad: + release_frozen_weight_columnwise((fc2_weight,)) + # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM # to prevent them from persisting on the all-gathered buffer. # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs @@ -1557,6 +1561,9 @@ def fc2_wgrad_gemm( bulk_overlap=ctx.ub_bulk_dgrad, ) + if not ctx.fc1_weight_requires_grad: + release_frozen_weight_columnwise((fc1_weight,)) + # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM # to prevent them from persisting on the all-gathered buffer. # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 78a4d31852..edd32bf335 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -25,6 +25,7 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, + release_frozen_weight_columnwise, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -1036,6 +1037,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. bulk_overlap=bwd_args.ub_bulk_dgrad, ) nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + if not bwd_args.requires_wgrad: + release_frozen_weight_columnwise((weight_fp8,)) # FSDP2 only handles deallocation all-gathered weights that it allocates. # Columnwise data is derived from rowwise data after allgather for fp8