PaperlayerView exact revision on arXiv

Source-linked AI summary

FP8-RL: A Practical and Stable Low-Precision Stack for LLM Reinforcement Learning

Zhaopeng Qiu, Shuang Yu, Jingqi Zhang, Shuai Zhang, Xue Huang, Jingyi Yang, Junjie Lai

arXiv:2601.18150v2cs.LGcs.CL

TL;DR

LLM reinforcement learning is bottlenecked by rollout cost and faces synchronization and train–inference mismatch challenges when using FP8. FP8-RL addresses these issues with dynamic quantization, FP8 linear layers and KV caches, and importance-sampling correction, delivering up to 44% rollout improvement while preserving learning behavior comparable to BF16.

  • Problem

    LLM RL rollouts are costly, while changing policy weights and lower-precision generation create synchronization and train–inference mismatch challenges that can destabilize training.

  • Method

    The report builds a practical FP8 rollout stack with dynamic weight synchronization, blockwise W8A8 linear quantization, per-step QKV-recalibrated KV-cache quantization, and importance-sampling-based rollout correction.

  • Results

    Up to 44% combined rollout improvement is achieved across dense and MoE models while learning behavior remains comparable to BF16 baselines.

  • Takeaways & Limitations

    FP8 can accelerate long-context LLM RL rollouts while retaining BF16-comparable learning behavior when mismatch correction is used.

  • Takeaways & Limitations

    For MoE models, FP8 rollout-only configurations can crash, and token-level TIS alone may be insufficient to control mismatch.

Abstract

from arXiv · show

Reinforcement learning (RL) for large language models (LLMs) is increasingly bottlenecked by rollout (generation), where long output sequence lengths make attention and KV-cache memory dominate end-to-end step time. FP8 offers an attractive lever for accelerating RL by reducing compute cost and memory traffic during rollout, but applying FP8 in RL introduces unique engineering and algorithmic challenges: policy weights change every step (requiring repeated quantization and weight synchronization into the inference engine) and low-precision rollouts can deviate from the higher-precision policy assumed by the trainer, causing train-inference mismatch and potential instability. This report presents a practical FP8 rollout stack for LLM RL, implemented in the veRL ecosystem with support for common training backends (e.g., FSDP/Megatron-LM) and inference engines (e.g., vLLM/SGLang). We (i) enable FP8 W8A8 linear-layer rollout using blockwise FP8 quantization, (ii) extend FP8 to KV-cache to remove long-context memory bottlenecks via per-step QKV scale recalibration, and (iii) mitigate mismatch using importance-sampling-based rollout correction (token-level TIS/MIS variants). Across dense and MoE models, these techniques deliver up to 44% rollout throughput gains while preserving learning behavior comparable to BF16 baselines.

1 Introduction

The report targets rollout as a dominant LLM-RL bottleneck and develops a practical FP8 workflow for dynamic weight synchronization, quantized rollout, and broad backend support. It combines linear-layer and KV-cache quantization with rollout correction, achieving substantial throughput gains while preserving learning behavior comparable to BF16 baselines.

  • Rollout can consume approximately 80% of total iteration time in synchronous LLM RL, making it a high-leverage target for acceleration.
  • FP8 rollouts must repeatedly synchronize and requantize changing policy weights, while low-precision generation can create train–inference mismatch and destabilize RL.
  • The proposed production-ready workflow supports dynamic weight synchronization across FSDP/Megatron-LM training backends and vLLM/SGLang inference engines.
  • Blockwise W8A8 linear quantization reduces rollout compute cost, while per-step QKV scale recalibration enables FP8 KV-cache quantization for greater long-context capacity and concurrency.
  • Across dense and MoE models, the approach delivers up to 44% rollout speedup while maintaining learning behavior comparable to BF16 baselines with standard rollout correction.

2 Our Solution: FP8 Rollout for LLM RL

