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

TL;DR

This work tackles the rollout bottleneck in on-policy RL for large language models by identifying training-rollout precision mismatch as a key source of instability when using FP8 in rollout. It introduces Jet-RL, a unified FP8 precision flow that aligns training and rollout, along with a mixed granularity FP8 GEMM design and implementation that leverages FP8 kernels to accelerate end-to-end RL training. The approach yields robust convergence across multiple models and tasks, achieving up to 33% rollout speedup, 41% training speedup, and 16% end-to-end speedup with minimal accuracy loss compared to BF16 baselines. Overall, Jet-RL demonstrates that on-policy FP8 RL can be both faster and more stable, enabling scalable, efficient reasoning in large language models.

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.

Figures (6)

• Figure 1: Overview of RL training different between JetRL and other methods. JetRL proposes unified precision flow for FP8 RL training accommodate both performance and throughput .
• Figure 2: Rollout Generation dominates the RL training latency. When the rollout length is larger than 8k, rollout will take $>75\%$ of the total latency, making it the primary bottleneck.
• Figure 3: Naive BF16 Train + FP8 Rollout fails when rollout context length increases. We observe that although the BF16-Train-FP8-Rollout method might exhibit similar performance compared with BF16 training, its performance quickly degrades as we extend the rollout length to more than 8k. We use the Qwen3-8B-Base model for this experiment, and evaluate and train on MATH.
• Figure 4: We conduct training using Qwen3-8B (reasoning) and Qwen3-8B-Base on MATH. When the model is trained on a task that is easy for itself, BF16-Train-FP8-Rollout is less likely to degrade. When the task is hard for the model we trained on, degradation is likely to happen.
• Figure 5: Overall FP8 precision flow of Jet-RL. The inference graph is a subgraph of the training graph.