From d34be8de62fd5fb961093291a10d0e0f17fb2d83 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Fri, 31 Jul 2026 15:47:18 -0700 Subject: [PATCH 1/2] [PyTorch] MXFP4 weight QAT on MXFP8 and FP8 block-scaling recipes Weight-only MXFP4 quantization-aware training: weights are projected onto the MXFP4 grid (E2M1 payloads, 1x32 power-of-two UE8M0 scales) before the host recipe quantizes them, so training sees exactly the deployment weight while activations and gradients keep the base recipe untouched. The only new numerics is the fused projection bf16/fp32 -> MXFP4 grid -> bf16/fp32 (tex.mxfp4_fake_quantize): per 1x32 block, scale 2^clamp(ceil(log2(amax/6)), -126, 125) derived from the amax bit pattern, RTNE onto E2M1, straight-through-estimator gradient. All MXFP8 and 128x128 blockwise FP8 encoding/decoding of the projected weight runs through the existing quantizer pipeline unchanged. Since every projected value is a grid point p * 2^e, the existing encodings are decoded-value exact: MXFP8 rowwise unconditionally, 128x128 blockwise for tile scale spreads up to 2^14. Recipes: MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling) and MXFP4QATFloat8BlockScaling(Float8BlockScaling); QAT can also be enabled on the stock host recipes via the mxfp4_qat_weights field (NVTE_MXFP4_QAT=1). backward_override keeps its base-recipe semantics. Also fixes ptx::exp2f(e8m0) to decode UE8M0 code 0 as 2^-127 and code 255 as NaN. Known limitations: for amax > 6*2^125 the fake-dequant saturates at the cap while packed deployment moves to scale 2^126; fp16 weights and fp16 activation/dequantize dtypes are rejected (the MXFP4 grid exceeds fp16 range); distributed end-to-end verification with a real optimizer master step and a fused projection+quantize kernel are future work. --- .../fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py | 385 +++++++ tests/pytorch/distributed/test_torch_fsdp2.py | 28 +- .../pytorch/references/mxfp4_qat_reference.py | 48 + tests/pytorch/test_mxfp4_qat.py | 981 ++++++++++++++++++ .../test_ops_grouped_linear_mxfp4_qat.py | 167 +++ .../pytorch/test_ops_grouped_mlp_mxfp4_qat.py | 143 +++ transformer_engine/common/cast/cast.cu | 33 + .../common/cast/mxfp4/fake_quantize_mxfp4.cuh | 127 +++ .../common/include/transformer_engine/cast.h | 9 + transformer_engine/common/recipe/__init__.py | 44 +- transformer_engine/common/util/ptx.cuh | 2 + transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 18 + .../pytorch/csrc/extensions/pybind.cpp | 3 + transformer_engine/pytorch/module/base.py | 107 +- .../pytorch/module/grouped_linear.py | 63 +- .../pytorch/module/layernorm_linear.py | 60 +- .../pytorch/module/layernorm_mlp.py | 25 +- transformer_engine/pytorch/module/linear.py | 33 +- transformer_engine/pytorch/mxfp4_qat.py | 65 ++ .../pytorch/ops/basic/basic_linear.py | 28 + .../pytorch/ops/basic/grouped_linear.py | 121 ++- .../pytorch/ops/fused/grouped_mlp.py | 80 +- .../ops/fused/userbuffers_backward_linear.py | 24 + .../ops/fused/userbuffers_forward_linear.py | 32 + transformer_engine/pytorch/ops/fuser.py | 21 +- transformer_engine/pytorch/ops/op.py | 36 + 27 files changed, 2630 insertions(+), 55 deletions(-) create mode 100644 tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py create mode 100644 tests/pytorch/references/mxfp4_qat_reference.py create mode 100644 tests/pytorch/test_mxfp4_qat.py create mode 100644 tests/pytorch/test_ops_grouped_linear_mxfp4_qat.py create mode 100644 tests/pytorch/test_ops_grouped_mlp_mxfp4_qat.py create mode 100644 transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh create mode 100644 transformer_engine/pytorch/mxfp4_qat.py diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py new file mode 100644 index 0000000000..fe09cdb635 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py @@ -0,0 +1,385 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FSDP2 coverage for MXFP4-QAT fusible operations. + +Run with: + torchrun --nproc_per_node=2 -m pytest -v -s --tb=short + +Covers only configurations that pure MXFP8 already supports: + +* BasicLinear/Linear: all backward override modes. +* GroupedLinear, independent high-precision weights: ``None``. +* Fused GroupedMLP, independent high-precision weights: ``None``, when the CuTe + DSL fusion is available. +""" + +from collections.abc import Callable, Sequence + +import pytest +import torch +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor +from torch.distributed._composable.fsdp import fully_shard + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common.recipe import ( + MXFP4QATMXFP8BlockScaling, + MXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize + +_DTYPE = torch.bfloat16 + + +def _skip_if_unsupported() -> None: + supported, reason = fp8.check_mxfp8_support() + if not supported: + pytest.skip(reason) + + +def _device() -> torch.device: + return torch.device("cuda", torch.cuda.current_device()) + + +def _named_weights(module: torch.nn.Module) -> list[tuple[str, torch.nn.Parameter]]: + weights = [ + (name, param) + for name, param in module.named_parameters() + if name.rsplit(".", 1)[-1].startswith("weight") + ] + assert weights, f"{type(module).__name__} has no weight parameters" + return weights + + +def _make_master_weight( + out_features: int, + in_features: int, + *, + device: torch.device, +) -> torch.Tensor: + """Construct BF16 values that MXFP8 preserves but MXFP4 changes.""" + pattern = torch.tensor( + [ + 1.0, + 0.3125, + -0.3125, + 0.6875, + -0.6875, + 0.15625, + -0.15625, + 0.8125, + ], + dtype=_DTYPE, + device=device, + ) + repeats = (in_features + pattern.numel() - 1) // pattern.numel() + row = pattern.repeat(repeats)[:in_features] + return row.unsqueeze(0).expand(out_features, -1).contiguous() + + +def _initialize_master_weights(module: torch.nn.Module) -> dict[str, torch.Tensor]: + master_weights = {} + with torch.no_grad(): + for name, param in _named_weights(module): + weight = _make_master_weight( + param.shape[0], + param.shape[1], + device=param.device, + ) + param.copy_(weight) + master_weights[name] = weight + return master_weights + + +def _project_weights( + master_weights: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + projected_weights = { + name: mxfp4_fake_quantize(weight).detach() for name, weight in master_weights.items() + } + for name in master_weights: + assert not torch.equal( + master_weights[name], projected_weights[name] + ), f"Test weight {name} unexpectedly lies entirely on the MXFP4 grid" + return projected_weights + + +@torch.no_grad() +def _load_weights(module: torch.nn.Module, weights: dict[str, torch.Tensor]) -> None: + named_weights = _named_weights(module) + assert [name for name, _ in named_weights] == list(weights) + for name, param in named_weights: + param.copy_(weights[name]) + + +def _fsdp2_shard(module: torch.nn.Module) -> torch.nn.Module: + world_size = torch.distributed.get_world_size() + mesh = DeviceMesh("cuda", list(range(world_size))) + fully_shard(module, mesh=mesh) + return module + + +def _run_reference( + factory: Callable[[], torch.nn.Module], + weights: dict[str, torch.Tensor], + x: torch.Tensor, + dy: torch.Tensor, + extra_args: Sequence[torch.Tensor], + *, + override: str | None, +) -> tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: + module = factory() + _load_weights(module, weights) + x_ref = x.detach().clone().requires_grad_(True) + recipe = MXFP8BlockScaling(backward_override=override) + with te.autocast(enabled=True, recipe=recipe): + y_ref = module(x_ref, *extra_args) + y_ref.backward(dy) + assert x_ref.grad is not None + return module, y_ref.detach(), x_ref.grad.detach() + + +def _assert_projected_forward( + y_qat: torch.Tensor, + y_projected: torch.Tensor, + y_master: torch.Tensor, +) -> None: + torch.testing.assert_close(y_qat, y_projected, rtol=0, atol=0) + assert not torch.equal( + y_qat, y_master + ), "QAT forward matched the unprojected master-weight path" + + +def _assert_fsdp_weight_grads( + fsdp_module: torch.nn.Module, + reference_module: torch.nn.Module, + *, + rtol: float = 0, + atol: float = 0, +) -> None: + fsdp_weights = _named_weights(fsdp_module) + reference_weights = _named_weights(reference_module) + assert [name for name, _ in fsdp_weights] == [name for name, _ in reference_weights] + for (name, fsdp_weight), (_, reference_weight) in zip(fsdp_weights, reference_weights): + assert isinstance(fsdp_weight, DTensor), f"{name} was not FSDP2-sharded" + assert fsdp_weight.grad is not None, f"STE did not return a gradient for {name}" + assert isinstance(fsdp_weight.grad, DTensor), f"{name}.grad is not sharded" + assert reference_weight.grad is not None + grad = fsdp_weight.grad.full_tensor() + assert torch.isfinite(grad).all(), f"{name}.grad contains non-finite values" + torch.testing.assert_close( + grad, + reference_weight.grad, + rtol=rtol, + atol=atol, + msg=lambda msg: (f"{name}: FSDP2 QAT wgrad did not reach the master shard\n{msg}"), + ) + + +@pytest.mark.parametrize("op_kind", ("BasicLinear", "Linear")) +@pytest.mark.parametrize( + "override", + ( + pytest.param(None, id="none"), + pytest.param("dequantized", id="dequantized"), + pytest.param("high_precision", id="high_precision"), + ), +) +def test_fsdp2_linear_qat_backward_overrides(op_kind: str, override: str | None) -> None: + """BasicLinear/Linear preserve all pure-MXFP8 override modes under FSDP2.""" + _skip_if_unsupported() + device = _device() + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + op_cls = getattr(te_ops, op_kind) + kwargs = {"device": device, "dtype": _DTYPE} + if op_kind == "Linear": + kwargs["bias"] = False + return op_cls(in_features, out_features, **kwargs) + + qat_module = factory() + master_weights = _initialize_master_weights(qat_module) + projected_weights = _project_weights(master_weights) + _fsdp2_shard(qat_module) + + # Identity inputs make dgrad and wgrad comparisons strict and expose + # per-element differences between master and projected weights. + x = torch.eye(in_features, dtype=_DTYPE, device=device).requires_grad_(True) + dy = torch.eye(out_features, dtype=_DTYPE, device=device) + qat_recipe = MXFP4QATMXFP8BlockScaling(backward_override=override) + with te.autocast(enabled=True, recipe=qat_recipe): + y_qat = qat_module(x) + y_qat.backward(dy) + assert x.grad is not None + + projected_ref, y_projected, dx_projected = _run_reference( + factory, + projected_weights, + x, + dy, + (), + override=override, + ) + master_ref, y_master, dx_master = _run_reference( + factory, + master_weights, + x, + dy, + (), + override=override, + ) + + _assert_projected_forward(y_qat.detach(), y_projected, y_master) + expected_dx = dx_master if override == "high_precision" else dx_projected + torch.testing.assert_close(x.grad, expected_dx, rtol=0, atol=0) + expected_wgrad_ref = master_ref if override == "high_precision" else projected_ref + _assert_fsdp_weight_grads(qat_module, expected_wgrad_ref) + + +def test_fsdp2_grouped_linear_qat_none() -> None: + """Independent GroupedLinear weights support QAT mode None under FSDP2.""" + _skip_if_unsupported() + device = _device() + group_size = 2 + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + return te_ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=False, + device=device, + dtype=_DTYPE, + single_grouped_weight=False, + ) + + qat_module = factory() + master_weights = _initialize_master_weights(qat_module) + projected_weights = _project_weights(master_weights) + _fsdp2_shard(qat_module) + + split_sizes = torch.tensor([128, 128], dtype=torch.int32, device=device) + eye = torch.eye(in_features, dtype=_DTYPE, device=device) + x = torch.cat((eye, eye)).requires_grad_(True) + dy = torch.cat((eye, eye)) + qat_recipe = MXFP4QATMXFP8BlockScaling(backward_override=None) + with te.autocast(enabled=True, recipe=qat_recipe): + y_qat = qat_module(x, split_sizes) + y_qat.backward(dy) + assert x.grad is not None + + projected_ref, y_projected, dx_projected = _run_reference( + factory, + projected_weights, + x, + dy, + (split_sizes,), + override=None, + ) + _, y_master, _ = _run_reference( + factory, + master_weights, + x, + dy, + (split_sizes,), + override=None, + ) + + _assert_projected_forward(y_qat.detach(), y_projected, y_master) + torch.testing.assert_close(x.grad, dx_projected, rtol=0, atol=0) + _assert_fsdp_weight_grads(qat_module, projected_ref) + + +def test_fsdp2_fused_grouped_mlp_qat_none() -> None: + """Fused GroupedMLP mode None preserves QAT weights through FSDP2 backward.""" + _skip_if_unsupported() + fused_cls = te_ops.fused.GroupedMLP_CuTeGEMMGLU + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused GroupedMLP is not supported on this system") + + device = _device() + group_size = 2 + hidden_size = 128 + + def factory() -> torch.nn.Module: + fc1 = te_ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=_DTYPE, + single_grouped_weight=False, + ) + fc2 = te_ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=_DTYPE, + single_grouped_weight=False, + ) + return te_ops.Sequential( + fc1, + te_ops.ScaledSwiGLU(glu_interleave_size=32), + fc2, + ) + + qat_module = factory() + master_weights = _initialize_master_weights(qat_module) + projected_weights = _project_weights(master_weights) + _fsdp2_shard(qat_module) + + split_sizes = torch.tensor([128, 128], dtype=torch.int32, device=device) + num_tokens = int(split_sizes.sum().item()) + values = torch.linspace( + -0.25, + 0.25, + num_tokens * hidden_size, + dtype=torch.float32, + device=device, + ) + x = values.reshape(num_tokens, hidden_size).to(_DTYPE).requires_grad_(True) + probs = torch.linspace(0.5, 1.0, num_tokens, dtype=torch.float32, device=device).to(_DTYPE) + dy = values.flip(0).reshape(num_tokens, hidden_size).to(_DTYPE) + extra_args = (split_sizes, probs, split_sizes) + + qat_recipe = MXFP4QATMXFP8BlockScaling(backward_override=None) + with te.autocast(enabled=True, recipe=qat_recipe): + y_qat = qat_module(x, *extra_args) + forward_ops = qat_module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 and isinstance(forward_ops[0][0], fused_cls) + y_qat.backward(dy) + assert x.grad is not None + + projected_ref, y_projected, dx_projected = _run_reference( + factory, + projected_weights, + x, + dy, + extra_args, + override=None, + ) + _, y_master, _ = _run_reference( + factory, + master_weights, + x, + dy, + extra_args, + override=None, + ) + + _assert_projected_forward(y_qat.detach(), y_projected, y_master) + torch.testing.assert_close(x.grad, dx_projected, rtol=0, atol=0) + # NCCL reduction of nontrivial BF16 grouped-MLP wgrads need not be + # bitwise identical to the single-rank reference. + _assert_fsdp_weight_grads(qat_module, projected_ref, rtol=2e-2, atol=2e-3) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index f386659b6c..6f5dd14088 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -5,7 +5,6 @@ import os import sys import subprocess -import sys from pathlib import Path sys.path.append(str(Path(__file__).resolve().parent.parent)) @@ -101,6 +100,33 @@ def test_fsdp2_mem_leak_tests(): assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_mxfp4_qat_ops_tests(): + """MXFP4-QAT parity with pure-MXFP8 fusible-op support under FSDP2.""" + mxfp8_available, reason = te.is_mxfp8_available(return_reason=True) + if not mxfp8_available: + pytest.skip(reason) + test_path = _FSDP2_DIR / "run_fsdp2_mxfp4_qat_ops.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + valid_returncodes=(0,), + env=os.environ, + timeout=600, + ) + + @pytest.mark.skipif(NUM_PROCS < 4, reason="Requires 4+ GPUs for DP4→DP2 resharding test") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") @pytest.mark.parametrize("recipe", _parametrize_recipes()) diff --git a/tests/pytorch/references/mxfp4_qat_reference.py b/tests/pytorch/references/mxfp4_qat_reference.py new file mode 100644 index 0000000000..f8be5c3e67 --- /dev/null +++ b/tests/pytorch/references/mxfp4_qat_reference.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pure-PyTorch reference for MXFP4 weight fake-quantization. + +Bit-identical specification of transformer_engine.pytorch.mxfp4_qat +(tex.mxfp4_fake_quantize): per 1x32 block, scale = 2^clamp(ceil(log2(amax/6)), -126, +125) derived from the amax bit pattern, RTNE onto the E2M1 grid, saturate at 6, +multiply back. Non-finite values NaN-poison their whole block, -0 keeps its sign, and +the scale floor/cap keep every grid value representable in bf16/fp32. +""" +import torch + +_E2M1_MAX = 6.0 +_MXFP4_BLOCK = 32 + + +def round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" + fine = torch.round(y * 2.0) * 0.5 + mid = torch.round(y) + coarse = torch.round(y * 0.5) * 2.0 + return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) + + +def mxfp4_fake_quantize_reference(weight: torch.Tensor) -> torch.Tensor: + """Numerical oracle for the mxfp4_fake_quantize kernel (composite torch ops).""" + rows, cols = weight.shape + w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) + amax = w32.abs().amax(dim=-1, keepdim=True) + nonfinite = ~torch.isfinite(amax) + + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + exp = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)) + exp = exp.clamp(min=-126, max=125) + scale = torch.ldexp(torch.ones_like(amax), exp) + scale = torch.where(amax > 0, scale, torch.ones_like(scale)) + + y = (w32 / scale).abs().clamp(max=_E2M1_MAX) + q = round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) + q = torch.copysign(q, w32) + out = q * scale + out = torch.where(nonfinite.expand_as(out), torch.full_like(out, float("nan")), out) + return out.view(rows, cols).to(weight.dtype) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py new file mode 100644 index 0000000000..2710bab113 --- /dev/null +++ b/tests/pytorch/test_mxfp4_qat.py @@ -0,0 +1,981 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-step and end-to-end tests for the MXFP4 weight-QAT recipes.""" +import ctypes +import os + +for _n in ("libnvrtc.so.13", "libnvrtc.so.12", "libnvrtc.so"): + try: + ctypes.CDLL(_n, mode=ctypes.RTLD_GLOBAL) + break + except OSError: + continue + +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import ( + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8_autocast +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize +from references.mxfp4_qat_reference import mxfp4_fake_quantize_reference +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + +torch.manual_seed(1234) +DEV = "cuda" +RUN_PERF = os.environ.get("NVTE_MXFP4_QAT_PERF", "0") == "1" +E2M1_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=DEV) +SHAPES = [(128, 256), (192, 448), (512, 1024)] + +FP8_TOL_1OP = 0.045 +FP8_TOL_2OP = 0.065 +BF16_TOL = 0.005 + + +def make_weight(m, n, dtype=torch.bfloat16, zero_blocks=0, outliers=0): + w = (torch.randn(m, n, device=DEV, dtype=torch.float32) * 0.02).to(dtype) + if outliers: + idx = torch.randint(0, m * n, (outliers,), device=DEV) + w.view(-1)[idx] *= 500.0 + if zero_blocks: + for i in range(zero_blocks): + r = torch.randint(0, m, (1,)).item() + k = torch.randint(0, n // 32, (1,)).item() + w[r, k * 32 : (k + 1) * 32] = 0.0 + return w + + +def rel_err(got, ref): + got, ref = got.to(torch.float64), ref.to(torch.float64) + return ((got - ref).norm() / ref.norm().clamp(min=1e-30)).item() + + +def fake_quant_ref_fp64(w): + """Independent fp64 reference: nearest-with-ties-to-even onto the E2M1 grid.""" + m, n = w.shape + w64 = w.to(torch.float64).view(m, n // 32, 32) + amax = w64.abs().amax(dim=-1, keepdim=True) + bits = amax.to(torch.float32).view(torch.int32) + exp_field = bits >> 23 + mant = bits & 0x7FFFFF + exp = exp_field - 129 + (mant > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)).clamp(-126, 125) + exp = torch.where(amax > 0, exp, torch.zeros_like(exp)) + scale = torch.ldexp(torch.ones_like(amax), exp) + y = (w64 / scale).clamp(min=-6.0, max=6.0) + grid = E2M1_GRID.to(torch.float64) + cand = torch.cat([-grid.flip(0)[:-1], grid]) + d = (y.unsqueeze(-1) - cand).abs() + near = d.argmin(dim=-1) + even = torch.tensor([abs(v) in (0.0, 1.0, 2.0, 4.0) for v in cand.tolist()], device=w.device) + dmin = d.gather(-1, near.unsqueeze(-1)).squeeze(-1) + tie_even = d.eq(dmin.unsqueeze(-1)) & even + has_even_tie = tie_even.any(dim=-1) + near_even = torch.where(has_even_tie, tie_even.to(torch.uint8).argmax(dim=-1), near) + q = cand[near_even] + return (q * scale).view(m, n).to(torch.float64) + + +def test_fake_quant_grid_scale_rtne(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w_hat = mxfp4_fake_quantize(w) + assert w_hat.dtype == w.dtype + + wh32 = w_hat.to(torch.float32).view(m, n // 32, 32) + amax = w.to(torch.float32).abs().view(m, n // 32, 32).amax(dim=-1, keepdim=True) + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mant = bits & 0x7FFFFF + exp = exp_field - 129 + (mant > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)).clamp(-126, 125) + scale = torch.ldexp(torch.ones_like(amax), exp) + scale = torch.where(amax > 0, scale, torch.ones_like(scale)) + + vals = (wh32 / scale).abs() + assert torch.isin(vals, E2M1_GRID).all(), "values off the E2M1 grid" + nz = (amax.double() > 6.0 * 2.0**-126) & (amax.double() <= 6.0 * 2.0**125) + assert ((scale.double() * 6.0)[nz] >= amax.double()[nz]).all() + assert (((scale.double() / 2.0) * 6.0)[nz] < amax.double()[nz]).all() + ref = fake_quant_ref_fp64(w) + assert torch.equal(w_hat.to(torch.float64), ref), f"RTNE mismatch {m}x{n}" + assert torch.equal(mxfp4_fake_quantize(w_hat), w_hat), "not idempotent" + + +def test_D_lossless_in_bf16(): + for m, n in SHAPES: + w = make_weight(m, n, outliers=4) + w_hat_bf16 = mxfp4_fake_quantize(w) + w_hat_fp32 = mxfp4_fake_quantize(w.to(torch.float32)) + assert torch.equal( + w_hat_bf16.to(torch.float32), w_hat_fp32 + ), "MXFP4-grid weight is not exact in bf16" + + +def test_E_mxfp8_rowwise_lossless(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w[0, 0] = 2.0**-128 + w_hat = mxfp4_fake_quantize(w) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + wh32 = w_hat.to(torch.float32) + + data = q._rowwise_data.view(torch.uint8)[:m, :n] + codes = q._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) + raw_vals = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) + raw_scale = torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127) + raw_dq = (raw_vals * raw_scale.unsqueeze(-1)).view(m, n) + assert torch.equal( + raw_dq, wh32 + ), f"RAW MXFP8 encoding of the MXFP4-grid weight is not lossless {m}x{n}" + + assert torch.equal( + dq, wh32 + ), f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" + + +def test_dequantize_e8m0_extreme_codes(): + """Real MXFP8 software dequantize with planted UE8M0 codes 0/255.""" + w = make_weight(32, 64) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize( + mxfp4_fake_quantize(w) + ) + data = q._rowwise_data.view(torch.uint8) + codes = q._rowwise_scale_inv.view(torch.uint8) + data[0, :32] = 56 + codes[0, 0] = 0 + data[1, :32] = 56 + codes[1, 0] = 255 + dq = q.dequantize(dtype=torch.float32) + assert (dq[0, :32] == 2.0**-127).all(), "UE8M0 code 0 must dequantize to 2^-127" + assert torch.isnan(dq[1, :32]).all(), "UE8M0 code 255 must dequantize to NaN" + + +def test_G_blockwise_lossless_and_transpose(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4) + w_hat = mxfp4_fake_quantize(w) + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, + ).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + assert torch.equal( + dq, w_hat.to(torch.float32) + ), f"128x128 blockwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" + q.update_usage(rowwise_usage=False, columnwise_usage=True) + dqc = q.dequantize(dtype=torch.float32) + assert torch.equal( + dqc, w_hat.to(torch.float32) + ), f"columnwise 128x128 dequant differs from the MXFP4-grid weight {m}x{n}" + + +def test_lossless_dequant_to_bf16(): + """Both weight encodings dequantize to bf16 bit-exactly (bwd override target dtype).""" + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w_hat = mxfp4_fake_quantize(w) + + q8 = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) + dq8 = q8.dequantize(dtype=torch.bfloat16) + assert torch.equal(dq8, w_hat), f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" + + w2_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4)) + qb = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, + ).quantize(w2_hat) + assert torch.equal( + qb.dequantize(dtype=torch.bfloat16), w2_hat + ), f"blockwise 128x128 dequant to bf16 not lossless {m}x{n}" + + +def test_C_mxfp8_colwise_bound(): + for m, n in SHAPES: + w_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4, outliers=4)) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3).quantize(w_hat) + q.update_usage(rowwise_usage=False, columnwise_usage=True) + dq = q.dequantize(dtype=torch.float32) + ref = w_hat.to(torch.float32) + blk = ref.view(m // 32, 32, n) + bound = blk.abs().amax(dim=1, keepdim=True) * 2.0**-3 + err = (dq.view(m // 32, 32, n) - blk).abs() + assert (err <= bound + 1e-30).all(), f"colwise 32x1 requant error above bound {m}x{n}" + + +_INTREE_KERNELS = {} + + +def _intree_kernel(fast_math=False): + """JIT-compile the in-tree kernel; fast_math=True rebuilds it with --use_fast_math.""" + if fast_math not in _INTREE_KERNELS: + import os + from torch.utils.cpp_extension import load_inline + + src_dir = os.environ.get( + "NVTE_MXFP4_QAT_TEST_SRC", + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "transformer_engine", + "common", + ), + ) + cuda_src = r""" +#include +#include +#include +#include "cast/mxfp4/fake_quantize_mxfp4.cuh" + +torch::Tensor mxfp4_fake_quantize_intree(torch::Tensor w) { + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.numel() % 32 == 0); + auto out = torch::empty_like(w); + auto stream = at::cuda::getCurrentCUDAStream(); + if (w.scalar_type() == at::kBFloat16) { + transformer_engine::dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(w.data_ptr()), + reinterpret_cast(out.data_ptr()), w.numel(), stream); + } else if (w.scalar_type() == at::kFloat) { + transformer_engine::dispatch::mxfp4::fake_quantize_mxfp4_launch( + w.data_ptr(), out.data_ptr(), w.numel(), stream); + } else { + TORCH_CHECK(false, "bf16/fp32 only"); + } + return out; +} +""" + _INTREE_KERNELS[fast_math] = load_inline( + name="nvte_mxfp4_qat_intree_test" + ("_fastmath" if fast_math else ""), + cpp_sources="torch::Tensor mxfp4_fake_quantize_intree(torch::Tensor w);", + cuda_sources=cuda_src, + functions=["mxfp4_fake_quantize_intree"], + extra_include_paths=[src_dir], + extra_cuda_cflags=[ + "-O3", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ] + + (["--use_fast_math"] if fast_math else []), + verbose=False, + ) + return _INTREE_KERNELS[fast_math] + + +def _both_paths(w): + got_kernel = _intree_kernel().mxfp4_fake_quantize_intree(w.contiguous()) + got_torch = mxfp4_fake_quantize_reference(w) + return got_kernel, got_torch + + +def _assert_bits_equal(a, b, msg): + assert a.dtype == b.dtype, f"{msg}: dtype mismatch" + ib = torch.int16 if a.dtype == torch.bfloat16 else torch.int32 + assert torch.equal(a.view(ib), b.view(ib)), f"{msg}: raw bits differ" + + +def _assert_same_with_nan(a, b, msg): + an, bn = torch.isnan(a), torch.isnan(b) + assert torch.equal(an, bn), f"{msg}: NaN masks differ" + _assert_bits_equal( + torch.where(an, torch.zeros_like(a), a), + torch.where(bn, torch.zeros_like(b), b), + msg, + ) + + +PARITY_SHAPES = SHAPES + [(3, 416), (1, 32)] + + +def test_kernel_matches_torch_reference(): + for m, n in PARITY_SHAPES: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, f"kernel != torch reference {m}x{n} {dtype}") + + +def test_edge_value_domains(): + m, n = 64, 128 + w = make_weight(m, n) + + z = torch.zeros(m, n, device=DEV, dtype=torch.bfloat16) + gk, gt = _both_paths(z) + assert torch.equal(gk, z) and torch.equal(gt, z) + + wz = w.clone() + wz[0, :32] = -0.0 + gk, gt = _both_paths(wz) + assert (gk[0, :32] == 0).all() and torch.signbit(gk[0, :32]).all(), "-0.0 sign lost" + _assert_bits_equal(gk, gt, "-0 block") + + ws = w.clone() + ws[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + gk, gt = _both_paths(ws) + assert torch.isfinite(gk[1, :32]).all() and torch.equal(gk, gt) + ref = fake_quant_ref_fp64(ws) + assert torch.equal(gt.to(torch.float64), ref), "subnormal block deviates from fp64 ref" + + wp = w.clone() + wp[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + gk, gt = _both_paths(wp) + assert (gk[5, :32].to(torch.float32) == 2.0**-127).all(), "2^-127 grid point lost" + _assert_bits_equal(gk, gt, "2^-127 grid point") + ref = fake_quant_ref_fp64(wp) + assert torch.equal(gt.to(torch.float64), ref), "2^-127 block deviates from fp64 ref" + + wh = w.clone() + wh[2, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + wh[2, 0] = -3.0e38 + gk, gt = _both_paths(wh) + assert torch.isfinite(gk[2, :32]).all() and torch.equal(gk, gt) + ref = fake_quant_ref_fp64(wh) + assert torch.equal(gt.to(torch.float64), ref), "huge block deviates from fp64 ref" + + for bad in (float("inf"), float("-inf"), float("nan")): + wb = w.clone() + wb[3, 40] = bad + gk, gt = _both_paths(wb) + _assert_same_with_nan(gk, gt, f"bad={bad}") + assert torch.isnan(gk[3, 32:64]).all(), f"{bad}: block not NaN-poisoned" + assert torch.isfinite(gk[3, :32]).all(), f"{bad}: neighbor block affected" + clean_rows = torch.ones(m, dtype=torch.bool) + clean_rows[3] = False + assert torch.isfinite(gk[clean_rows]).all() + + try: + mxfp4_fake_quantize(torch.zeros(32, 32, device=DEV, dtype=torch.float16)) + raise SystemExit("fp16 should be rejected") + except ValueError: + pass + + +def test_blockwise_recipe_is_128x128(): + r = MXFP4QATFloat8BlockScaling() + assert r.x_block_scaling_dim == 1 and r.w_block_scaling_dim == 2 + assert r.grad_block_scaling_dim == 1 + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, + ) + assert q.block_len == 128 + w_hat = mxfp4_fake_quantize(make_weight(256, 512)) + qt = q.quantize(w_hat) + assert tuple(qt._rowwise_scale_inv.shape)[0] == 256 // 128 + + +def test_blockwise_over_ratio_is_bounded_not_lossless(): + w = make_weight(128, 128) + w[0, :32] = torch.tensor(2.0**-14, dtype=torch.bfloat16) + w[64, 0] = 100.0 + w_hat = mxfp4_fake_quantize(w) + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, + ).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + assert torch.isfinite(dq).all() + err = (dq - w_hat.to(torch.float32)).abs().max().item() + tile_amax = w_hat.to(torch.float32).abs().max().item() + assert err <= tile_amax * 2.0**-9, "over-ratio loss larger than FP8 headroom" + + +def test_kernel_perf(): + if not RUN_PERF: + print(" SKIP (set NVTE_MXFP4_QAT_PERF=1 to run benchmarks)") + return + rows, cols = 8192, 8192 + w = make_weight(rows, cols) + k = _intree_kernel() + for _ in range(3): + k.mxfp4_fake_quantize_intree(w) + mxfp4_fake_quantize_reference(w) + torch.cuda.synchronize() + + def timeit(fn, iters=20): + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters * 1e3 + + t_kernel = timeit(lambda: k.mxfp4_fake_quantize_intree(w)) + t_torch = timeit(lambda: mxfp4_fake_quantize_reference(w)) + moved = rows * cols * 2 * 2 + print( + f" kernel: {t_kernel:.1f} us ({moved / t_kernel / 1e3:.0f} GB/s), " + f"torch: {t_torch:.1f} us ({moved / t_torch / 1e3:.0f} GB/s)" + ) + assert t_kernel < t_torch, "kernel slower than the torch reference" + + +def test_kernel_fast_math_immune(): + """The kernel stays bit-identical when compiled with --use_fast_math.""" + k = _intree_kernel(fast_math=True) + for m, n in PARITY_SHAPES: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) + got = k.mxfp4_fake_quantize_intree(w.contiguous()) + ref = mxfp4_fake_quantize_reference(w) + _assert_bits_equal(got, ref, f"fast-math build diverged {m}x{n} {dtype}") + + w = make_weight(64, 128) + w[0, :32] = 0.0 + w[1, 5] = float("inf") + w[2, 9] = float("nan") + w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + w[4, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + w[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + got = k.mxfp4_fake_quantize_intree(w.contiguous()) + ref = mxfp4_fake_quantize_reference(w) + _assert_same_with_nan(got, ref, "fast-math edge") + assert ( + got[5, :32].to(torch.float32) == 2.0**-127 + ).all(), "fast-math build flushed the 2^-127 grid point" + + +def test_exp2f_e8m0_all_codes(): + """All 256 UE8M0 codes through exp2f/exp2f_rcp extracted verbatim from ptx.cuh.""" + import re as _re + from torch.utils.cpp_extension import load_inline + + src_dir = os.environ.get( + "NVTE_MXFP4_QAT_TEST_SRC", + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "..", "transformer_engine", "common" + ), + ) + ptx_src = open(os.path.join(src_dir, "util", "ptx.cuh")).read() + m_exp = _re.search( + r"__device__ __forceinline__ float exp2f\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S + ) + m_rcp = _re.search( + r"template <>\n__device__ __forceinline__ float exp2f_rcp\(e8m0_t biased_exp\)" + r" \{.*?\n\}", + ptx_src, + _re.S, + ) + assert m_exp and m_rcp, "could not extract exp2f/exp2f_rcp from ptx.cuh" + cuda_src = ( + "#include \n#include \n" + "#include \n" + "namespace teptx {\n" + "typedef unsigned char e8m0_t;\n" + "constexpr unsigned FP32_MANTISSA_BITS = 23;\n" + "template __device__ __forceinline__ T exp2f_rcp(e8m0_t);\n" + + m_rcp.group(0) + + "\n" + + m_exp.group(0) + + "\n}\n" + "__global__ void all_codes_kernel(float *e, float *r) {\n" + " const int b = threadIdx.x;\n" + " e[b] = teptx::exp2f(static_cast(b));\n" + " r[b] = teptx::exp2f_rcp(static_cast(b));\n" + "}\n" + "void run_all_codes(torch::Tensor e, torch::Tensor r) {\n" + " all_codes_kernel<<<1, 256, 0, at::cuda::getCurrentCUDAStream()>>>(\n" + " e.data_ptr(), r.data_ptr());\n" + "}\n" + ) + mod = load_inline( + name="nvte_mxfp4_qat_exp2f_test", + cpp_sources="void run_all_codes(torch::Tensor e, torch::Tensor r);", + cuda_sources=cuda_src, + functions=["run_all_codes"], + verbose=False, + ) + e = torch.empty(256, dtype=torch.float32, device=DEV) + r = torch.empty_like(e) + mod.run_all_codes(e, r) + torch.cuda.synchronize() + codes = torch.arange(256, dtype=torch.int32) + exp_ref = torch.ldexp(torch.ones(256), codes - 127) + rcp_ref = torch.ldexp(torch.ones(256), 127 - codes) + assert torch.isnan(e[255]) and torch.isnan(r[255]), "code 255 must be NaN" + assert e.view(torch.int32)[255].item() == 0x7FFFFFFF, "code 255 NaN payload" + assert r.view(torch.int32)[255].item() == 0x7FFFFFFF, "code 255 rcp NaN payload" + assert torch.equal(e[:255].cpu(), exp_ref[:255]), "exp2f wrong (code 0 == 2^-127?)" + assert torch.equal(r[:255].cpu(), rcp_ref[:255]), "exp2f_rcp wrong" + + +def test_ste_identity_gradient(): + """mxfp4_fake_quantize carries an identity (STE) gradient and matches the reference forward.""" + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(64, 128, dtype=dtype).requires_grad_(True) + up = torch.randn_like(w) + out = mxfp4_fake_quantize(w) + out.backward(up) + assert torch.equal(w.grad, up), f"STE gradient is not identity ({dtype})" + _assert_bits_equal(out.detach(), mxfp4_fake_quantize_reference(w.detach()), "STE fwd") + + +def test_misaligned_and_noncontiguous(): + """Misaligned storage-offset views and non-contiguous inputs match an aligned copy bitwise.""" + for dtype, off in ((torch.bfloat16, 4), (torch.float32, 2)): + m, n = 64, 128 + flat = (torch.randn(m * n + off, device=DEV, dtype=torch.float32) * 0.02).to(dtype) + v = flat[off:].view(m, n) + assert v.is_contiguous() and v.data_ptr() % 16 != 0 + got = mxfp4_fake_quantize(v) + ref = mxfp4_fake_quantize(v.clone()) + _assert_bits_equal(got, ref, f"misaligned view {dtype}") + + nc = (torch.randn(64, 256, device=DEV, dtype=torch.float32) * 0.02)[:, ::2] + assert not nc.is_contiguous() + got = mxfp4_fake_quantize_reference(nc) + assert torch.equal(got, mxfp4_fake_quantize_reference(nc.contiguous())) + + +def test_bf16_exhaustive_bit_patterns(): + """All 65536 bf16 bit patterns: kernel/torch parity, fp64 oracle on finite blocks.""" + bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) + w = bits.view(torch.bfloat16).view(2048, 32) + gk, gt = _both_paths(w) + _assert_same_with_nan(gk, gt, "bf16 exhaustive") + finite = torch.isfinite(w.to(torch.float32)).all(dim=-1) + ref = fake_quant_ref_fp64(w[finite]) + assert torch.equal(gt[finite].to(torch.float64), ref), "bf16 exhaustive vs fp64 ref" + + +def test_fp32_bit_fuzz(): + """Random raw fp32 bit patterns (subnormals/NaN/Inf included).""" + g = torch.Generator().manual_seed(99) + for _ in range(4): + bits = torch.randint(-(2**31), 2**31 - 1, (4096, 32), generator=g, dtype=torch.int64) + w = bits.to(torch.int32).cuda().view(torch.float32) + gk, gt = _both_paths(w) + _assert_same_with_nan(gk, gt, "fp32 fuzz") + finite = torch.isfinite(w).all(dim=-1) + if finite.any(): + ref = fake_quant_ref_fp64(w[finite]) + assert torch.equal(gt[finite].to(torch.float64), ref), "fp32 fuzz vs fp64 ref" + + +def test_scale_threshold_and_rtne_midpoints(): + """amax = 1.5*2^k ulp triples and all E2M1 RTNE midpoints vs explicit expected values.""" + rows = [] + for k in (-100, -10, 0, 42, 100): + a = torch.tensor(1.5 * 2.0**k, dtype=torch.float32) + ab = a.view(torch.int32) + for delta in (-1, 0, 1): + row = torch.zeros(32, dtype=torch.float32) + row[0] = (ab + delta).view(torch.float32) + rows.append(row) + w = torch.stack(rows).cuda() + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, "threshold parity") + assert torch.equal(gt.to(torch.float64), fake_quant_ref_fp64(w)), "threshold vs fp64 ref" + + mids = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]) + expected = torch.tensor([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0]) + for t in (-60, 0, 60): + sc = 2.0**t + block = torch.zeros(32) + block[:7] = mids * sc + block[7] = 6.0 * sc + w = block.view(1, 32).to(torch.float32).cuda() + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, f"midpoints t={t}") + exp_row = torch.zeros(32) + exp_row[:7] = expected * sc + exp_row[7] = 6.0 * sc + assert torch.equal( + gt[0].cpu().to(torch.float32), exp_row.to(torch.float32) + ), f"midpoint values t={t}" + + +def test_deployment_top_cap_domain(): + """amax <= 6*2^125 is deployment-bitwise; above it we saturate at the cap while + TileKernels moves to scale 2^126.""" + s125 = 6.0 * 2.0**125 + w = torch.zeros(2, 32, dtype=torch.float32) + w[0, 0] = s125 + just_above = (torch.tensor(s125, dtype=torch.float32).view(torch.int32) + 1).view(torch.float32) + w[1, 0] = just_above + wc = w.cuda() + gk, gt = _both_paths(wc) + _assert_bits_equal(gk, gt, "top-cap parity") + assert gk[0, 0].item() == s125, "boundary block must be exact" + assert gk[1, 0].item() == s125, "above-boundary block must satfinite at 6*2^125" + assert torch.equal(gt.to(torch.float64), fake_quant_ref_fp64(wc)), "top-cap vs fp64 ref" + + +def test_blockwise_feasibility_enumeration(): + """Per-element E4M3-lattice prediction == TE 128x128 quantizer for exponent spreads d=0..16.""" + payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + boundary = None + for d in range(17): + w = torch.zeros(128, 128, dtype=torch.float32) + w[0, 0] = 6.0 + vals = payloads * 2.0 ** (-d) + w[1, :7] = vals + w[1, 7:14] = -vals + wc = w.cuda() + assert torch.equal(mxfp4_fake_quantize(wc), wc), f"tile not on MXFP4 grid (d={d})" + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + force_pow_2_scales=True, + block_scaling_dim=2, + ).quantize(wc) + S = q._rowwise_scale_inv.reshape(-1)[0].item() + dq = q.dequantize(dtype=torch.float32) + pred = ((wc / S).to(torch.float8_e4m3fn).to(torch.float32) * S) == wc + got = dq == wc + assert torch.equal(pred, got), f"E4M3-lattice prediction mismatch at spread d={d}" + if boundary is None and not bool(pred.all()): + boundary = d + assert ( + boundary is not None and boundary >= 12 + ), f"full-payload exactness should hold at least through spread 2^11, got {boundary}" + print( + f" full-grid exact through spread d={boundary - 1}; first loss at d={boundary} " + "(E4M3 subnormal lattice bound)" + ) + + +def test_fp16_pipeline_rejected(): + """fp16 weights and fp16 activation dtypes must fail loudly under MXFP4 QAT.""" + mod = te.Linear(256, 128, bias=False, params_dtype=torch.float16, device=DEV) + import transformer_engine.pytorch.ops as te_ops + + ops_mod = te_ops.Linear(256, 128, bias=False, dtype=torch.float16, device=DEV) + x = torch.randn(64, 256, device=DEV, dtype=torch.float16) + for qat_mod in (mod, ops_mod): + for recipe in (MXFP4QATMXFP8BlockScaling(), MXFP4QATFloat8BlockScaling()): + try: + with fp8_autocast(enabled=True, fp8_recipe=recipe): + qat_mod(x) + raise SystemExit(f"fp16 pipeline must be rejected ({recipe.__class__.__name__})") + except (NotImplementedError, ValueError): + pass + + +def test_recipe_switch_invalidates_cache(): + """base<->QAT recipe switches must requantize instead of reusing cached weight workspaces.""" + import warnings as _warnings + from transformer_engine.common.recipe import MXFP8BlockScaling + + torch.manual_seed(3) + x = torch.randn(64, 256, device=DEV, dtype=torch.bfloat16) + mod = te.Linear(256, 128, bias=False, params_dtype=torch.bfloat16, device=DEV) + base, qat = MXFP8BlockScaling(), MXFP4QATMXFP8BlockScaling() + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + with fp8_autocast(enabled=True, fp8_recipe=base): + mod(x, is_first_microbatch=True) + with fp8_autocast(enabled=True, fp8_recipe=qat): + y_switch = mod(x, is_first_microbatch=False) + y_fresh = mod(x, is_first_microbatch=True) + assert torch.equal( + y_switch, y_fresh + ), "base->QAT switch reused an UNPROJECTED cached weight" + with fp8_autocast(enabled=True, fp8_recipe=base): + y2_switch = mod(x, is_first_microbatch=False) + y2_fresh = mod(x, is_first_microbatch=True) + assert torch.equal(y2_switch, y2_fresh), "QAT->base switch reused a PROJECTED cached weight" + + +def test_recipe_field_controls_qat(): + """mxfp4_qat_weights field (env NVTE_MXFP4_QAT) turns base recipes into QAT recipes.""" + import warnings as _warnings + from transformer_engine.common.recipe import Float8BlockScaling, MXFP8BlockScaling + + base = MXFP8BlockScaling(mxfp4_qat_weights=True) + assert base.mxfp4_qat() and MXFP4QATMXFP8BlockScaling().mxfp4_qat() + assert not MXFP8BlockScaling().mxfp4_qat() + + torch.manual_seed(9) + x = torch.randn(64, 256, device=DEV, dtype=torch.bfloat16) + mod = te.Linear(256, 128, bias=False, params_dtype=torch.bfloat16, device=DEV) + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + with fp8_autocast(enabled=True, fp8_recipe=base): + y_field = mod(x, is_first_microbatch=True) + with fp8_autocast(enabled=True, fp8_recipe=MXFP4QATMXFP8BlockScaling()): + y_class = mod(x, is_first_microbatch=True) + assert torch.equal(y_field, y_class), "field-driven QAT differs from class-driven QAT" + + assert Float8BlockScaling(mxfp4_qat_weights=True).mxfp4_qat() + blockwise_1d = Float8BlockScaling( + mxfp4_qat_weights=True, + w_block_scaling_dim=1, + ) + assert blockwise_1d.mxfp4_qat() + assert blockwise_1d.w_block_scaling_dim == 1 + + +def test_ops_linear_backward_overrides(): + """te.ops linear paths apply weight QAT and preserve each backward policy.""" + import transformer_engine.pytorch.ops as te_ops + + for op_kind in ("basic", "linear"): + for recipe_cls in (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling): + for override in (None, "dequantized", "high_precision"): + torch.manual_seed(13) + recipe = recipe_cls(backward_override=override) + if op_kind == "basic": + op = te_ops.BasicLinear( + 256, + 128, + device=DEV, + dtype=torch.bfloat16, + ) + else: + # Exercise the BasicLinear+Bias forward fusion used by te.ops.Linear. + op = te_ops.Linear( + 256, + 128, + bias=True, + device=DEV, + dtype=torch.bfloat16, + ) + with torch.no_grad(): + op.bias.zero_() + + with torch.no_grad(): + op.weight.copy_( + (torch.randn_like(op.weight, dtype=torch.float32) * 0.02).to( + op.weight.dtype + ) + ) + w_master = op.weight.detach().to(torch.float32) + w_hat = mxfp4_fake_quantize(op.weight.detach()).to(torch.float32) + + x = (torch.randn(128, 256, device=DEV, dtype=torch.float32) * 0.5).to( + torch.bfloat16 + ) + x.requires_grad_(True) + with fp8_autocast(enabled=True, fp8_recipe=recipe): + out = op(x) + + ref_out_hat = x.detach().to(torch.float32) @ w_hat.t() + ref_out_master = x.detach().to(torch.float32) @ w_master.t() + e_hat = rel_err(out, ref_out_hat) + e_master = rel_err(out, ref_out_master) + assert ( + e_hat < FP8_TOL_1OP + ), f"{op_kind} {recipe_cls.__name__} {override}: fwd err {e_hat}" + assert e_hat < e_master, ( + f"{op_kind} {recipe_cls.__name__} {override}: " + "forward did not use the MXFP4-projected weight" + ) + + grad = (torch.randn_like(out, dtype=torch.float32) * 0.1).to(torch.bfloat16) + out.backward(grad) + + ref_dx_hat = grad.to(torch.float32) @ w_hat + ref_dx_master = grad.to(torch.float32) @ w_master + e_hat = rel_err(x.grad, ref_dx_hat) + e_master = rel_err(x.grad, ref_dx_master) + if override == "high_precision": + assert ( + e_master < BF16_TOL + ), f"{op_kind} {recipe_cls.__name__}: high-precision dgrad err {e_master}" + assert e_master < e_hat, "high-precision dgrad did not use the master weight" + elif override == "dequantized": + assert ( + e_hat < BF16_TOL + ), f"{op_kind} {recipe_cls.__name__}: dequantized dgrad err {e_hat}" + assert e_hat < e_master, "dequantized dgrad did not use the QAT weight" + else: + assert ( + e_hat < FP8_TOL_1OP + ), f"{op_kind} {recipe_cls.__name__}: quantized dgrad err {e_hat}" + assert e_hat < e_master, "quantized dgrad did not use the QAT weight" + + assert op.weight.grad is not None, "STE did not return wgrad to the master weight" + ref_dw = grad.to(torch.float32).t() @ x.detach().to(torch.float32) + wgrad_tol = BF16_TOL if override == "high_precision" else FP8_TOL_2OP + assert rel_err(op.weight.grad, ref_dw) < wgrad_tol + + +def _run_module(module_kind, recipe, override, fuse): + torch.manual_seed(7) + m_splits = [192, 320] + tokens, in_f, out_f = sum(m_splits), 256, 512 + recipe.backward_override = override + + if module_kind == "linear": + mod = te.Linear( + in_f, + out_f, + bias=False, + params_dtype=torch.bfloat16, + device=DEV, + fuse_wgrad_accumulation=fuse, + ) + params = [mod.weight] + else: + mod = te.GroupedLinear( + 2, + in_f, + out_f, + bias=False, + params_dtype=torch.bfloat16, + device=DEV, + fuse_wgrad_accumulation=fuse, + ) + params = [mod.weight0, mod.weight1] + with torch.no_grad(): + for p in params: + p.copy_((torch.randn_like(p, dtype=torch.float32) * 0.02).to(p.dtype)) + w_hats = [mxfp4_fake_quantize(p.detach()) for p in params] + if fuse: + for p in params: + p.main_grad = torch.zeros_like(p, dtype=torch.float32) + p.grad_added_to_main_grad = False + + x = (torch.randn(tokens, in_f, device=DEV, dtype=torch.float32) * 0.5).to(torch.bfloat16) + x.requires_grad_(True) + with fp8_autocast(enabled=True, fp8_recipe=recipe): + if module_kind == "linear": + out = mod(x) + else: + out = mod(x, m_splits) + + offs = [0, m_splits[0], tokens] if module_kind == "grouped" else [0, tokens] + nsplit = len(params) + + def cat_ref(weights, mat, transpose): + return torch.cat( + [ + mat[offs[i] : offs[i + 1]].to(torch.float32) + @ ( + weights[i].detach().to(torch.float32).t() + if transpose + else weights[i].detach().to(torch.float32) + ) + for i in range(nsplit) + ] + ) + + ref_out_hat = cat_ref(w_hats, x, True) + ref_out_master = cat_ref(params, x, True) + e_hat, e_master = rel_err(out, ref_out_hat), rel_err(out, ref_out_master) + assert e_hat < FP8_TOL_1OP, f"fwd err {e_hat}" + assert e_hat < e_master, "fwd closer to master than to the MXFP4-grid weight" + + g = (torch.randn_like(out, dtype=torch.float32) * 0.1).to(out.dtype) + out.backward(g) + + ref_dx_hat = cat_ref(w_hats, g, False) + ref_dx_master = cat_ref(params, g, False) + e_h, e_m = rel_err(x.grad, ref_dx_hat), rel_err(x.grad, ref_dx_master) + if override == "high_precision": + assert e_m < BF16_TOL, f"hp dgrad err vs master {e_m}" + assert e_m < e_h, "hp dgrad must use the unquantized master (B)" + elif override == "dequantized": + assert e_h < BF16_TOL, f"dequantized dgrad err vs w_hat {e_h}" + assert e_h < e_m, "dequantized dgrad must use the MXFP4-grid weight (A)" + else: + assert e_h < FP8_TOL_1OP, f"mode-None dgrad err {e_h}" + assert e_h < e_m, "mode-None dgrad must use the MXFP4-grid weight (C)" + + if module_kind == "linear": + with fp8_autocast(enabled=True, fp8_recipe=recipe): + out_c1 = mod(x.detach(), is_first_microbatch=True) + out_c2 = mod(x.detach(), is_first_microbatch=False) + assert rel_err(out_c2, ref_out_hat) < FP8_TOL_1OP, "cached-weight fwd drifted" + assert rel_err(out_c1, out_c2) < 1e-6, "workspace update changed the weights" + + wgrad_tol = BF16_TOL if override == "high_precision" else FP8_TOL_2OP + for i, p in enumerate(params): + ref_wg = g[offs[i] : offs[i + 1]].to(torch.float32).t() @ x[offs[i] : offs[i + 1]].to( + torch.float32 + ) + if fuse: + assert p.grad_added_to_main_grad is True, "fused-wgrad flag not set (G)" + e = rel_err(p.main_grad, ref_wg) + assert e < wgrad_tol, f"main_grad err {e} (G)" + else: + assert p.grad is not None, "STE did not deliver a weight grad (G)" + e = rel_err(p.grad, ref_wg) + assert e < wgrad_tol, f"wgrad err {e} (G)" + + +def test_e2e_matrix(): + for module_kind in ("linear", "grouped"): + for recipe_cls in (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling): + for override in (None, "dequantized", "high_precision"): + for fuse in (False, True): + _run_module(module_kind, recipe_cls(), override, fuse) + print( + f" e2e OK {module_kind} {recipe_cls.__name__} " + f"override={override} fuse={fuse}" + ) + + +TESTS = [ + test_fake_quant_grid_scale_rtne, + test_D_lossless_in_bf16, + test_E_mxfp8_rowwise_lossless, + test_dequantize_e8m0_extreme_codes, + test_G_blockwise_lossless_and_transpose, + test_C_mxfp8_colwise_bound, + test_lossless_dequant_to_bf16, + test_kernel_matches_torch_reference, + test_edge_value_domains, + test_blockwise_recipe_is_128x128, + test_blockwise_over_ratio_is_bounded_not_lossless, + test_exp2f_e8m0_all_codes, + test_ste_identity_gradient, + test_misaligned_and_noncontiguous, + test_bf16_exhaustive_bit_patterns, + test_fp32_bit_fuzz, + test_scale_threshold_and_rtne_midpoints, + test_deployment_top_cap_domain, + test_blockwise_feasibility_enumeration, + test_recipe_switch_invalidates_cache, + test_recipe_field_controls_qat, + test_ops_linear_backward_overrides, + test_fp16_pipeline_rejected, + test_kernel_perf, + test_kernel_fast_math_immune, + test_e2e_matrix, +] + +if __name__ == "__main__": + import sys + import traceback + + failed = 0 + for t in TESTS: + try: + t() + torch.cuda.synchronize() + print(f"PASS {t.__name__}") + except Exception: + failed += 1 + print(f"FAIL {t.__name__}") + traceback.print_exc() + print(f"{len(TESTS) - failed}/{len(TESTS)} passed") + sys.exit(1 if failed else 0) diff --git a/tests/pytorch/test_ops_grouped_linear_mxfp4_qat.py b/tests/pytorch/test_ops_grouped_linear_mxfp4_qat.py new file mode 100644 index 0000000000..9aeaee0546 --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_mxfp4_qat.py @@ -0,0 +1,167 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP4-QAT coverage for the fusible ops.GroupedLinear.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops.basic.grouped_linear as grouped_linear_impl +from transformer_engine.common.recipe import ( + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8_autocast +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize + + +def _skip_if_recipe_unavailable(recipe) -> None: + if recipe.mxfp8(): + available, reason = te.is_mxfp8_available(return_reason=True) + else: + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not available: + pytest.skip(reason) + + +def _relative_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + actual = actual.to(torch.float64) + expected = expected.to(torch.float64) + return ((actual - expected).norm() / expected.norm().clamp(min=1e-30)).item() + + +@pytest.mark.parametrize( + "recipe_cls", + (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling), +) +@pytest.mark.parametrize("path", ("graph_safe", "legacy")) +def test_ops_grouped_linear_qat_none_forward_saved_weight_and_ste( + recipe_cls, + path, + monkeypatch, +): + """Project each expert once, save host-FP8 for dgrad, and return wgrad to the master.""" + recipe = recipe_cls(backward_override=None) + _skip_if_recipe_unavailable(recipe) + if path == "legacy": + monkeypatch.setattr( + grouped_linear_impl.GroupedLinear, + "_is_graph_safe_path_supported", + staticmethod(lambda **kwargs: False), + ) + else: + original_path_check = grouped_linear_impl.GroupedLinear._is_graph_safe_path_supported + + def require_graph_safe_path(**kwargs): + supported = original_path_check(**kwargs) + if not supported: + pytest.skip("graph-safe grouped-tensor path is unavailable") + return True + + monkeypatch.setattr( + grouped_linear_impl.GroupedLinear, + "_is_graph_safe_path_supported", + staticmethod(require_graph_safe_path), + ) + + projection_calls = 0 + + def tracked_fake_quantize(weight): + nonlocal projection_calls + projection_calls += 1 + return mxfp4_fake_quantize(weight) + + monkeypatch.setattr(grouped_linear_impl, "mxfp4_fake_quantize", tracked_fake_quantize) + + torch.manual_seed(1234) + num_groups, in_features, out_features = 2, 256, 128 + split_sizes = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + op = te.ops.GroupedLinear( + num_groups, + in_features, + out_features, + bias=False, + device="cuda", + dtype=torch.bfloat16, + ) + with torch.no_grad(): + for idx in range(num_groups): + weight = getattr(op, f"weight{idx}") + weight.copy_((0.05 * torch.randn_like(weight, dtype=torch.float32)).to(weight.dtype)) + projected_weights = [ + mxfp4_fake_quantize(getattr(op, f"weight{idx}").detach()).to(torch.float32) + for idx in range(num_groups) + ] + + x = torch.randn(256, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) + grad_output = torch.randn(256, out_features, device="cuda", dtype=torch.bfloat16) + with fp8_autocast(enabled=True, fp8_recipe=recipe): + output = op(x, split_sizes) + output.backward(grad_output) + + # The forward host-QTensor is saved for mode None, so backward must not replay QAT. + assert projection_calls == num_groups + + x_splits = torch.split(x.detach().to(torch.float32), split_sizes.tolist()) + dy_splits = torch.split(grad_output.to(torch.float32), split_sizes.tolist()) + expected_output = torch.cat( + [inp @ weight.t() for inp, weight in zip(x_splits, projected_weights)] + ) + expected_dgrad = torch.cat([dy @ weight for dy, weight in zip(dy_splits, projected_weights)]) + assert _relative_error(output, expected_output) < 0.05 + assert _relative_error(x.grad, expected_dgrad) < 0.05 + + for idx, (inp, dy) in enumerate(zip(x_splits, dy_splits)): + weight = getattr(op, f"weight{idx}") + assert weight.grad is not None + expected_wgrad = dy.t() @ inp + assert _relative_error(weight.grad, expected_wgrad) < 0.07 + + +@pytest.mark.parametrize( + "recipe_cls", + (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling), +) +@pytest.mark.parametrize("backward_override", ("dequantized", "high_precision")) +def test_ops_grouped_linear_qat_rejects_unsupported_backward_overrides( + recipe_cls, + backward_override, +): + """QAT support must not exceed the pure GroupedLinear override range.""" + recipe = recipe_cls(backward_override=backward_override) + _skip_if_recipe_unavailable(recipe) + op = te.ops.GroupedLinear( + 2, + 256, + 128, + bias=False, + device="cuda", + dtype=torch.bfloat16, + ) + x = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + split_sizes = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + with pytest.raises(NotImplementedError, match="only supports backward_override=None"): + with fp8_autocast(enabled=True, fp8_recipe=recipe): + op(x, split_sizes) + + +@pytest.mark.parametrize( + "recipe_cls", + (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling), +) +def test_ops_grouped_linear_qat_rejects_primary_quantized_weights(recipe_cls): + """A packed primary weight cannot provide the master required by QAT.""" + recipe = recipe_cls(backward_override=None) + _skip_if_recipe_unavailable(recipe) + with pytest.raises(NotImplementedError, match="high-precision master is required"): + with te.quantized_model_init(enabled=True, recipe=recipe): + te.ops.GroupedLinear( + 2, + 256, + 128, + bias=False, + device="cuda", + dtype=torch.bfloat16, + ) diff --git a/tests/pytorch/test_ops_grouped_mlp_mxfp4_qat.py b/tests/pytorch/test_ops_grouped_mlp_mxfp4_qat.py new file mode 100644 index 0000000000..69b4b38254 --- /dev/null +++ b/tests/pytorch/test_ops_grouped_mlp_mxfp4_qat.py @@ -0,0 +1,143 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP4-QAT coverage for the fused ops grouped-MLP path.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_impl +from transformer_engine.common.recipe import ( + MXFP4QATMXFP8BlockScaling, + MXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8_autocast +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize + + +def _require_fused_mxfp8_grouped_mlp() -> None: + available, reason = te.is_mxfp8_available(return_reason=True) + if not available: + pytest.skip(reason) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("CuTe DSL fused MXFP8 grouped-MLP is not available") + + +def _make_grouped_mlp(): + fc1 = te.ops.GroupedLinear( + 2, + 128, + 256, + bias=False, + device="cuda", + dtype=torch.bfloat16, + single_grouped_weight=False, + ) + activation = te.ops.ScaledSwiGLU(glu_interleave_size=32) + fc2 = te.ops.GroupedLinear( + 2, + 128, + 128, + bias=False, + device="cuda", + dtype=torch.bfloat16, + single_grouped_weight=False, + ) + return te.ops.Sequential(fc1, activation, fc2), fc1, fc2 + + +def _relative_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + actual = actual.to(torch.float64) + expected = expected.to(torch.float64) + return ((actual - expected).norm() / expected.norm().clamp(min=1e-30)).item() + + +def test_fused_grouped_mlp_qat_none_matches_projected_pure_mxfp8(monkeypatch): + """QAT None projects once in forward, saves host FP8, and returns master wgrads.""" + _require_fused_mxfp8_grouped_mlp() + + torch.manual_seed(1234) + qat_module, qat_fc1, qat_fc2 = _make_grouped_mlp() + ref_module, ref_fc1, ref_fc2 = _make_grouped_mlp() + + qat_weights = [] + ref_weights = [] + with torch.no_grad(): + for qat_fc, ref_fc in ((qat_fc1, ref_fc1), (qat_fc2, ref_fc2)): + for group_idx in range(2): + qat_weight = getattr(qat_fc, f"weight{group_idx}") + ref_weight = getattr(ref_fc, f"weight{group_idx}") + qat_weight.copy_( + (0.05 * torch.randn_like(qat_weight, dtype=torch.float32)).to(qat_weight.dtype) + ) + ref_weight.copy_(mxfp4_fake_quantize(qat_weight.detach())) + qat_weights.append(qat_weight) + ref_weights.append(ref_weight) + + projection_calls = 0 + real_fake_quantize = grouped_mlp_impl.mxfp4_fake_quantize + + def tracked_fake_quantize(weight): + nonlocal projection_calls + projection_calls += 1 + return real_fake_quantize(weight) + + monkeypatch.setattr(grouped_mlp_impl, "mxfp4_fake_quantize", tracked_fake_quantize) + + split_sizes = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + x_qat = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + x_ref = x_qat.detach().clone().requires_grad_(True) + probs_qat = torch.rand(256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + probs_ref = probs_qat.detach().clone().requires_grad_(True) + + with fp8_autocast( + enabled=True, + fp8_recipe=MXFP4QATMXFP8BlockScaling(backward_override=None), + ): + out_qat = qat_module(x_qat, split_sizes, probs_qat, split_sizes) + with fp8_autocast( + enabled=True, + fp8_recipe=MXFP8BlockScaling(backward_override=None), + ): + out_ref = ref_module(x_ref, split_sizes, probs_ref, split_sizes) + + qat_fused_op = qat_module._module_groups[0]._forward_ops[0][0] + ref_fused_op = ref_module._module_groups[0]._forward_ops[0][0] + assert isinstance(qat_fused_op, te.ops.fused.GroupedMLP_CuTeGEMMGLU) + assert isinstance(ref_fused_op, te.ops.fused.GroupedMLP_CuTeGEMMGLU) + assert projection_calls == 4 + assert _relative_error(out_qat, out_ref) < 1e-6 + + grad_output = torch.randn_like(out_qat) + out_qat.backward(grad_output) + out_ref.backward(grad_output) + + # The fused None path must reuse the forward host-QTensor, not replay QAT. + assert projection_calls == 4 + assert _relative_error(x_qat.grad, x_ref.grad) < 1e-6 + assert _relative_error(probs_qat.grad, probs_ref.grad) < 1e-6 + for qat_weight, ref_weight in zip(qat_weights, ref_weights): + assert qat_weight.grad is not None + assert ref_weight.grad is not None + assert _relative_error(qat_weight.grad, ref_weight.grad) < 1e-6 + + +@pytest.mark.parametrize("backward_override", ("dequantized", "high_precision")) +def test_fused_grouped_mlp_qat_rejects_unsupported_overrides(backward_override): + """Do not extend QAT beyond the pure fused grouped-MLP override range.""" + _require_fused_mxfp8_grouped_mlp() + module, _, _ = _make_grouped_mlp() + split_sizes = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + probs = torch.rand(256, device="cuda", dtype=torch.bfloat16) + + with pytest.raises(NotImplementedError, match="fused grouped-MLP ops support only"): + with fp8_autocast( + enabled=True, + fp8_recipe=MXFP4QATMXFP8BlockScaling( + backward_override=backward_override, + ), + ): + module(x, split_sizes, probs, split_sizes) diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 1e3c04573b..bf44e3bffb 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -16,6 +16,7 @@ #include "../utils.cuh" #include "dispatch/dequantize.cuh" #include "dispatch/quantize.cuh" +#include "mxfp4/fake_quantize_mxfp4.cuh" #include "transformer_engine/transpose.h" void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { @@ -47,6 +48,38 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, dispatch::quantize_fwd_helper(input, output, quant_config, stream); } +void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp4_fake_quantize); + using namespace transformer_engine; + const Tensor &input_tensor = *convertNVTETensorCheck(input); + Tensor &output_tensor = *convertNVTETensorCheck(output); + NVTE_CHECK(input_tensor.dtype() == DType::kBFloat16 || input_tensor.dtype() == DType::kFloat32, + "MXFP4 fake-quantize supports BF16/FP32 inputs, got ", + to_string(input_tensor.dtype())); + NVTE_CHECK(output_tensor.dtype() == input_tensor.dtype(), + "MXFP4 fake-quantize output dtype must match the input dtype."); + NVTE_CHECK(input_tensor.data.shape == output_tensor.data.shape, + "MXFP4 fake-quantize output shape must match the input shape."); + const size_t numel = input_tensor.numel(); + NVTE_CHECK(numel % 32 == 0, "MXFP4 fake-quantize needs the element count divisible by 32."); + NVTE_CHECK(!input_tensor.data.shape.empty() && input_tensor.data.shape.back() % 32 == 0, + "MXFP4 fake-quantize needs the innermost dimension divisible by 32 " + "(1x32 blocks must not cross rows)."); + if (numel == 0) { + return; + } + if (input_tensor.dtype() == DType::kBFloat16) { + dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(output_tensor.data.dptr), numel, stream); + } else { + dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(output_tensor.data.dptr), numel, stream); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dequantize); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh new file mode 100644 index 0000000000..6b4c437221 --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh @@ -0,0 +1,127 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ +#define TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ + +#include +#include + +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp4 { +namespace fake_quantize_kernel { + +constexpr size_t MXFP4_BLOCK_SIZE = 32; +constexpr size_t THREADS_PER_CHUNK = 256; +constexpr float E2M1_MAX = 6.0f; +constexpr int SCALE_EXP_MIN = -126; +constexpr int SCALE_EXP_MAX = 125; + +__device__ __forceinline__ float round_e2m1_magnitude(const float y) { + if (y <= 2.0f) { + return rintf(y * 2.0f) * 0.5f; + } + if (y <= 4.0f) { + return rintf(y); + } + return rintf(y * 0.5f) * 2.0f; +} + +__device__ __forceinline__ int block_scale_exponent_from_bits(const int bits) { + const int exp_field = bits >> 23; + const int mantissa = bits & 0x7FFFFF; + int e = exp_field - 129 + (mantissa > 0x400000 ? 1 : 0); + if (exp_field == 0) { + e = SCALE_EXP_MIN; + } + return max(SCALE_EXP_MIN, min(SCALE_EXP_MAX, e)); +} + +__device__ __forceinline__ float exp2_from_int(const int e) { + return __int_as_float((e + 127) << 23); +} + +__device__ __forceinline__ float mul_rn_no_ftz(const float a, const float b) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + fake_quantize_mxfp4_kernel(const IType *__restrict__ input, IType *__restrict__ output, + const size_t num_vectors) { + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + constexpr int LANES_PER_BLOCK = MXFP4_BLOCK_SIZE / VEC_ELTS; + + const size_t stride = static_cast(gridDim.x) * blockDim.x; + for (size_t vec_idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + vec_idx < num_vectors; vec_idx += stride) { + Vec in_vec; + in_vec.load_from(input, vec_idx); + + float elts[VEC_ELTS]; + int amax_bits = 0; +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + elts[i] = static_cast(in_vec.data.elt[i]); + amax_bits = max(amax_bits, __float_as_int(elts[i]) & 0x7fffffff); + } + const unsigned lane = threadIdx.x & 31u; + const unsigned group_mask = (LANES_PER_BLOCK == 32) ? 0xffffffffu + : (((1u << LANES_PER_BLOCK) - 1u) + << (lane & ~(LANES_PER_BLOCK - 1u))); +#pragma unroll + for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) { + amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset)); + } + const int nonfinite = amax_bits >= 0x7f800000; + + Vec out_vec; + if (nonfinite) { + const float qnan = __int_as_float(0x7fc00000); +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + out_vec.data.elt[i] = static_cast(qnan); + } + } else { + const int e = (amax_bits > 0) ? block_scale_exponent_from_bits(amax_bits) : 0; + const float scale = exp2_from_int(e); + const float inv_scale = exp2_from_int(-e); +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + const float av = __int_as_float(__float_as_int(elts[i]) & 0x7fffffff); + const float y = fminf(mul_rn_no_ftz(av, inv_scale), E2M1_MAX); + const float q = copysignf(round_e2m1_magnitude(y), elts[i]); + out_vec.data.elt[i] = static_cast(mul_rn_no_ftz(q, scale)); + } + } + out_vec.store_to(output, vec_idx); + } +} + +} // namespace fake_quantize_kernel + +template +void fake_quantize_mxfp4_launch(const IType *input, IType *output, const size_t numel, + cudaStream_t stream) { + using namespace fake_quantize_kernel; + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + const size_t num_vectors = numel / VEC_ELTS; + const size_t blocks_needed = (num_vectors + THREADS_PER_CHUNK - 1) / THREADS_PER_CHUNK; + const int grid = static_cast(blocks_needed < 65535 ? blocks_needed : 65535); + fake_quantize_mxfp4_kernel + <<>>(input, output, num_vectors); +} + +} // namespace mxfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..e19f40d3fd 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -418,6 +418,15 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, */ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Project a BF16/FP32 tensor onto the MXFP4 (E2M1) grid with 1x32 + * power-of-two block scales, writing back dequantized values. + * + * \param[in] input Input tensor (BF16/FP32, innermost dim divisible by 32). + * \param[in,out] output Output tensor with the same dtype and shape. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); + /*! \brief Casts input grouped tensor from reduced to higher precision. * In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise * data of the output tensor, regardless of whether the row- or columnwise scaling is used. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8c209ace15..576ba451fe 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -8,10 +8,8 @@ import os from enum import Enum from typing import Any, Literal, Optional, Union, Callable, NamedTuple -from dataclasses import field from pydantic.dataclasses import dataclass - _BACKWARD_OVERRIDES = (None, "high_precision", "dequantized") _NVFP4_4OVER6_SCOPES = ("none", "weights", "activations", "all") _NVFP4_4OVER6_ERR_MODES = ("MAE", "MSE") @@ -168,6 +166,12 @@ def custom(cls): """Whether the given recipe is custom.""" return issubclass(cls, CustomRecipe) + def mxfp4_qat(self): + """Whether the recipe applies MXFP4 weight QAT on top of its base recipe.""" + return bool(getattr(self, "mxfp4_qat_weights", False)) or isinstance( + self, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling) + ) + @dataclass(repr=False) class DelayedScaling(Recipe): @@ -368,6 +372,7 @@ class MXFP8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -437,6 +442,7 @@ class Float8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" def __post_init__(self) -> None: assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" @@ -461,6 +467,12 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + if self.mxfp4_qat_weights: + if not self.fp8_quant_fwd_weight.power_2_scale: + raise ValueError( + "MXFP4 QAT requires power-of-two weight scales; " + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is incompatible." + ) def _make_repr(self) -> str: return ( @@ -481,6 +493,34 @@ def _make_repr(self) -> str: ) +@dataclass(repr=False) +class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): + """ + MXFP8 recipe with MXFP4 weight quantization-aware training. + + Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid before + the regular MXFP8 weight quantization, whose rowwise encoding of the projected + weight is lossless. Everything else follows the base recipe. bf16/fp32 only. + """ + + mxfp4_qat_weights: bool = True + + +@dataclass(repr=False) +class MXFP4QATFloat8BlockScaling(Float8BlockScaling): + """ + Float8 block-scaling recipe with MXFP4 weight quantization-aware training. + + Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid before + regular blockwise FP8 weight quantization, lossless within each block's FP8 + dynamic-range headroom. Weight blocks stay 128x128 by default; set + ``w_block_scaling_dim=1`` for 1x128 rowwise and 128x1 columnwise blocks. + Everything else follows the base recipe. bf16/fp32 only. + """ + + mxfp4_qat_weights: bool = True + + @dataclass(repr=False) class NVFP4BlockScaling(Recipe): """ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..94a82990e3 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -380,6 +380,8 @@ __device__ __forceinline__ bf16 exp2f_rcp(e8m0_t biased_exp) { } __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { + if (biased_exp == 0) return __int_as_float(0x00400000); + if (biased_exp == 255) return __int_as_float(0x7fffffff); return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..2a77895ead 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -347,6 +347,8 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz py::object dequantize(const py::handle &input, DType otype); +at::Tensor mxfp4_fake_quantize(const at::Tensor &input); + py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..41dcd03e7f 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -556,6 +556,24 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, py::cast(std::move(dbias_torch))); } +at::Tensor mxfp4_fake_quantize(const at::Tensor &input) { + TORCH_CHECK(input.is_cuda(), "mxfp4_fake_quantize: input must be a CUDA tensor."); + TORCH_CHECK(input.dim() == 2, "mxfp4_fake_quantize: input must be 2D, got dim=", input.dim()); + TORCH_CHECK(input.size(1) % 32 == 0, + "mxfp4_fake_quantize: innermost dimension must be divisible by 32, got ", + input.size(1)); + const c10::cuda::CUDAGuard device_guard(input.device()); + auto input_contiguous = input.contiguous(); + if (reinterpret_cast(input_contiguous.data_ptr()) % 16 != 0) { + input_contiguous = input_contiguous.clone(); + } + auto output = at::empty_like(input_contiguous); + auto input_cpp = makeTransformerEngineTensor(input_contiguous); + auto output_cpp = makeTransformerEngineTensor(output); + nvte_mxfp4_fake_quantize(input_cpp.data(), output_cpp.data(), at::cuda::getCurrentCUDAStream()); + return output; +} + py::object dequantize(const py::handle &input, transformer_engine::DType otype) { init_extension(); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7e9d114be8..facbb0c489 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -199,6 +199,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); + m.def("mxfp4_fake_quantize", &transformer_engine::pytorch::mxfp4_fake_quantize, + "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", py::arg("input"), + py::call_guard()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); m.def("create_empty_quantized_tensor", diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 7c11e635c6..e69f4c22d7 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -46,6 +46,7 @@ _fsdp_gather_tensors, ) from ..constants import dist_group_type +from ..mxfp4_qat import mxfp4_fake_quantize from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer @@ -730,6 +731,23 @@ def fill_userbuffers_buffer_for_all_gather( raise ValueError(f"Unsupported quantizer for Userbuffers ({quantizer})") +def _mxfp4_qat_block_recipe_signature(recipe: Optional[Recipe]) -> Optional[tuple]: + """Immutable QAT blockwise configuration for detecting recipe mutation/switches.""" + if recipe is None or not recipe.float8_block_scaling() or not recipe.mxfp4_qat(): + return None + return ( + recipe.fp8_format, + recipe.fp8_quant_fwd_inp, + recipe.fp8_quant_fwd_weight, + recipe.fp8_quant_bwd_grad, + recipe.x_block_scaling_dim, + recipe.w_block_scaling_dim, + recipe.grad_block_scaling_dim, + recipe.backward_override, + recipe.mxfp4_qat(), + ) + + def _is_weight_workspace_valid( workspace: QuantizedTensorStorage, quantizer: Quantizer, @@ -747,6 +765,33 @@ def _is_weight_workspace_valid( return False if quantizer.columnwise_usage and workspace._columnwise_data is None: return False + elif isinstance(workspace, Float8BlockwiseQTensorStorage): + if not isinstance(quantizer, Float8BlockQuantizer): + return False + workspace_quantizer = workspace._quantizer + if not isinstance(workspace_quantizer, Float8BlockQuantizer): + return False + if workspace._is_2D_scaled != (quantizer.block_scaling_dim == 2): + return False + if ( + workspace._fp8_dtype != quantizer.dtype + or workspace_quantizer.block_scaling_dim != quantizer.block_scaling_dim + or workspace_quantizer.force_pow_2_scales != quantizer.force_pow_2_scales + or workspace_quantizer.amax_epsilon != quantizer.amax_epsilon + ): + return False + if quantizer.rowwise_usage and ( + not workspace_quantizer.rowwise_usage + or workspace._rowwise_data is None + or workspace._rowwise_scale_inv is None + ): + return False + if quantizer.columnwise_usage and ( + not workspace_quantizer.columnwise_usage + or workspace._columnwise_data is None + or workspace._columnwise_scale_inv is None + ): + return False elif isinstance(workspace, NVFP4TensorStorage): if quantizer.rowwise_usage and workspace._rowwise_data is None: return False @@ -800,8 +845,26 @@ def quantize_weight( ``_fp8_workspaces``. """ + _mxfp4_qat_active = ( + FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + if FP8GlobalStateManager.is_fp8_enabled() + else False + ) + if _mxfp4_qat_active and workspace_dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "the MXFP4 grid (values up to 6*2^125) exceeds fp16 range. Use bf16 " + "or fp32." + ) + # Already-quantized weight (primary FP8 parameters) if isinstance(tensor, QuantizedTensor): + if _mxfp4_qat_active: + raise NotImplementedError( + "MXFP4 QAT recipes do not support primary quantized weights: the " + "high-precision master weight is required to project onto the " + "MXFP4 grid." + ) update_rowwise = True if quantizer.rowwise_usage else None update_columnwise = True if quantizer.columnwise_usage else None tensor.update_usage( @@ -833,6 +896,8 @@ def quantize_weight( if update_workspace: if tensor is None: raise ValueError("tensor kwarg must be provided to update FP8 workspace") + if _mxfp4_qat_active: + tensor = mxfp4_fake_quantize(tensor) if hasattr(workspace, "quantize_"): workspace.quantize_(tensor, noop_flag=skip_update_flag) else: @@ -842,6 +907,8 @@ def quantize_weight( # Cache miss — create new workspace if tensor is None or quantizer is None: raise ValueError("tensor and quantizer kwargs must be provided to construct FP8 workspace") + if _mxfp4_qat_active: + tensor = mxfp4_fake_quantize(tensor) if cache: # Ensure the tensor in the cache is an instance of torch.Tensor, # as it persists beyond a single forward pass. @@ -1061,7 +1128,10 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_block_scaling() and isinstance( recipe_state, Float8BlockScalingRecipeState ): - return + # Block dims and scale policy are recipe fields, so a 1D/2D switch + # within float8_block_scaling still needs fresh quantizers. + if not recipe.mxfp4_qat() or recipe == recipe_state.recipe: + return if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): @@ -1473,13 +1543,29 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: meta["fp8_checkpoint"] = fp8_enabled _original_recipe = None + _current_recipe = FP8GlobalStateManager.get_fp8_recipe() + _original_block_signature = meta.get("_mxfp4_qat_block_recipe_signature") + _current_block_signature = _mxfp4_qat_block_recipe_signature(_current_recipe) + _block_recipe_config_changed = False if fp8_parameters or fp8_enabled: _original_recipe = meta.get("recipe", None) - if self.fp8_initialized and FP8GlobalStateManager.get_fp8_recipe() == _original_recipe: + _block_recipe_config_changed = ( + _original_recipe is not None + and _original_block_signature != _current_block_signature + ) + if ( + self.fp8_initialized + and _current_recipe == _original_recipe + and not _block_recipe_config_changed + ): # FP8 init has already been run and recipe is the same, don't do anything. return - meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + meta["recipe"] = _current_recipe + if _block_recipe_config_changed: + # RecipeState holds the same recipe object, so in-place mutation of it + # is invisible to ``==``; hence the separate immutable signature. + self.fast_setattr("fp8_meta_tensors_initialized", False) else: # If fp8 isn't enabled, turn off and return. self.fast_setattr("fp8_initialized", False) @@ -1503,12 +1589,19 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: self.init_fp8_meta_tensors(meta["recipe"]) self.fast_setattr("fp8_initialized", True) - meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + meta["recipe"] = _current_recipe _current_recipe = meta["recipe"] - if _original_recipe is not None and not ( - issubclass(_current_recipe.__class__, _original_recipe.__class__) - or issubclass(_original_recipe.__class__, _current_recipe.__class__) + meta["_mxfp4_qat_block_recipe_signature"] = _current_block_signature + if _block_recipe_config_changed: + # Cached workspaces own buffers for the old block layout; not updatable in place. + self._fp8_workspaces.clear() + if _original_recipe is not None and ( + not ( + issubclass(_current_recipe.__class__, _original_recipe.__class__) + or issubclass(_original_recipe.__class__, _current_recipe.__class__) + ) + or _current_recipe.mxfp4_qat() != _original_recipe.mxfp4_qat() ): warnings.warn( f"Recipe type changed from {_original_recipe.__class__.__name__} " diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 143fee7ba2..85e0fe5c53 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -52,6 +52,7 @@ materialize_weight_for_backward, finalize_weight_grads, ) +from ..mxfp4_qat import mxfp4_fake_quantize from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -104,6 +105,7 @@ def _is_grouped_tensor_path_supported( cpu_offloading: bool, backward_override: Optional[str], save_original_input: bool, + is_dist_weight: bool, activation_dtype: torch.dtype, input_quantizers: List[Optional[Quantizer]], output_quantizers: List[Optional[Quantizer]], @@ -144,6 +146,9 @@ def _is_grouped_tensor_path_supported( or save_original_input ): return False + # The fused backward has no materialize/finalize support for GTP weights. + if is_dist_weight: + return False # 3. Filter by compute capability and cuBLAS version device_capability = get_device_compute_capability() if not (9, 0) <= device_capability <= (11, 0): @@ -533,6 +538,18 @@ def forward( origin_weights = weights is_dist_weight = is_distributed_weight(weights[0]) if is_dist_weight: + if ( + is_grad_enabled + and weight_requires_grad + and wgrad_store is not None + and wgrad_store.delay_wgrad_compute() + ): + # finalize_weight_grads would reduce-scatter an uninitialized buffer + # before the delayed wgrad GEMM runs. + raise NotImplementedError( + "GroupedLinear does not support delay_wgrad_compute with" + " distributed weights." + ) weights = materialize_weight_for_forward(weights) # Configure quantizers @@ -597,6 +614,7 @@ def forward( cpu_offloading=cpu_offloading, backward_override=backward_override, save_original_input=save_original_input, + is_dist_weight=is_dist_weight, activation_dtype=activation_dtype, input_quantizers=input_quantizers, output_quantizers=output_quantizers, @@ -762,7 +780,8 @@ def forward( ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers - ctx.weights_requires_grad = weights[0].requires_grad + # GTP-materialized ``weights`` may not preserve ``requires_grad``. + ctx.weights_requires_grad = weight_requires_grad if fuse_wgrad_accumulation and ctx.weights_requires_grad: # Keep weakrefs to weights to preserve attributes like main_grad # when we need to modify the weight python objects @@ -789,6 +808,8 @@ def forward( ctx.activation_dtype = activation_dtype ctx.fp8 = fp8 ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + # Backward usually runs outside the quantization autocast. + ctx.mxfp4_qat = fp8 and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() ctx.backward_override = backward_override ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation ctx.cpu_offloading = cpu_offloading @@ -1135,23 +1156,42 @@ def backward( ctx.device, ) weights_for_dgrad = weights + if ( + is_dist_weight + and (ctx.fp8 or ctx.backward_override == "dequantized") + and ctx.weight_quantizers[0] is not None + ): + # Re-gathered GTP masters are high-precision; replay the forward's + # projection + quantization. ``dequantized`` forces ``ctx.fp8`` False. + weights_for_dgrad = [] + for weight, weight_quantizer in zip(weights, ctx.weight_quantizers): + if not isinstance(weight, QuantizedTensorStorage): + if ctx.mxfp4_qat: + weight = mxfp4_fake_quantize(weight) + weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = weight_quantizer(weight) + weights_for_dgrad.append(weight) if ctx.backward_override == "dequantized": + # Dequantize from the fprop quantized layout (the replay above, for GTP). weights_for_dgrad = [ ( weight.dequantize(dtype=ctx.activation_dtype) if isinstance(weight, QuantizedTensorStorage) else cast_if_needed(weight, ctx.activation_dtype) ) - for weight in weights + for weight in weights_for_dgrad ] elif ctx.backward_override == "high_precision": + # GTP: ``saved_weights`` are only the local shards, so use the re-gathered + # masters. They stay un-projected: the quantize branch above skips this. + hp_weights = weights if is_dist_weight else saved_weights weights_for_dgrad = [ ( weight.dequantize(dtype=ctx.activation_dtype) if isinstance(weight, QuantizedTensorStorage) else cast_if_needed(weight, ctx.activation_dtype) ) - for weight in saved_weights + for weight in hp_weights ] # Make sure weights are available in column-wise format # for dgrad computation. @@ -1867,21 +1907,20 @@ def forward( f"does not match number of GEMMs ({num_gemms})." ) - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = ( - FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor - ) - else: - skip_fp8_weight_update = None - if skip_fp8_weight_update is not None: - is_first_microbatch = False - # Preprocess input tensor if isinstance(inp, QuantizedTensorStorage): raise TypeError("GroupedLinear doesn't support input tensor in FP8.") inp = self.prepare_forward(inp, num_gemms=self.num_gemms) try: + grouped_weight = getattr(self, "weight", None) + if grouped_weight is not None and is_distributed_weight(grouped_weight): + # The per-GEMM member views lose the distributed-weight marker. + raise NotImplementedError( + "single_grouped_weight does not support distributed (GTP)" + " weights: the per-GEMM member tensors bypass the" + " materialize/finalize protocol." + ) weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a588e21a0c..b6c2a29f6b 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -32,6 +32,7 @@ _2X_ACC_WGRAD, ) from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..mxfp4_qat import mxfp4_fake_quantize from ..utils import ( assert_dim_for_fp8_exec, cast_if_needed, @@ -318,9 +319,22 @@ def forward( # ------------------------------------------------------ origin_weight = weight is_dist_weight = is_distributed_weight(origin_weight) + if ( + is_dist_weight + and is_grad_enabled + and weight_requires_grad + and wgrad_store is not None + and wgrad_store.delay_wgrad_compute() + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported with distributed" + " (GTP) weights: the delayed wgrad closure would skip" + " finalize_weight_grads" + ) if is_dist_weight: weight = materialize_weight_for_forward(weight)[0] out_features = weight.shape[0] + # The materialized tensor may drop ``requires_grad``; use ``weight_requires_grad``. new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -417,7 +431,7 @@ def forward( # ------------------------------------------------------ # Deallocate GEMM input tensor if no longer needed - if not weight.requires_grad and not return_layernorm_output: + if not weight_requires_grad and not return_layernorm_output: clear_tensor_data(ln_out, ln_out_total) ln_out = ln_out_total = None elif with_input_all_gather and not return_layernorm_output_gathered: @@ -464,7 +478,7 @@ def forward( ln_out_to_save = ln_out_hp ctx.weight_quantizer = weight_quantizer ctx.ln_out_needs_gather = ( - weight.requires_grad and parallel_mode == "column" and sequence_parallel + weight_requires_grad and parallel_mode == "column" and sequence_parallel ) # Input with column-wise usage is needed for wgrad GEMM. @@ -492,7 +506,7 @@ def forward( mu, rsigma, weightmat if fp8 and not is_weight_param_quantized else None, - ln_out_to_save if weight.requires_grad else None, + ln_out_to_save if weight_requires_grad else None, ) nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") @@ -529,27 +543,26 @@ def forward( ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects ctx.requires_dgrad = inp_requires_grad - ctx.requires_wgrad = weight.requires_grad + ctx.requires_wgrad = weight_requires_grad ctx.is_weight_param_quantized = is_weight_param_quantized ctx.is_fsdp2 = is_fsdp2 - if fuse_wgrad_accumulation and weight.requires_grad: - # Keep weakref to weight to preserve attributes like main_grad - # when we need to modify the weight python object - ctx.origin_weight_ref = weakref.ref(weight) + if fuse_wgrad_accumulation and weight_requires_grad: + # Weakref the origin param (not the GTP-materialized ``weight``) to keep main_grad. + ctx.origin_weight_ref = weakref.ref(origin_weight) # Save overwrite_main_grad flag now while we have access to weight object ctx.origin_weight_overwrites_main_grad = getattr( - weight, "overwrite_main_grad", False + origin_weight, "overwrite_main_grad", False ) # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates # the main_grad buffer lazily before backprop - if hasattr(weight, "__fsdp_param__"): + if hasattr(origin_weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward - ctx.main_grad_func = weight.get_main_grad + ctx.main_grad_func = origin_weight.get_main_grad elif is_dist_weight: ctx.main_grad_func = origin_weight.grad_buffer else: - ctx.main_grad_func = lambda: weight.main_grad + ctx.main_grad_func = lambda: origin_weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer ctx.grad_weight_quantizer = grad_weight_quantizer ctx.grad_output_quantizer = grad_output_quantizer @@ -810,8 +823,25 @@ def backward( # fsdp2 quantized-tensor hooks when workspace was not saved. weight = saved_weight elif ctx.weight_quantizer is not None: + # Replay the forward's projection + quantization. Keyed on the recipe captured + # in the forward; ``saved_weight`` must stay un-projected for high_precision. + w_for_encode = saved_weight + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + w_for_encode = mxfp4_fake_quantize(w_for_encode) ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) - weight = ctx.weight_quantizer(saved_weight) + weight = ctx.weight_quantizer(w_for_encode) + elif ( + is_dist_weight + and (ctx.fp8 or ctx.backward_override == "dequantized") + and ctx.weight_quantizer is not None + and not isinstance(weight, QuantizedTensorStorage) + ): + # Distributed weight re-gathered a BF16 weight: replay the forward's projection + # and quantization. ``dequantized`` lands here too: it forces ``ctx.fp8`` False. + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + weight = mxfp4_fake_quantize(weight) + ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = ctx.weight_quantizer(weight) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): @@ -854,7 +884,9 @@ def backward( else: weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) elif ctx.backward_override == "high_precision": - weight_for_dgrad = saved_weight + # GTP: ``saved_weight`` is only the local shard, so use the re-gathered master. + # It is still un-projected: the quantize branch above skips high_precision. + weight_for_dgrad = weight if is_dist_weight else saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 19e20d775c..99f535ec6c 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -33,6 +33,7 @@ _2X_ACC_WGRAD, ) from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..mxfp4_qat import mxfp4_fake_quantize from ..jit import ( bias_gelu_fused, bgrad_dgelu_fused, @@ -63,6 +64,7 @@ _get_cuda_rng_state, _set_cuda_rng_state, ) +from ..distributed_weight import is_distributed_weight from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -247,11 +249,20 @@ def _forward( backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override else: backward_override = None - assert backward_override is None, ( - "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." - " Replace LayerNormMLP with LayerNormLinear + Linear to enable" - " high_precision/dequantized backward." - ) + # A no-grad forward saves no backward state, so the override is moot there. + if backward_override is not None and is_grad_enabled: + raise NotImplementedError( + "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in" + " LayerNormMLP. Replace LayerNormMLP with LayerNormLinear + Linear to enable" + " high_precision/dequantized backward." + ) + # No materialize/finalize wiring here: a shard would be consumed as a full weight. + if is_distributed_weight(fc1_weight) or is_distributed_weight(fc2_weight): + raise NotImplementedError( + "LayerNormMLP does not support distributed (GTP) weights. Replace" + " LayerNormMLP with LayerNormLinear + Linear to enable distributed" + " weights." + ) # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take if is_grad_enabled and not recompute_for_bwd: @@ -1228,6 +1239,8 @@ def backward( if isinstance(origin_fc2_weight, QuantizedTensorStorage): fc2_weight = origin_fc2_weight elif ctx.fc2_weight_quantizer is not None: + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + origin_fc2_weight = mxfp4_fake_quantize(origin_fc2_weight) ctx.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) fc2_weight = ctx.fc2_weight_quantizer(origin_fc2_weight) @@ -1522,6 +1535,8 @@ def fc2_wgrad_gemm( if isinstance(origin_fc1_weight, QuantizedTensorStorage): fc1_weight = origin_fc1_weight elif ctx.fc1_weight_quantizer is not None: + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + origin_fc1_weight = mxfp4_fake_quantize(origin_fc1_weight) ctx.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) fc1_weight = ctx.fc1_weight_quantizer(origin_fc1_weight) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 6b0f941bc4..8d64fc5caa 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -15,6 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from ..mxfp4_qat import mxfp4_fake_quantize from transformer_engine.pytorch.torch_version import torch_version from .base import ( @@ -196,6 +197,9 @@ class LinearBwdArgs: wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD backward_override: Optional[str] = None is_weight_param_quantized: bool = False + # Forward projected the weight onto the MXFP4 grid before quantizing; backward must + # replay that, and can't infer it (``fp8`` may be forced False, and the autocast has exited). + mxfp4_qat: bool = False custom: bool = False debug: bool = False @@ -272,6 +276,18 @@ def _linear_forward_impl( weight = args.weight is_dist_weight = is_distributed_weight(args.weight) + if ( + is_dist_weight + and args.is_grad_enabled + and weight.requires_grad + and args.wgrad_store is not None + and args.wgrad_store.delay_wgrad_compute() + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported with distributed" + " (GTP) weights: the delayed wgrad closure would skip" + " finalize_weight_grads" + ) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -689,6 +705,7 @@ def _linear_setup_ctx( bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator bwd_args.backward_override = backward_override + bwd_args.mxfp4_qat = fp8 and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() bwd_args.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) bwd_args.custom = fwd_args.custom bwd_args.debug = fwd_args.debug @@ -1003,16 +1020,24 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # fsdp2 quantized-tensor hooks when workspace was not saved. weight_fp8 = saved_weight elif bwd_args.weight_quantizer is not None: + # Replay the forward's projection + quantization. Keyed on the recipe captured + # in the forward; ``saved_weight`` must stay un-projected for high_precision. + w_for_encode = saved_weight + if bwd_args.mxfp4_qat: + w_for_encode = mxfp4_fake_quantize(w_for_encode) bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) - weight_fp8 = bwd_args.weight_quantizer(saved_weight) + weight_fp8 = bwd_args.weight_quantizer(w_for_encode) elif ( is_dist_weight - and bwd_args.fp8 + and (bwd_args.fp8 or bwd_args.backward_override == "dequantized") and bwd_args.weight_quantizer is not None and not isinstance(weight_fp8, QuantizedTensorStorage) ): # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer # so the dgrad operand isn't cast by the delayed recipe. + # ``dequantized`` lands here too: it forces ``bwd_args.fp8`` False. + if bwd_args.mxfp4_qat: + weight_fp8 = mxfp4_fake_quantize(weight_fp8) bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(weight_fp8) @@ -1056,7 +1081,9 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. else: weight_for_dgrad = cast_if_needed(weight_for_dgrad, bwd_args.activation_dtype) elif bwd_args.backward_override == "high_precision": - weight_for_dgrad = saved_weight + # GTP: ``saved_weight`` is only the local shard, so use the re-gathered master. + # It is still un-projected: the quantize branch above skips high_precision. + weight_for_dgrad = weight_fp8 if is_dist_weight else saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py new file mode 100644 index 0000000000..dee82ce223 --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -0,0 +1,65 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP4 weight fake-quantization for quantization-aware training. + +Projects weights onto the MXFP4 (E2M1) grid with 1x32 power-of-two block scales, +returning dequantized values that are exact in bf16/fp32. Re-quantizing them with +the host recipe matches a direct MXFP4 converter in decoded values, but not in +bytes: the host encoder re-canonicalizes (payload, scale) per block. Scale floor +2^-126 / cap 2^125 (deployment contract); non-finite values NaN-poison their 1x32 +block; fp16 is rejected. Specification: +tests/pytorch/references/mxfp4_qat_reference.py. +""" +import torch + +__all__ = ["mxfp4_fake_quantize"] + +_MXFP4_BLOCK = 32 + + +class _MXFP4FakeQuantizeSTE(torch.autograd.Function): + """Straight-through estimator: identity gradient through the projection.""" + + @staticmethod + def forward(ctx, weight): # pylint: disable=arguments-differ + import transformer_engine_torch as tex + + if not hasattr(tex, "mxfp4_fake_quantize"): + raise RuntimeError( + "transformer_engine_torch was built without mxfp4_fake_quantize; " + "rebuild Transformer Engine." + ) + w = weight.contiguous() + if w.data_ptr() % 16 != 0: + # ``contiguous()`` keeps storage offsets; the kernel's vector loads need 16B alignment. + w = w.clone() + return tex.mxfp4_fake_quantize(w) + + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return grad_output + + +def mxfp4_fake_quantize(weight: torch.Tensor) -> torch.Tensor: + """Project ``weight`` onto the MXFP4 grid and return it in the input dtype. + + Per 1x32 block: scale = 2^clamp(ceil(log2(amax / 6)), -126, 125), RTNE onto E2M1, + saturate at 6, multiply back. Straight-through (identity) gradient. + """ + if not weight.is_cuda: + raise ValueError("MXFP4 QAT expects a CUDA weight tensor") + if weight.dim() != 2: + raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") + if weight.dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " + "(fp16 cannot represent the full MXFP4 scale range)" + ) + if weight.shape[1] % _MXFP4_BLOCK != 0: + raise ValueError( + "MXFP4 QAT needs the weight inner dim divisible by " + f"{_MXFP4_BLOCK}, got {weight.shape[1]}" + ) + return _MXFP4FakeQuantizeSTE.apply(weight) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4..8f3ecd70e3 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -19,6 +19,7 @@ gather_along_first_dim, reduce_scatter_along_first_dim, ) +from ...distributed_weight import is_distributed_weight from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...module.base import ( _2X_ACC_FPROP, @@ -26,6 +27,7 @@ _2X_ACC_WGRAD, ) from ...module._common import set_quantizer_amax_reduction_group +from ...mxfp4_qat import mxfp4_fake_quantize from ...tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage @@ -558,6 +560,30 @@ def _functional_forward( # Check weight tensor w = weight + # No materialize/finalize protocol here, so the wrapper would silently compute with + # the local shard. The Userbuffers fused op bypasses this and guards itself. + if is_distributed_weight(w): + raise NotImplementedError( + "te.ops.Linear/BasicLinear does not support distributed (GTP) " + "weights. Use the te.Linear module for distributed weights." + ) + mxfp4_qat = ( + with_quantized_compute + and FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ) + if mxfp4_qat: + if is_quantized_tensor(w): + raise NotImplementedError( + "MXFP4 QAT recipes do not support primary quantized weights: the " + "high-precision master weight is required to project onto the " + "MXFP4 grid." + ) + if dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "the MXFP4 grid exceeds fp16 range. Use bf16 or fp32." + ) if not with_quantized_compute: w = maybe_dequantize(w, dtype) elif with_quantized_compute and not is_quantized_tensor(w): @@ -567,6 +593,8 @@ def _functional_forward( rowwise=True, columnwise=input_requires_grad and backward_override is None, ) + if mxfp4_qat: + w = mxfp4_fake_quantize(w) w = weight_quantizer(w) # Check output tensor diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index d3e2dc0a13..c5eb892093 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -23,6 +23,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) +from ...mxfp4_qat import mxfp4_fake_quantize from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage @@ -65,7 +66,6 @@ compute_grouped_dbias_dscales, ) - # Keys for passing caller-provided output and grad-input buffers to a grouped # linear (or fused grouped MLP) through Sequential's ``op_kwargs``. OUTPUT_BUFFER_KEY = "output" @@ -481,6 +481,50 @@ def _make_grouped_biases_from_packed(self, packed_biases: torch.Tensor) -> None: for group_idx in range(self.num_groups): self.register_parameter(f"bias{group_idx}", None) + @staticmethod + def _mxfp4_qat_recipe() -> Optional[Recipe]: + """Return the active MXFP4-QAT recipe, if any.""" + if ( + FP8GlobalStateManager.is_fp8_enabled() or FP8GlobalStateManager.with_fp8_parameters() + ) and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + return FP8GlobalStateManager.get_fp8_recipe() + return None + + @staticmethod + def _reject_backward_override(context: str, *, requires_backward: bool) -> None: + """Reject backward-override recipes that this op cannot honor. + + Both backends replay the quantized weights in the backward dgrad GEMM, so the + override would be silently ignored. ``requires_backward`` comes from the fuser, + since ``torch.is_grad_enabled()`` is always False in fuser_forward. + """ + if not requires_backward: + return + if not FP8GlobalStateManager.is_fp8_enabled(): + return + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + if backward_override is not None: + raise NotImplementedError( + "ops.GroupedLinear only supports backward_override=None, got " + f"{backward_override!r} ({context}): the grouped backward " + "replays the quantized weights and would silently ignore the " + "override." + ) + + @classmethod + def _validate_mxfp4_qat_discrete_weights(cls, weights: Sequence[torch.Tensor]) -> bool: + """Validate the supported QAT subset and report whether projection is needed.""" + qat_recipe = cls._mxfp4_qat_recipe() + if qat_recipe is None: + return False + if any(is_quantized_tensor(weight) for weight in weights): + raise NotImplementedError( + "MXFP4 QAT recipes do not support quantized primary weights in the " + "operations API: the high-precision master is required for the QAT " + "projection." + ) + return True + def _quantize_weights( self, weights: Sequence[torch.Tensor], @@ -488,6 +532,13 @@ def _quantize_weights( ) -> Sequence[torch.Tensor]: """Construct quantized weight tensors.""" + if self._mxfp4_qat_recipe() is not None: + raise NotImplementedError( + "MXFP4 QAT recipes do not support quantized primary weights in the " + "operations API: the high-precision master is required for the QAT " + "projection." + ) + # Manually construct MXFP8 weights if isinstance(quantizers[0], MXFP8Quantizer): return self._quantize_weights_mxfp8(weights, quantizers) @@ -839,10 +890,35 @@ def _get_grouped_weight_for_gemm( columnwise_usage: bool, with_quantized_compute: bool, dtype: torch.dtype, + requires_backward: bool, ) -> GroupedTensorStorage: """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. Supports MXFP8/BF16/FP16 compute paths. """ + # Consumed in place here, re-materialized in the backward: forward/backward would differ. + if is_distributed_weight(weight_param): + raise NotImplementedError( + "GroupedLinear with single_grouped_weight=True does not support " + "distributed (GTP) weights. Use discrete per-expert weights." + ) + mxfp4_qat = False + if with_quantized_compute: + self._reject_backward_override( + "single_grouped_weight path", requires_backward=requires_backward + ) + mxfp4_qat = self._mxfp4_qat_recipe() is not None + if mxfp4_qat: + if weight_param.quantizer is not None: + raise NotImplementedError( + "MXFP4 QAT recipes do not support quantized primary weights in the " + "operations API: the high-precision master is required for the QAT " + "projection." + ) + if dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "use bf16 or fp32." + ) num_groups = self.num_groups is_weight_quantized = weight_param.quantizer is not None if is_weight_quantized and with_quantized_compute: @@ -881,8 +957,12 @@ def _get_grouped_weight_for_gemm( # Quantized compute path, use the fused group quantize kernel. weight_quantizer = weight_quantizers[0] weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + weight_data = weight_param.rowwise_data.view(weight_param.logical_shape) + if mxfp4_qat: + # Packed == per-expert: rows never straddle experts, 1x32 blocks span the inner dim. + weight_data = mxfp4_fake_quantize(weight_data) return tex.group_quantize( - weight_param.rowwise_data.view(weight_param.logical_shape), + weight_data, weight_quantizer, num_groups, None, @@ -895,16 +975,31 @@ def _get_discrete_weights_for_gemm( columnwise_usage: bool, with_quantized_compute: bool, dtype: torch.dtype, + requires_backward: bool, ) -> list[torch.Tensor]: """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. Returns a Python list, which dispatches the GEMM to ``discrete_in`` mode. """ + with_mxfp4_qat = False + if with_quantized_compute: + self._reject_backward_override( + "discrete-weight path", requires_backward=requires_backward + ) + with_mxfp4_qat = self._validate_mxfp4_qat_discrete_weights(weight_params) + if with_mxfp4_qat and dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "use bf16 or fp32." + ) out: list[torch.Tensor] = [] for w, quantizer in zip(weight_params, weight_quantizers): if not with_quantized_compute: w = maybe_dequantize(w, dtype) elif with_quantized_compute and not is_quantized_tensor(w): quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + if with_mxfp4_qat: + # Identity STE: the manually computed wgrad lands on the original param. + w = mxfp4_fake_quantize(w) w = quantizer(w) out.append(w) return out @@ -1003,6 +1098,11 @@ def fuser_forward( weight_quantizers = [None] * num_groups with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() if with_quantized_compute: + if self._mxfp4_qat_recipe() is not None and self._is_distributed_weight(): + raise NotImplementedError( + "MXFP4 QAT in ops.GroupedLinear does not support DistributedWeight " + "materialize/replay. FSDP2-unsharded high-precision weights are supported." + ) for group_idx in range(num_groups): input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) @@ -1188,6 +1288,12 @@ def _fuser_forward_split_quantize( # Extract params if self.single_grouped_weight: + # Same restriction as ``_get_grouped_weight_for_gemm``; QAT runs on the split views. + if self._is_distributed_weight(): + raise NotImplementedError( + "GroupedLinear with single_grouped_weight=True does not support " + "distributed (GTP) weights. Use discrete per-expert weights." + ) weights = self.weight.quantized_tensors if weights is None: weights = self.weight.split_into_quantized_tensors() @@ -1203,6 +1309,7 @@ def _fuser_forward_split_quantize( columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, dtype=dtype, + requires_backward=input_requires_grad or weight_requires_grad, ) # Split input tensor and convert dtypes if needed @@ -1337,6 +1444,7 @@ def _fuser_forward_grouped_tensor( columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, dtype=dtype, + requires_backward=input_requires_grad or weight_requires_grad, ) else: # Discrete weights @@ -1346,6 +1454,7 @@ def _fuser_forward_grouped_tensor( columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, dtype=dtype, + requires_backward=input_requires_grad or weight_requires_grad, ) # Allocate output buffer and wrap as a GroupedTensor view. @@ -1594,7 +1703,9 @@ def _fuser_backward_split_quantize( # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. # Return discarded (see finalize_weight_grads); the dummy is returned below instead. if ctx.weight_requires_grad and is_dist_weight: - assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + # Explicit raise, not an assert: ``python -O`` strips asserts. + if delay_wgrad: + raise RuntimeError("Delayed wgrad compute is unsupported with distributed weights.") finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace @@ -1818,7 +1929,9 @@ def _fuser_backward_grouped_tensor( # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. # Return discarded (see finalize_weight_grads); the dummy is returned below instead. if ctx.weight_requires_grad and is_dist_weight: - assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + # Explicit raise, not an assert: ``python -O`` strips asserts. + if delay_wgrad: + raise RuntimeError("Delayed wgrad compute is unsupported with distributed weights.") finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3d..7af8f219c7 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -26,6 +26,8 @@ finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD +from ...fp8 import FP8GlobalStateManager +from ...mxfp4_qat import mxfp4_fake_quantize from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor @@ -960,6 +962,45 @@ def fuser_forward( fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) + fp8_recipe = ( + FP8GlobalStateManager.get_fp8_recipe() + if FP8GlobalStateManager.is_fp8_enabled() + else None + ) + mxfp4_qat = fp8_recipe is not None and fp8_recipe.mxfp4_qat() + if ( + fp8_recipe is not None + and fp8_recipe.backward_override is not None + and requires_grad + ): + # The fused backward replays the quantized weights, so an override would be + # silently ignored. ``requires_grad`` is the fuser's signal; grad mode is off here. + raise NotImplementedError( + "Pure and MXFP4-QAT fused grouped-MLP ops support only " + f"backward_override=None, got {fp8_recipe.backward_override!r}." + ) + if mxfp4_qat: + if dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "use bf16 or fp32." + ) + if fc1_is_dist: + raise NotImplementedError( + "MXFP4 QAT fused grouped-MLP ops do not support DistributedWeight/GTP. " + "FSDP2 high-precision parameters are supported after FSDP2 all-gather." + ) + if fc1_is_dist: + # High-precision GTP weights would be quantized per-shard before + # ``materialize_weight_for_forward``, so the all-gather would never happen. + fc1_weight_list = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] + fc2_weight_list = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] + if any(not is_quantized_tensor(w) for w in fc1_weight_list + fc2_weight_list): + raise NotImplementedError( + "Fused grouped-MLP ops do not support high-precision " + "distributed (GTP) weights; distributed weights must be " + "natively quantized." + ) # Extract split sizes from extra input fc1_split_sizes = basic_op_extra_inputs[0][0] @@ -1004,6 +1045,11 @@ def fuser_forward( "FC1 expected GroupedTensor weight with single_grouped_weight=True." ) if fc1_op.weight.quantizer is not None: + if mxfp4_qat: + raise NotImplementedError( + "MXFP4 QAT fused grouped-MLP ops require high-precision " + "master weights, not primary quantized weights." + ) fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) fc1_op.weight.quantizer = fc1_weight_quantizer grouped_fc1_weight = fc1_op.weight @@ -1011,8 +1057,13 @@ def fuser_forward( if fc1_op.weight.rowwise_data is None: raise RuntimeError("FC1 grouped weight has no rowwise_data to quantize.") fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc1_weight_data = fc1_op.weight.rowwise_data.view(fc1_op.weight.logical_shape) + if mxfp4_qat: + # Packed == per-expert: rows never straddle experts, 1x32 blocks + # span the inner dim. + fc1_weight_data = mxfp4_fake_quantize(fc1_weight_data) grouped_fc1_weight = _group_quantize_for_grouped_mlp( - fc1_op.weight.rowwise_data.view(fc1_op.weight.logical_shape), + fc1_weight_data, fc1_weight_quantizer, num_groups, None, @@ -1024,8 +1075,14 @@ def fuser_forward( quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) if not is_quantized_tensor(weight): quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc1_weights.append(quantizer(weight)) + weight_for_quantize = mxfp4_fake_quantize(weight) if mxfp4_qat else weight + quantized_fc1_weights.append(quantizer(weight_for_quantize)) else: + if mxfp4_qat: + raise NotImplementedError( + "MXFP4 QAT fused grouped-MLP ops require high-precision " + "master weights, not primary quantized weights." + ) quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights if fc1_is_dist: @@ -1038,6 +1095,11 @@ def fuser_forward( "FC2 expected GroupedTensor weight with single_grouped_weight=True." ) if fc2_op.weight.quantizer is not None: + if mxfp4_qat: + raise NotImplementedError( + "MXFP4 QAT fused grouped-MLP ops require high-precision " + "master weights, not primary quantized weights." + ) fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) fc2_op.weight.quantizer = fc2_weight_quantizer grouped_fc2_weight = fc2_op.weight @@ -1045,8 +1107,12 @@ def fuser_forward( if fc2_op.weight.rowwise_data is None: raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc2_weight_data = fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape) + if mxfp4_qat: + # See the FC1 packed-projection note. + fc2_weight_data = mxfp4_fake_quantize(fc2_weight_data) grouped_fc2_weight = _group_quantize_for_grouped_mlp( - fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), + fc2_weight_data, fc2_weight_quantizer, num_groups, None, @@ -1058,8 +1124,14 @@ def fuser_forward( quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) if not is_quantized_tensor(weight): quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - quantized_fc2_weights.append(quantizer(weight)) + weight_for_quantize = mxfp4_fake_quantize(weight) if mxfp4_qat else weight + quantized_fc2_weights.append(quantizer(weight_for_quantize)) else: + if mxfp4_qat: + raise NotImplementedError( + "MXFP4 QAT fused grouped-MLP ops require high-precision " + "master weights, not primary quantized weights." + ) quantized_fc2_weights.append(weight) grouped_fc2_weight = quantized_fc2_weights diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index c1070e38a6..2df85f4501 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -20,6 +20,8 @@ get_ub, using_cublasmp_backend, ) +from ...mxfp4_qat import mxfp4_fake_quantize +from ...quantization import FP8GlobalStateManager from ...quantized_tensor import Quantizer from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...utils import canonicalize_device, canonicalize_dtype, clear_tensor_data @@ -315,6 +317,13 @@ def _functional_backward( if with_quantized_compute: if not is_quantized_tensor(w): weight_quantizer.set_usage(columnwise=True) + if ( + FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ): + # Only direct ``_functional_backward`` callers land here; the fuser + # path saves the already-projected weight. + w = mxfp4_fake_quantize(w) w = weight_quantizer(w) else: w = maybe_dequantize(w, dtype) @@ -659,6 +668,18 @@ def fuse_backward_ops( """ + # Disable Userbuffers for backward overrides, in sync with the forward op: the + # fuser fuses forward/backward separately, so a lone UB backward reads unset ctx. + recipe = unused.get("recipe", None) + if recipe is not None: + backward_override = recipe.backward_override + elif FP8GlobalStateManager.is_fp8_enabled(): + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + if backward_override is not None: + return ops + # Return immediately if environment is not distributed if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: return ops @@ -697,6 +718,9 @@ def fuse_backward_ops( continue if linear.tensor_parallel_size == 1: continue + if not linear.sequence_parallel: + # Userbuffers requires SP; don't fuse rather than fail at runtime. + continue if linear.tensor_parallel_mode == "row" and bias is not None: continue else: diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 3a8ff5438d..e397034036 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -14,12 +14,14 @@ from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import get_distributed_world_size +from ...distributed_weight import is_distributed_weight from ...quantization import FP8GlobalStateManager from ...module.base import ( fill_userbuffers_buffer_for_all_gather, get_ub, _2X_ACC_FPROP, ) +from ...mxfp4_qat import mxfp4_fake_quantize from ...quantized_tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage @@ -224,10 +226,36 @@ def _functional_forward( # Initialize weight tensor w = weight + if is_distributed_weight(w): + # No materialize/finalize protocol here: the local shard would be + # consumed as the full weight. + raise NotImplementedError( + "Userbuffers linear ops do not support distributed (GTP) weights." + ) + mxfp4_qat = ( + with_quantized_compute + and FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ) + if mxfp4_qat: + if is_quantized_tensor(w): + raise NotImplementedError( + "MXFP4 QAT recipes do not support primary quantized weights: the " + "high-precision master weight is required to project onto the " + "MXFP4 grid." + ) + if dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "the MXFP4 grid exceeds fp16 range. Use bf16 or fp32." + ) if not with_quantized_compute: w = maybe_dequantize(w, dtype) elif with_quantized_compute and not is_quantized_tensor(w): weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + if mxfp4_qat: + # The quantized weight is also the one saved for backward. + w = mxfp4_fake_quantize(w) w = weight_quantizer(w) # Construct output tensor if needed @@ -319,6 +347,7 @@ def fuser_forward( raise RuntimeError( f"Unsupported recipe for Userbuffers ({recipe.__class__.__name__})" ) + # MXFP4 QAT is handled in _functional_forward, which projects before quantizing. # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -441,6 +470,9 @@ def fuse_forward_ops( continue if linear.tensor_parallel_size == 1: continue + if not linear.sequence_parallel: + # Userbuffers requires SP; don't fuse rather than fail at runtime. + continue if linear.tensor_parallel_mode == "row" and bias is not None: continue else: diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..172a688292 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -18,6 +18,7 @@ FusibleOperation, FusedOperation, OperationContext, + _recipe_fusion_signature, ) @@ -352,6 +353,8 @@ def __init__( self.recipe_type = None self.first_op_requiring_backward = 0 self.backward_override = None + # Recipe field snapshot; recipe_type alone misses same-class field changes. + self._recipe_signature = None self._last_amax_history_len = 0 # Flatten list of parameters @@ -440,13 +443,20 @@ def maybe_fuse_ops( need_reset = False recipe_type = type(recipe) backward_override = recipe.backward_override if recipe is not None else None - fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + recipe_signature = _recipe_fusion_signature(recipe) + fusion_params = ( + recipe_type, + first_op_requiring_backward, + backward_override, + recipe_signature, + ) if fusion_params != ( self.recipe_type, self.first_op_requiring_backward, self.backward_override, + self._recipe_signature, ): - # Recipe type, backward override, or grad requirements have changed + # Recipe type/config, backward override, or grad requirements have changed. need_reset = True elif ( recipe is not None @@ -493,7 +503,12 @@ def maybe_fuse_ops( ) # Save current fusion params - self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params + ( + self.recipe_type, + self.first_op_requiring_backward, + self.backward_override, + self._recipe_signature, + ) = fusion_params # Save amax history length if isinstance(recipe, DelayedScaling): diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..fd6b080d68 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -24,6 +24,34 @@ from ..tensor import Quantizer +def _recipe_fusion_signature(recipe: Optional[Recipe]) -> Optional[tuple]: + """Snapshot of the recipe fields that get baked into quantizers. + + Recipe class alone is not enough: same-class fields can change between forward + passes, including by in-place mutation. DelayedScaling amax parameters are + excluded, since they are handled separately to preserve amax history. + """ + if recipe is None: + return None + fp8_format = getattr(recipe, "fp8_format", None) + return ( + type(recipe), + recipe.mxfp4_qat(), + getattr(recipe, "backward_override", None), + fp8_format.name if fp8_format is not None else None, + # QParams are frozen dataclasses, so they compare by value. + getattr(recipe, "fp8_quant_fwd_inp", None), + getattr(recipe, "fp8_quant_fwd_weight", None), + getattr(recipe, "fp8_quant_bwd_grad", None), + getattr(recipe, "fp4_quant_fwd_inp", None), + getattr(recipe, "fp4_quant_fwd_weight", None), + getattr(recipe, "fp4_quant_bwd_grad", None), + getattr(recipe, "x_block_scaling_dim", None), + getattr(recipe, "w_block_scaling_dim", None), + getattr(recipe, "grad_block_scaling_dim", None), + ) + + @dataclasses.dataclass class OperationContext: """State needed to apply an operation @@ -190,6 +218,8 @@ def __init__(self) -> None: # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + # Recipe signature the quantization state was built with. + self._recipe_state_signature: Optional[tuple] = None @property def is_fused_op(self) -> bool: @@ -243,6 +273,7 @@ def reset_recipe_state( if recipe is None: self._fp8_metas = None self._quantizers = None + self._recipe_state_signature = None return # Communication group for FP8 amax reductions @@ -252,8 +283,12 @@ def reset_recipe_state( # This could happen for example if calling BasicOperation.forward directly, as in that # case, the OperationFuser is not persistent, or when loading from a checkpoint need_to_reset_recipe_state = False + recipe_signature = _recipe_fusion_signature(recipe) if self._fp8_metas is None or self._quantizers is None: need_to_reset_recipe_state = True + elif recipe_signature != self._recipe_state_signature: + # Same-class field changes are baked into the quantizers too. + need_to_reset_recipe_state = True else: for mode in ("forward", "backward"): fp8_meta_key = FP8GlobalStateManager.get_meta_tensor_key( @@ -270,6 +305,7 @@ def reset_recipe_state( # Construct quantization recipe states self._fp8_metas = {"forward": None, "backward": None} self._quantizers = {"forward": [], "backward": []} + self._recipe_state_signature = recipe_signature for mode in ("forward", "backward"): num_quantizers = self.num_quantizers(mode) if num_quantizers == 0: From f599c3d6ae2fb3b23ffab7322b7166142a3a83d8 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 1 Aug 2026 02:11:44 -0700 Subject: [PATCH 2/2] update Signed-off-by: zhihaow6 --- .../run_fsdp2_mxfp4_qat_modules.py | 352 ++++++++++++++ .../fsdp2_tests/run_gtp_mxfp4_qat.py | 458 ++++++++++++++++++ tests/pytorch/distributed/test_torch_fsdp2.py | 54 +++ .../pytorch/ops/basic/grouped_linear.py | 4 + 4 files changed, 868 insertions(+) create mode 100644 tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_modules.py create mode 100644 tests/pytorch/distributed/fsdp2_tests/run_gtp_mxfp4_qat.py diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_modules.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_modules.py new file mode 100644 index 0000000000..960465ce50 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_modules.py @@ -0,0 +1,352 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FSDP2 coverage for MXFP4-QAT module-level APIs. + +Run with: + torchrun --nproc_per_node=2 -m pytest -v -s --tb=short + +Module counterpart of ``run_fsdp2_mxfp4_qat_ops.py`` (same reference scheme: +master vs MXFP4-projected weights discriminate which operand each backward +override used). Covers only configurations that pure recipes already support: + +* te.Linear / te.LayerNormLinear: all backward override modes. +* te.LayerNormMLP: mode ``None`` (the module rejects overrides). +* te.GroupedLinear, independent per-expert weights: all override modes. + +Both QAT host recipes (MXFP8 and 128x128 blockwise FP8) are parametrized. +""" + +from collections.abc import Callable, Sequence + +import pytest +import torch +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor +from torch.distributed._composable.fsdp import fully_shard + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import ( + Float8BlockScaling, + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, + MXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize + +_DTYPE = torch.bfloat16 + +_RECIPES = [ + pytest.param( + (MXFP4QATMXFP8BlockScaling, MXFP8BlockScaling, fp8.check_mxfp8_support), + id="mxfp8", + ), + pytest.param( + (MXFP4QATFloat8BlockScaling, Float8BlockScaling, fp8.check_fp8_block_scaling_support), + id="blockwise", + ), +] + +_OVERRIDES = ( + pytest.param(None, id="none"), + pytest.param("dequantized", id="dequantized"), + pytest.param("high_precision", id="high_precision"), +) + + +def _skip_if_unsupported(check_fn) -> None: + supported, reason = check_fn() + if not supported: + pytest.skip(reason) + + +def _device() -> torch.device: + return torch.device("cuda", torch.cuda.current_device()) + + +def _named_weights(module: torch.nn.Module) -> list[tuple[str, torch.nn.Parameter]]: + """GEMM weight parameters (weight, weightN, fc1_weight, ...), not layernorm affines.""" + weights = [ + (name, param) + for name, param in module.named_parameters() + if (lambda last: last.startswith("weight") or last.endswith("_weight"))( + name.rsplit(".", 1)[-1] + ) + and "layer_norm" not in name + ] + assert weights, f"{type(module).__name__} has no weight parameters" + return weights + + +def _make_master_weight( + out_features: int, + in_features: int, + *, + device: torch.device, +) -> torch.Tensor: + """Construct BF16 values that MXFP8/blockwise FP8 preserve but MXFP4 changes.""" + pattern = torch.tensor( + [ + 1.0, + 0.3125, + -0.3125, + 0.6875, + -0.6875, + 0.15625, + -0.15625, + 0.8125, + ], + dtype=_DTYPE, + device=device, + ) + repeats = (in_features + pattern.numel() - 1) // pattern.numel() + row = pattern.repeat(repeats)[:in_features] + return row.unsqueeze(0).expand(out_features, -1).contiguous() + + +def _initialize_master_weights(module: torch.nn.Module) -> dict[str, torch.Tensor]: + master_weights = {} + with torch.no_grad(): + for name, param in _named_weights(module): + weight = _make_master_weight( + param.shape[0], + param.shape[1], + device=param.device, + ) + param.copy_(weight) + master_weights[name] = weight + return master_weights + + +def _project_weights( + master_weights: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + projected_weights = { + name: mxfp4_fake_quantize(weight).detach() for name, weight in master_weights.items() + } + for name in master_weights: + assert not torch.equal( + master_weights[name], projected_weights[name] + ), f"Test weight {name} unexpectedly lies entirely on the MXFP4 grid" + return projected_weights + + +@torch.no_grad() +def _load_weights(module: torch.nn.Module, weights: dict[str, torch.Tensor]) -> None: + named_weights = _named_weights(module) + assert [name for name, _ in named_weights] == list(weights) + for name, param in named_weights: + param.copy_(weights[name]) + + +def _fsdp2_shard(module: torch.nn.Module) -> torch.nn.Module: + world_size = torch.distributed.get_world_size() + mesh = DeviceMesh("cuda", list(range(world_size))) + fully_shard(module, mesh=mesh) + return module + + +def _run_reference( + factory: Callable[[], torch.nn.Module], + weights: dict[str, torch.Tensor], + x: torch.Tensor, + dy: torch.Tensor, + extra_args: Sequence, + *, + pure_cls: type, + override: str | None, +) -> tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: + module = factory() + _load_weights(module, weights) + x_ref = x.detach().clone().requires_grad_(True) + recipe = pure_cls(backward_override=override) + with te.autocast(enabled=True, recipe=recipe): + y_ref = module(x_ref, *extra_args) + y_ref.backward(dy) + assert x_ref.grad is not None + return module, y_ref.detach(), x_ref.grad.detach() + + +def _assert_projected_forward( + y_qat: torch.Tensor, + y_projected: torch.Tensor, + y_master: torch.Tensor, +) -> None: + torch.testing.assert_close(y_qat, y_projected, rtol=0, atol=0) + assert not torch.equal( + y_qat, y_master + ), "QAT forward matched the unprojected master-weight path" + + +def _assert_fsdp_weight_grads( + fsdp_module: torch.nn.Module, + reference_module: torch.nn.Module, + *, + rtol: float = 0, + atol: float = 0, +) -> None: + fsdp_weights = _named_weights(fsdp_module) + reference_weights = _named_weights(reference_module) + assert [name for name, _ in fsdp_weights] == [name for name, _ in reference_weights] + for (name, fsdp_weight), (_, reference_weight) in zip(fsdp_weights, reference_weights): + assert isinstance(fsdp_weight, DTensor), f"{name} was not FSDP2-sharded" + assert fsdp_weight.grad is not None, f"STE did not return a gradient for {name}" + assert isinstance(fsdp_weight.grad, DTensor), f"{name}.grad is not sharded" + assert reference_weight.grad is not None + grad = fsdp_weight.grad.full_tensor() + assert torch.isfinite(grad).all(), f"{name}.grad contains non-finite values" + torch.testing.assert_close( + grad, + reference_weight.grad, + rtol=rtol, + atol=atol, + msg=lambda msg: (f"{name}: FSDP2 QAT wgrad did not reach the master shard\n{msg}"), + ) + + +def _make_input(tokens: int, features: int, device: torch.device) -> torch.Tensor: + values = torch.linspace(-0.25, 0.25, tokens * features, dtype=torch.float32, device=device) + return values.reshape(tokens, features).to(_DTYPE) + + +def _run_module_case( + factory: Callable[[], torch.nn.Module], + x: torch.Tensor, + extra_args: Sequence, + *, + qat_cls: type, + pure_cls: type, + override: str | None, + wgrad_rtol: float = 0, + wgrad_atol: float = 0, +) -> None: + """Shared body: run FSDP2 QAT vs projected/master references, assert the discriminators.""" + qat_module = factory() + master_weights = _initialize_master_weights(qat_module) + projected_weights = _project_weights(master_weights) + _fsdp2_shard(qat_module) + + x = x.detach().clone().requires_grad_(True) + qat_recipe = qat_cls(backward_override=override) + with te.autocast(enabled=True, recipe=qat_recipe): + y_qat = qat_module(x, *extra_args) + dy = _make_input(y_qat.shape[0], y_qat.shape[1], y_qat.device).flip(0) + y_qat.backward(dy) + assert x.grad is not None + + projected_ref, y_projected, dx_projected = _run_reference( + factory, projected_weights, x, dy, extra_args, pure_cls=pure_cls, override=override + ) + master_ref, y_master, dx_master = _run_reference( + factory, master_weights, x, dy, extra_args, pure_cls=pure_cls, override=override + ) + + _assert_projected_forward(y_qat.detach(), y_projected, y_master) + expected_dx = dx_master if override == "high_precision" else dx_projected + torch.testing.assert_close(x.grad, expected_dx, rtol=0, atol=0) + expected_wgrad_ref = master_ref if override == "high_precision" else projected_ref + _assert_fsdp_weight_grads( + qat_module, expected_wgrad_ref, rtol=wgrad_rtol, atol=wgrad_atol + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_fsdp2_module_linear_qat(recipes, override) -> None: + """te.Linear preserves all pure override modes under FSDP2.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + return te.Linear( + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(256, in_features, device) + _run_module_case( + factory, x, (), qat_cls=qat_cls, pure_cls=pure_cls, override=override + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_fsdp2_module_layernorm_linear_qat(recipes, override) -> None: + """te.LayerNormLinear preserves all pure override modes under FSDP2.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + return te.LayerNormLinear( + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(256, in_features, device) + _run_module_case( + factory, x, (), qat_cls=qat_cls, pure_cls=pure_cls, override=override + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +def test_fsdp2_module_layernorm_mlp_qat_none(recipes) -> None: + """te.LayerNormMLP mode None under FSDP2 (the module rejects overrides).""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + hidden_size = 128 + + def factory() -> torch.nn.Module: + return te.LayerNormMLP( + hidden_size, + 4 * hidden_size, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(256, hidden_size, device) + _run_module_case( + factory, x, (), qat_cls=qat_cls, pure_cls=pure_cls, override=None + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_fsdp2_module_grouped_linear_qat(recipes, override) -> None: + """te.GroupedLinear (independent expert weights) preserves override modes under FSDP2.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + num_gemms = 2 + in_features = out_features = 128 + m_splits = [128, 128] + + def factory() -> torch.nn.Module: + return te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(sum(m_splits), in_features, device) + _run_module_case( + factory, x, (m_splits,), qat_cls=qat_cls, pure_cls=pure_cls, override=override + ) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_gtp_mxfp4_qat.py b/tests/pytorch/distributed/fsdp2_tests/run_gtp_mxfp4_qat.py new file mode 100644 index 0000000000..0f47610065 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/run_gtp_mxfp4_qat.py @@ -0,0 +1,458 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-rank DistributedWeight (GTP) coverage for MXFP4-QAT. + +Run with: + torchrun --nproc_per_node=2 -m pytest -v -s --tb=short + +TE ships no DistributedWeight implementer, so this file provides a real one +(``_GTPShardedWeight``: dim-0 shards, all-gather materialize, reduce-scatter +finalize) and drives the module APIs through it end to end. The reference +scheme mirrors ``run_fsdp2_mxfp4_qat_ops.py``: master vs MXFP4-projected full +weights discriminate which operand each backward override used. + +Covers: + +* te.Linear / te.LayerNormLinear / te.GroupedLinear: all override modes + (None replays projection + host quantization on the re-gathered master; + dequantized dequantizes that replay; high_precision reads the un-projected + re-gathered master). +* Loud rejection of the unsupported GTP surfaces (ops API, single grouped + weight, LayerNormMLP, delayed wgrad). + +Shares ``conftest.py`` (dist init) with the FSDP2 scripts in this directory. +""" + +from collections.abc import Callable, Sequence + +import pytest +import torch +import torch.distributed as dist + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common.recipe import ( + Float8BlockScaling, + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, + MXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize + +_DTYPE = torch.bfloat16 + +_RECIPES = [ + pytest.param( + (MXFP4QATMXFP8BlockScaling, MXFP8BlockScaling, fp8.check_mxfp8_support), + id="mxfp8", + ), + pytest.param( + (MXFP4QATFloat8BlockScaling, Float8BlockScaling, fp8.check_fp8_block_scaling_support), + id="blockwise", + ), +] + +_OVERRIDES = ( + pytest.param(None, id="none"), + pytest.param("dequantized", id="dequantized"), + pytest.param("high_precision", id="high_precision"), +) + + +def _skip_if_unsupported(check_fn) -> None: + supported, reason = check_fn() + if not supported: + pytest.skip(reason) + + +def _device() -> torch.device: + return torch.device("cuda", torch.cuda.current_device()) + + +class _GTPShardedWeight(torch.nn.Parameter): + """Real DistributedWeight implementer: dim-0 shards over the default group. + + materialize all-gathers every member's full weight; finalize reduce-scatters + each full wgrad back to its shard with AVG (the test feeds every rank the + same batch, so AVG keeps single-rank references bit-comparable). + """ + + is_distributed_weight = True + + @staticmethod + def wrap(shard: torch.Tensor, full_shape: torch.Size) -> "_GTPShardedWeight": + weight = _GTPShardedWeight(shard.detach().clone()) + weight._full_shape = tuple(full_shape) + weight._group = None # set for grouped members; None -> single weight + weight.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + return weight + + def _members(self) -> list["_GTPShardedWeight"]: + return self._group if self._group is not None else [self] + + def _materialize(self) -> list[torch.Tensor]: + full_weights = [] + for member in self._members(): + full = torch.empty( + member._full_shape, dtype=member.dtype, device=member.device + ) + dist.all_gather_into_tensor(full, member.data.contiguous()) + full_weights.append(full) + return full_weights + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + out = self._materialize() + return out if self._group is not None else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + out = self._materialize() + return out if self._group is not None else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wgrad_list = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + shard_grads = [] + for member, wgrad in zip(self._members(), wgrad_list): + shard_grad = torch.empty_like(member.data, dtype=wgrad.dtype) + dist.reduce_scatter_tensor( + shard_grad, wgrad.contiguous(), op=dist.ReduceOp.AVG + ) + shard_grads.append(shard_grad) + return shard_grads if self._group is not None else shard_grads[0] + + def grad_buffer(self) -> torch.Tensor: + if not hasattr(self, "_grad_buffer"): + self._grad_buffer = torch.zeros( + self._full_shape, dtype=torch.float32, device=self.device + ) + return self._grad_buffer + + +def _shard(full: torch.Tensor) -> torch.Tensor: + world_size = dist.get_world_size() + rank = dist.get_rank() + assert full.shape[0] % world_size == 0 + return full.chunk(world_size, dim=0)[rank].contiguous() + + +def _make_master_weight( + out_features: int, + in_features: int, + *, + device: torch.device, +) -> torch.Tensor: + """Construct BF16 values that MXFP8/blockwise FP8 preserve but MXFP4 changes.""" + pattern = torch.tensor( + [ + 1.0, + 0.3125, + -0.3125, + 0.6875, + -0.6875, + 0.15625, + -0.15625, + 0.8125, + ], + dtype=_DTYPE, + device=device, + ) + repeats = (in_features + pattern.numel() - 1) // pattern.numel() + row = pattern.repeat(repeats)[:in_features] + return row.unsqueeze(0).expand(out_features, -1).contiguous() + + +def _gtp_wrap_weights( + module: torch.nn.Module, weight_names: Sequence[str] +) -> list[_GTPShardedWeight]: + """Replace each named full-weight Parameter with its rank's dim-0 shard wrapper.""" + wrappers = [] + for name in weight_names: + full = getattr(module, name) + wrapper = _GTPShardedWeight.wrap(_shard(full.data), full.shape) + setattr(module, name, wrapper) + wrappers.append(wrapper) + if len(wrappers) > 1: + wrappers[0]._group = wrappers + return wrappers + + +def _load_weights(module: torch.nn.Module, weights: dict[str, torch.Tensor]) -> None: + with torch.no_grad(): + for name, value in weights.items(): + getattr(module, name).copy_(value) + + +def _make_input(tokens: int, features: int, device: torch.device) -> torch.Tensor: + values = torch.linspace(-0.25, 0.25, tokens * features, dtype=torch.float32, device=device) + return values.reshape(tokens, features).to(_DTYPE) + + +def _run_reference( + factory: Callable[[], torch.nn.Module], + weights: dict[str, torch.Tensor], + x: torch.Tensor, + dy: torch.Tensor, + extra_args: Sequence, + *, + pure_cls: type, + override: str | None, +) -> tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: + module = factory() + _load_weights(module, weights) + x_ref = x.detach().clone().requires_grad_(True) + recipe = pure_cls(backward_override=override) + with te.autocast(enabled=True, recipe=recipe): + y_ref = module(x_ref, *extra_args) + y_ref.backward(dy) + assert x_ref.grad is not None + return module, y_ref.detach(), x_ref.grad.detach() + + +def _run_gtp_case( + factory: Callable[[], torch.nn.Module], + weight_names: Sequence[str], + x: torch.Tensor, + extra_args: Sequence, + *, + qat_cls: type, + pure_cls: type, + override: str | None, +) -> None: + """Shared body: GTP-sharded QAT module vs full-weight projected/master references.""" + device = _device() + gtp_module = factory() + master_weights = { + name: _make_master_weight( + getattr(gtp_module, name).shape[0], + getattr(gtp_module, name).shape[1], + device=device, + ) + for name in weight_names + } + _load_weights(gtp_module, master_weights) + projected_weights = { + name: mxfp4_fake_quantize(weight).detach() for name, weight in master_weights.items() + } + for name in weight_names: + assert not torch.equal(master_weights[name], projected_weights[name]) + wrappers = _gtp_wrap_weights(gtp_module, weight_names) + + x = x.detach().clone().requires_grad_(True) + qat_recipe = qat_cls(backward_override=override) + with te.autocast(enabled=True, recipe=qat_recipe): + y_qat = gtp_module(x, *extra_args) + dy = _make_input(y_qat.shape[0], y_qat.shape[1], device).flip(0) + y_qat.backward(dy) + assert x.grad is not None + + leader = wrappers[0] + assert leader.calls["fwd"] > 0, "materialize_group_for_forward was bypassed" + assert leader.calls["bwd"] > 0, "materialize_group_for_backward was bypassed" + assert leader.calls["finalize"] > 0, "finalize_group_grads was bypassed" + + projected_ref, y_projected, dx_projected = _run_reference( + factory, projected_weights, x, dy, extra_args, pure_cls=pure_cls, override=override + ) + master_ref, y_master, dx_master = _run_reference( + factory, master_weights, x, dy, extra_args, pure_cls=pure_cls, override=override + ) + + # Forward used the projected full weight, not the raw master. + torch.testing.assert_close(y_qat.detach(), y_projected, rtol=0, atol=0) + assert not torch.equal(y_qat.detach(), y_master) + + # Dgrad operand per override (see module docstring). + expected_dx = dx_master if override == "high_precision" else dx_projected + torch.testing.assert_close(x.grad, expected_dx, rtol=0, atol=0) + + # Wgrad: finalize reduce-scattered (AVG of identical per-rank batches) the + # full wgrad back to this rank's shard of the master parameter. + expected_ref = master_ref if override == "high_precision" else projected_ref + for name, wrapper in zip(weight_names, wrappers): + assert wrapper.grad is not None, f"{name}: STE did not deliver a shard gradient" + assert wrapper.grad.shape == wrapper.shape, f"{name}: gradient is not shard-shaped" + ref_grad = getattr(expected_ref, name).grad + assert ref_grad is not None + torch.testing.assert_close( + wrapper.grad, + _shard(ref_grad), + rtol=0, + atol=0, + msg=lambda msg: f"{name}: GTP QAT wgrad did not reach the master shard\n{msg}", + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_gtp_linear_qat_backward_overrides(recipes, override) -> None: + """te.Linear + GTP preserves all pure override modes.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + return te.Linear( + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(256, in_features, device) + _run_gtp_case( + factory, ("weight",), x, (), qat_cls=qat_cls, pure_cls=pure_cls, override=override + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_gtp_layernorm_linear_qat_backward_overrides(recipes, override) -> None: + """te.LayerNormLinear + GTP preserves all pure override modes.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + in_features = out_features = 128 + + def factory() -> torch.nn.Module: + return te.LayerNormLinear( + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(256, in_features, device) + _run_gtp_case( + factory, ("weight",), x, (), qat_cls=qat_cls, pure_cls=pure_cls, override=override + ) + + +@pytest.mark.parametrize("recipes", _RECIPES) +@pytest.mark.parametrize("override", _OVERRIDES) +def test_gtp_grouped_linear_qat_backward_overrides(recipes, override) -> None: + """te.GroupedLinear + GTP replays projection per expert in every override mode.""" + qat_cls, pure_cls, check_fn = recipes + _skip_if_unsupported(check_fn) + device = _device() + num_gemms = 2 + in_features = out_features = 128 + m_splits = [128, 128] + + def factory() -> torch.nn.Module: + return te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=_DTYPE, + device=device, + ) + + x = _make_input(sum(m_splits), in_features, device) + _run_gtp_case( + factory, + ("weight0", "weight1"), + x, + (m_splits,), + qat_cls=qat_cls, + pure_cls=pure_cls, + override=override, + ) + + +# ── Loud rejection of unsupported GTP surfaces ─────────────────────── + + +def _qat_autocast(): + return te.autocast(enabled=True, recipe=MXFP4QATMXFP8BlockScaling()) + + +def test_gtp_ops_grouped_linear_qat_rejected() -> None: + """ops.GroupedLinear rejects QAT + DistributedWeight instead of computing on shards.""" + _skip_if_unsupported(fp8.check_mxfp8_support) + device = _device() + op = te_ops.GroupedLinear(2, 128, 128, bias=False, device=device, dtype=_DTYPE) + _gtp_wrap_weights(op, ("weight0", "weight1")) + split_sizes = torch.tensor([128, 128], dtype=torch.int64, device=device) + x = _make_input(256, 128, device).requires_grad_(True) + with pytest.raises(NotImplementedError, match="does not support DistributedWeight"): + with _qat_autocast(): + op(x, split_sizes) + + +def test_gtp_ops_linear_rejected() -> None: + """te.ops.Linear has no GTP protocol and must fail fast (pure and QAT alike).""" + _skip_if_unsupported(fp8.check_mxfp8_support) + device = _device() + op = te_ops.Linear(128, 128, bias=False, device=device, dtype=_DTYPE) + _gtp_wrap_weights(op, ("weight",)) + x = _make_input(256, 128, device).requires_grad_(True) + with pytest.raises(NotImplementedError, match="does not support distributed"): + with _qat_autocast(): + op(x) + + +def test_gtp_layernorm_mlp_rejected() -> None: + """LayerNormMLP has no GTP wiring and must fail fast.""" + _skip_if_unsupported(fp8.check_mxfp8_support) + device = _device() + module = te.LayerNormMLP(128, 512, bias=False, params_dtype=_DTYPE, device=device) + _gtp_wrap_weights(module, ("fc1_weight",)) + x = _make_input(256, 128, device).requires_grad_(True) + with pytest.raises(NotImplementedError, match="LayerNormMLP does not support distributed"): + with _qat_autocast(): + module(x) + + +def test_gtp_single_grouped_weight_rejected() -> None: + """single_grouped_weight splits members and must reject distributed weights.""" + _skip_if_unsupported(fp8.check_mxfp8_support) + device = _device() + module = te.GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=_DTYPE, + device=device, + single_grouped_weight=True, + ) + module.make_grouped_weights() + module.weight.is_distributed_weight = True + x = _make_input(256, 128, device).requires_grad_(True) + with pytest.raises( + NotImplementedError, match="single_grouped_weight does not support distributed" + ): + with _qat_autocast(): + module(x, [128, 128]) + + +def test_gtp_delayed_wgrad_rejected() -> None: + """delay_wgrad_compute would skip finalize_weight_grads and must fail fast.""" + _skip_if_unsupported(fp8.check_mxfp8_support) + device = _device() + module = te.Linear( + 128, + 128, + bias=False, + params_dtype=_DTYPE, + device=device, + delay_wgrad_compute=True, + ) + _gtp_wrap_weights(module, ("weight",)) + x = _make_input(256, 128, device).requires_grad_(True) + with pytest.raises( + NotImplementedError, match="Delayed weight grad computation is not supported" + ): + with _qat_autocast(): + module(x) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 6f5dd14088..fca9ad1967 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -127,6 +127,60 @@ def test_fsdp2_mxfp4_qat_ops_tests(): ) +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_mxfp4_qat_module_tests(): + """MXFP4-QAT parity with pure-recipe module-level support under FSDP2.""" + mxfp8_available, reason = te.is_mxfp8_available(return_reason=True) + if not mxfp8_available: + pytest.skip(reason) + test_path = _FSDP2_DIR / "run_fsdp2_mxfp4_qat_modules.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + valid_returncodes=(0,), + env=os.environ, + timeout=600, + ) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_gtp_mxfp4_qat_tests(): + """MXFP4-QAT through a real multi-rank DistributedWeight (GTP) implementer.""" + mxfp8_available, reason = te.is_mxfp8_available(return_reason=True) + if not mxfp8_available: + pytest.skip(reason) + test_path = _FSDP2_DIR / "run_gtp_mxfp4_qat.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + valid_returncodes=(0,), + env=os.environ, + timeout=600, + ) + + @pytest.mark.skipif(NUM_PROCS < 4, reason="Requires 4+ GPUs for DP4→DP2 resharding test") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") @pytest.mark.parametrize("recipe", _parametrize_recipes()) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index c5eb892093..80ebe8e711 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -854,6 +854,10 @@ def _is_graph_safe_path_supported( """ if not (9, 0) <= get_device_compute_capability() <= (11, 0): return False + # Every graph-safe flow runs the grouped-GEMM setup workspace query, which + # requires cuBLAS 13.3+; fall back to the split-quantize flow on older cuBLAS. + if tex.get_cublasLt_version() < 130300: + return False if with_quantized_compute: # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM # path; the compute-capability range was already checked above. On Hopper it