FP8-RL: 대규모 언어 모델 강화 학습을 위한 실용적이고 안정적인 저정밀도 스택

초록

대규모 언어 모델(LLM)의 강화 학습(RL)은 롤아웃(생성) 과정에서 점점 더 병목 현상을 겪고 있으며, 긴 출력 시퀀스 길이로 인해 어텐션과 KV 캐시 메모리가 종단 간 단계 시간을 지배하게 됩니다. FP8은 롤아웃 중 컴퓨팅 비용과 메모리 트래픽을 줄여 RL 가속화에 매력적인 수단을 제공하지만, FP8을 RL에 적용할 때는 고유한 엔지니어링 및 알고리즘적 과제가 발생합니다: 정책 가중치가 매 단계마다 변경되어(추론 엔진으로의 반복적인 양자화 및 가중치 동기화 필요) 저정밀도 롤아웃이 학습기가 가정한 고정밀도 정책과 차이를 보여 학습-추론 불일치와 잠재적 불안정성을 초래할 수 있습니다. 본 보고서는 veRL 생태계 내에서 구현되고 일반적인 학습 백엔드(FSDP/Megatron-LM 등) 및 추론 엔진(vLLM/SGLang 등)을 지원하는 실용적인 LLM RL용 FP8 롤아웃 스택을 소개합니다. 우리는 (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