The FP8 rollout system addresses dynamic weight synchronization and train–inference mismatch while accelerating generation through linear-layer and KV-cache quantization. Across dense and MoE models, rollout correction preserves learning behavior comparable to BF16 while delivering substantial speedups.

  • The system combines dynamic weight synchronization with importance-sampling correction to support stable FP8 rollouts in LLM RL.It is implemented in verl and targets the mismatch between quantized rollout and higher-precision training policies.
  • Mismatch Mitigation: Token-level TIS keeps FP8 dense-model training aligned with BF16, whereas omitting TIS causes noticeable accuracy degradation.The correction keeps quantization-induced divergence stable and within an acceptable range.
  • Dense Model: 10–20% speedup: FP8 W8A8 rollout outperforms BF16 on the Qwen3-8B-Base dense model, with larger gains at longer sequence lengths.Reduced memory traffic becomes more beneficial as memory bandwidth constrains long-context generation.
  • MoE Model: 30–50% speedup: FP8 W8A8 rollout improves time-per-token over BF16 on the Qwen3-30B-A3B-Base MoE model.The improvement is reported as 2–3× larger than the 10–20% dense-model speedup.
  • MoE Router Precision: Excluding the MoE router from FP8 quantization is sufficient to maintain routing consistency, with BF16 and FP32 routers showing similarly lower mismatch KL than an FP8 router.The FP8 router reaches approximately 0.004 mismatch KL early in training, while BF16 and FP32 routers remain closer to the BF16 baseline.
  • End-to-End FP8: Full FP8 remains viable with token-level TIS, while FP8 rollout-only can develop a KL spike above 5 and crash training.End-to-end FP8 follows the BF16 accuracy trend despite residual mismatch from kernel-level numerical differences.
  • KV Cache Quantization: 44% speedup: combining linear-layer and KV-cache quantization outperforms the approximately 20% linear-only and 38% KV-cache-only gains.KV-cache quantization targets the memory-bound bottleneck of long-context generation, while both components provide complementary benefits.

3 Related Work

Prior work combines low-precision computation with optimized serving and distributed RL architectures, while training–inference mismatch remains a central stability concern. Importance sampling methods address this mismatch by correcting off-policy updates, with truncated or masked variants controlling variance.

  • FP8 research established E4M3 and E5M2 encodings, while later work applied FP8 across LLM weights, gradients, and optimizer states.
  • Inference frameworks such as vLLM, TensorRT-LLM, and SGLang combine reduced precision with optimized attention, KV-cache management, and runtime scheduling.
  • Distributed RL systems decouple model roles and integrate specialized training and inference backends to improve resource utilization and generation performance.
  • Even synchronized weights can leave rollout and training policies divergent because of differences in numerics, kernels, and parallelism, creating instability or collapse.
  • Importance sampling corrects off-policy updates with likelihood ratios, while TIS and MIS reduce variance by clipping extreme weights or masking unreliable tokens or sequences.

4 Summary and Future Work

FP8-RL combines quantization and rollout-correction techniques to accelerate LLM RL while preserving BF16-comparable learning behavior. The report also identifies broader validation and deployment boundaries for the approach.

  • Summary: Up to 44% combined improvement is achieved through W8A8 linear quantization and KV-cache quantization in FP8-RL.W8A8 provides 10–20% rollout speedup for dense models and 30–50% for MoE models; KV-cache quantization adds gains under memory pressure.
  • Summary: Importance-sampling-based correction remains essential for maintaining accuracy comparable to BF16 baselines across settings.The correction addresses distribution shift between quantized rollout and full-precision training.
  • Future Work: Future work includes NVFP4, larger models, and multiturn or agentic RL scenarios with longer contexts.The report notes instability from accumulated quantization error for NVFP4 and identifies broader validation as future work.
  • Implementation: FP8 rollout requires automatic weight conversion and loading into the inference backend during rollout.The configuration enables FP8 quantization through a rollout setting, with the system handling conversion and loading into vLLM.
  • Implementation: vLLM 0.11+ and SGLang 0.55+ enable DeepGEMM by default, while vLLM 0.10.x requires manual activation.CUDA 12.9 or higher is also required for the reported FP8 DeepGEMM performance recommendations.

B.1 verl Implementation (Inference-Side Calibration)

The verl implementation enables inference-side FP8 KV-cache quantization through dynamic QKV scale recalculation. It recalibrates scales before each rollout and can combine KV-cache quantization with FP8 linear W8A8.

  • B.1 verl Implementation: The verl implementation uses vLLM’s native dynamic QKV scale calculation for FP8 KV-cache quantization.The configuration specifies FP8 E4M3 KV-cache storage and enables scale calculation.
  • B.1 verl Implementation: Scale recalibration is forced before each rollout phase by resetting internal calculation flags.This supports updated scale computation at rollout-phase granularity.
  • B.1 verl Implementation: vLLM 0.11 or later is required for FP8 KV-cache support.The implementation automatically updates QKV scales per step without external intervention.
  • B.1 verl Implementation: FP8 KV-cache quantization can be combined with FP8 linear W8A8 quantization for maximum throughput.The two quantization levels target complementary rollout bottlenecks.

