The 4-bitter Lesson: NVFP4 in the RL Loop
Reading the following blog post:
- The 4-bitter Lesson: Balancing Stability and Performance in NVFP4 RL (Ziang Li and the humans& team).
The core challenge: how to run hardware-native 4-bit (NVFP4) quantization inside an asynchronous reinforcement learning (RL) loop without triggering policy divergence or catastrophic gradient instability. The pretraining playbook for low precision does not transfer to RL, and the reason is a bias–variance argument.
1. Why Pretraining Quantization Fails in RL
In standard pretraining, gradient signals are dense and the quantization goal is mostly to avoid bias — usually solved with stochastic rounding, as in the 10-trillion-token NVFP4 pretraining recipe.
RL lives under a different bias–variance tradeoff. The policy-gradient estimator is already extremely noisy: off-policy staleness, advantage estimation, sampled rollouts. Layering quantization noise on top degrades the per-update true policy-gradient signal past the point where the optimizer can correct, which accelerates policy drift and collapses reward.
Takeaway: an RL quantization recipe must prioritize absolute error minimization over merely staying unbiased.
2. Baseline: Per-Token Scales & Targeted MoE Quantization
To bank the memory win while preserving stability, the team quantizes only the Mixture-of-Experts (MoE) layers — roughly 97% of parameters in a DeepSeek-V3-style architecture.
For activations, a global FP32 scale creates two problems:
- Batch dependency: a token’s quantization varies with the other tokens in the batch.
- Temporal leakage: future tokens in the same sequence can influence the shared scale, corrupting the causal structure of past tokens.
Fix — per-token activation scaling. Each token computes its own FP32 activation scale across the hidden dimension. Quantization stays local, and there is no separate calibration step. This scale computation, along with the FP8 E4M3 block scales, is fused directly into the activation quantization kernel to keep overhead low.
3. Gradient Stability: The Chain-Rule Inconsistency
A conservative start is NVFP4 forward, higher-precision (BF16) backward. But this introduces a chain-rule violation that produces severe gradient-norm spikes.
Let $Q(\cdot)$ be the quantization function (non-differentiable clipping + rounding).
- Forward actually executed: $y = x \cdot Q(w_{bf16})$
- Naive backward evaluates: $y = x \cdot w_{bf16}$
Because the backward pass is unaware of the clipping and rounding decisions made in the forward pass, the gradients misalign with what was actually computed, and training diverges.
Solution — Dequantized Backward. The backward pass uses the BF16-dequantized value of the exact quantized tensor from the forward pass. With $DQ(\cdot)$ as dequantization:
\[y = x \cdot DQ(Q(w_{bf16}))\]This re-aligns forward and backward.
NVFP4’s coarse grid still raises gradient variance, but Adam’s momentum and adaptive second moment absorb it. Implemented in NVIDIA’s TransformerEngine, this avoids storing extra quantized tensor copies, cutting peak linear memory by 70%.
4. The “Four-Over-Six” (4/6) Optimization
NVFP4 normally maps the block max to $\pm 6$. But the FP4 grid is sparse near the top (step size 2 between 4 and 6), so a value near $5/6$ of range can incur error up to $1/6$ of range. (I covered the standalone algorithm in Adaptive NVFP4 Quantization; here it is pushed into the RL loop.)
Intuition: adaptively restrict the max representation to $\pm 4$ when it lowers error, dropping the worst-case error to $1/8$ of range. Prior pretraining work applied 4/6 only to activations; this recipe applies it to both weights and activations — essential in RL where pretrained weights are highly sensitive to quantization error.
The bottleneck. Deciding between 4 and 6 requires an adjusted candidate:
candidate = fp4_value * fp8_scale * amax / (E2M1_MAXVAL * E4M3_4OVER6_MAXVAL)
Converting FP4 → FP32, computing blockwise error, and comparing in scalar FP32 causes severe hardware stalls.
The optimization. Skip FP32; do the casting losslessly in FP16 — not BF16, which lacks the mantissa bits to represent the E2M1 × E4M3 product precisely — using hardware-native PTX (cvt.rn.f16x2, mul.rn.f16x2). Rather than multiplying the candidate up by fractional scales twice, scale the target down by its inverse. Result: 2.8× speedup, >99.97% match to the strict FP32 reference.
Stability note: to avoid off-policy drift between trainer and sampler, 4/6 error accumulation must be bit-exact and deterministic across both stacks (TransformerEngine and FlashInfer) — the same training–inference numeric parity and reproducible-numerics discipline that stabilizes RL more broadly.
5. Strategic High Precision: Spend Where It Matters
Not all weights survive 4-bit. The final recipe keeps BF16 in two places:
- The final 15% of layers (as in NVIDIA’s pretraining paper).
- Shared MoE experts — active for every token, so the compute-to-impact ratio strongly favors keeping them high precision.
6. The “One-Step Trap”
Combining Dequantized Backward + FP16 4/6 + selective precision yields a stable reward curve with gradient spikes eliminated. A useful side effect: the same pipeline enables online post-training quantization — serving FP8/BF16 models in NVFP4 without a separate calibration pipeline (compare quantization-aware distillation for NVFP4).
The guiding idea (from Rich Sutton): the one-step trap. Modeling a complex dynamical system by optimizing isolated components — e.g., standardizing an NVFP4 pretraining recipe and assuming it works for RL — is a trap. The real challenge is co-designing components to handle the interactions between samplers and trainers.
Future direction: MXFP8 activations with NVFP4 weights. Activations are harder to quantize than weights; a mixed-precision instruction set (MXFP8 × MXFP4) at the hardware level could supply exactly the precision needed without giving up the throughput of 4-bit weights.