Jet-RL: Permitindo Aprendizagem por Reforço FP8 On-Policy com um Fluxo de Precisão Unificado para Treinamento e Rollout

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

January 20, 2026

Autores: Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

cs.AI

Resumo

O aprendizado por reforço (RL) é essencial para aprimorar as capacidades de raciocínio complexo de modelos de linguagem de grande escala (LLMs). No entanto, os *pipelines* de treinamento de RL existentes são computacionalmente ineficientes e intensivos em recursos, sendo que a fase de *rollout* responde por mais de 70% do tempo total de treinamento. O treinamento de RL quantizado, particularmente usando precisão FP8, oferece uma abordagem promissora para mitigar esse gargalo. Uma estratégia comumente adotada aplica a precisão FP8 durante o *rollout*, mantendo a precisão BF16 para o treinamento. Neste trabalho, apresentamos o primeiro estudo abrangente sobre o treinamento de RL com FP8 e demonstramos que a estratégia amplamente utilizada de "treinamento em BF16 + *rollout* em FP8" sofre de severa instabilidade de treinamento e colapso catastrófico de precisão em *rollouts* de longo horizonte e tarefas desafiadoras. Nossa análise mostra que essas falhas decorrem da natureza *off-policy* da abordagem, que introduz uma incompatibilidade numérica substancial entre o treinamento e a inferência. Motivados por essas observações, propomos o Jet-RL, um *framework* de treinamento de RL em FP8 que permite uma otimização de RL robusta e estável. A ideia central é adotar um fluxo de precisão FP8 unificado para treinamento e *rollout*, minimizando assim as discrepâncias numéricas e eliminando a necessidade de calibração inter-etapas ineficiente. Experimentos extensivos validam a eficácia do Jet-RL: nosso método alcança uma aceleração de até 33% na fase de *rollout*, até 41% na fase de treinamento e um ganho de velocidade de 16% de ponta a ponta em relação ao treinamento BF16, mantendo uma convergência estável em todas as configurações e incorrendo em uma degradação de precisão insignificante.

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.