TRACE: Rollout-Guided FP4 Quantization Makes MoE RL Training Fast and Stable, 5.4× Speedup

Large language model reinforcement learning (RL) is increasingly compute-intensive, with the rollout phase consuming a disproportionate share of resources. Reducing rollout precision from BF16 to FP4 offers significant speed and memory savings, but introduces a critical instability: training and rollout paths, even when starting from nearly identical high-precision values, can quantize to different FP4 codebook entries due to discrete rounding boundaries. This train–rollout discrepancy amplifies tiny numerical differences into policy divergence, causing reward collapse in MoE models where router decisions and expert selection are also affected.

Problem: FP4 Rounding Boundaries Amplify Tiny Differences

The paper illustrates the issue with a concrete example: a training activation of 60.24 and a rollout activation of 59.76 differ by only 0.48 in BF16. After applying the same global and block scales, their normalized values become 2.51 and 2.49 — straddling a rounding boundary. One quantizes to FP4 code 3, the other to code 2, expanding the final difference to 24. In MoE models, such shifts alter router logits and expert choices, further decoupling the training policy from the generation policy. Experiments on Qwen3.5-35B-A3B show that standard FP4 QAT, QaRL, and QUADS all suffer late-training reward degradation or collapse under joint FP4 W/A and FP4 KV cache rollout.

TRACE Core Idea: Align Training Rounding to Rollout Reality

Instead of minimizing each path's quantization error independently, TRACE directly minimizes the discrepancy between the two low-precision execution paths. Its first design, Rollout-Guided Quantization-Aware Training, records the actual FP4 quantization results (routed expert activations and KV states) during rollout. In subsequent training steps, these recorded results guide the fake-quantization rounding decisions on the training side. For each value to be quantized, training still chooses between the two adjacent FP4 E2M1 codebook entries, but prefers the entry closer to the rollout-side quantized value rather than applying standard round-to-nearest.

Compact Rollout Information Caching

Transmitting full rollout quantization results across GPU→CPU→disk→GPU incurs prohibitive storage and communication overhead, especially for long trajectories and asynchronous RL. The authors observe that over 99% of training/rollout quantization mismatches differ by only a single adjacent codebook entry. Since the training-side activation already narrows the candidate set to two neighbors, the rollout side need only transmit a few mantissa bits plus the scale to indicate the correct rounding direction. This compact cache reduces communication to near-native FP4 rollout throughput.

Experimental Results

Qwen3.5-35B-A3B RL average score: TRACE 75.3 vs BF16 rollout 74.9, FP4 QAT 59.6, QaRL 58.6, QUADS 68.8.

HMMT25 benchmark: TRACE 70.0 vs QUADS 59.0.

Decoding throughput (8K input, 128K output): up to 5.4× higher than BF16 rollout with compact caching; full-information caching severely degrades throughput.

End-to-end RL step time overhead: ~7.4% with compact cache.

Larger MoE models: Qwen3.5-122B-A10B, Qwen3.8-Flash-Next, Qwen3.8-2.4T-A95B all recover BF16-level performance under joint FP4 W/A + FP4 KV rollout.

Model Adapts to FP4 During RL Training

An unexpected phenomenon emerges: early in training, FP4 execution incurs a performance penalty, but once the train–rollout discrepancy is controlled, the policy gradually adapts to the low-precision environment. On Qwen3.5-35B-A3B HMMT25, TRACE catches up to BF16 after ~120 RL steps; on Qwen3.8-Flash-Next, after ~240 steps. This suggests that consistent, learnable low-precision execution lets the policy internalize numerical constraints rather than treating them as static noise.

Superiority Over Post-Training Quantization

When BF16-trained policies are quantized post-hoc, Vanilla NVFP4 scores 70.4, 4over6 scores 71.0, and H-Scale scores 71.4. TRACE, which trains with joint FP4 rollout from the start, achieves 75.3 — demonstrating that co-training with the target quantization yields a fundamentally more robust FP4 policy.

Paper: TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models (arXiv:2610.07767)