Jet-RL: Реализация обучения с подкреплением на FP8 с он-политикой с использованием единого потока точности для фазы обучения и исполнения политики

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

January 20, 2026

Авторы: Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

cs.AI

Аннотация

Обучение с подкреплением (RL) играет ключевую роль в повышении способностей крупных языковых моделей (LLM) к сложным рассуждениям. Однако существующие конвейеры RL-обучения вычислительно неэффективны и ресурсоемки, причем фаза прокрутки (rollout) занимает более 70% общего времени обучения. Квантованное RL-обучение, в частности с использованием точности FP8, предлагает перспективный подход к устранению этого узкого места. Распространенная стратегия предполагает применение точности FP8 на этапе прокрутки при сохранении точности BF16 для обучения. В данной работе мы представляем первое комплексное исследование RL-обучения с FP8 и показываем, что широко используемая стратегия «BF16-обучение + FP8-прокрутка» страдает от серьезной нестабильности обучения и катастрофического коллапса точности при длительных прокрутках и сложных задачах. Наш анализ показывает, что эти сбои вызваны внеполитическим (off-policy) характером подхода, который вносит существенное численное рассогласование между обучением и выводом. Мотивированные этими наблюдениями, мы предлагаем Jet-RL — фреймворк для FP8 RL-обучения, который обеспечивает robustную и стабильную RL-оптимизацию. Ключевая идея заключается в использовании единого потока точности FP8 как для обучения, так и для прокрутки, что минимизирует численные расхождения и устраняет необходимость в неэффективной межшаговой калибровке. Многочисленные эксперименты подтверждают эффективность Jet-RL: наш метод обеспечивает ускорение до 33% на фазе прокрутки, до 41% на фазе обучения и 16% сквозное ускорение по сравнению с обучением на BF16, сохраняя при этом стабильную сходимость во всех настройках и демонстрируя пренебрежимо малую деградацию точности.

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.