Jet-RL: Het mogelijk maken van on-policy FP8-versterkingsleren met een uniforme precisieflow voor training en rollout

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

Samenvatting

Versterkend leren (RL) is essentieel voor het verbeteren van de complexe redeneervaardigheden van grote taalmodellen (LLM's). Bestaande RL-trainingspijplijnen zijn echter rekenkundig inefficiënt en resource-intensief, waarbij de rollout-fase meer dan 70% van de totale trainingstijd in beslag neemt. Gekwantiseerde RL-training, met name het gebruik van FP8-precisie, biedt een veelbelovende aanpak om dit knelpunt te verlichten. Een veelgebruikte strategie past FP8-precisie toe tijdens de rollout, maar behoudt BF16-precisie voor de training. In dit werk presenteren we de eerste uitgebreide studie naar FP8 RL-training en tonen we aan dat de veelgebruikte BF16-training + FP8-rollout-strategie lijdt onder ernstige trainingsinstabiliteit en een catastrofale nauwkeurigheidsdaling bij lange rollouts en uitdagende taken. Onze analyse toont aan dat deze mislukkingen voortkomen uit het off-policy karakter van de aanpak, die een aanzienlijke numerieke mismatch tussen training en inferentie introduceert. Gemotiveerd door deze observaties stellen we Jet-RL voor, een FP8 RL-trainingsraamwerk dat robuuste en stabiele RL-optimalisatie mogelijk maakt. Het kernidee is om een uniforme FP8-precisiestroom aan te nemen voor zowel training als rollout, waardoor numerieke discrepanties worden geminimaliseerd en de behoefte aan inefficiënte inter-step-kalibratie wordt geëlimineerd. Uitgebreide experimenten valideren de effectiviteit van Jet-RL: onze methode bereikt een versnelling tot 33% in de rollout-fase, tot 41% versnelling in de trainingsfase en een end-to-end versnelling van 16% ten opzichte van BF16-training, terwijl stabiele convergentie in alle instellingen wordt behouden en met een verwaarloosbare nauwkeurigheidsafname.

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.