TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
TRACE aligns FP4 train and rollout paths for MoE RL, matching BF16 quality with up to 5.4x rollout speedup.
TRACE is an FP4 quantization framework for reinforcement-learning post-training of Mixture-of-Experts language models. It uses rollout-guided quantization-aware training so training-side FP4 rounding follows rollout-side outcomes, reducing train-rollout discrepancy, and caches selected mantissa and scale information from deeper layers. Across four large MoE models and reasoning, coding, and long-horizon RL tasks, joint FP4 weight, activation, and KV-cache rollout matches BF16 rollout performance. The method reports up to a 5.4x rollout speedup and stronger final FP4 results than post-hoc quantization of BF16-trained policies.
- TRACE aligns training-side FP4 rounding with rollout-side quantization outcomes.
- A cache keeps selected mantissa and scale data from deeper layers.
- Evaluated on four large MoE models for reasoning, coding, and long-horizon RL.
- Joint FP4 weights, activations, and KV cache match BF16 rollout quality with up to 5.4x speedup.
From the original article
1 min read at huggingface.co
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.
By Xin Wang, Hao Yu, Zhengyang Zhuge · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.