Jet-RL: 統一された学習とロールアウトの精度フローによるオンポリシーFP8強化学習の実現

要旨

強化学習（RL）は大規模言語モデル（LLM）の複雑な推論能力を高める上で不可欠である。しかし、既存のRLトレーニングパイプラインは計算効率が低くリソース集約的であり、ロールアウト段階が全体のトレーニング時間の70%以上を占めている。量子化RLトレーニング、特にFP8精度の使用は、このボトルネックを軽減する有望なアプローチを提供する。一般的に採用されている戦略は、ロールアウト中にFP8精度を適用し、トレーニングにはBF16精度を維持するというものである。本研究では、FP8 RLトレーニングに関する初の包括的な研究を提示し、広く使用されている「BF16トレーニング＋FP8ロールアウト」戦略が、長期的なロールアウトや困難なタスクにおいて、深刻なトレーニング不安定性と壊滅的な精度低下に陥ることを実証する。我々の分析によれば、これらの失敗はアプローチのオフポリシー性に起因し、トレーニングと推論の間に数値的不一致を大幅に導入している。これらの観察に動機付けられ、我々は堅牢で安定したRL最適化を可能にするFP8 RLトレーニングフレームワーク「Jet-RL」を提案する。核心となるアイデアは、トレーニングとロールアウトの両方に統一されたFP8精度フローを採用することで、数値的不一致を最小限に抑え、非効率的なステップ間キャリブレーションの必要性を排除することである。大規模な実験によりJet-RLの有効性が検証された：本手法は、ロールアウト段階で最大33%、トレーニング段階で最大41%、BF16トレーニングと比較してエンドツーエンドで最大16%の高速化を達成し、全ての設定で安定した収束を維持し、無視できる程度の精度劣落しか生じない。

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.

Jet-RL: 統一された学習とロールアウトの精度フローによるオンポリシーFP8強化学習の実現

Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow

要旨

Support