Jet-RL : Activation de l'Apprentissage par Renforcement FP8 sur Politique avec un Flux de Précision Unifié pour l'Entraînement et le Déploiement
Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow
January 20, 2026
Auteurs: Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu
cs.AI
Résumé
L'apprentissage par renforcement (RL) est essentiel pour améliorer les capacités de raisonnement complexe des grands modèles de langage (LLM). Cependant, les pipelines d'entraînement par RL existants sont inefficaces en calcul et gourmands en ressources, la phase de déploiement (« rollout ») représentant plus de 70 % du temps d'entraînement total. L'entraînement par RL quantifié, en particulier avec la précision FP8, offre une approche prometteuse pour atténuer ce goulot d'étranglement. Une stratégie couramment adoptée applique la précision FP8 pendant le déploiement tout en conservant la précision BF16 pour l'entraînement. Dans ce travail, nous présentons la première étude complète sur l'entraînement RL en FP8 et démontrons que la stratégie largement utilisée « entraînement-BF16 + déploiement-FP8 » souffre d'une instabilité d'entraînement sévère et d'un effondrement catastrophique de la précision dans des scénarios de déploiement à long horizon et sur des tâches difficiles. Notre analyse montre que ces échecs découlent de la nature hors politique (« off-policy ») de l'approche, qui introduit un décalage numérique substantiel entre l'entraînement et l'inférence. Motivés par ces observations, nous proposons Jet-RL, un cadre d'entraînement RL en FP8 qui permet une optimisation RL robuste et stable. L'idée clé est d'adopter un flux de précision FP8 unifié à la fois pour l'entraînement et le déploiement, minimisant ainsi les écarts numériques et éliminant le besoin d'une calibration inter-étapes inefficace. Des expériences approfondies valident l'efficacité de Jet-RL : notre méthode permet d'atteindre jusqu'à 33 % d'accélération dans la phase de déploiement, jusqu'à 41 % d'accélération dans la phase d'entraînement, et une accélération de bout en bout de 16 % par rapport à l'entraînement BF16, tout en maintenant une convergence stable dans tous les paramètres et en n'entraînant qu'une dégradation de précision négligeable.
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.