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_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/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 f386659b6c..fca9ad1967 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,87 @@ 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 < 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/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..80ebe8e711 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) @@ -803,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 @@ -839,10 +894,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 +961,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 +979,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 +1102,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 +1292,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 +1313,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 +1448,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 +1458,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 +1707,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 +1933,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: