Reinforcement learning (RL) for post-training large language models carries heavy computation and memory costs, much of it concentrated in rollout generation. That burden has motivated low-precision rollouts, but existing FP4 RL approaches share a key weakness: they optimize quantization accuracy on the training and rollout paths separately instead of directly shrinking the discrepancy between the two quantized execution paths.
To address this, researchers propose TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), an FP4 quantization framework designed for RL training of Mixture-of-Experts (MoE) language models.
TRACE's central mechanism is rollout-guided quantization-aware training, which feeds quantization outcomes from the rollout side back into the training side's FP4 rounding decisions. This directly reduces the train-rollout discrepancy rather than treating the two paths in isolation.
The framework also introduces an efficient quantization-information caching scheme that selectively retains mantissa and scale information from deeper layers, limiting the storage and communication overhead that rollout guidance would otherwise add.
According to the paper, TRACE was evaluated on four large-scale MoE language models spanning reasoning, coding, and long-horizon RL tasks. The results report joint FP4 weight/activation and FP4 KV-cache rollout with RL performance comparable to BF16 rollout, up to 5.4x rollout speedup, and strong final FP4 performance relative to post-hoc FP4 quantization of BF16-trained policies.