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

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

Source-linked AI summary

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

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

arXiv:2601.14243v2Jan 20, 2026cs.LGcs.CL

TL;DR

RL training is inefficient because long rollouts dominate runtime, while BF16-training plus FP8-rollout can suffer instability from training–rollout mismatch. Jet-RL uses a unified FP8 precision flow for training and rollout, delivering substantial speedups with stable convergence and near-BF16 accuracy.

Problem
Long autoregressive rollouts account for over 70% of RL training time, while BF16-training plus FP8-rollout can fail on long rollouts and challenging tasks because training and rollout precision differ.
Method
Jet-RL enforces an identical FP8 precision flow for training and inference rollout to reduce policy mismatch and remove inter-step calibration.
Results
Jet-RL robustly converges across evaluated settings, with up to 1.33× rollout speedup, 1.41× training speedup, and 1.16× end-to-end speedup while remaining close to BF16 performance.
Takeaways & Limitations
Unified FP8 precision provides a reliable path for accelerating low-precision RL training without sacrificing model accuracy.
Takeaways & Limitations
BF16-training plus FP8-rollout remains fragile and tends to fail depending on dataset and model settings, particularly for long rollouts and challenging tasks.

Abstract

from arXiv · show

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.

1. Introduction

RL training is bottlenecked by long autoregressive rollouts, while BF16-training plus FP8-rollout can become unstable on long or challenging tasks. Jet-RL addresses this mismatch with unified FP8 precision across training and rollout and reports substantial speedups with near-BF16 performance.

Motivation: Rollout generation exceeds 70% of total RL training time when sequences are longer than 8K tokens, making it the primary efficiency bottleneck.Profiling shows rollout generation alone accounts for over 70% of training time beyond 8K rollout length.
Motivation: BF16-training plus FP8-rollout is widely used to accelerate rollouts while retaining BF16 training for expected stability and accuracy.The strategy is adopted in modern RL frameworks including VeRL, SLIME, NeMo-RL, and OpenRLHF.
Failure modes: The mixed-precision strategy rapidly loses accuracy beyond 8K-token rollouts and can collapse at 16K, especially in long-generation settings.The reported Qwen3-8B-Base experiment compares 4K, 8K, and 16K generations against BF16 training.
Failure modes: Accumulated numerical discrepancies between training and rollout amplify off-policy effects, causing trajectory divergence and unstable RL optimization.The mismatch is negligible for short sequences but progressively increases with rollout length.
Failure modes: BF16-training plus FP8-rollout is fragile on harder reasoning tasks or weaker base models, where its training curve quickly diverges from BF16.The strategy may show no degradation when models have strong task priors from pretraining.
Jet-RL: Jet-RL enforces an identical FP8 precision flow for training and inference, eliminating policy mismatch and inter-step calibration.The framework uses mixed per-group and per-block quantization with FP8 GEMM kernels.
Results: Jet-RL achieves up to 1.33× rollout speedup, 1.41× training speedup, and 1.16× end-to-end speedup while maintaining convergence and near-BF16 accuracy.The reported BF16-to-FP8 performance gap is reduced from typically more than 5% to approximately 1%.

2. Background

The paper introduces FP8 quantization, linear-layer GEMM workloads, and the phases and computational roles of a standard RL training pipeline. Rollout generation is identified as the dominant performance bottleneck as response length increases.

