fp4-train

fp4-train is an NVFP4 quantizer with stochastic rounding and error feedback, the primitives for stable 4-bit training, loadable through kernels. Low-bit training is limited less by the matmul than by the quantizer: round-to-nearest is biased, so the quantization error correlates with the signal and the optimizer drifts. This kernel quantizes to NVFP4 (e2m1 values with an e4m3 scale per 16-element block) with stochastic rounding, so the quantization is unbiased, and carries a per-element error-feedback residual so sub-ULP information accumulates instead of vanishing. The packed output is in the two-codes-per-byte layout the Blackwell FP4 tensor cores consume.

Usage

import torch
from kernels import get_kernel

fp4 = get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True)

# NVFP4-stored-weight training: the weight lives in 4-bit and is updated in place.
# Error feedback carries the sub-ULP part of each update, so round-to-nearest no
# longer stalls; this is the regime where SR/EF matter.
ef = fp4.ErrorFeedback(stochastic=False)                      # RTN + error feedback
W = W0.clone()                                                # the NVFP4 weight (bf16 values)
for step in range(num_steps):
    grad = compute_grad(W)                                    # your backward
    W = ef.qdq("layer0.weight", W - lr * grad, step=step)     # apply the update in NVFP4
# stochastic=True instead gives unbiased stochastic rounding (a noisier estimator)

# quantization primitives
packed, scale = fp4.quantize(W, stochastic=True, seed=0, step=global_step)  # NVFP4 + e4m3
Wq = fp4.dequantize(packed, scale, *W.shape)

# quantized-operand linear (full-precision master weights kept by the caller)
x = torch.randn(8, 512, 2048, device="cuda")
Wm = torch.randn(6144, 2048, device="cuda", requires_grad=True)
y = fp4.fp4_linear(x, Wm, stochastic=True, step=global_step)  # NVFP4 operands, STE grad

# native NVFP4 tensor-core matmul on Blackwell (a @ b^T, K a multiple of 16)
a = torch.randn(2048, 4096, device="cuda"); b = torch.randn(4096, 4096, device="cuda")
y = fp4.fp4_mm(a, b)                                          # hardware FP4 path, bf16 out

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

How it works

Each 16-element block gets a scale of block_absmax / 6 encoded to e4m3 (so the largest element maps near the e2m1 maximum of 6). Every element is divided by the decoded scale and encoded to the e2m1 grid {0, .5, 1, 1.5, 2, 3, 4, 6} with a sign bit. Rounding is stochastic: for a value between two grid points, the higher is chosen with probability equal to the fractional distance, using a counter-based (seed, step, global-index) RNG, so the expected quantized value equals the input. With error feedback, the per-element residual x - dequant(quantize(x)) is carried and added into the next quantization, so the time-averaged quantized value equals the true value even under round-to-nearest. Codes are packed two per byte; the block scales are e4m3, the layout NVFP4 tensor cores read.

Correctness

Measured on L4, H200, and RTX PRO 6000 through get_kernel:

• Exact encoding. The quantize/dequantize matches an independent reference implementation to 0 (max abs diff = 0).
• Stochastic rounding is unbiased. For a value 2.3 sitting between the grid points 2 and 3, the stochastic-rounding mean is 2.300 against a true 2.3, where round-to-nearest is biased to 2.000.
• Error feedback cancels the drift. With round-to-nearest plus error feedback the time-averaged quantized value is 2.297, tracking the true 2.3, where plain round-to-nearest stays stuck at 2.000.
• Round-to-nearest stalls low-bit training; stochastic rounding and error feedback do not. With the weights stored in NVFP4 and updated in place, so each step's update is smaller than a quantization step, round-to-nearest discards every update and training stalls, while stochastic rounding (unbiased) and error feedback (residual-carrying) apply the updates and recover near-bf16 loss. A pure accumulation of a sub-ULP increment reaches a true 1.50 under SR (1.50) and EF (1.50) while RTN stays at the starting 0.50. A linear model trained with NVFP4-stored weights reaches MSE 1.0 (EF) and 13 (SR) against a bf16 0.0, where RTN sits at 77; a small MLP reaches loss 0.159 (EF) and 0.171 (SR) against a bf16 0.158, where RTN stalls at 1.03. Error feedback recovers essentially the full-precision loss; stochastic rounding, being a noisier estimator, trails it but still beats RTN by a wide margin.

Scope

The verified contribution is the quantizer that makes 4-bit training stable: correct NVFP4 encoding, unbiased stochastic rounding, and error feedback that preserves sub-ULP information. Their value shows up when the low precision is in the training state itself, meaning weights stored and updated in NVFP4, where the per-step update is below a quantization step: there round-to-nearest stalls and stochastic rounding or error feedback recover near-bf16 loss, as measured above. When instead only the forward matmul operands are quantized and a full-precision master weight is kept, the choice of rounding is second order, since the master carries the sub-ULP information; that regime does not need this kernel's rounding modes and is not where the gain is. Two matmuls are bundled: fp4_matmul takes a numerically exact FP4 path (dequantize, bf16 accumulate) that is portable to every compute-capability-8.0+ GPU, and fp4_mm runs the operands on the Blackwell NVFP4 tensor cores by swizzling the e4m3 block scales into the layout torch._scaled_mm consumes. On an RTX PRO 6000 the hardware path returns a bit-identical result (relative error 0) and is about 10-14x faster than the dequantize path and 1.5-1.7x faster than a bf16 matmul.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ (e4m3 scale via software conversion).
• K a multiple of 16; float32 or bfloat16 input.

References

NVFP4 and OCP microscaling (MXFP) formats; Gupta et al., "Deep Learning with Limited Numerical Precision" (stochastic rounding, 2015); Seide et al., "1-bit SGD" (error feedback, 2014); low-precision training with unbiased quantization.

License

Apache-2.0.