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

2026-01-20

AI Summary

This work addresses the computational inefficiency of existing reinforcement learning (RL) training pipelines and the numerical instability arising from mixed-precision strategies—specifically, BF16 training coupled with FP8 inference—which cause precision collapse and training divergence in long-horizon tasks. To resolve this, we propose Jet-RL, the first end-to-end fully FP8 on-policy RL framework that unifies numerical precision across both training and rollout phases, thereby eliminating the mismatch inherent in hybrid-precision approaches without requiring calibration. Experimental results demonstrate that Jet-RL achieves a 33% speedup in the rollout phase and a 41% acceleration in training, yielding a 16% end-to-end performance gain, all while maintaining stable convergence and incurring only negligible accuracy degradation.

Technology Category

• Machine Learning: Reinforcement Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

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.

Problem

• Reinforcement Learning
• FP8 Quantization
• Training Instability
• Numerical Mismatch
• Rollout Efficiency

Innovation

• FP8 quantization
• on-policy reinforcement learning
• unified precision training
• numerical stability
• efficient LLM alignment