2.1. Quantization Basis: FP8 quantization maps high-precision tensors to lower-precision representations to accelerate computation and reduce memory usage.The paper focuses on FP8 E4M3, whose maximum representable value is Δmax = 448.
2.2. Quantization of a Linear Layer: Linear layers contain three GEMMs: FProp computes outputs, WGrad computes weight gradients, and DGrad computes activation gradients.The corresponding matrix equations define the forward and backward computations.
2.2. Quantization of a Linear Layer: FP8 GEMMs require specific row-wise and column-wise operand layouts summarized in Table 1.The first operand uses row-wise storage and the second uses column-wise storage for current FP8 TensorCore hardware.
2.3. Workload of Reinforcement Learning: A PPO-style RL pipeline uses Actor, Reference, Reward, and Critic models for policy optimization, regularization, reward computation, and value estimation.The Actor generates responses, while the other models evaluate them and the Actor is subsequently updated.
2.3. Workload of Reinforcement Learning: RL training comprises Rollout, Evaluation, and Update phases, with rollout handled by inference engines and evaluation/update handled by training frameworks.Rollout performs autoregressive decoding; Evaluation runs model forward passes; Update performs the Actor’s training step.
2.3. Workload of Reinforcement Learning: When rollout lengths increase, autoregressive generation dominates RL training latency and becomes the pipeline’s most expensive component.The paper reports that rollout latency scales with generated-response length.

3. Motivation

The paper motivates Jet-RL by showing that BF16-train-FP8-rollout can become unstable or inaccurate, especially with long rollouts and challenging tasks. These failures are linked to numerical mismatch between rollout and training, motivating a more reliable approach.

3.1. Rollout is the Bottleneck in RL Training: When rollout length exceeds 8K tokens, rollout generation alone accounts for over 70% of total RL training time.The measurement varies maximum rollout length from 1K to 16K on Qwen3-8B-Base using GSM8K and MATH.
3.1. Rollout is the Bottleneck in RL Training: FP8 is motivated as a rollout-acceleration technique because it can offer an ideal 2× speedup over BF16 and is widely reported not to degrade inference performance.The paper also notes that FP8 can be integrated into existing RL training pipelines.
3.2. BF16-Train-FP8-Rollout with Calibration is Slow: Calibration-based FP8 approaches are impractical for frequently synchronized RL actors because data-dependent calibration can take tens of minutes even for small 8B models.Repeating calibration at every synchronization step conflicts with the goal of acceleration.
3.3. BF16-Train-FP8-Rollout without Calibration is Unstable: BF16-train-FP8-rollout breaks the on-policy consistency requirement because rollout and training can produce different logits for identical prompts.The resulting mismatch can cause unstable optimization and degraded performance.
3.3. BF16-Train-FP8-Rollout without Calibration is Unstable: The paper identifies long-rollout generation and challenging tasks as the two settings where BF16-train-FP8-rollout tends to fail.The reported conclusion is described as fragile and dependent on dataset and model settings.
3.2. BF16-Train-FP8-Rollout with Calibration is Slow: With 16K-token rollouts, FP8 rollout accuracy collapses after only 20 training steps, whereas performance is on par with BF16 at rollout lengths below 4K.The FP8 rollout begins diverging from BF16 when rollout length reaches 8K.
3.3. BF16-Train-FP8-Rollout without Calibration is Unstable: BF16-train-FP8-rollout is more likely to degrade on difficult tasks or weaker base models than when the model has strong task priors.The Qwen3-8B-Base setting falls behind BF16 on GSM8K, while Qwen3-8B reasoning tracks BF16 more closely.

4. Jet-RL: Enabling on-policy FP8 RL Training

Jet-RL aligns training and rollout through a unified FP8 precision flow, making the process on-policy and reducing numerical mismatch. Its design combines graph-consistent quantization, FP8 GEMMs, and finer-grained quantization choices for stable acceleration.

Unified FP8 Precision Flow: Jet-RL makes the inference graph a precision-consistent subgraph of the training graph, converting BF16-train-FP8-rollout from an off-policy process into an on-policy one.The forward pass in training and inference shares quantization behavior while retaining a higher-precision BF16 master weight for training stability.
Precision-Flow Modeling: The precision-propagation graph represents operators and weights as nodes, with edges encoding tensor precision and quantization granularity.Training and inference graphs share operator topology, while forward, backward, and activation-save connections describe the training flow.
Unified FP8 Precision Flow: Jet-RL stores backward-pass activations in FP8 because downstream GEMMs access the quantized tensors produced by the forward pass.This choice is supported by prior evidence that FP8 activation storage can maintain training stability.
Backward Computation: Gradients transported between operators remain in BF16, while DGrad and WGrad GEMMs use FP8 to preserve accuracy while accelerating computation.Quantizing transported gradients can cause underflow or quantization noise that degrades convergence.
GEMM Quantization Granularity: Jet-RL uses 128 × 128 per-block weight quantization and 1 × 128 per-group activation and gradient quantization instead of unstable per-tensor quantization.FProp uses an (1 × 128) × (128 × 128) configuration, while WGrad uses a finer-grained 1 × 128 by 128 × 1 configuration.

