Jet-RL: Ermöglichung von On-Policy-FP8-Reinforcement-Learning mit einheitlicher Präzisionssteuerung für Training und Rollout

papers.abstract

Reinforcement Learning (RL) ist entscheidend für die Verbesserung der komplexen Denkfähigkeiten großer Sprachmodelle (LLMs). Allerdings sind bestehende RL-Trainingspipelines rechnerisch ineffizient und ressourcenintensiv, wobei die Rollout-Phase über 70 % der gesamten Trainingszeit ausmacht. Quantisiertes RL-Training, insbesondere unter Verwendung von FP8-Präzision, bietet einen vielversprechenden Ansatz zur Minderung dieses Engpasses. Eine häufig angewandte Strategie nutzt FP8-Präzision während des Rollouts, behält jedoch BF16-Präzision für das Training bei. In dieser Arbeit präsentieren wir die erste umfassende Studie zu FP8-RL-Training und zeigen, dass die weit verbreitete BF16-Training- + FP8-Rollout-Strategie unter langen Rollout-Horizonten und anspruchsvollen Aufgaben unter schwerer Trainingsinstabilität und katastrophalem Genauigkeitsverlust leidet. Unsere Analyse zeigt, dass diese Fehler auf den Off-Policy-Charakter des Ansatzes zurückzuführen sind, der eine erhebliche numerische Diskrepanz zwischen Training und Inferenz einführt. Aufbauend auf diesen Beobachtungen schlagen wir Jet-RL vor, ein FP8-RL-Trainingsframework, das eine robuste und stabile RL-Optimierung ermöglicht. Die Kernidee besteht darin, einen einheitlichen FP8-Präzisionsfluss sowohl für das Training als auch für den Rollout zu verwenden, um so numerische Abweichungen zu minimieren und die Notwendigkeit ineffizienter kalibrierung zwischen den Schritten zu beseitigen. Umfangreiche Experimente validieren die Wirksamkeit von Jet-RL: Unsere Methode erreicht eine Beschleunigung von bis zu 33 % in der Rollout-Phase, bis zu 41 % in der Trainingsphase und eine 16 %ige Ende-zu-Ende-Beschleunigung gegenüber BF16-Training, bei gleichzeitig stabiler Konvergenz über alle Einstellungen hinweg und nur vernachlässigbarer Genauigkeitseinbuße.

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: Ermöglichung von On-Policy-FP8-Reinforcement-Learning mit einheitlicher Präzisionssteuerung für Training und Rollout

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

papers.abstract

Support