FP8-RL: Una Pila Práctica y Estable de Baja Precisión para el Aprendizaje por Refuerzo de Modelos de Lenguaje Grandes

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

January 26, 2026

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

cs.AI

Resumen

El aprendizaje por refuerzo (RL) para modelos de lenguaje grandes (LLM) se ve cada vez más limitado por la generación (rollout), donde las largas secuencias de salida hacen que la atención y la memoria de la caché KV dominen el tiempo total de paso. FP8 ofrece una palanca atractiva para acelerar el RL al reducir el coste computacional y el tráfico de memoria durante la generación, pero aplicar FP8 en el RL introduce desafíos únicos de ingeniería y algoritmos: los pesos de la política cambian en cada paso (lo que requiere una cuantificación repetida y una sincronización de pesos en el motor de inferencia) y las generaciones de baja precisión pueden desviarse de la política de mayor precisión asumida por el entrenador, causando una discrepancia entrenamiento-inferencia y una posible inestabilidad. Este informe presenta una pila práctica de generación FP8 para RL de LLM, implementada en el ecosistema veRL con soporte para backends de entrenamiento comunes (por ejemplo, FSDP/Megatron-LM) y motores de inferencia (por ejemplo, vLLM/SGLang). Nosotros (i) habilitamos la generación en capas lineales W8A8 FP8 utilizando cuantificación FP8 por bloques, (ii) extendemos FP8 a la caché KV para eliminar los cuellos de botella de memoria en contextos largos mediante el reajuste de la escala QKV por paso, y (iii) mitigamos la discrepancia utilizando una corrección de la generación basada en muestreo por importancia (variantes a nivel de token TIS/MIS). En modelos densos y MoE, estas técnicas ofrecen ganancias de rendimiento en la generación de hasta el 44%, preservando un comportamiento de aprendizaje comparable al de las líneas base 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.