TRACE：面向 MoE 语言模型 FP4 强化学习的 rollout 引导量化感知训练

arXiv:2610.07767v1 公告类型：新 摘要：对大语言模型（LLMs）进行训练后强化学习（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 相当；与对 BF16 训练策略进行事后 FP4 量化相比，TRACE 可实现最高 5.4 倍的 rollout 加速，并取得出色的最终 FP4 性能。