用FP8加速大模型强化学习推理，提升44%吞吐量且不损失训练效果
FP8-RL: A Practical and Stable Low-Precision Stack for LLM Reinforcement Learning
- 采用块级量化实现FP8线性层推理，降低内存占用
- 通过每步重校准实现KV缓存的FP8化，解决长序列瓶颈
- 用重要性采样修正精度损失，避免训练与推理不一致
大语言模型强化学习的瓶颈在于生成阶段，长输出序列导致注意力和键值缓存内存占用过高。FP8可降低计算成本和内存流量，但应用于强化学习时面临独特挑战：策略权重每步变化需重复量化与同步，且低精度生成可能偏离高精度策略，引发训练-推理不匹配与不稳定。本文在veRL生态中构建实用的FP8推理栈，支持常见训练后端（如FSDP/Megatron-LM）和推理引擎（如vLLM/SGLang）。提出三项技术：(i) 块级FP8量化实现W8A8线性层推理；(ii) 通过每步QKV缩放重校准将FP8扩展至键值缓存，消除长上下文内存瓶颈；(iii) 使用基于重要性采样的生成修正（令牌级TIS/MIS变体）缓解偏差。在密集模型与MoE模型上，该方案实现最高44%的生成吞吐提升，同时保持与BF16基线相当的学习行为。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。