[vincentzed]

Description

NVFP4BlockScaling supports NVFP4 backward GEMMs, or higher-precision backward GEMMs through backward_override.
Some RL post-training setups (i.e a possible recipe in Miles) use the latter to keep gradient quantization error below the noise of small policy grad updates.

We already know that the error contrib is not equal:

• dgrad (dX = dY W) error propagates into every earlier layer.
• wgrad (dW = dY^T X) error stays in one weight gradient.

Meeting this requirement in the middle, we present:

• MXFP8 operands for dgrad, with NVFP4 operands for fprop and wgrad. It also lets
  gradient quantizers use 4over6 (we found this to be beneficial), which is currently limited to weights and activations.

• Both options are recipe fields with env vars. Frameworks that construct a default NVFP4BlockScaling, such as
  Megatron-Core, can opt in from the environment, the same way they use NVTE_NVFP4_ROW_SCALED_ACTIVATION and
  NVTE_BACKWARD_OVERRIDE. Megatron-Core cannot pass a CustomRecipe yet (Add format-neutral TE custom recipe support (--custom-recipe) Megatron-LM#7176).

The code style should be relatively aligned w/ the TE requirements.

Codex below>

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• NVFP4BlockScaling.dgrad_mxfp8 / NVTE_NVFP4_DGRAD_MXFP8.
  • NVFP4BlockScalingRecipeState.make_quantizers returns a HybridQuantizer for two slots:
    • weight: NVFP4 rowwise (fprop) plus MXFP8 columnwise (dgrad), requantized from the dequantized NVFP4 weight
      (columnwise_source="rowwise_dequantized").
    • grad_output: MXFP8 rowwise (dgrad) plus NVFP4 columnwise (wgrad).
  • Input quantizers are unchanged, including row-scaled activations.
  • The NVFP4 wgrad operand keeps the recipe's gradient RHT, stochastic rounding, and nvfp4_4over6_grad settings.
• NVFP4BlockScaling.nvfp4_4over6_grad / NVTE_NVFP4_4OVER6_GRAD.
  • Gradient quantizers use 4over6.
  • Like the existing 4over6 scopes, the gradient quantizers must have RHT and stochastic rounding disabled.
  • The E4M3 bound is 256 when nvfp4_4over6_e4m3_use_256 is "all" and 448 otherwise. "activations" keeps
    meaning forward activations.
  • nvfp4_4over6="all" still leaves gradients unchanged unless this flag is set.
• Validation.
  • dgrad_mxfp8 with backward_override fails in __post_init__.
  • quantized_model_init rejects dgrad_mxfp8 recipes, because distributed master-weight updates reject hybrid
    weights whose columnwise data is derived from rowwise data.
  • Quantized NVFP4 primary weights used under a dgrad_mxfp8 autocast are rejected in the module's weight/recipe
    check, because they carry no MXFP8 columnwise data for dgrad.
• Paths that assume uniform NVFP4 layouts skip this mode:
  • the module and op-fuser graph-safe grouped GEMM predicates;
  • the fused grouped MLP.
  • Both GroupedLinear implementations use their per-group paths, as they already do for row_scaled_activation.
• Hybrid gradient quantizers in existing paths. These also affect CustomRecipe hybrids:
  • LayerNormMLP takes the unfused bias-gradient path for HybridQuantizer, since tex.bgrad_quantize does not
    accept it.
  • te.ops.GroupedLinear backward splits and quantizes grad_output with the hybrid-aware grouped split helper
    instead of tex.split_quantize, which cannot convert hybrid quantizers.
  • The cuBLASMp wgrad re-gather in Linear and LayerNormLinear quantizes the original gradient for
    HybridQuantizer, as Linear already does for MXFP8. Previously it requantized the columnwise operand from
    dequantized rowwise data.
• Docs:
  • docs/envvars.rst entries for both env vars;
  • the NVTE_NVFP4_4OVER6_E4M3_USE_256 description now covers gradients;
  • recipe and make_quantizers docstrings.
• Tests: tests/pytorch/nvfp4/test_nvfp4_backward_options.py has 81 tests.
  • Env-var parsing (in a subprocess) and the backward_override conflict. These run without FP4 hardware.
  • Gradient-4over6 RHT/SR rejection and quantizer layouts for both options.
  • Independent references. te.Linear is compared against FP32 GEMMs on dequantized operands, with tolerances
    set by BF16 output rounding:
    • dgrad_mxfp8: fprop and wgrad on NVFP4 operands; dgrad on MXFP8 of the dequantized NVFP4 weight.
      Negative controls asserted: NVFP4 dgrad, and an MXFP8 weight quantized from the original weight.
    • nvfp4_4over6_grad: dgrad and wgrad on 4over6 gradients. Negative control asserted: standard NVFP4
      gradients.
  • Dispatch parity. Bitwise equality with the equivalent CustomRecipe for:
    • Linear, LayerNormLinear and GroupedLinear, with and without bias and row-scaled activations;
    • fuse_wgrad_accumulation (main_grad), delay_wgrad_compute, save_original_input,
      is_first_microbatch weight caching, and GroupedLinear(single_grouped_weight=True) for both the module and the op;
    • te.ops.BasicLinear and te.ops.GroupedLinear.
  • Against full NVFP4:
    • the forward and wgrad are bitwise unchanged (the LayerNormLinear norm-weight gradient depends on dgrad, so it
      is excluded);
    • the dgrad error against the backward_override="dequantized" gradient is under half.
  • LayerNormMLP (gelu, swiglu, with and without bias):
    • the forward is bitwise unchanged;
    • FC2 gradients are bitwise unchanged;
    • all gradients are finite.
  • Both quantized-primary-weight rejections.
  • All-zero gradients, as RL losses produce for loss-masked prompt and padding tokens and for zero-advantage
    micro-batches:
    • patterns: leading zero rows, sparse nonzero rows, zero 16-element blocks in both directions, and an all-zero
      grad_output;
    • run on Linear and GroupedLinear with five recipe configurations, plus all-zero input rows;
    • checks: gradients are finite, dgrad is exactly zero on zero rows, and weight gradients are exactly zero for an
      all-zero grad_output.

Numerics on RL gradients

These numbers are not part of the test suite.

Setup:

• Model: Qwen3-0.6B after 300 GRPO steps on Countdown, with every decoder linear replaced by te.Linear.
• Batch: one GRPO batch of 256 sequences (32 prompts x 8 samples, group-normalized ±1 rewards).
• Forward: row-scaled NVFP4 activations and 1D NVFP4 weights, with no RHT/SR/2D.
• Error: distance of the full-batch weight gradient from the backward_override="dequantized" gradient.
• R: that error divided by the policy gradient's sampling noise (half-batch split). Sampling noise is 0.85 of the
  gradient norm for the RL loss and 0.096 for an SFT loss on the same tokens.
• Run-to-run spread is about 2%, because the BF16 attention backward is nondeterministic.

backward	RL rel. err	RL R	SFT rel. err	SFT R
BF16 on original operands (backward_override="high_precision")	0.266	0.31	0.154	1.60
MXFP8 dgrad + wgrad (zoo nvfp4_row_scaled_fwd_mxfp8_bwd_factory)	0.112	0.13	0.046	0.47
dgrad_mxfp8=True	0.206	0.24	0.107	1.11
NVFP4 dgrad + wgrad	0.454	0.53	0.230	2.39
nvfp4_4over6_grad=True (bound 448 / 256)	0.440 / 0.441	0.52	0.235 / 0.236	2.44
dgrad_mxfp8=True, nvfp4_4over6_grad=True	0.208	0.25	0.118	1.23
For reference: NVFP4 with row-scaled grad_output (not in this PR)	0.455	0.54	0.230	2.39

• dgrad_mxfp8 halves the NVFP4 backward error.
• Gradient 4over6 is roughly neutral:
  • the RL error is 3% lower, which is close to run-to-run spread;
  • the SFT error is 2–3% higher;
  • it adds nothing on top of dgrad_mxfp8.
• Row scaling of grad_output does not change the error. The gradient error comes from 4-bit element resolution
  within 16-element blocks, not from the per-tensor scale range.

Performance

These numbers are not part of the test suite.

Setup:

• Measurement: kernel time of quantize(dY) + dgrad + wgrad per linear per step.
• Hardware and software: 1x B300 (sm_103), TE 2.19 with this change, CUDA 13.0, PyTorch 2.13.
• Timing: flashinfer.testing.bench_gpu_time_with_cupti with CUDA graphs and cold L2; median of 50 iterations after
  10 warm-up iterations.
• Operands: random normal (x: std 0.5, W: std 0.02, dY: std 1e-3); RL recipe settings (no RHT/SR/2D).
• Excluded:
  • X columnwise quantization, which happens in the forward;
  • W quantization, which happens once per optimizer step.
• Script: in the collapsed section below.

Times are in ms.

Benchmark script

"""Backward cost of NVFP4BlockScaling(dgrad_mxfp8=True) vs full NVFP4, MXFP8 and BF16 backward on B300.

Per linear and per step, the backward runs quantize(dY) + dgrad + wgrad. The quantizers come from the recipe
(this PR; RL settings: no RHT/SR/2D, row-scaled activations):
  bf16        dgrad = dY W, wgrad = dY^T X in BF16
  mxfp8       Composer 2 / TE zoo: MXFP8 dY (both directions), MXFP8 dgrad + wgrad
  nvfp4       NVFP4 dY (both directions), NVFP4 dgrad + wgrad
  dgrad_mxfp8 hybrid dY (MXFP8 rowwise + NVFP4 columnwise), MXFP8 dgrad, NVFP4 wgrad
X columnwise (for wgrad) is quantized in the forward pass and W columnwise once per optimizer step; both are
reported separately. Timing: flashinfer bench_gpu_time_with_cupti, CUDA graph, cold L2.
Shapes: Qwen3-8B dense MLP at 16k tokens; Qwen3-30B-A3B routed expert at 1k and 8k tokens per expert;
DeepSeek-V3 routed expert at 8k tokens per expert.
Writes results/bench_dgrad_mxfp8.json.
"""

import json
import statistics
from pathlib import Path

import torch  # noqa: I001
import transformer_engine.pytorch  # noqa: F401
import transformer_engine_torch as tex
from flashinfer.testing.utils import bench_gpu_time_with_cupti
from transformer_engine.common.recipe import NVFP4BlockScaling
from transformer_engine.pytorch.cpp_extensions import general_gemm
from transformer_engine.pytorch.quantization import RecipeState
from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer

ROOT = Path(__file__).resolve().parents[1]
SHAPES = {  # name: (tokens, in_features, out_features)
    "qwen3-8b gate_up 16k": (16384, 4096, 24576),
    "qwen3-8b down 16k": (16384, 12288, 4096),
    "qwen3-30b-a3b expert fc1 1k": (1024, 2048, 1536),
    "qwen3-30b-a3b expert fc2 1k": (1024, 768, 2048),
    "qwen3-30b-a3b expert fc1 8k": (8192, 2048, 1536),
    "qwen3-30b-a3b expert fc2 8k": (8192, 768, 2048),
    "deepseek-v3 expert fc1 8k": (8192, 7168, 4096),
    "deepseek-v3 expert fc2 8k": (8192, 2048, 7168),
}
RL = dict(disable_rht=True, disable_stochastic_rounding=True, disable_2d_quantization=True, row_scaled_activation=True)


def time_ms(fn):
    return statistics.median(bench_gpu_time_with_cupti(fn, use_cuda_graph=True, cold_l2_cache=True,
                                                       repeat_iters=50, dry_run_iters=10))


def recipe_quantizers(recipe):
    inp, weight, _ = RecipeState.create(recipe, mode="forward", num_quantizers=3).make_quantizers()
    grad_output, _ = RecipeState.create(recipe, mode="backward", num_quantizers=2).make_quantizers()
    return inp, weight, grad_output


def main():
    torch.manual_seed(0)
    out = {}
    for sname, (t, k, n) in SHAPES.items():
        x = (torch.randn(t, k, device="cuda") * 0.5).to(torch.bfloat16)
        w = (torch.randn(n, k, device="cuda") * 0.02).to(torch.bfloat16)
        dy = (torch.randn(t, n, device="cuda") * 1e-3).to(torch.bfloat16)
        dx_out = torch.empty(t, k, device="cuda", dtype=torch.bfloat16)
        dw_out = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
        res = {"bf16": {"dgrad_ms": time_ms(lambda: general_gemm(w, dy, out_dtype=torch.bfloat16, layout="NN", grad=True, out=dx_out)),
                        "wgrad_ms": time_ms(lambda: general_gemm(x, dy, out_dtype=torch.bfloat16, layout="NT", grad=True, out=dw_out))}}

        mq = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)
        wq, xq, dyq = mq(w), mq(x), mq(dy)
        res["mxfp8"] = {"quant_dy_ms": time_ms(lambda: mq.update_quantized(dy, dyq)),
                        "quant_x_col_ms": time_ms(lambda: mq.update_quantized(x, xq)),
                        "dgrad_ms": time_ms(lambda: general_gemm(wq, dyq, out_dtype=torch.bfloat16, layout="NN", grad=True, out=dx_out)),
                        "wgrad_ms": time_ms(lambda: general_gemm(xq, dyq, out_dtype=torch.bfloat16, layout="NT", grad=True, out=dw_out))}

        _, wq4, gq4 = recipe_quantizers(NVFP4BlockScaling(**RL))
        xq4 = type(gq4)(rowwise=True, columnwise=True, with_rht=False, with_post_rht_amax=False,
                        with_2d_quantization=False, stochastic_rounding=False)
        w4, x4, dy4 = wq4(w), xq4(x), gq4(dy)
        res["nvfp4"] = {"quant_dy_ms": time_ms(lambda: gq4.update_quantized(dy, dy4)),
                        "quant_x_col_ms": time_ms(lambda: xq4.update_quantized(x, x4)),
                        "dgrad_ms": time_ms(lambda: general_gemm(w4, dy4, out_dtype=torch.bfloat16, layout="NN", grad=True, out=dx_out)),
                        "wgrad_ms": time_ms(lambda: general_gemm(x4, dy4, out_dtype=torch.bfloat16, layout="NT", grad=True, out=dw_out))}

        _, wqh, gqh = recipe_quantizers(NVFP4BlockScaling(dgrad_mxfp8=True, **RL))
        wh, dyh = wqh(w), gqh(dy)
        w_col, dy_row, dy_col = wh.columnwise_sub_storage, dyh.rowwise_sub_storage, dyh.columnwise_sub_storage
        res["dgrad_mxfp8"] = {"quant_dy_ms": time_ms(lambda: gqh.update_quantized(dy, dyh)),
                              "quant_x_col_ms": res["nvfp4"]["quant_x_col_ms"],
                              "quant_w_step_ms": time_ms(lambda: wqh.update_quantized(w, wh)),
                              "dgrad_ms": time_ms(lambda: general_gemm(w_col, dy_row, out_dtype=torch.bfloat16, layout="NN", grad=True, out=dx_out)),
                              "wgrad_ms": time_ms(lambda: general_gemm(x4, dy_col, out_dtype=torch.bfloat16, layout="NT", grad=True, out=dw_out))}
        res["nvfp4"]["quant_w_step_ms"] = time_ms(lambda: wq4.update_quantized(w, w4))
        for cfg, r in res.items():
            r["bwd_ms"] = r.get("quant_dy_ms", 0.0) + r["dgrad_ms"] + r["wgrad_ms"]
        out[sname] = res
        print(sname, {c: round(r["bwd_ms"], 4) for c, r in res.items()},
              {c: {kk: round(v, 4) for kk, v in r.items()} for c, r in res.items()}, flush=True)
    (ROOT / "results/bench_dgrad_mxfp8.json").write_text(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()

shape (tokens, in, out)	BF16	MXFP8 bwd	NVFP4 bwd	dgrad_mxfp8
Qwen3-8B MLP gate_up (16384, 4096, 24576)	4.62	2.84	1.50	2.38
Qwen3-8B MLP down (16384, 12288, 4096)	2.33	1.34	0.76	1.07
DeepSeek-V3 expert fc1 (8192, 7168, 4096)	0.739	0.420	0.253	0.351
DeepSeek-V3 expert fc2 (8192, 2048, 7168)	0.399	0.265	0.182	0.246
Qwen3-30B-A3B expert fc1 (8192, 2048, 1536)	0.068	0.094	0.062	0.083
Qwen3-30B-A3B expert fc2 (8192, 768, 2048)	0.035	0.060	0.054	0.064

• At the larger shapes, dgrad_mxfp8 is 1.1–1.25x faster than an MXFP8 backward and 1.6–2.2x faster than BF16.
• At Qwen3-30B-A3B expert size, every quantized backward is slower than BF16, because quantization launches and
  small-GEMM latency dominate.
• The hybrid grad_output quantization is two casts that each read dY: 0.53 ms vs 0.35 ms for the single NVFP4 cast
  at gate_up. A fused MXFP8-rowwise + NVFP4-columnwise cast would remove that overhead; it is left as follow-up.

Testing

B300 (sm_103), CUDA 13.0, PyTorch 2.13, cuDNN 9.20, driver 615.71.09, single GPU.

main (9c33764 + this change): editable build with NVTE_CUDA_ARCHS="100a;103a" and NVTE_WITH_NCCL_EP=0; the
container's NCCL 2.29 lacks the NCCL EP APIs.

suite	with this change	unmodified main
tests/pytorch/nvfp4/test_nvfp4_backward_options.py	81 passed	–
tests/pytorch/nvfp4	9384 passed	–
test_backward_override.py	5141 passed	–
test_fusible_ops.py	2388 passed	–
test_sanity.py	17320 passed	–
test_grouped_linear.py	1397 passed	–
test_identity_quantizer.py	96 passed, 2 xfailed	–
test_recipe.py	145 passed	–
test_weight_swizzle_in_layers.py	52 passed	–
test_grouped_mlp.py -k 'not prequantized'	457 passed, 8 failed	same 8 failures
test_grouped_mlp.py (full)	an illegal memory access in test_grouped_linear_prequantized_mxfp8_grad cascades into 387 errors	identical
test_custom_recipe.py + test_hybrid_quantization.py	359 passed, 4 MXFP8-DPA failures (no backend on this system)	same failures
test_checkpoint.py	11 failed: checkpoint artifacts not present	same
qa/L0_license/test.sh	OK
qa/L0_pytorch_lint/test.sh (PYTHON_ONLY=1)	10.00/10; the two remaining messages are in fused_mla_q_uproj.py

v2.19 backport: installed TE 2.19 with the patched Python files.

• The new tests pass (81), tests/pytorch/nvfp4 passes (9380), and test_grouped_mlp -k 'not prequantized' passes
  (438).
• test_backward_override, test_grouped_linear and test_recipe pass.
• The same 5 pre-existing MXFP8-DPA and debug-tool failures occur with and without the change.
• test_fusible_ops.py does not import in that image (libtorch_cuda.so: undefined symbol ncclCommResume).

Limitations

• Userbuffers all-gather of grad_output is not supported. fill_userbuffers_buffer_for_all_gather handles only
  FP8 and MXFP8 quantizers, which is already the case for NVFP4 gradient quantizers. The failure is an explicit
  Unsupported quantizer for Userbuffers error.
• The cuBLASMp re-gather change is untested. It needs a cuBLASMp build and multiple GPUs.
• Not exercised: FSDP2, CPU offloading, CUDA graphs, debug mode, and multi-GPU tensor/sequence parallelism.
• Quantized primary weights are not supported with dgrad_mxfp8; this is an explicit error. This includes Megatron-Core's
  TEFusedMLPWithGroupedLinear dense MLP, which creates its weights under quantized_model_init.
• Runtime toggles do not take effect on an already-built module. Changing dgrad_mxfp8 or nvfp4_4over6_grad
  between steps on a module that already built its quantizers does not rebuild them. The same holds for every other
  NVFP4BlockScaling field; [PyTorch] Support dynamic recipe updates during training #3343 addresses recipe changes in general.
• The per-token recipe of Add NVFP4 per-token quantization recipe #3045 is not covered. That recipe excludes 4over6 and has its own GEMM path.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete (see Limitations)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes (the failures listed under Testing also occur on unmodified main)

[negvet]

Rebased NVIDIA/Megatron-LM#7176 and opened it for review (it is functional).
Even without it, you can probably run something in mcore already now. --fp4/8-recipe flag supports custom. Just make sure to also set --fp4/8-format to e.g. e2m1 so that the path is activated. 7176 is doing it in a better way, though.