FP8-RL: Uma Pilha Prática e Estável de Baixa Precisão para Aprendizagem por Reforço de LLM

Resumo

O aprendizado por reforço (RL) para grandes modelos de linguagem (LLMs) está cada vez mais limitado pela etapa de *rollout* (geração), na qual sequências de saída longas fazem com que a atenção e a memória da *KV-cache* dominem o tempo total da etapa. O formato FP8 oferece uma alavanca atraente para acelerar o RL, reduzindo o custo computacional e o tráfego de memória durante o *rollout*. No entanto, a aplicação do FP8 no RL introduz desafios únicos de engenharia e algoritmo: os pesos da política mudam a cada passo (exigindo quantização repetida e sincronização dos pesos no motor de inferência) e os *rollouts* de baixa precisão podem se desviar da política de maior precisão assumida pelo treinador, causando incompatibilidade entre treino e inferência e potencial instabilidade. Este relatório apresenta uma *stack* prática de *rollout* em FP8 para RL de LLMs, implementada no ecossistema veRL com suporte para *backends* de treino comuns (por exemplo, FSDP/Megatron-LM) e motores de inferência (por exemplo, vLLM/SGLang). Nós (i) habilitamos o *rollout* de camadas lineares W8A8 em FP8 usando quantização em blocos (*blockwise*), (ii) estendemos o FP8 para a *KV-cache* para remover os gargalos de memória de contexto longo via recalibração da escala QKV por passo, e (iii) mitigamos a incompatibilidade usando correção de *rollout* baseada em amostragem por importância (variantes TIS/MIS a nível de *token*). Em modelos densos e MoE, essas técnicas proporcionam ganhos de até 44% na taxa de transferência (*throughput*) do *rollout*, preservando um comportamento de aprendizado comparável às baselines em 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.

FP8-RL: Uma Pilha Prática e Estável de Baixa Precisão para Aprendizagem por Reforço de LLM

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

Resumo

Support