TRACE：让FP4强化学习训练与Rollout真正对齐

强化学习已经成为大语言模型后训练的重要环节，但它的成本并不只来自参数更新。训练过程中，模型需要反复生成Rollout样本，而生成阶段同时受到计算量、显存占用和KV缓存规模的限制。将Rollout压缩到FP4，可以显著降低硬件负担；然而，低精度并非简单地把BF16模型在训练完成后直接量化即可解决。

问题：训练和Rollout并不是同一个量化路径

现有FP4强化学习方案通常分别优化训练路径和Rollout路径的量化精度。这样做能够降低各自的局部误差，却未必能保证两条路径产生相近的计算结果。对于依赖当前策略持续采样的强化学习而言，这种差异尤其敏感：训练阶段看到的模型行为与Rollout阶段实际产生的数据分布可能发生偏移，进而影响策略优化。

TRACE的核心判断是，真正需要优化的目标不只是“每条路径各自量化得足够好”，还包括“训练与Rollout之间尽可能一致”。论文将这一思路概括为Train-Rollout Quantization Alignment，并面向混合专家（MoE）语言模型设计了对应框架。

核心方法

• Rollout引导的量化感知训练：TRACE将Rollout侧的量化结果带回训练过程，用于指导训练侧的FP4舍入决策。换言之，训练侧不再独立选择量化结果，而是尽量贴近部署时Rollout真正使用的低精度行为。
• 面向整条生成链路的FP4支持：框架覆盖权重和激活，并进一步支持FP4 KV缓存Rollout。这意味着优化目标不仅是模型参数，还包括长序列生成中占用大量存储的缓存部分。
• 选择性缓存量化信息：Rollout指导会引入额外的信息存储与通信成本。TRACE并不保存所有中间信息，而是有选择地保留更深层网络中的尾数与缩放信息，在提供指导的同时压缩开销。

实验结果与解读

论文在四个大规模MoE语言模型上进行了评估，任务覆盖推理、代码和长时程强化学习。结果表明，联合使用FP4权重、激活和KV缓存进行Rollout时，TRACE能够获得与BF16 Rollout相当的强化学习表现；与对BF16训练策略进行事后FP4量化相比，它也展现出更强的最终FP4效果。论文报告的最高Rollout加速为5.4倍。

这项工作的价值在于，它把量化问题从“如何尽量还原高精度模型”转向“如何让训练过程适应实际的低精度执行路径”。对于MoE模型而言，Rollout规模大、生成频繁，且缓存与通信成本容易成为瓶颈，因此这种路径级对齐思路具有较强的工程针对性。

不过，现有材料主要给出了方法概览和总体结果，尚未披露不同模型、任务及硬件配置下的详细对比数据。因此，5.4倍加速应理解为论文实验中的最高结果，而不是所有场景下的固定收益。总体来看，TRACE说明：在强化学习低比特化中，训练与部署的一致性可能比单独追求量化误差更关键。

来源：Hugging Face Daily Papers