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

TL;DR TRACE 通过 rollout 引导的量化感知训练，对齐训练与 rollout 的 FP4 量化路径，使 MoE 语言模型 RL 训练能用 FP4 权重/激活/KV-cache 达到 BF16 同等效果，提速最高 5.4 倍。

问题

问题背景

强化学习（RL）后训练已成为提升大语言模型（LLMs）推理、编程等复杂能力的关键范式。在 MoE 架构上，rollout 生成阶段的计算与存储开销尤为突出，促使业界探索 FP4 等低精度 rollout 以降低资源消耗、提升训练吞吐。

现有方法局限

现有 FP4 RL 方法普遍将训练路径和 rollout 路径的量化视为两个独立优化问题：分别最小化各自路径上的量化误差，或对两条路径采用相同的静态 FP4 格式。这种做法忽略了训练与 rollout 之间量化行为不一致带来的系统性偏差。具体而言，即使在每条路径上量化精度都足够高，两条路径在 FP4 舍入决策、溢出处理、缩放因子计算等方面的微小差异，也会导致 rollout 生成的样本分布与训练更新时使用的梯度信号方向不一致。这种不一致在多次 RL 迭代中累积，最终使得优化目标偏离真实的 on-policy 梯度，损害策略性能。

为什么这个问题难/重要

RL 训练要求 rollout 数据与当前策略严格对齐，低精度 rollout 引入的非线性量化噪声会在不同层、不同专家路由上表现出异质影响。MoE 模型的专家稀疏激活进一步放大了这种敏感性：量化误差可能改变路由决策，导致 rollout 轨迹分布偏移。要直接对齐训练和 rollout 的量化过程，需要跨路径传递量化信息（如舍入结果、缩放因子），但这会带来显著的通信与存储开销，在分布式训练环境中难以承受。业界对低成本 RL 训练需求强烈，FP4 rollout 可带来高达 5.4 倍加速，但若策略质量大幅下降则得不偿失。

行业类比

类似在自动驾驶仿真中，若模拟环境的传感器模型与实车部署的低精度推理不一致，训练出的控制策略可能在真实道路上出现性能退化；在 LLM 的 RLHF 场景下，rollout 端 FP4 与训练端 BF16 的行为 gap 相当于用有偏样本估计梯度，最终模型在真实 BF16 部署时可能表现出对话质量下降或推理错误。

核心洞察

• **范式转变：从独立量化到执行路径对齐。** 现有 FP4 RL 方法通常独立优化训练和 rollout 路径的量化精度，而 TRACE 的核心创新在于利用 rollout 侧的实际量化结果指导训练侧的 FP4 舍入决策，直接最小化两条执行路径之间的差异。这一视角的独特之处在于：RL 训练中 rollout 的量化误差会通过策略梯度信号传递回训练侧，若两条路径的量化行为不一致，训练得到的策略可能在低精度推理时表现下降。TRACE 通过对齐量化行为，从根本上避免了这种“训练-推理不一致”导致的性能损失，对从业者的启示是：在 RL 低精度化中，执行路径一致性比单独追求量化误差最小化更重要。
• **选择性信息缓存：以层敏感度平衡对齐与开销。** TRACE 仅对较深层缓存 mantissa 和 scale 信息，而非全量传递 rollout 量化信息，从而在不牺牲对齐效果的前提下大幅降低存储和通信开销。这一设计揭示了 MoE 模型中不同层对 train-rollout 量化差异的敏感度不同：深层网络由于涉及更复杂的路由与专家组合，其量化误差对最终输出影响更大。相比现有方法中对所有层一视同仁的处理，TRACE 通过层选择性缓存实现了更高效的资源利用。对于大规模分布式 RL 训练，这种按层重要性分配量化对齐信息的策略可显著减轻通信瓶颈，同时保持与 BF16 rollout 相当的最终性能，值得工程实践借鉴。

方法

输入与目标

RL 后训练中，rollout 生成消耗大量计算与存储，TRACE 以 BF16 策略模型为起点，目标是在 rollout 阶段启用 FP4 权重/激活与 FP4 KV-cache，同时保持训练与 rollout 量化路径一致。

关键模块

• Rollout-Guided Quantization-Aware Training：训练过程中，训练侧的 FP4 舍入决策参考 rollout 侧的实际量化结果（mantissa 与 scale），使两条路径的量化误差分布对齐，而非各自独立优化量化精度。
• Mantissa-Only Train-Rollout Communication：为降低由 rollout 指导引入的存储与通信开销，TRACE 仅从深层网络选择性缓存 mantissa 和 scale 信息，浅层不传递或减少传递，实现紧凑指导。

输出

训练得到的策略模型支持 FP4 权重/激活和 FP4 KV-cache rollout，在 reasoning、coding、long-horizon RL 任务上性能与 BF16 rollout 相当，并实现最高 5.4× rollout 加速；最终 FP4 性能优于对 BF16 训练策略做 post-hoc FP4 量化。

与同类方法的差异：TRACE 直接减少 train-rollout 量化路径间的 discrepancy，而非仅独立优化两条路径的量化精度。

