FP8-RL：大規模言語モデルの強化学習における実用的で安定した低精度スタック

要旨

大規模言語モデル(LLM)における強化学習(RL)は、ロールアウト(生成)処理がますますボトルネックとなっている。長い出力シーケンスにより、アテンションとKVキャッシュのメモリ使用量がエンドツーエンドのステップ時間を支配するようになっている。FP8演算は、ロールアウト時の計算コストとメモリ転送量を削減することでRLを加速する有力な手段となるが、RLへのFP8適用には特有の工学的・算法的課題が存在する：ポリシー重みがステップ毎に変化（推論エンジンへの繰り返し量子化と重み同期が必要）、低精度ロールアウトが学習器が想定する高精度ポリシーから乖離（学習-推論間の不一致や不安定性の原因）。本報告では、veRLエコシステム上に実装した実用的なLLM RL向けFP8ロールアウトスタックを提案する。これは一般的な学習バックエンド(FSDP/Megatron-LM等)と推論エンジン(vLLM/SGLang等)をサポートする。(i)ブロック単位FP8量子化によるFP8 W8A8線形層ロールアウトの実現、(ii)ステップ毎のQKVスケール再調整による長文脈メモリボトルネック解消のためのKVキャッシュへのFP8拡張、(iii)重要度抽様に基づくロールアウト補正（トークンレベルTIS/MIS変種）による不一致緩和を実施。密モデル・MoEモデル双方で、これらの技術によりBF16ベースラインと同等の学習動作を維持しつつ、最大44%のロールアウトスループット向上を達成した。

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.

FP8-RL：大規模言語モデルの強化学習における実用的で安定した低精度スタック

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

要旨

Support