arXiv cs.CL· Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, Yizhong Cao, Mi Zhang, Dayiheng Liu, Jianwei Zhang·· 6 天前AI 评分22
TRACE：面向 MoE 语言模型 FP4 强化学习的 rollout 引导量化感知训练
TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
AI 导读
针对现有 FP4 RL 方法分别优化训练与 rollout 路径量化精度、未直接缩小两条量化执行路径差异的问题，研究者提出 TRACE 框架，用 rollout 侧量化结果引导训练侧 FP4 舍入决策，并缓存深层尾数与 scale 信息以降低存储和通信开销。
正文 · 原文
Authors:Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, Yizhong Cao, Mi Zhang, Dayiheng Liu, Jianwei Zhang
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.
| Subjects: | Machine Learning (cs.LG); Computation and Language (cs.CL) |
| Cite as: | arXiv:2610.07767 [cs.LG] |
| (or arXiv:2610.07767v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.07767 arXiv-issued DOI via DataCite |
Submission history
From: Xin Wang [view email]
[v1]
Tue, 6 Oct 2026 05:03:47 UTC (638 KB)
[v2]
Fri, 9 Oct 2026 04:13:24 UTC (636 KB)
来源：arXiv cs.CL · arxiv.org