实验

实验设计

TRACE 在四款大规模 MoE 模型上评估，覆盖推理、代码、长程 RL 任务（具体数据集名称未在摘要中披露）。训练阶段使用 rollout 引导的量化感知训练：仅将更深层的 mantissa 和 scale 信息从 rollout 侧传回训练侧，以降低通信开销。实验对比 BF16 rollout 基线和事后 FP4 量化方案。

关键发现

• rollout 加速：联合 FP4 权重/激活与 FP4 KV-cache 的 rollout 最高达到 5.4× 加速，相比 BF16 rollout。
• RL 性能：TRACE 的 FP4 rollout 在 RL 性能上与 BF16 rollout 相当。
• 最终 FP4 表现：相比对 BF16 训练策略的事后 FP4 量化，TRACE 取得更优的最终 FP4 性能。

与基线对比解读

现有 FP4 RL 方法分别优化训练路径与 rollout 路径的量化误差，未直接对齐两者，导致 rollout 与训练执行路径存在差异。TRACE 通过 rollout 侧量化结果引导训练侧 FP4 舍入决策，直接减小 train-rollout discrepancy，这是与独立量化优化的核心差异。与事后 FP4 量化相比，TRACE 在训练过程中保持量化一致性，避免策略训练完成后量化导致的性能退化，因此在最终 FP4 策略上表现更强。

行业影响

落地场景

TRACE 主要针对大规模 MoE 语言模型的 RL 后训练阶段，可将 rollout 生成过程从 BF16 压缩到 FP4，显著降低计算与显存开销。适用于需要高频、长轨迹 rollout 的场景：

• 推理 / 代码生成 RL 训练：如数学推理、代码生成任务的策略优化，rollout 吞吐提升 5.4 倍。
• 长时程智能体 RL：多轮交互、长文本生成的训练，KV-cache 也可 FP4，显存压力大幅缓解。
• 在线策略服务：需要低延迟、低成本的推理服务，可直接使用 FP4 量化策略而无需额外后量化。

具体用例：内容平台的个性化推荐模型训练中，基于用户反馈的 RL 微调需要大量 candidate generation rollout，TRACE 可将 rollout 放在 FP4 硬件上执行，训练服务器则保持高精度梯度更新，整体迭代速度提升数倍。

商业价值

主要落在降本与体验提升两条线：

• 降本：FP4 rollout 将 rollout 阶段的 GPU 算力、显存、网络通信消耗大幅降低。尤其 MoE 模型参数量大，低精度 KV-cache 可支持更长上下文，减少因 OOM 导致的训练中断。
• 体验提升：与其他 FP4 RL 方法不同，TRACE 直接对齐 train-rollout 量化路径，使 RL 最终策略在 FP4 推理下性能与 BF16 接近，避免后量化带来的精度崩溃，保障线上服务质量。

与现有产品 / 工作流接口

TRACE 可作为训练框架插件集成到现有 RLHF / GRPO 流程中：

1. 在 rollout 执行器中启用 FP4 量化权重 / 激活 / KV-cache，参考论文中的 mantissa-only 通信 协议降低节点间传输量。
2. 训练侧使用 rollout 侧反馈的 quantization outcome 进行 QAT，可复用现有 FSDP / Megatron 分布式训练栈。
3. 缓存方案仅需改动数据加载与通信模块，无需重构 MoE 路由。

对现有工作的差异：对比单一路径 FP4 量化或后量化方案，TRACE 通过显式 train-rollout 对齐，避免了 rollout 低精度导致训练目标偏离的难题，更适合直接替换 BF16 rollout 的生产环境。

局限

• **泛化性受限**：TRACE 仅在 MoE 架构上验证，未涉及稠密模型或其他稀疏结构（如细粒度 MoE）。其 rollout 引导与缓存策略可能依赖 MoE 特有的路由与专家计算模式，迁移到稠密 LLM 时效果未知。此外，缓存方案仅保留深层 mantissa 和 scale，浅层完全采用独立量化，未见对浅层量化误差影响的系统分析和消融，可能导致部分任务下量化敏感层被忽略。
• **训练稳定性风险**：rollout 侧量化结果作为引导信号，在训练早期策略尚未收敛时可能包含较大噪声或不稳定舍入决策，进而误导训练侧 FP4 舍入方向。论文未展示不同训练阶段（如 RL 初始阶段 vs 后期）的引导信号分布或训练 loss 动态，也未讨论是否需要 warm-up 或引导信号平滑机制来缓解早期波动。
• **对比基线有限**：实验对比了 post-hoc FP4 量化与独立路径量化感知训练，但未与更先进的低比特训练方案（如 FP8+INT4 混合精度、基于梯度量化的 RL 训练）或显式校准 train-rollout 分布偏移的方法（如 score centering、量化误差正则）进行系统性比较。同时，硬件实现细节（如 FP4 矩阵乘法在主流 GPU 上的实际利用率、缓存通信的 PCIe 开销）未深入分析，实际部署加速比可能受限于特定硬件对 FP4 算子的支持程度。