Jet-RL: Abilitazione dell'Apprendimento per Rinforzo FP8 On-Policy con un Flusso di Precisione Unificato per Addestramento e Rollout

Abstract

L'apprendimento per rinforzo (RL) è fondamentale per potenziare le capacità di ragionamento complesso dei grandi modelli linguistici (LLM). Tuttavia, le pipeline di addestramento RL esistenti sono computazionalmente inefficienti e ad alto consumo di risorse, con la fase di rollout che rappresenta oltre il 70% del tempo totale di addestramento. L'addestramento RL quantizzato, in particolare l'utilizzo della precisione FP8, offre un approccio promettente per mitigare questo collo di bottiglia. Una strategia comunemente adottata applica la precisione FP8 durante il rollout mantenendo la precisione BF16 per l'addestramento. In questo lavoro, presentiamo il primo studio completo sull'addestramento RL in FP8 e dimostriamo che la strategia ampiamente utilizzata di addestramento in BF16 + rollout in FP8 soffre di una grave instabilità di addestramento e di un collasso catastrofico dell'accuratezza in scenari con rollout a lungo orizzonte e task complessi. La nostra analisi mostra che questi fallimenti originano dalla natura *off-policy* dell'approccio, che introduce una sostanziale discrepanza numerica tra addestramento e inferenza. Sulla base di queste osservazioni, proponiamo Jet-RL, un framework di addestramento RL in FP8 che abilita un'ottimizzazione RL robusta e stabile. L'idea chiave è adottare un flusso di precisione FP8 unificato sia per l'addestramento che per il rollout, minimizzando così le discrepanze numeriche ed eliminando la necessità di una inefficiente calibrazione inter-step. Esperimenti estensivi convalidano l'efficacia di Jet-RL: il nostro metodo raggiunge un incremento di velocità fino al 33% nella fase di rollout, fino al 41% nella fase di addestramento e un incremento di velocità end-to-end del 16% rispetto all'addestramento BF16, mantenendo al contempo una convergenza stabile in tutti gli scenari e subendo una degradazione dell'accuratezza trascurabile.

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: Abilitazione dell'Apprendimento per Rinforzo FP8 On-Policy con un Flusso di Precisione Unificato per Addestramento e Rollout

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

Abstract

Support