Jet-RL：通过统一训练与推理精度流实现基于策略的FP8强化学习
机器学习
2026-01-27 v2 计算与语言
摘要
强化学习（RL）对于增强大型语言模型（LLM）的复杂推理能力至关重要。然而，现有的RL训练流程计算效率低下且资源密集，其中推理阶段占总训练时间的70%以上。量化RL训练，特别是使用FP8精度，为缓解这一瓶颈提供了一种有前景的方法。一种常用的策略是在推理阶段应用FP8精度，同时在训练阶段保留BF16精度。在这项工作中，我们首次对FP8 RL训练进行了全面研究，并证明了广泛使用的BF16训练+FP8推理策略在长程推理和具有挑战性的任务下会遭受严重的训练不稳定性和灾难性的精度崩溃。我们的分析表明，这些失败源于该方法的离策略性质，它在训练和推理之间引入了显著的数值不匹配。受这些观察的启发，我们提出了Jet-RL，一个FP8 RL训练框架，能够实现稳健且稳定的RL优化。其关键思想是采用统一的FP8精度流进行训练和推理，从而最小化数值差异，并消除低效的步间校准需求。大量实验验证了Jet-RL的有效性：我们的方法在推理阶段实现了高达33%的加速，在训练阶段实现了高达41%的加速，并且与BF16训练相比实现了16%的端到端加速，同时在所有设置下保持稳定收敛，且精度下降可忽略不计。
引用
@article{arxiv.2601.14243,
title = {Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow},
author = {Haocheng Xi and Charlie Ruan and Peiyuan Liao and Yujun Lin and Han Cai and Yilong Zhao and Shuo Yang and Kurt Keutzer and Song Han and Ligeng Zhu},
journal= {arXiv preprint arXiv:2601.14243},
year = {2026}
}
备注
11 pages, 6 figures, 4 tables