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(