Jet-RL: Habilitación del Aprendizaje por Refuerzo FP8 en Política con un Flujo de Precisión Unificado para Entrenamiento y Ejecución

Resumen

El aprendizaje por refuerzo (RL) es fundamental para mejorar las capacidades de razonamiento complejo de los modelos de lenguaje grandes (LLM). Sin embargo, las canalizaciones de entrenamiento de RL existentes son computacionalmente ineficientes y requieren muchos recursos, siendo la fase de ejecución (rollout) la responsable de más del 70% del tiempo total de entrenamiento. El entrenamiento de RL cuantizado, particularmente el uso de precisión FP8, ofrece un enfoque prometedor para mitigar este cuello de botella. Una estrategia comúnmente adoptada aplica precisión FP8 durante la ejecución, manteniendo la precisión BF16 para el entrenamiento. En este trabajo, presentamos el primer estudio integral del entrenamiento de RL con FP8 y demostramos que la estrategia ampliamente utilizada de "entrenamiento-BF16 + ejecución-FP8" sufre de una grave inestabilidad en el entrenamiento y un colapso catastrófico de la precisión bajo ejecuciones de horizonte largo y tareas desafiantes. Nuestro análisis muestra que estos fallos se originan en la naturaleza fuera de política (off-policy) del enfoque, que introduce una discrepancia numérica sustancial entre el entrenamiento y la inferencia. Motivados por estas observaciones, proponemos Jet-RL, un marco de entrenamiento de RL con FP8 que permite una optimización de RL robusta y estable. La idea clave es adoptar un flujo de precisión FP8 unificado tanto para el entrenamiento como para la ejecución, minimizando así las discrepancias numéricas y eliminando la necesidad de una calibración entre pasos ineficiente. Experimentos exhaustivos validan la efectividad de Jet-RL: nuestro método logra hasta un 33% de aceleración en la fase de ejecución, hasta un 41% de aceleración en la fase de entrenamiento, y una aceleración integral del 16% sobre el entrenamiento BF16, manteniendo una convergencia estable en todos los escenarios y con una degradación de precisión insignificante.

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: Habilitación del Aprendizaje por Refuerzo FP8 en Política con un Flujo de Precisión Unificado para Entrenamiento y Ejecución

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

Resumen

Support