TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
作者: Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, Yizhong Cao, Mi Zhang, Dayiheng Liu, Jianwei Zhang
分类: cs.LG, cs.CL
发布日期: 2026-10-06
💡 一句话要点
提出TRACE以解决FP4强化学习中的量化精度问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 量化训练 强化学习 大语言模型 Mixture-of-Experts FP4 回滚引导 计算效率
📋 核心要点
- 现有FP4强化学习方法在训练和回滚路径的量化精度优化上存在独立性，导致性能差异。
- TRACE通过回滚引导的量化感知训练，利用回滚侧的量化结果来优化训练过程，减少路径间差异。
- 实验表明TRACE在多个任务上实现了与BF16回滚相当的性能，同时回滚速度提升了5.4倍。
📝 摘要（中文）
强化学习（RL）在后训练大语言模型（LLMs）中生成回滚时会产生大量计算和内存开销，这促使低精度回滚的需求。然而，现有的FP4 RL方法存在一个关键限制：它们主要独立优化训练和回滚路径的量化精度，而不是直接减少两条量化执行路径之间的差异。为此，本文提出了TRACE（Train-Rollout Quantization Alignment via Compact GuidancE），一个针对Mixture-of-Experts（MoE）语言模型的FP4量化框架。TRACE结合了回滚引导的量化感知训练，利用回滚侧的量化结果来指导训练侧的FP4舍入决策，从而直接减少训练与回滚之间的差异。实验结果表明，TRACE在多个大规模MoE语言模型上实现了与BF16回滚相当的RL性能，同时在回滚速度上提升了5.4倍。
🔬 方法详解
问题定义：本文旨在解决现有FP4强化学习方法在训练和回滚路径量化精度优化上的独立性问题，导致的性能差异。
核心思路：TRACE通过回滚引导的量化感知训练，利用回滚侧的量化结果来指导训练侧的FP4舍入决策，从而直接减少训练与回滚之间的差异。
技术框架：TRACE的整体架构包括回滚引导的量化训练模块和高效的量化信息缓存机制，后者选择性保留来自深层的尾数和尺度信息，以降低存储和通信开销。
关键创新：TRACE的主要创新在于将回滚侧的量化结果与训练过程紧密结合，直接优化训练和回滚路径之间的量化一致性，这与现有方法的独立优化策略形成鲜明对比。
关键设计：TRACE采用高效的量化信息缓存方案，选择性地保留重要的量化信息，并在损失函数中引入了新的约束，以确保训练过程中的量化决策能够有效指导回滚过程。
🖼️ 关键图片
📊 实验亮点
TRACE在四个大规模MoE语言模型上进行评估，结果显示其在推理、编码和长时间跨度的RL任务中，能够实现与BF16回滚相当的性能，同时回滚速度提升高达5.4倍，展现出强大的FP4性能。
🎯 应用场景
该研究的潜在应用领域包括大规模语言模型的训练和优化，尤其是在需要高效计算和内存管理的场景中，如自然语言处理、对话系统和代码生成等。TRACE的设计理念和方法可以为未来的量化训练技术提供新的思路，推动更高效的模型部署与应用。
📄 摘要（原文）
Reinforcement learning (RL) for post-training large language models (LLMs) incurs substantial computation and memory overhead during rollout generation, which motivates low-precision rollout for efficient RL training. However, existing FP4 RL methods suffer from a key limitation: they primarily optimize quantization accuracy on the training and rollout paths independently rather than directly reducing the discrepancy between the two quantized execution paths. In this work, we propose TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), an FP4 quantization framework for RL training of Mixture-of-Experts (MoE) language models that addresses the limitation of existing FP4 RL methods. TRACE incorporates rollout-guided quantization-aware training that uses rollout-side quantization outcomes to guide training-side FP4 rounding decisions, directly reducing train-rollout discrepancy. Moreover, TRACE adopts an efficient quantization-information caching scheme that selectively retains mantissa and scale information from deeper layers to reduce the storage and communication overhead introduced by rollout guidance. We evaluate TRACE on four large-scale MoE language models across reasoning, coding, and long-horizon RL tasks. Our results demonstrate that TRACE enables joint FP4 weight/activation and FP4 KV-cache rollout with RL performance comparable to BF16 rollout, while achieving up to 5.4xrollout speedup and strong final FP4 performance compared with post-hoc FP4 quantization of BF16-trained policies.