B.2 NeMo-RL Implementation (Trainer-Side Calibration)

The NeMo-RL implementation performs QKV scale recalibration on the training side using updated policy weights and selected training data. New scales are synchronized to the inference engine for subsequent rollouts.

  • B.2 NeMo-RL Implementation: NeMo-RL recalibrates QKV scales at the end of each training step using updated policy weights.This places scale recalibration on the trainer side rather than relying solely on inference-side calculation.
  • B.2 NeMo-RL Implementation: Calibration uses a subset of prompts and generated responses from the training data.The implementation provides control over the calibration data distribution.
  • B.2 NeMo-RL Implementation: New QKV scales are synchronized to the inference engine for the subsequent rollout.The process keeps inference-side cache quantization aligned with recalibrated trainer-side scales.
  • B.2 NeMo-RL Implementation: Calibration overhead is approximately 2-3% of total step time.The reported overhead is measured for the trainer-side recalibration procedure.

B.3 NeMo-RL Experimental Validation

The trainer-side calibration approach was validated on Qwen3-8B-Base using the NeMo-RL framework.

  • B.3 NeMo-RL Experimental Validation: The NeMo-RL framework validates trainer-side QKV scale calibration on Qwen3-8B-Base.This passage reports the validation target and framework but does not state a performance result.

B.3.1 Experimental Setup

The setup compares BF16, linear-layer W8A8, and full FP8 quantization, with token-level TIS applied to all FP8 configurations for rollout correction.

  • Configurations evaluated: The evaluation compares a BF16 baseline, Linear W8A8, and Full FP8 configurations.Full FP8 quantizes linear layers, KV-cache storage, and attention computations.
  • Rollout correction: All FP8 configurations use token-level TIS with C=2 for rollout correction.

B.3.2 Training Effectiveness

Token-level TIS restores validation-accuracy alignment across quantization configurations, although jointly quantizing KV-cache and attention increases mismatch KL divergence.

  • Accuracy alignment: After enabling token-level TIS, Full FP8 validation accuracy aligns with the BF16 baseline and Linear W8A8.
  • Mismatch KL: Mismatch KL divergence is higher when both KV-cache and attention are quantized to FP8 than with Linear W8A8.The passage attributes this to compounded quantization errors across multiple components.
  • Calibration comparison: Trainer-Side calibration produces training-stability results largely consistent with Inference-Side calibration.Both calibration paradigms are reported to achieve comparable effectiveness under FP8 KV-cache quantization.

B.3.3 Performance Results

Adding FP8 quantization to KV-cache and attention substantially improves rollout performance, especially for longer responses, while calibration adds only a small step-time cost.

  • Speedup magnitude: An additional ~30% speedup from FP8 KV-cache and attention yields an overall ~48% speedup over the BF16 baseline.
  • Length dependency: Performance gains become more pronounced at longer response lengths as attention represents a larger workload share.
  • Calibration overhead: QKV scale recalibration consumes approximately 2-3% of total step time.The reported overhead is minor relative to the rollout acceleration.
  • Calibration comparison: Trainer-Side calibration achieves comparable training effectiveness and performance gains to Inference-Side calibration.

C.1 Experimental Setup

The end-to-end evaluation compares BF16 and FP8 training-rollout combinations on Qwen3-8B-Base, using figures that track training curves and rollout performance across these precision settings.

  • Configurations: The experiment compares BF16 training with BF16 rollout, BF16 training with FP8 rollout, and FP8 training with FP8 rollout.
  • Rollout performance: The evaluation uses Qwen3-8B-Base with Trainer-Side calibration for rollout-performance comparison.
  • Training effectiveness: Training curves compare FP8 training plus FP8 rollout, BF16 training plus FP8 rollout, and BF16 training plus BF16 rollout.

C.2 Results

Across the reported training dynamics, end-to-end FP8 closely tracks BF16 while reducing training–inference mismatch relative to FP8 rollout-only, though residual mismatch remains.

  • FP8 training + rollout closely tracks BF16 training + rollout across response length, reward, and validation accuracy.The comparison indicates preserved learning dynamics and final quality under this setup.
  • FP8 training + rollout exhibits lower training–inference mismatch KL than FP8 rollout-only.Aligning precision between training and rollout reduces distribution drift.
  • FP8 training + rollout still shows slightly higher mismatch than BF16 training + rollout.Precision alignment helps but does not fully remove all sources of mismatch.
Loading 2601.18150v2…