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 信息以降低存储和通信开销。
正文 · AI 翻译
作者：Xin Wang、Hao Yu、Zhengyang Zhuge、Bochao Mao、Zheng Li、Junda Feng、Yuyan Luo、Yi Zhang、Yizhong Cao、Mi Zhang、Dayiheng Liu、Jianwei Zhang
摘要：用于大语言模型（LLM）后训练的强化学习（RL）在 rollout 生成过程中会产生大量的计算和内存开销，这促使人们采用低精度 rollout 来实现高效的 RL 训练。然而，现有的 FP4 RL 方法存在一个关键局限：它们主要独立地优化训练路径和 rollout 路径上的量化精度，而不是直接减小两条量化执行路径之间的差异。在这项工作中，我们提出了 TRACE（Train-Rollout Quantization Alignment via Compact GuidancE），这是一个用于混合专家（MoE）语言模型 RL 训练的 FP4 量化框架，旨在解决现有 FP4 RL 方法的局限。TRACE 引入了由 rollout 引导的量化感知训练，利用 rollout 侧的量化结果来指导训练侧的 FP4 舍入决策，从而直接减小训练与 rollout 之间的差异。此外，TRACE 采用了一种高效的量化信息缓存方案，有选择地保留更深层的尾数和缩放信息，以降低 rollout 引导带来的存储和通信开销。我们在四个大规模 MoE 语言模型上，针对推理、编码和长时程 RL 任务对 TRACE 进行了评估。我们的结果表明，TRACE 能够实现 FP4 权重/激活与 FP4 KV-cache 的联合 rollout，其 RL 性能与 BF16 rollout 相当，同时实现了最高 5.4 倍的 rollout 加速，并且与对 BF16 训练策略进行事后 FP4 量化相比，最终 FP4 性能更强。
| 主题： | 机器学习（cs.LG）；计算与语言（cs.CL） |
| 引用为： | arXiv:2610.07767 [cs.LG] |
| （或使用 arXiv:2610.07767v2 [cs.LG] 查看此版本） | |
| https://doi.org/10.48550/arXiv.2610.07767 由 DataCite 提供的 arXiv 颁发 DOI |
提交历史
来自：Xin Wang [查看电子邮件]
[v1]
2026年10月6日星期二 05:03:47 UTC（638 KB）
[v2]
2026年10月9日星期五 04:13:24 UTC（636 KB）
来源：arXiv cs.CL · arxiv.org