Resumen Técnico: Jet-RL

Declaración del Problema

El Aprendizaje por Refuerzo (RL) se ha vuelto fundamental para mejorar las capacidades de razonamiento complejo de los Grandes Modelos de Lenguaje (LLMs), particularmente a través de la generación de Cadena de Pensamiento (CoT). Sin embargo, los procesos de entrenamiento de RL son computacionalmente ineficientes, con la fase de rollout (generación de respuestas) consumiendo más del 70% del tiempo total de entrenamiento, especialmente cuando la longitud de la generación supera los 8,000 tokens.

Para abordar este cuello de botella, una estrategia común consiste en utilizar cuantización FP8 durante la fase de rollout, manteniendo al mismo tiempo la precisión BF16 para la fase de entrenamiento (BF16-train + FP8-rollout). Aunque este enfoque es ampliamente adoptado en marcos como VeRL y OpenRLHF, este artículo identifica dos modos de falla críticos:

1. Inestabilidad en Rollouts Largos: A medida que la longitud de generación aumenta (por ejemplo, >8k tokens), el desajuste numérico entre la distribución de entrenamiento BF16 y la distribución de rollout FP8 se acumula. Esto conduce a una inestabilidad severa en el entrenamiento y a un colapso catastrófico de la precisión.
2. Falla en Tareas Desafiantes: La estrategia falla cuando se aplica a tareas de razonamiento difíciles o modelos base más débiles. En estos escenarios, la menor confianza del modelo amplifica los errores de cuantización, causando que la curva de entrenamiento diverja rápidamente en comparación con las líneas base de precisión completa.

La causa raíz se identifica como una naturaleza off-policy: el actor aprende de datos generados por una política (rollout FP8) que difiere numéricamente de la política que se está actualizando (entrenamiento BF16), violando el supuesto on-policy esencial para un RL estable.

Metodología: Jet-RL

Los autores proponen Jet-RL, un marco diseñado para permitir un entrenamiento de RL en FP8 robusto y on-policy, forzando un flujo de precisión unificado entre el entrenamiento y el rollout.

Arquitectura Central

Jet-RL elimina el desajuste de precisión asegurando que tanto el paso hacia adelante (forward pass) del entrenamiento como el rollout de inferencia utilicen el mismo flujo de cuantización FP8.

• Grafo de Precisión Unificado: El marco modela el modelo como un grafo dirigido donde las aristas representan la propagación de tensores. Jet-L asegura que el grafo de inferencia ($\mathcal{G}_{infer}$) sea un subgrafo del grafo forward de entrenamiento ($\mathcal{G}^{fwd}_{train}$), compartiendo la misma granularidad de cuantización y precisión.
• Pesos Maestros (Master Weights): Mientras que el paso forward y el rollout operan en FP8, los pesos maestros se mantienen en BF16 para estabilizar las actualizaciones de entrenamiento, de manera similar al entrenamiento de precisión mixta estándar.

Esquema de Cuantización

Para mantener la estabilidad mientras se acelera la computación, Jet-RL emplea una estrategia de cuantización mixta para las Capas Lineales (GEMMs):

• Activaciones: Cuantizadas utilizando una cuantización 1×128 por grupo.
• Pesos: Cuantizados utilizando una cuantización 128×128 por bloque.
• Paso hacia atrás (Backward Pass):
  • Las activaciones guardadas para el paso hacia atrás se almacenan en FP8 (fusionadas con el operador anterior).
  • Los gradientes transportados entre operadores se mantienen en BF16 para prevenir el subdesbordamiento (underflow) de gradientes y la degradación de la convergencia.
  • Las operaciones de GEMM hacia atrás (DGrad y WGrad) se realizan en FP8 para aceleración, utilizando configuraciones de kernel específicas (por ejemplo, 1×128 y 128×1 para WGrad) para alinearse con los requisitos del hardware.

Implementación

Jet-RL está implementado utilizando vLLM para el motor de inferencia y VeRL para el marco de entrenamiento de RL. Aprovecha kernels de GEMM FP8 de última generación de DeepGEMM y kernels personalizados de Triton para cuantización, transposición y operaciones fusionadas de activación/RMSNorm.

Contribuciones Clave

1. Identificación del Problema: El artículo demuestra que la estrategia estándar de BF16-train + FP8-rollout es fundamentalmente defectuosa para rollouts largos y tareas desafiantes debido a desajustes numéricos off-policy, lo que conduce al colapso del entrenamiento.
2. Marco Jet-RL: La propuesta de un flujo de precisión FP8 unificado que impone consistencia on-policy entre el entrenamiento y el rollout, eliminando la necesidad de calibración entre pasos y resolviendo el desajuste de la política.
3. Desempeño y Estabilidad: La demostración de que Jet-RL logra aceleraciones significativas manteniendo la convergencia y la precisión cercana a las líneas base de BF16 en diversos modelos y tareas.

Resultados Experimentales

Los autores evaluaron Jet-RL en modelos incluyendo Llama3.1-8B, Qwen2.5-7B y Qwen3-8B-Base en conjuntos de datos como GSM8K, MATH y DeepMATH.

Precisión y Convergencia

• Estabilidad: A diferencia de la línea base de BF16-train + FP8-rollout, que no logró converger en Qwen2.5-7B y mostró una degradación significativa en rollouts de 16k, Jet-RL convergió con éxito en todas las configuraciones.
• Brecha de Precisión: Jet-RL redujo la brecha de rendimiento respecto al entrenamiento BF16 a aproximadamente un 1% (por ejemplo, 1.1% de degradación en Qwen3-8B-Base frente al 2.9% de la línea base). En algunos casos (Llama3.1-8B), Jet-RL incluso superó a la línea base de BF16 por un 2.0%.
• Robustez: Jet-RL mantuvo el desempeño en tareas difíciles (DeepMATH) y contextos largos (16k) donde la estrategia de rollout FP8 ingenua colapsó.

Eficiencia

• Aceleración de Rollout: La cuantización FP8 proporcionó una aceleración de 1.07× a 1.33× en la fase de rollout (tokens/s), siendo los modelos más grandes (32B) los que obtuvieron mayores ganancias.
• Aceleración de Entrenamiento: El flujo unificado de FP8 aceleró la fase de actualización del actor en 1.54× y la inferencia del modelo de referencia en 1.80×, resultando en una aceleración de 1.41× para la fase de entrenamiento.
• End-to-End: El marco logró una aceleración end-to-end de 1.16× para el entrenamiento de modelos de 8B, con expectativas de mayores ganancias para modelos más grandes.

Significado y Reivindicaciones

El artículo afirma que Jet-RL proporciona un camino confiable y eficiente para aplicar la computación FP8 para acelerar el entrenamiento de RL a gran escala. Al abordar la inestabilidad fundamental causada por los desajustes de precisión, Jet-RL permite el uso de FP8 no solo para la inferencia, sino para todo el bucle de entrenamiento de RL. Esto permite una aceleración significativa (hasta un 33% de aceleración en el rollout y un 41% en el entrenamiento) sin sacrificar la estabilidad de la convergencia o la precisión requerida para entrenar modelos de razonamiento complejos. El trabajo establece que un flujo de precisión unificado es esencial para escalar el entrenamiento de RL con hardware de baja precisión.