5. Evaluation

Jet-RL is evaluated across multiple models, datasets, rollout lengths, benchmarks, and tensor-parallel settings. It remains more stable than BF16-train-FP8-rollout and delivers rollout, training, and end-to-end speedups over BF16.

Accuracy Evaluation: Across 8K-rollout experiments, BF16-train-FP8-rollout is unstable, while Jet-RL converges in all reported model settings and stays close to BF16 performance.The evaluation uses Llama3.1-8B, Qwen2.5-7B, and Qwen3-8B-Base on GSM8K + MATH, with results reported against BF16 baselines.
Rollout Efficiency: 1.33× is the maximum FP8 rollout speedup over BF16, with larger models benefiting more and higher tensor parallelism reducing the observed gain.The reported speedup range is 1.07× to 1.33×; communication overhead explains the lower benefit at greater tensor parallelism.
End-to-End Efficiency: 1.41× training-phase throughput and 1.16× end-to-end step-time speedup are achieved for Qwen3-8B with an 8K rollout.Actor update reaches 1.54× and reference-model inference 1.80×, contributing to the training-phase improvement.

6. Related Works

Related work spans low-precision inference and training, RL methods for alignment and reasoning, systems optimization, trajectory control, and quantization–adaptation co-design. Jet-RL is positioned within efforts to reduce the computational cost of long on-policy rollouts and low-precision training.

Low-Precision Training and Inference: Post-training quantization compresses pretrained LLMs without fine-tuning, while fully quantized training performs low-precision computations during training.Prior work includes activation-aware, INT8, and FP8 approaches targeting memory, computation, and accuracy trade-offs.
Reinforcement Learning: RL research has progressed from human-feedback alignment methods toward reasoning-oriented algorithms using automatically verifiable outcomes.Examples include PPO, DPO, GRPO, and GSPO-related approaches.
RL Efficiency Challenges: Long autoregressive rollouts and on-policy data generation make RL for reasoning computationally intensive and can reduce hardware utilization.This motivates both system-level and algorithmic efficiency improvements.
Systems and Algorithmic Optimizations: Asynchronous RL systems improve utilization by decoupling rollout generation from training, whereas trajectory methods address redundant or excessive reasoning.AReaL and ReaLHF target parallelism, while NoThinking targets overthinking and earlier termination.
Co-Design and Off-Policy Mitigation: QeRL combines NVFP4 quantization with LoRA, while truncated importance sampling addresses off-policy effects through clipped importance ratios.These approaches represent co-design and algorithmic alternatives for reducing rollout cost or mitigating policy mismatch.

7. Conclusion

Jet-RL addresses the training–rollout mismatch underlying instability in BF16-training plus FP8-rollout RL by using an identical FP8 precision flow for training and inference. It converges robustly while preserving near-BF16 performance and delivering substantial speedups across RL pipeline stages.

Jet-RL replaces the flawed BF16-train + FP8-rollout strategy with identical FP8 precision flows for training forward passes and inference rollouts.This design targets the mismatch that causes training instability and catastrophic performance collapse.
Jet-RL robustly converges across all evaluated models and benchmark settings, with performance usually within 1% of the BF16 RL baseline.
1.33× rollout phase speedup, 1.41× training phase speedup, and 1.16× end-to-end speedup are achieved without sacrificing model accuracy.