Technical Summary: Jet-RL

Problem Statement

Reinforcement Learning (RL) has become pivotal for enhancing the complex reasoning capabilities of Large Language Models (LLMs), particularly through Chain-of-Thought (CoT) generation. However, RL training pipelines are computationally inefficient, with the rollout phase (generating responses) consuming over 70% of total training time, especially when generation lengths exceed 8,000 tokens.

To address this bottleneck, a common strategy involves using FP8 quantization during the rollout phase while retaining BF16 precision for the training phase (BF16-train + FP8-rollout). While this approach is widely adopted in frameworks like VeRL and OpenRLHF, this paper identifies two critical failure modes:

1. Long-Rollout Instability: As generation length increases (e.g., >8k tokens), the numerical mismatch between the BF16 training distribution and the FP8 rollout distribution accumulates. This leads to severe training instability and catastrophic accuracy collapse.
2. Challenging Task Failure: The strategy fails when applied to difficult reasoning tasks or weaker base models. In these scenarios, the model's lower confidence amplifies quantization errors, causing the training curve to diverge rapidly compared to full-precision baselines.

The root cause is identified as an off-policy nature: the actor learns from data generated by a policy (FP8 rollout) that differs numerically from the policy being updated (BF16 training), violating the on-policy assumption essential for stable RL.

Methodology: Jet-RL

The authors propose Jet-RL, a framework designed to enable robust, on-policy FP8 RL training by enforcing a unified precision flow between training and rollout.

Core Architecture

Jet-RL eliminates the precision mismatch by ensuring that both the training forward pass and the inference rollout utilize the identical FP8 quantization flow.

• Unified Precision Graph: The framework models the model as a directed graph where edges represent tensor propagation. Jet-RL ensures the inference graph ($\mathcal{G}_{infer}$) is a subgraph of the training forward graph ($\mathcal{G}^{fwd}_{train}$), sharing the same quantization granularity and precision.
• Master Weights: While the forward pass and rollout operate in FP8, the master weights are retained in BF16 to stabilize training updates, similar to standard mixed-precision training.

Quantization Scheme

To maintain stability while accelerating computation, Jet-RL employs a mixed quantization strategy for Linear Layers (GEMMs):

• Activations: Quantized using 1×128 per-group quantization.
• Weights: Quantized using 128×128 per-block quantization.
• Backward Pass:
  • Activations saved for the backward pass are stored in FP8 (fused with the previous operator).
  • Gradients transported between operators are kept in BF16 to prevent gradient underflow and convergence degradation.
  • The backward GEMM operations (DGrad and WGrad) are performed in FP8 for acceleration, utilizing specific kernel configurations (e.g., 1×128 and 128×1 quantization for WGrad) to align with hardware requirements.

Implementation

Jet-RL is implemented using vLLM for the inference engine and VeRL for the RL training framework. It leverages state-of-the-art FP8 GEMM kernels from DeepGEMM and custom Triton kernels for quantization, transposition, and fused activation/RMSNorm operations.

Key Contributions

1. Problem Identification: The paper demonstrates that the standard BF16-train + FP8-rollout strategy is fundamentally flawed for long-rollout and challenging tasks due to off-policy numerical mismatches, leading to training collapse.
2. Jet-RL Framework: The proposal of a unified FP8 precision flow that enforces on-policy consistency between training and rollout, eliminating the need for inter-step calibration and resolving policy mismatch.
3. Performance and Stability: The demonstration that Jet-RL achieves significant speedups while maintaining convergence and accuracy close to BF16 baselines across diverse models and tasks.

Experimental Results

The authors evaluated Jet-RL on models including Llama3.1-8B, Qwen2.5-7B, and Qwen3-8B-Base across datasets like GSM8K, MATH, and DeepMATH.

Accuracy and Convergence

• Stability: Unlike the BF16-train + FP8-rollout baseline, which failed to converge on Qwen2.5-7B and showed significant degradation on 16k rollouts, Jet-RL converged successfully in all settings.
• Accuracy Gap: Jet-RL reduced the performance gap relative to BF16 training to approximately 1% (e.g., 1.1% degradation on Qwen3-8B-Base vs. 2.9% for the baseline). In some cases (Llama3.1-8B), Jet-RL even outperformed the BF16 baseline by 2.0%.
• Robustness: Jet-RL maintained performance on difficult tasks (DeepMATH) and long contexts (16k) where the naive FP8 rollout strategy collapsed.

Efficiency

• Rollout Speedup: FP8 quantization provided a 1.07× to 1.33× speedup in the rollout phase (tokens/s), with larger models (32B) seeing the highest gains.
• Training Speedup: The unified FP8 flow accelerated the actor update phase by 1.54× and reference model inference by 1.80×, resulting in a 1.41× speedup for the training phase.
• End-to-End: The framework achieved a 1.16× end-to-end speedup for 8B model training, with expectations of higher gains for larger models.

Significance and Claims

The paper claims that Jet-RL provides a reliable and efficient path for applying FP8 computation to accelerate large-scale RL training. By addressing the fundamental instability caused by precision mismatches, Jet-RL enables the use of FP8 not just for inference, but for the entire RL training loop. This allows for significant acceleration (up to 33% rollout speedup and 41% training speedup) without sacrificing the convergence stability or accuracy required for training complex reasoning models. The work establishes that a unified precision flow is essential for scaling RL training with low-precision hardware.