TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang and 4 others
Alibaba Group · Ohio State University
Published October 6, 2026 · on the feed October 7, 2026 · ▲ 88 · arXiv 2610.07767
The abstract
Reinforcement learning (RL) for post-training large language models (LLMs) incurs substantial computation and memory overhead during rollout generation, which motivates low-precision rollout for efficient RL training. However, existing FP4 RL methods suffer from a key limitation: they primarily optimize quantization accuracy on the training and rollout paths independently rather than directly reducing the discrepancy between the two quantized execution paths. In this work, we propose TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), an FP4 quantization framework for RL training of Mixture-of-Experts (MoE) language models that addresses the limitation of existing FP4 RL methods. TRACE incorporates rollout-guided quantization-aware training that uses rollout-side quantization outcomes to guide training-side FP4 rounding decisions, directly reducing train-rollout discrepancy. Moreover, TRACE adopts an efficient quantization-information caching scheme that selectively retains mantissa and scale information from deeper layers to reduce the storage and communication overhead introduced by rollout guidance. We evaluate TRACE on four large-scale MoE language models across reasoning, coding, and long-horizon RL tasks. Our results demonstrate that TRACE enables joint FP4 weight/activation and FP4 KV-cache rollout with RL performance comparable to BF16 rollout, while achieving up to 5.4xrollout speedup and strong final FP4 performance compared with post-hoc FP4 quantization of BF16-trained policies.
The Brief
2:38A short podcast summary of the paper. Written and voiced by AI.
Transcript
TRACE. Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models. By Xin Wang and eleven others, from Alibaba Group and the Ohio State University. Published on the sixth of October. No code repository is linked to it.
Reinforcement learning for post-training large language models is expensive in both computation and memory, and much of that cost sits in generating rollouts. That is the motivation for running rollouts at low precision. Existing FP4 methods for reinforcement learning have a key limitation, however. They optimise quantization accuracy on the training path and on the rollout path separately. What they do not do is directly shrink the discrepancy between those two quantized paths, which is the gap that matters when one path generates the data the other learns from.
TRACE, short for Train-Rollout Quantization Alignment via Compact GuidancE, is an FP4 quantization framework for reinforcement learning of Mixture-of-Experts language models. Its central move is rollout-guided quantization-aware training: the outcomes of quantization on the rollout side are used to guide FP4 rounding decisions on the training side. Rather than making each path accurate on its own, it aims the two paths at each other, reducing the train-rollout discrepancy directly.
Put simply, when training has to decide how to round a value to FP4, it takes into account how the rollout side's quantization actually turned out. Guidance from the rollout side has a cost, though, because that information has to be stored and communicated. To keep it down, TRACE uses an efficient caching scheme for quantization information. It selectively retains mantissa and scale information from deeper layers, which reduces the storage and communication overhead the guidance would otherwise introduce.
TRACE is evaluated on four large-scale Mixture-of-Experts language models, across reasoning, coding, and long-horizon reinforcement learning tasks. It enables rollout with FP4 weights and activations together with an FP4 KV-cache, and reaches reinforcement learning performance comparable to BF16 rollout. Rollout speedup reaches up to five point four times.
There is a second comparison, about the policy that comes out at the end. One could simply train in BF16 and quantize the finished policy to FP4 afterwards, post hoc. Against that route, the final FP4 performance of TRACE is strong. So the gains are not only in the speed of training: the low-precision model that training produces is itself better suited to running in FP4.
Where earlier FP4 methods made each path accurate in isolation, TRACE targets the agreement between them, and builds training around that agreement, with a cache compact enough that the guidance stays affordable.
Paper Quest · 7 chapters · about 39 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
Four bits make rollouts fast. They also make two copies of the same model disagree.
You are training a large sparse language model with reinforcement learning, and most of the cost is generating rollouts. Running them in a 4-bit float format would be far faster.
The road ahead
- 1 · The problem ~6 min Rounding rollouts to 4 bits makes reinforcement learning cheaper, but training and rollout can round the same value to different codes, and the policies drift apart.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min Instead of making each path round accurately on its own, TRACE makes training round each value toward the code the rollout side actually produced, and keeps that record cheap.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Rollouts run in FP4 and leave a compact rounding record from deeper layers; training loads it, rounds to match, then updates the policy.
- 6 · The experiment ~6 min You’ll guess: Running rollouts in FP4 with TRACE instead of BF16, what was the largest rollout speedup reported?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.