Jet-RL: 통합 훈련 및 롤아웃 정밀도 흐름을 통한 온-폴리시 FP8 강화 학습 가능화

초록

강화학습(RL)은 대규모 언어 모델(LLM)의 복잡한 추론 능력 향상에 필수적입니다. 그러나 기존 RL 학습 파이프라인은 계산 효율성이 낮고 리소스 집약적이며, 롤아웃 단계가 전체 학습 시간의 70% 이상을 차지합니다. 양자화된 RL 학습, 특히 FP8 정밀도를 사용하는 방식은 이러한 병목 현상을 완화할 수 있는 유망한 접근법입니다. 일반적으로 채택되는 전략은 롤아웃 시에는 FP8 정밀도를 적용하고 학습 시에는 BF16 정밀도를 유지하는 것입니다. 본 연구에서는 FP8 RL 학습에 대한 최초의 포괄적인 연구를 제시하며, 널리 사용되는 BF16 학습 + FP8 롤아웃 전략이 장기 롤아웃 및 난이도 높은 작업에서 심각한 학습 불안정성과 치명적 정확도 붕괴를 겪는다는 것을 입증합니다. 우리의 분석에 따르면 이러한 실패는 접근법의 오프-폴리시 특성에서 비롯되며, 이는 학습과 추론 간에 상당한 수치적 불일치를 초래합니다. 이러한 관찰을 바탕으로 우리는 강력하고 안정적인 RL 최적화를 가능하게 하는 FP8 RL 학습 프레임워크인 Jet-RL을 제안합니다. 핵심 아이디어는 학습과 롤아웃 모두에 통합된 FP8 정밀도 흐름을 채택하여 수치적 불일치를 최소화하고 비효율적인 단계 간 보정의 필요성을 제거하는 것입니다. 광범위한 실험을 통해 Jet-RL의 효과를 검증했습니다. 우리의 방법은 BF16 학습 대비 롤아웃 단계에서 최대 33%, 학습 단계에서 최대 41%의 속도 향상과 16%의 종단 간 속도 향상을 달성하면서도 모든 설정에서 안정적인 수렴을 유지하고 미미한 정확도 저하만 발생시켰습니다.

English

Reinforcement learning (RL) is essential for enhancing the complex reasoning capabilities of large language models (LLMs). However, existing RL training pipelines are computationally inefficient and resource-intensive, with the rollout phase accounting for over 70% of total training time. Quantized RL training, particularly using FP8 precision, offers a promising approach to mitigating this bottleneck. A commonly adopted strategy applies FP8 precision during rollout while retaining BF16 precision for training. In this work, we present the first comprehensive study of FP8 RL training and demonstrate that the widely used BF16-training + FP8-rollout strategy suffers from severe training instability and catastrophic accuracy collapse under long-horizon rollouts and challenging tasks. Our analysis shows that these failures stem from the off-policy nature of the approach, which introduces substantial numerical mismatch between training and inference. Motivated by these observations, we propose Jet-RL, an FP8 RL training framework that enables robust and stable RL optimization. The key idea is to adopt a unified FP8 precision flow for both training and rollout, thereby minimizing numerical discrepancies and eliminating the need for inefficient inter-step calibration. Extensive experiments validate the effectiveness of Jet-RL: our method achieves up to 33% speedup in the rollout phase, up to 41% speedup in the training phase, and a 16% end-to-end speedup over BF16 training, while maintaining stable convergence across all settings and incurring negligible accuracy degradation.

Jet-RL: 통합 훈련 및 롤아웃 정밀도 흐름을 통한 온-폴리시 FP8 강화 학습 가능화

Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow

초록

Support