提出统一精度流的FP8强化学习框架,显著提速且稳定
Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow
- 训练与推理统一用FP8精度,减少数值偏差
- 滚动阶段提速33%,整体训练提速41%,端到端快16%
- 适合需要高效长程推理的LLM强化学习应用
强化学习对提升大语言模型的复杂推理能力至关重要。然而现有强化学习训练流程计算效率低、资源消耗大,其中滚动阶段占总训练时间超过70%。采用FP8精度的量化强化学习有望缓解这一瓶颈。当前常用策略是在滚动时使用FP8,而训练保留BF16精度。本文首次全面研究FP8强化学习训练,发现该方法在长时序滚动和复杂任务下存在严重训练不稳定性及灾难性精度下降。分析表明,问题源于该策略的离线性,导致训练与推理间存在显著数值不匹配。为此,我们提出Jet-RL框架,实现训练与滚动统一使用FP8精度流,从而最小化数值差异,无需低效的跨步校准。大量实验验证其有效性:相比BF16训练,滚动阶段最高提速33%,训练阶段最高提速41%,端到端提速16%,且在所有设置下保持稳定收敛,精度损失可忽略。
原文摘要 · Abstract (English)
Reinforcement learning (RL) is essential for enhancing the complex reasoning capabilities of large language models (LLMs). However, existing RL training pipelines are computationally inefficient and resource-intensive, with the rollout phase accounting for over 70% of total training time. Quantized RL training, particularly using FP8 precision, offers a promising approach to mitigating this bottleneck. A commonly adopted strategy applies FP8 precision during rollout while retaining BF16 precision for training. In this work, we present the first comprehensive study of FP8 RL training and demonstrate that the widely used BF16-training + FP8-rollout strategy suffers from severe training instability and catastrophic accuracy collapse under long-horizon rollouts and challenging tasks. Our analysis shows that these failures stem from the off-policy nature of the approach, which introduces substantial numerical mismatch between training and inference. Motivated by these observations, we propose Jet-RL, an FP8 RL training framework that enables robust and stable RL optimization. The key idea is to adopt a unified FP8 precision flow for both training and rollout, thereby minimizing numerical discrepancies and eliminating the need for inefficient inter-step calibration. Extensive experiments validate the effectiveness of Jet-RL: our method achieves up to 33% speedup in the rollout phase, up to 41% speedup in the training phase, and a 16% end-to-end speedup over BF16 training, while maintaining stable convergence across all settings and incurring negligible accuracy degradation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。