Paper review · October 8, 2026
TRACE: Rollout-Based QAT (Quantization-Aware Training) for FP4 Reinforcement Learning of MoE Language Models
TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
The method uses the rollout-side FP4 rounding results to guide rounding on the training side, which reduces train-rollout mismatch. In RL on MoE models, it achieved up to 5.4x faster rollouts at BF16-level performance.
Problem
Rollout generation is expensive in LLM reinforcement learning. Low-precision FP4 rollouts are attractive, but they widen the numerical mismatch between the training path and the rollout path, which causes policy mismatch and training collapse. In MoE models, this difference can even change expert routing. Existing methods only reduce the quantization error of each path separately. They do not directly reduce the difference between the two quantization paths.
Method
(1) Rollout-guided QAT: During rollout, the FP4 quantization results for routed-expert activations and the KV cache are recorded. During training, each value is rounded to whichever of its two adjacent FP4 codewords is closer to the rollout codeword. Compared with standard RTN rounding, this is guaranteed not to increase the local mismatch at each point. (2) Efficient caching: Most mismatches are a single step between adjacent codewords, and corrections are concentrated in deep layers. So only the mantissa and scale information of the later layers is stored and transmitted. Experiments were run in VeRL, Megatron and SGLang environments using GRPO and R3 routing replay.
Results
- The authors report that rollouts using both FP4 W/A and an FP4 KV cache kept RL performance and training behavior similar to BF16 rollouts. The evaluation used four Qwen-family MoE models on reasoning, coding (DeepSWE, Terminal-Bench) and long-horizon tasks (GDPval).
- Rollouts were up to 5.4x faster than BF16 rollouts.
- When training Qwen3.8-Flash-Next on Terminal-Bench 2.1, standard FP4 QAT saw its train-rollout log-prob gap grow and its score drop sharply. TRACE suppressed the mismatch, and its gap with BF16 gradually narrowed or reversed during training.
- Its final FP4 performance was higher than pipelines that train with BF16 rollouts and then apply post-training quantization to FP4 (vanilla NVFP4, 4over6, H-Scale).
- Over 99% of mismatched quantized values differed by one adjacent step in the FP4 codebook. This is the basis for the design that transmits only mantissa information. The authors also estimate that keeping the full information would require up to about 51 TB per RL step.
Limitations
The authors state that the guarantee is limited to local, per-point mismatch and does not mean that network-wide or policy-level differences decrease monotonically. It cannot remove activation differences caused by policy staleness in asynchronous training. Only routed experts and the KV cache are quantized, and the other modules stay in BF16. All evaluated models are Qwen-family MoE models, and the two large models were trained for only 50 steps.
Why it matters
The paper offers a way to run the rollout stage, the most expensive part of RL post-training for large MoE models, in FP4 while keeping training stable. The key idea is that the goal should be 'agreement between the two paths' rather than 'minimizing quantization error'.
Key terms
- Train-rollout mismatch
- A phenomenon in RL where the inference engine that generates responses and the training engine that computes gradients output different probabilities for the same input. It causes policy mismatch and training instability.
- Quantization-aware training (QAT)
- A technique that simulates low-precision quantization in the forward pass during training so the model adapts to quantization error.
- NVFP4 / FP4 E2M1
- A 4-bit floating-point format with 2 exponent bits and 1 mantissa bit. It is used with block scales and can represent very few values.
- Round-to-Nearest(RTN)
- The standard method of rounding a value to the nearest quantization codeword. Small differences near a boundary can be amplified into large differences.
Original
Authors: Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, et al. · Posted 2026-10-06 · License arXiv 비독점 배포 라이선스 · Based on the full text
This review is an AI-written summary of the paper. Please treat the original paper as the reference.