From 36e4e6f8e52eae15c0da4ce82394588c9fae9735 Mon Sep 17 00:00:00 2001 From: Vincent Zhong Date: Mon, 28 Sep 2026 12:56:32 +0000 Subject: [PATCH] [PyTorch] Add MXFP8 dgrad and gradient 4over6 NVFP4 options Add two NVFP4BlockScaling fields. Both can be set through env vars, so frameworks that construct a default recipe (e.g. Megatron-Core) can opt in without code changes. * dgrad_mxfp8 (NVTE_NVFP4_DGRAD_MXFP8): the dgrad GEMM uses MXFP8 operands while fprop and wgrad keep NVFP4 operands. The weight and grad_output slots become HybridQuantizers; the weight columnwise MXFP8 data is requantized from the NVFP4 fprop weight. Rejected together with backward_override and by quantized_model_init; quantized NVFP4 primary weights are rejected under this recipe. The graph-safe grouped GEMM path and the fused grouped MLP are skipped in this mode. * nvfp4_4over6_grad (NVTE_NVFP4_4OVER6_GRAD): 4over6 for gradient quantizers. Like the existing 4over6 scopes, it requires RHT and stochastic rounding to be disabled for those quantizers. Hybrid gradient quantizers now work in paths that assumed non-hybrid quantizers: * LayerNormMLP takes the unfused bgrad path, since tex.bgrad_quantize does not accept HybridQuantizer. * te.ops.GroupedLinear splits and quantizes grad_output with the hybrid-aware grouped split helper instead of tex.split_quantize. * The cuBLASMp wgrad re-gather in Linear and LayerNormLinear quantizes the original gradient for HybridQuantizer, as Linear already does for MXFP8, instead of requantizing dequantized rowwise data. Tests cover env var parsing, validation, quantizer layouts, te.Linear against GEMMs on dequantized operands for both options (with negative controls), bitwise parity with the equivalent CustomRecipe across modules, module options and te.ops, LayerNormMLP, the quantized primary weight rejections, the dgrad error reduction, and all-zero gradients and inputs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015LRZMCLfqm7g6gHe49iBPg Signed-off-by: Vincent Zhong --- docs/envvars.rst | 14 +- .../nvfp4/test_nvfp4_backward_options.py | 516 ++++++++++++++++++ transformer_engine/common/recipe/__init__.py | 21 + transformer_engine/pytorch/module/base.py | 6 + .../pytorch/module/grouped_linear.py | 1 + .../pytorch/module/layernorm_linear.py | 6 + .../pytorch/module/layernorm_mlp.py | 5 +- transformer_engine/pytorch/module/linear.py | 8 +- .../pytorch/ops/basic/grouped_linear.py | 4 +- .../pytorch/ops/fused/grouped_mlp.py | 3 + transformer_engine/pytorch/quantization.py | 45 +- 11 files changed, 615 insertions(+), 14 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_backward_options.py diff --git a/docs/envvars.rst b/docs/envvars.rst index 52b586ba1c5..aec1c3f351f 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -347,7 +347,7 @@ Kernel Configuration :Type: ``str`` (``none``, ``weights``, ``activations``, or ``all``) :Default: ``all`` - :Description: Select NVFP4 4over6 quantizers that use 256 instead of 448 as the global E4M3 scale bound. By default, all 4over6 quantizers use 256. Set the env var to ``none`` (or set ``NVFP4BlockScaling(nvfp4_4over6_e4m3_use_256="none")``) to use the standard NVFP4 448 bound for all 4over6 quantizers. This option is only meaningful for tensor roles that also enable :envvar:`NVTE_NVFP4_4OVER6`. + :Description: Select NVFP4 4over6 quantizers that use 256 instead of 448 as the global E4M3 scale bound. By default, all 4over6 quantizers use 256. Set the env var to ``none`` (or set ``NVFP4BlockScaling(nvfp4_4over6_e4m3_use_256="none")``) to use the standard NVFP4 448 bound for all 4over6 quantizers. This option is only meaningful for tensor roles that also enable :envvar:`NVTE_NVFP4_4OVER6` or :envvar:`NVTE_NVFP4_4OVER6_GRAD`. Gradient 4over6 quantizers use 256 only when this option is ``all``. .. envvar:: NVTE_NVFP4_4OVER6_ERR_MODE @@ -361,6 +361,18 @@ Kernel Configuration :Default: ``0`` :Description: Use the faster NVFP4 4over6 candidate error path that compares candidates in the E4M3-scaled domain after the E2M1 x E4M3 product is rounded to FP16. Error differences and accumulation remain FP32. By default, 4over6 error comparison uses the original input-domain path; ``NVTE_USE_FAST_MATH`` does not control this error-comparison path. +.. envvar:: NVTE_NVFP4_4OVER6_GRAD + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Apply 4over6 adaptive NVFP4 block scaling to gradient quantizers in the ``NVFP4BlockScaling`` recipe. When set to ``1`` (or when ``NVFP4BlockScaling(nvfp4_4over6_grad=True)`` is used), gradient blocks use the same map-to-4 versus map-to-6 candidate selection as :envvar:`NVTE_NVFP4_4OVER6`. Requires RHT and stochastic rounding to be disabled for gradient quantizers (``NVTE_NVFP4_DISABLE_RHT=1`` and ``NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1``, or ``disable_rht=True`` and ``disable_stochastic_rounding=True``). + +.. envvar:: NVTE_NVFP4_DGRAD_MXFP8 + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Run the dgrad GEMM of the ``NVFP4BlockScaling`` recipe with MXFP8 operands while the fprop and wgrad GEMMs keep NVFP4 operands. When set to ``1`` (or when ``NVFP4BlockScaling(dgrad_mxfp8=True)`` is used), the weight quantizer emits NVFP4 rowwise data for fprop and MXFP8 columnwise data for dgrad, and the ``grad_output`` quantizer emits MXFP8 rowwise data for dgrad and NVFP4 columnwise data for wgrad. Cannot be combined with ``NVTE_BACKWARD_OVERRIDE`` or with quantized primary weights. + Torch Compilation and Fusion ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/tests/pytorch/nvfp4/test_nvfp4_backward_options.py b/tests/pytorch/nvfp4/test_nvfp4_backward_options.py new file mode 100644 index 00000000000..a1047c32059 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_backward_options.py @@ -0,0 +1,516 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for NVFP4BlockScaling backward options (MXFP8 dgrad and gradient 4over6). + +Also covers zero-valued grad_output rows and blocks, which RL losses produce for loss-masked +prompt and padding tokens and for micro-batches whose advantages are all zero. +""" + +import os +import subprocess +import sys + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common.recipe import CustomRecipe, NVFP4BlockScaling +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.quantization import RecipeState +from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +requires_nvfp4 = pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +requires_nvfp4_and_mxfp8 = pytest.mark.skipif( + not (nvfp4_available and mxfp8_available), + reason=reason_for_no_nvfp4 if not nvfp4_available else reason_for_no_mxfp8, +) + +IN_FEATURES = 256 +OUT_FEATURES = 512 +M_SPLITS = [128, 256, 128] +NUM_TOKENS = sum(M_SPLITS) + +# References are FP32 GEMMs on dequantized operands; the tolerance covers BF16 output rounding. +BF16_OUTPUT_TOLS = {"rtol": 1.6e-2, "atol": 1e-5} + + +def make_recipe(**kwargs) -> NVFP4BlockScaling: + """NVFP4 recipe with RHT, stochastic rounding, and 2D weight quantization disabled by default.""" + kwargs = { + "disable_rht": True, + "disable_stochastic_rounding": True, + "disable_2d_quantization": True, + **kwargs, + } + return NVFP4BlockScaling(**kwargs) + + +def nvfp4_quantizer(row_scaled: bool = False, use_4over6: bool = False) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=256 if use_4over6 else 448, + ) + + +def mxfp8_quantizer() -> MXFP8Quantizer: + return MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3) + + +def dgrad_mxfp8_custom_recipe(row_scaled_activation: bool = False) -> CustomRecipe: + """CustomRecipe equivalent to ``make_recipe(dgrad_mxfp8=True)``.""" + + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return nvfp4_quantizer(row_scaled=row_scaled_activation) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer(), + columnwise_quantizer=mxfp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer(), + columnwise_quantizer=nvfp4_quantizer(), + ) + return nvfp4_quantizer() + + return CustomRecipe(qfactory=qfactory) + + +def make_quantizers(recipe, mode: str, num_quantizers: int) -> list: + return RecipeState.create(recipe, mode=mode, num_quantizers=num_quantizers).make_quantizers() + + +def make_module(kind: str, bias: bool = False, activation: str = "gelu") -> torch.nn.Module: + torch.manual_seed(1234) + kwargs = {"bias": bias, "params_dtype": torch.bfloat16} + if kind == "linear": + return te.Linear(IN_FEATURES, OUT_FEATURES, **kwargs) + if kind == "layernorm_linear": + return te.LayerNormLinear(IN_FEATURES, OUT_FEATURES, normalization="RMSNorm", **kwargs) + if kind == "layernorm_mlp": + return te.LayerNormMLP(IN_FEATURES, OUT_FEATURES, activation=activation, **kwargs) + if kind == "grouped_linear": + return te.GroupedLinear(len(M_SPLITS), IN_FEATURES, OUT_FEATURES, **kwargs) + raise ValueError(f"Unknown module kind: {kind}") + + +def make_inputs(out_features: int = OUT_FEATURES): + generator = torch.Generator(device="cuda").manual_seed(0) + x = torch.randn(NUM_TOKENS, IN_FEATURES, device="cuda", generator=generator) + dy = torch.randn(NUM_TOKENS, out_features, device="cuda", generator=generator) * 1e-3 + return x.to(torch.bfloat16), dy.to(torch.bfloat16) + + +def default_forward(module, x): + if isinstance(module, te.GroupedLinear): + return module(x, M_SPLITS) + if isinstance(module, te_ops.GroupedLinear): + return module(x, torch.tensor(M_SPLITS, device="cuda")) + return module(x) + + +def run_step(module, recipe, x, dy, forward=default_forward): + """One forward and backward step. Returns the output, dgrad, and parameter gradients. + + Parameters with ``main_grad`` (fused wgrad accumulation) report ``main_grad``. + """ + x = x.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + y = forward(module, x) + y.backward(dy) + if ( + getattr(module, "wgrad_store", None) is not None + and module.wgrad_store.delay_wgrad_compute() + ): + module.backward_dw() + grads = { + name: (p.main_grad if hasattr(p, "main_grad") else p.grad).detach().clone() + for name, p in module.named_parameters() + } + return y.detach(), x.grad.detach(), grads + + +def assert_steps_equal(step, step_ref): + y, dx, grads = step + y_ref, dx_ref, grads_ref = step_ref + torch.testing.assert_close(y, y_ref, rtol=0, atol=0) + torch.testing.assert_close(dx, dx_ref, rtol=0, atol=0) + assert grads.keys() == grads_ref.keys() + for name, grad in grads.items(): + torch.testing.assert_close(grad, grads_ref[name], rtol=0, atol=0) + + +def rel_err(x: torch.Tensor, ref: torch.Tensor) -> float: + return ((x.float() - ref.float()).norm() / ref.float().norm()).item() + + +def dequantize(quantizer, tensor: torch.Tensor) -> torch.Tensor: + """Rowwise-quantize and dequantize to FP32.""" + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer(tensor).dequantize(dtype=torch.float32) + + +def dequantize_columnwise(quantizer, tensor: torch.Tensor) -> torch.Tensor: + """Emulate columnwise quantization by quantizing the transpose rowwise; dequantize to FP32.""" + return dequantize(quantizer, tensor.t().contiguous()).t() + + +def test_recipe_env_vars(): + """Frameworks that construct a default NVFP4BlockScaling can opt in through env vars.""" + env = dict(os.environ, NVTE_NVFP4_DGRAD_MXFP8="1", NVTE_NVFP4_4OVER6_GRAD="1") + code = ( + "from transformer_engine.common.recipe import NVFP4BlockScaling\n" + "recipe = NVFP4BlockScaling()\n" + "assert recipe.dgrad_mxfp8 and recipe.nvfp4_4over6_grad, recipe\n" + ) + subprocess.run([sys.executable, "-c", code], env=env, check=True) + + +def test_dgrad_mxfp8_rejects_backward_override(): + # pydantic reports failed __post_init__ assertions as ValidationError, a ValueError. + with pytest.raises(ValueError, match="NVTE_NVFP4_DGRAD_MXFP8"): + make_recipe(dgrad_mxfp8=True, backward_override="high_precision") + + +@requires_nvfp4 +def test_4over6_grad_rejects_rht_and_stochastic_rounding(): + recipe = NVFP4BlockScaling( + nvfp4_4over6_grad=True, disable_rht=False, disable_stochastic_rounding=False + ) + with pytest.raises(ValueError, match="4over6"): + make_quantizers(recipe, "backward", 2) + + +@requires_nvfp4_and_mxfp8 +def test_dgrad_mxfp8_quantizers(): + recipe = make_recipe(dgrad_mxfp8=True, row_scaled_activation=True) + input_quantizer, weight_quantizer, _ = make_quantizers(recipe, "forward", 3) + assert isinstance(input_quantizer, NVFP4Quantizer) + assert input_quantizer.row_scaled_nvfp4 + assert isinstance(weight_quantizer, HybridQuantizer) + assert isinstance(weight_quantizer.rowwise_quantizer, NVFP4Quantizer) + assert isinstance(weight_quantizer.columnwise_quantizer, MXFP8Quantizer) + assert weight_quantizer.columnwise_source == "rowwise_dequantized" + + grad_output_quantizer, grad_input_quantizer = make_quantizers(recipe, "backward", 2) + assert isinstance(grad_output_quantizer, HybridQuantizer) + assert isinstance(grad_output_quantizer.rowwise_quantizer, MXFP8Quantizer) + assert isinstance(grad_output_quantizer.columnwise_quantizer, NVFP4Quantizer) + assert grad_output_quantizer.columnwise_source == "original" + assert not grad_output_quantizer.columnwise_quantizer.row_scaled_nvfp4 + assert isinstance(grad_input_quantizer, NVFP4Quantizer) + + # The NVFP4 wgrad operand keeps the recipe's gradient RHT and stochastic rounding. + recipe = NVFP4BlockScaling( + dgrad_mxfp8=True, disable_rht=False, disable_stochastic_rounding=False + ) + grad_output_quantizer, _ = make_quantizers(recipe, "backward", 2) + assert grad_output_quantizer.columnwise_quantizer.with_rht + assert grad_output_quantizer.columnwise_quantizer.stochastic_rounding + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("e4m3_use_256, e4m3_max", [("all", 256), ("activations", 448)]) +def test_4over6_grad_quantizers(e4m3_use_256, e4m3_max): + recipe = make_recipe(nvfp4_4over6_grad=True, nvfp4_4over6_e4m3_use_256=e4m3_use_256) + for quantizer in make_quantizers(recipe, "backward", 2): + assert quantizer.nvfp4_use_4over6 + assert quantizer.nvfp4_e4m3_max == e4m3_max + for quantizer in make_quantizers(recipe, "forward", 3): + assert not quantizer.nvfp4_use_4over6 + + # nvfp4_4over6="all" does not cover gradients. + grad_output_quantizer, _ = make_quantizers(make_recipe(nvfp4_4over6="all"), "backward", 2) + assert not grad_output_quantizer.nvfp4_use_4over6 + + # With MXFP8 dgrad, 4over6 applies to the NVFP4 wgrad operand. + recipe = make_recipe(nvfp4_4over6_grad=True, dgrad_mxfp8=True) + grad_output_quantizer, _ = make_quantizers(recipe, "backward", 2) + assert grad_output_quantizer.columnwise_quantizer.nvfp4_use_4over6 + + +@requires_nvfp4_and_mxfp8 +def test_dgrad_mxfp8_linear_reference(): + """te.Linear matches GEMMs on dequantized operands. + + fprop and wgrad use NVFP4 operands; dgrad uses MXFP8 operands, with the weight requantized + from its dequantized NVFP4 value. The NVFP4 dgrad and the MXFP8 weight quantized from the + original weight are negative controls. + """ + x, dy = make_inputs() + module = make_module("linear") + y, dx, grads = run_step(module, make_recipe(dgrad_mxfp8=True), x, dy) + + w = module.weight.detach() + w_nvfp4 = dequantize(nvfp4_quantizer(), w) + y_ref = dequantize(nvfp4_quantizer(), x) @ w_nvfp4.t() + dx_ref = dequantize(mxfp8_quantizer(), dy) @ dequantize_columnwise( + mxfp8_quantizer(), w_nvfp4.to(torch.bfloat16) + ) + dw_ref = dequantize_columnwise(nvfp4_quantizer(), dy).t() @ dequantize_columnwise( + nvfp4_quantizer(), x + ) + torch.testing.assert_close(y, y_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + torch.testing.assert_close(dx, dx_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + torch.testing.assert_close(grads["weight"], dw_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + + dx_nvfp4 = dequantize(nvfp4_quantizer(), dy) @ dequantize_columnwise(nvfp4_quantizer(), w) + dx_mxfp8_from_original = dequantize(mxfp8_quantizer(), dy) @ dequantize_columnwise( + mxfp8_quantizer(), w + ) + for dx_wrong in (dx_nvfp4, dx_mxfp8_from_original): + with pytest.raises(AssertionError): + torch.testing.assert_close(dx, dx_wrong.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + + +@requires_nvfp4 +def test_4over6_grad_linear_reference(): + """te.Linear with gradient 4over6 matches GEMMs on dequantized 4over6 gradients.""" + x, dy = make_inputs() + module = make_module("linear") + _, dx, grads = run_step(module, make_recipe(nvfp4_4over6_grad=True), x, dy) + + w = module.weight.detach() + for use_4over6 in (True, False): + dx_ref = dequantize(nvfp4_quantizer(use_4over6=use_4over6), dy) @ dequantize_columnwise( + nvfp4_quantizer(), w + ) + dy_col = dequantize_columnwise(nvfp4_quantizer(use_4over6=use_4over6), dy) + dw_ref = dy_col.t() @ dequantize_columnwise(nvfp4_quantizer(), x) + if use_4over6: + torch.testing.assert_close(dx, dx_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + torch.testing.assert_close( + grads["weight"], dw_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS + ) + else: # Negative control: standard NVFP4 gradients. + with pytest.raises(AssertionError): + torch.testing.assert_close(dx, dx_ref.to(torch.bfloat16), **BF16_OUTPUT_TOLS) + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("kind", ["linear", "layernorm_linear", "grouped_linear"]) +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_dgrad_mxfp8_matches_custom_recipe(kind, bias, row_scaled_activation): + """The built-in recipe dispatches exactly like the equivalent CustomRecipe.""" + x, dy = make_inputs() + recipe = make_recipe(dgrad_mxfp8=True, row_scaled_activation=row_scaled_activation) + custom_recipe = dgrad_mxfp8_custom_recipe(row_scaled_activation) + assert_steps_equal( + run_step(make_module(kind, bias), recipe, x, dy), + run_step(make_module(kind, bias), custom_recipe, x, dy), + ) + + +def _fused_wgrad_linear(): + module = te.Linear( + IN_FEATURES, + OUT_FEATURES, + bias=False, + params_dtype=torch.bfloat16, + fuse_wgrad_accumulation=True, + ) + module.weight.main_grad = torch.zeros(OUT_FEATURES, IN_FEATURES, device="cuda") + return module + + +def _microbatches(module, x): + return module(x, is_first_microbatch=True) + module(x, is_first_microbatch=False) + + +MODULE_OPTIONS = { + "fuse_wgrad_accumulation": (_fused_wgrad_linear, default_forward), + "delay_wgrad_compute": ( + lambda: te.Linear( + IN_FEATURES, + OUT_FEATURES, + bias=False, + params_dtype=torch.bfloat16, + delay_wgrad_compute=True, + ), + default_forward, + ), + "save_original_input": ( + lambda: te.Linear( + IN_FEATURES, + OUT_FEATURES, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ), + default_forward, + ), + "weight_caching": (lambda: make_module("linear"), _microbatches), + "single_grouped_weight": ( + lambda: te.GroupedLinear( + len(M_SPLITS), + IN_FEATURES, + OUT_FEATURES, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=True, + ), + default_forward, + ), + "ops_basic_linear": ( + lambda: te_ops.BasicLinear(IN_FEATURES, OUT_FEATURES, dtype=torch.bfloat16), + default_forward, + ), + "ops_grouped_linear": ( + lambda: te_ops.GroupedLinear( + len(M_SPLITS), IN_FEATURES, OUT_FEATURES, bias=False, dtype=torch.bfloat16 + ), + default_forward, + ), + "ops_grouped_linear_single_grouped_weight": ( + lambda: te_ops.GroupedLinear( + len(M_SPLITS), + IN_FEATURES, + OUT_FEATURES, + bias=False, + dtype=torch.bfloat16, + single_grouped_weight=True, + ), + default_forward, + ), +} + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("option", list(MODULE_OPTIONS)) +def test_dgrad_mxfp8_module_options(option): + build, forward = MODULE_OPTIONS[option] + x, dy = make_inputs() + torch.manual_seed(1234) + step = run_step(build(), make_recipe(dgrad_mxfp8=True), x, dy, forward) + torch.manual_seed(1234) + step_ref = run_step(build(), dgrad_mxfp8_custom_recipe(), x, dy, forward) + assert_steps_equal(step, step_ref) + assert torch.isfinite(step[1]).all() + assert all(torch.isfinite(grad).all() for grad in step[2].values()) + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("kind", ["linear", "layernorm_linear", "grouped_linear"]) +def test_dgrad_mxfp8_error(kind): + """MXFP8 dgrad keeps the NVFP4 forward and wgrad and reduces the dgrad error.""" + x, dy = make_inputs() + _, dx_ref, _ = run_step(make_module(kind), make_recipe(backward_override="dequantized"), x, dy) + y_nvfp4, dx_nvfp4, grads_nvfp4 = run_step(make_module(kind), make_recipe(), x, dy) + y, dx, grads = run_step(make_module(kind), make_recipe(dgrad_mxfp8=True), x, dy) + torch.testing.assert_close(y, y_nvfp4, rtol=0, atol=0) + for name, grad in grads.items(): + if not name.startswith("layer_norm"): # The norm weight gradient depends on dgrad. + torch.testing.assert_close(grad, grads_nvfp4[name], rtol=0, atol=0) + assert rel_err(dx, dx_ref) < 0.5 * rel_err(dx_nvfp4, dx_ref) + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("activation", ["gelu", "swiglu"]) +@pytest.mark.parametrize("bias", [False, True]) +def test_dgrad_mxfp8_layernorm_mlp(activation, bias): + x, dy = make_inputs(out_features=IN_FEATURES) + module = make_module("layernorm_mlp", bias, activation) + y_nvfp4, _, grads_nvfp4 = run_step(module, make_recipe(), x, dy) + module = make_module("layernorm_mlp", bias, activation) + y, dx, grads = run_step(module, make_recipe(dgrad_mxfp8=True), x, dy) + torch.testing.assert_close(y, y_nvfp4, rtol=0, atol=0) + assert torch.isfinite(dx).all() + for name, grad in grads.items(): + assert torch.isfinite(grad).all() + if name.startswith("fc2"): # FC2 gradients do not depend on dgrad. + torch.testing.assert_close(grad, grads_nvfp4[name], rtol=0, atol=0) + + +@requires_nvfp4_and_mxfp8 +def test_dgrad_mxfp8_rejects_quantized_primary_weights(): + with pytest.raises(ValueError, match="dgrad_mxfp8"): + with te.quantized_model_init(enabled=True, recipe=make_recipe(dgrad_mxfp8=True)): + make_module("linear") + + # NVFP4 primary weights provide no MXFP8 columnwise data for dgrad. + with te.quantized_model_init(enabled=True, recipe=make_recipe()): + module = make_module("linear") + x, _ = make_inputs() + with pytest.raises(RuntimeError, match="dgrad_mxfp8"): + with te.autocast(enabled=True, recipe=make_recipe(dgrad_mxfp8=True)): + module(x) + + +BACKWARD_OPTIONS = { + "nvfp4": {}, + "row_scaled_activation": {"row_scaled_activation": True}, + "4over6_grad": {"nvfp4_4over6_grad": True}, + "dgrad_mxfp8": {"dgrad_mxfp8": True}, + "dgrad_mxfp8_4over6_grad_row_scaled_activation": { + "dgrad_mxfp8": True, + "nvfp4_4over6_grad": True, + "row_scaled_activation": True, + }, +} + + +def zero_mask(pattern: str) -> torch.Tensor: + mask = torch.zeros(NUM_TOKENS, OUT_FEATURES, dtype=torch.bool, device="cuda") + if pattern == "prompt_rows": # Loss-masked prompt tokens at the start of each sequence + for start in range(0, NUM_TOKENS, 128): + mask[start : start + 96] = True + elif pattern == "sparse_rows": # Few tokens carry gradient + mask[:] = True + mask[::37] = False + elif pattern == "blocks": # All-zero 16-element blocks in both directions + mask[:, 16:48] = True + mask[32:64, :] = True + elif pattern == "all": # Micro-batch with zero advantages + mask[:] = True + else: + raise ValueError(f"Unknown zero-mask pattern: {pattern}") + return mask + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("options", list(BACKWARD_OPTIONS)) +@pytest.mark.parametrize("pattern", ["prompt_rows", "sparse_rows", "blocks", "all"]) +@pytest.mark.parametrize("kind", ["linear", "grouped_linear"]) +def test_zero_grad_output(options, pattern, kind): + x, dy = make_inputs() + mask = zero_mask(pattern) + dy = dy.masked_fill(mask, 0) + _, dx, grads = run_step(make_module(kind), make_recipe(**BACKWARD_OPTIONS[options]), x, dy) + assert torch.isfinite(dx).all() + assert all(torch.isfinite(grad).all() for grad in grads.values()) + zero_rows = mask.all(dim=1) + assert (dx[zero_rows] == 0).all() + if pattern == "all": + assert all((grad == 0).all() for grad in grads.values()) + else: + assert (dx[~zero_rows] != 0).any() + + +@requires_nvfp4_and_mxfp8 +@pytest.mark.parametrize("options", list(BACKWARD_OPTIONS)) +def test_zero_input_rows(options): + """Padding tokens produce all-zero input rows.""" + x, dy = make_inputs() + x[::3] = 0 + y, dx, grads = run_step(make_module("linear"), make_recipe(**BACKWARD_OPTIONS[options]), x, dy) + assert torch.isfinite(y).all() + assert torch.isfinite(dx).all() + assert all(torch.isfinite(grad).all() for grad in grads.values()) + assert (y[::3] == 0).all() diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 128e8280bbe..8471fc0e169 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -547,6 +547,20 @@ class NVFP4BlockScaling(Recipe): to keep the standard NVFP4 448 bound for 4over6 tensors. nvfp4_4over6_err_mode : {'MAE', 'MSE'}, default = 'MAE' Error metric used by NVFP4 4over6 candidate selection. + nvfp4_4over6_grad : bool, default = False + If set to `True`, gradient quantizers also use 4over6 adaptive block + scaling. Requires `disable_rht=True` and + `disable_stochastic_rounding=True`. Gradient 4over6 tensors use 256 as + the global E4M3 scale bound only when ``nvfp4_4over6_e4m3_use_256`` is + ``'all'``; its ``'activations'`` scope covers forward activations only. + dgrad_mxfp8 : bool, default = False + If set to `True`, the dgrad GEMM uses MXFP8 operands while the fprop + and wgrad GEMMs keep NVFP4 operands. The weight quantizer emits NVFP4 + rowwise data for fprop and MXFP8 columnwise data for dgrad, requantized + from the dequantized NVFP4 weight. The ``grad_output`` quantizer emits + MXFP8 rowwise data for dgrad and NVFP4 columnwise data for wgrad. + Cannot be combined with `backward_override` or with quantized primary + weights (`quantized_model_init`). backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, `high_precision` keeps original high-precision operands for backward, @@ -564,6 +578,8 @@ class NVFP4BlockScaling(Recipe): nvfp4_4over6: str = os.getenv("NVTE_NVFP4_4OVER6", "none") nvfp4_4over6_e4m3_use_256: str = os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") nvfp4_4over6_err_mode: str = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").upper() + nvfp4_4over6_grad: bool = os.getenv("NVTE_NVFP4_4OVER6_GRAD", "0") == "1" + dgrad_mxfp8: bool = os.getenv("NVTE_NVFP4_DGRAD_MXFP8", "0") == "1" fp4_format: Format = Format.E2M1 fp8_format: Format = Format.E4M3 @@ -588,6 +604,9 @@ def __post_init__(self) -> None: assert ( self.nvfp4_4over6_err_mode in _NVFP4_4OVER6_ERR_MODES ), "NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'." + assert ( + not self.dgrad_mxfp8 or self.backward_override is None + ), "NVTE_NVFP4_DGRAD_MXFP8 cannot be combined with NVTE_BACKWARD_OVERRIDE." # Quantization params # Note: RHT is currently only applied to column-wise usage so that @@ -620,6 +639,8 @@ def _make_repr(self) -> str: f"nvfp4_4over6={self.nvfp4_4over6}, " f"nvfp4_4over6_e4m3_use_256={self.nvfp4_4over6_e4m3_use_256}, " f"nvfp4_4over6_err_mode={self.nvfp4_4over6_err_mode}, " + f"nvfp4_4over6_grad={self.nvfp4_4over6_grad}, " + f"dgrad_mxfp8={self.dgrad_mxfp8}, " f"fp4_quant_fwd_inp={self.fp4_quant_fwd_inp}, " f"fp4_quant_fwd_weight={self.fp4_quant_fwd_weight}, " f"fp4_quant_bwd_grad={self.fp4_quant_bwd_grad}, " diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 0b5ab7bd92b..0b2ab8041bb 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -2062,6 +2062,12 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] + if recipe.nvfp4() and recipe.dgrad_mxfp8: + # Quantized primary weights provide no MXFP8 columnwise data for dgrad. + raise RuntimeError( + "NVFP4BlockScaling(dgrad_mxfp8=True) does not support quantized primary weights" + " from quantized_model_init()." + ) weight_tensors = [getattr(self, name) for name in self.weight_names] for i, tensor in enumerate(weight_tensors): if isinstance(tensor, QuantizedTensorStorage): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e403630fd0a..156ab1fa523 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -148,6 +148,7 @@ def is_module_grouped_tensor_path_supported( device_capability >= (10, 0) and not recipe.disable_rht and not recipe.row_scaled_activation + and not recipe.dgrad_mxfp8 ) return False diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index e2a7fa921a6..ce456e8401a 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1009,6 +1009,8 @@ def _layernorm_linear_backward_impl( # Prepare grad output tensor # Note: Cast to expected dtype and perform tensor-parallel communication + # Keep the original gradient for the cuBLASMp wgrad re-gather. + grad_output_arg = grad_output nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") ( grad_output, @@ -1198,6 +1200,10 @@ def _layernorm_linear_backward_impl( ): if args.grad_output_quantizer is not None: set_quantizer_usage_for_wgrad_all_gather(args.grad_output_quantizer) + if isinstance(args.grad_output_quantizer, HybridQuantizer): + # Quantize the original gradient so the columnwise operand is not + # requantized from the rowwise data. + grad_output = grad_output_arg.reshape(-1, grad_output_arg.shape[-1]).contiguous() grad_output, _ = gather_along_first_dim( grad_output, args.tp_group, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ad5864d4a60..293be5ee878 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1712,11 +1712,12 @@ def fc2_wgrad_gemm( ) # activation in high precision if args.fp8: - # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now + # TODO: float8 blockwise current scaling, identity, hybrid, and custom + # quantizers have no bgrad fusion for now if ( isinstance( args.fc1_grad_output_quantizer, - (Float8BlockQuantizer, IdentityQuantizer), + (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer), ) or args.fp8_recipe.custom() ): diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7f00e19644a..b0d5abe807a 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -91,6 +91,7 @@ is_value_opaque_quantizer, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer +from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom from ..export import is_in_onnx_export_mode, assert_warmed_up @@ -1471,8 +1472,9 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # through the same FP8-aware all-gather as the non-overlap path in # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing the # grad_output quantizer. Per-tensor FP8 can reconstruct columnwise - # data from the gathered rowwise data; MXFP8 must instead quantize - # the original gradient columnwise to avoid double quantization. + # data from the gathered rowwise data; MXFP8 and hybrid quantizers + # must instead quantize the original gradient columnwise to avoid + # double quantization. if ( bwd_args.requires_wgrad and bwd_args.ub_overlap_ag @@ -1481,7 +1483,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None ): if grad_output_quantizer is not None: set_quantizer_usage_for_wgrad_all_gather(grad_output_quantizer) - if isinstance(grad_output_quantizer, MXFP8Quantizer): + if isinstance(grad_output_quantizer, (MXFP8Quantizer, HybridQuantizer)): grad_output = grad_output_arg.reshape(-1, grad_output_arg.shape[-1]).contiguous() grad_output, _ = gather_along_first_dim( grad_output, diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 95516500450..7a0f87c9d09 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -19,6 +19,7 @@ from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore +from ...module._split_quantization import _split_quantize from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -128,6 +129,7 @@ def is_op_fuser_grouped_tensor_path_supported( device_capability >= (10, 0) and not recipe.disable_rht and not recipe.row_scaled_activation + and not recipe.dgrad_mxfp8 ) return False @@ -1548,7 +1550,7 @@ def _fuser_backward_split_quantize( rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad, ) - dys = tex.split_quantize(dy, split_sizes_int, ctx.grad_output_quantizers) + dys, _ = _split_quantize(dy, split_sizes_int, ctx.grad_output_quantizers, ctx.dtype) else: dys = torch.split(dy, split_sizes_int) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 5462b0d2198..6db670d3e53 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -892,6 +892,9 @@ def fuse_grouped_mlp_ops( # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. if recipe.nvfp4() and recipe.disable_rht: return ops + # The fused NVFP4 grouped MLP does not support hybrid MXFP8 dgrad operands. + if recipe.nvfp4() and recipe.dgrad_mxfp8: + return ops # The fused MXFP8 backward reinterprets the grad output's storage as E4M3, so an E5M2 # backward format would have its gradients misread rather than converted. This declines # MXFP8 with Format.HYBRID. fp8_format does not describe NVFP4 gradients, so NVFP4 is diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 34ae20b4984..e09db4fc771 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -923,12 +923,18 @@ def quantized_model_init( This functionality is *EXPERIMENTAL*. """ + recipe = get_default_fp8_recipe() if recipe is None else recipe + if enabled and recipe.nvfp4() and recipe.dgrad_mxfp8: + raise ValueError( + "quantized_model_init does not support NVFP4BlockScaling(dgrad_mxfp8=True)." + ) + qstate = FP8GlobalStateManager.quantization_state _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe _high_precision_init_val = qstate.high_precision_init_val qstate.fp8_parameters = enabled - qstate.fp8_recipe = get_default_fp8_recipe() if recipe is None else recipe + qstate.fp8_recipe = recipe qstate.high_precision_init_val = preserve_high_precision_init_val try: yield @@ -1668,7 +1674,7 @@ def __init__( device = torch.device("cuda") def make_quantizers(self) -> list: - """Build one ``NVFP4Quantizer`` per slot, dispatched by tensor type. + """Build one quantizer per slot, dispatched by tensor type. Per-slot behavior, resolved via :meth:`RecipeState._slot_tensor_type`: @@ -1677,9 +1683,11 @@ def make_quantizers(self) -> list: ``recipe.fp4_quant_fwd_inp``. * ``"grad_output"`` / ``"grad_input"`` -> ``recipe.fp4_quant_bwd_grad``. * NVFP4 4over6 is applied to non-gradient slots selected by - ``recipe.nvfp4_4over6``. Gradient slots always use standard NVFP4, - which lets gradient RHT and stochastic rounding follow the base - recipe. + ``recipe.nvfp4_4over6`` and to gradient slots if + ``recipe.nvfp4_4over6_grad`` is set. + * If ``recipe.dgrad_mxfp8`` is set, the ``"weight"`` and + ``"grad_output"`` slots use ``HybridQuantizer`` instances whose dgrad + operands (weight columnwise, grad_output rowwise) are MXFP8. When the owning module/op provides a role list via ``get_quantizer_roles``, the per-slot ``tensor_type`` drives dispatch. @@ -1688,6 +1696,8 @@ def make_quantizers(self) -> list: layout slot ``idx % 3 == 1`` is always weight and the rest fall into the input config, matching the legacy index-based behavior. """ + from .tensor.hybrid_tensor import HybridQuantizer + from .tensor.mxfp8_tensor import MXFP8Quantizer from .tensor.nvfp4_tensor import NVFP4Quantizer def _qparams(tensor_type: str): @@ -1707,6 +1717,8 @@ def _make(tensor_type: str) -> NVFP4Quantizer: nvfp4_use_4over6 = tensor_type == "weight" elif self.recipe.nvfp4_4over6 == "activations": nvfp4_use_4over6 = tensor_type != "weight" + else: + nvfp4_use_4over6 = self.recipe.nvfp4_4over6_grad nvfp4_e4m3_max = 448 if nvfp4_use_4over6: # Current 4over6 kernels target RL and post-training quantization paths. @@ -1723,7 +1735,7 @@ def _make(tensor_type: str) -> NVFP4Quantizer: if tensor_type == "weight": nvfp4_e4m3_max = 256 elif self.recipe.nvfp4_4over6_e4m3_use_256 == "activations": - if tensor_type != "weight": + if tensor_type not in ("weight", "grad_output", "grad_input"): nvfp4_e4m3_max = 256 elif self.recipe.nvfp4_4over6_e4m3_use_256 == "none": nvfp4_e4m3_max = 448 @@ -1745,10 +1757,29 @@ def _make(tensor_type: str) -> NVFP4Quantizer: nvfp4_4over6_err_mode=self.recipe.nvfp4_4over6_err_mode, ) + def _make_slot(tensor_type: str) -> Union[NVFP4Quantizer, HybridQuantizer]: + quantizer = _make(tensor_type) + if not self.recipe.dgrad_mxfp8: + return quantizer + if tensor_type == "weight": + # Derive the MXFP8 dgrad weight from the quantized NVFP4 fprop weight. + return HybridQuantizer( + rowwise_quantizer=quantizer, + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), + columnwise_source="rowwise_dequantized", + ) + if tensor_type == "grad_output": + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), + columnwise_quantizer=quantizer, + columnwise_source="original", + ) + return quantizer + if self.mode not in ("forward", "backward"): raise RuntimeError(f"Unexpected recipe mode ({self.mode})") - return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] + return [_make_slot(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] def _handle_delayed_scaling_requests(