论文解读 · 2026年10月8日
TRACE:面向MoE语言模型FP4强化学习的基于Rollout的QAT(量化感知训练)
TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
以Rollout侧的FP4舍入结果引导训练侧的舍入,减少了训练-Rollout不一致,在MoE模型的RL中以BF16水平的性能实现了最高5.4倍的Rollout加速。
问题
在LLM强化学习中,Rollout生成成本很高。FP4低精度Rollout很有吸引力,但训练路径与Rollout路径之间的数值不一致会增大,引发策略不一致和训练崩溃。在MoE中,这种差异甚至可能改变专家路由。现有方法只是分别降低各路径的量化误差,并不直接缩小两条量化路径之间的差异。
方法
(1) Rollout引导的QAT:在Rollout过程中记录路由专家激活值和KV缓存的FP4量化结果。在训练阶段,在相邻的两个FP4码字中舍入到更接近Rollout码字的一方。与普通RTN舍入相比,可保证在每个位置上局部不一致不会增加。(2) 高效缓存:不一致大多仅相差一个相邻码字,且修正集中在深层。因此只存储和传输后段层的尾数(mantissa)与缩放信息。实验在VeRL、Megatron、SGLang环境中使用GRPO和R3路由重放进行。
结果
- 作者报告,在同时使用FP4 W/A和FP4 KV缓存的Rollout中,保持了与BF16 Rollout相近的RL性能和训练态势。评估使用4个Qwen系列MoE模型,任务包括推理、编程(DeepSWE、Terminal-Bench)和长程任务(GDPval)。
- 与BF16 Rollout相比,Rollout速度最高快至5.4倍。
- 在Qwen3.8-Flash-Next的Terminal-Bench 2.1训练中,普通FP4 QAT的训练-Rollout log-prob差异不断扩大,分数骤降。TRACE抑制了不一致,与BF16的差距在训练中逐渐缩小甚至反超。
- 最终FP4性能高于先用BF16 Rollout训练、再进行FP4训练后量化的流水线(vanilla NVFP4、4over6、H-Scale)。
- 发生不一致的量化值中99%以上在FP4码本中仅相差一个相邻格。这正是只传输尾数信息这一设计的依据。同时给出的估算显示,若保存全部信息,每个RL步最多需要约51TB。
局限
作者指出,该保证仅限于逐位置的局部不一致,并不意味着整个网络或策略层面的差异会单调减小。异步训练中策略staleness导致的激活值差异本身无法消除。量化对象仅为路由专家和KV缓存,其余模块保持BF16。评估模型均为Qwen系列MoE,两个大型模型仅训练了50步。
意义
提出了一种在大规模MoE的RL后训练中,将成本最高的Rollout阶段以FP4运行的同时保持训练稳定性的方法。核心观点是,目标应是“两条路径的一致”,而非“量化误差最小化”。
关键术语
- 训练-Rollout不一致(Train-Rollout Mismatch)
- 在RL中,生成响应的推理引擎与计算梯度的训练引擎对相同输入给出不同概率的现象。会引发策略不一致和训练不稳定。
- 量化感知训练(QAT)
- 在前向传播中模拟低精度量化进行训练,使模型适应量化误差的技术。
- NVFP4 / FP4 E2M1
- 由2位指数、1位尾数组成的4位浮点格式。与块缩放配合使用,可表示的值非常少。
- Round-to-Nearest(RTN)
- 将值舍入到最近量化码字的标准方式。边界附近的微小差异可能被放大为较大差异。
原文
作者: Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, 等 · 发布 2026-10-06 · 许可 arXiv 비독점 배포 라이선스 · 基于全文分析
本解读是 AI 阅读论文后撰写的摘要,请以论文原文为准。