FP8-RL: Практичный и стабильный стек низкой точности для обучения с подкреплением больших языковых моделей

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

January 26, 2026

Авторы: Zhaopeng Qiu, Shuang Yu, Jingqi Zhang, Shuai Zhang, Xue Huang, Jingyi Yang, Junjie Lai

cs.AI

Аннотация

Обучение с подкреплением (RL) для больших языковых моделей (LLM) всё чаще упирается в ограничение, связанное с генерацией (rollout), когда длинные выходные последовательности приводят к тому, что внимание и кэш ключей-значений (KV-cache) становятся доминирующими факторами во времени выполнения шага. Формат FP8 предлагает привлекательный способ ускорения RL за счёт снижения вычислительных затрат и объёма передаваемых данных в памяти в процессе генерации. Однако применение FP8 в RL порождает уникальные инженерные и алгоритмические проблемы: веса политики меняются на каждом шаге (что требует повторного квантования и синхронизации весов в механизм вывода), а генерация с низкой точностью может отклоняться от политики с более высокой точностью, предполагаемой обучающим модулем, вызывая рассогласование между обучением и выводом и потенциальную нестабильность. В данном отчёте представлен практический стек генерации на FP8 для RL LLM, реализованный в экосистеме veRL с поддержкой распространённых бэкендов обучения (например, FSDP/Megatron-LM) и механизмов вывода (например, vLLM/SGLang). Мы (i) реализуем генерацию с линейными слоями в формате W8A8 FP8, используя поблочное квантование в FP8, (ii) расширяем использование FP8 на KV-cache для устранения узких мест памяти при длинных контекстах за счёт перекалибровки масштабов QKV на каждом шаге, и (iii) смягчаем рассогласование с помощью коррекции генерации на основе важности сэмплирования (вариации TIS/MIS на уровне токенов). Применение этих методов к плотным и MoE-моделям позволяет достичь повышения пропускной способности генерации до 44% при сохранении поведения обучения, сопоставимого с базовыми уровнями на BF16.

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.