[PyTorch] MXFP4 weight QAT on MXFP8 and FP8 block-scaling recipes

[xiuhu17]

Description

Weight-only MXFP4 quantization-aware training: weights are projected onto the MXFP4 grid (E2M1 payloads, 1x32 power-of-two UE8M0 scales) before the host recipe quantizes them, so training sees exactly the deployment weight while activations and gradients keep the base recipe untouched.

The only new numerics is the fused projection bf16/fp32 -> MXFP4 grid -> bf16/fp32 (tex.mxfp4_fake_quantize: per 1x32 block, scale 2^clamp(ceil(log2(amax/6)), -126, 125) derived from the amax bit pattern, RTNE onto E2M1, straight-through-estimator gradient). All MXFP8 and 128x128 blockwise FP8 encoding/decoding of the projected weight runs through the existing quantizer pipeline unchanged. Since every projected value is a grid point p * 2^e, the existing encodings are decoded-value exact: MXFP8 rowwise unconditionally, 128x128 blockwise for tile scale spreads up to 2^14.

Recipes: MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling) and MXFP4QATFloat8BlockScaling(Float8BlockScaling); QAT can also be enabled on the stock host recipes via the mxfp4_qat_weights field (NVTE_MXFP4_QAT=1). backward_override keeps its base-recipe semantics. Also fixes ptx::exp2f(e8m0) to decode UE8M0 code 0 as 2^-127 and code 255 as NaN.

Test coverage

All suites run on B200 (CUDA 13.0 build, cuBLAS 13.0; 2 ranks for the distributed suites). Skips are environment gates (perf benchmark opt-in; graph-safe grouped GEMM and the CuTe DSL fused grouped MLP need cuBLAS 13.3+).

Level	Suite	What it verifies	Result
Kernel precision	tests/pytorch/test_mxfp4_qat.py	Fused bf16/fp32→MXFP4-grid→bf16/fp32 kernel bitwise vs composite-torch reference and an independent fp64 RTNE oracle: all 65536 bf16 patterns, 524k fp32 bit-fuzz, directed edges (±0, subnormal ladder 2^-127…2^-133, 2^-128 RTNE tie, every midpoint, 1.5·2^k±1ulp, satfinite, scale cap, NaN/Inf block poisoning, misaligned views); --use_fast_math rebuild bit-identical; UE8M0 code 0→2^-127 / 255→NaN through real dequantize; MXFP8 rowwise losslessness + columnwise bound; 128×128 blockwise exactness boundary; STE identity; input rejections (CPU/non-2D/inner%32/fp16)	pass
Single GPU e2e	same file	te.Linear/te.GroupedLinear × both QAT recipes × 3 backward overrides × fused/unfused wgrad, dgrad-operand discriminators (projected vs master); weight-workspace caching and microbatch lifecycle; base↔QAT recipe-switch cache invalidation	pass (69 total, 6 env skips across the three single-GPU files)
Single GPU ops	test_ops_grouped_linear_mxfp4_qat.py, test_ops_grouped_mlp_mxfp4_qat.py	te.ops.GroupedLinear QAT None on both backends with a projection-count contract; fused grouped-MLP QAT None bitwise vs projected pure-MXFP8 reference; loud rejection of unsupported overrides and primary-quantized weights	pass
FSDP2 ops	distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py	BasicLinear/Linear all 3 overrides, ops.GroupedLinear None, fused grouped-MLP None under fully_shard; forward/dgrad bitwise vs projected reference, sharded wgrad bitwise vs the reference module	7 pass, 1 env skip
FSDP2 modules	run_fsdp2_mxfp4_qat_modules.py	te.Linear/te.LayerNormLinear/te.GroupedLinear × 3 overrides and te.LayerNormMLP (None) × both QAT recipes under fully_shard, same discriminator scheme	20 pass
Single-GPU vs FSDP2 parity	run_fsdp2_mxfp4_qat_parity.py	Real data parallelism: each rank trains its slice of a deterministic global batch; forward/dgrad must match the full-batch single-GPU run bitwise, wgrad must match a split-sum single-GPU reference (isolates summation-order effects from the distributed implementation)	18 pass
GTP (DistributedWeight)	run_gtp_mxfp4_qat.py	Real multi-rank implementer (dim-0 shards, all-gather materialize, reduce-scatter finalize): te.Linear/te.LayerNormLinear/te.GroupedLinear × 3 overrides × both recipes with protocol-hook call assertions, plus loud rejection of the unsupported GTP surfaces (ops API, single grouped weight, LayerNormMLP, delayed wgrad)	23 pass

Launchers for the four distributed suites live in distributed/test_torch_fsdp2.py (qa L1); the three single-GPU files are wired into qa L0.

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• 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

[greptile-apps]

Greptile Summary

Adds weight-only MXFP4 quantization-aware training on MXFP8 and block-scaled FP8 recipes.

• Adds the fused MXFP4 fake-quantization kernel and PyTorch binding.
• Integrates projected-weight handling into fused modules and fusible operations.
• Adds distributed, grouped-operation, numerical-reference, and recipe coverage.
• Corrects UE8M0 exponent decoding at boundary codes.

Confidence Score: 5/5

The PR appears safe to merge based on the currently established evidence.

No blocking failure remains established; the previously reported cached-weight issue cannot be classified from the available backend evidence.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Integrates MXFP4 projection with cached weight workspaces and refreshes recipe-dependent quantization state.
transformer_engine/pytorch/mxfp4_qat.py	Defines the validated PyTorch MXFP4 fake-quantization API and identity-gradient estimator.
transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh	Implements the fused blockwise MXFP4 projection numerics.
transformer_engine/common/recipe/init.py	Adds MXFP4-QAT recipe variants and recipe configuration support.
transformer_engine/pytorch/ops/basic/basic_linear.py	Applies projected QAT weights in the fusible basic-linear execution path.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  W["BF16/FP32 master weight"] --> QAT["MXFP4 fake projection<br/>E2M1 + 1x32 UE8M0 scale"]
  QAT --> HOST["Existing MXFP8 or FP8 block-scaling quantizer"]
  HOST --> GEMM["Quantized GEMM"]
  GEMM --> BW["Backward pass"]
  BW --> STE["Straight-through gradient"]
  STE --> W

Loading

_{Reviews (8): Last reviewed commit: "[PyTorch] Restyle MXFP4-QAT tests and ad..." | Re-trigger Greptile}

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.