A medida que los modelos de lenguaje (LLMs) avanzan de la simple generación de texto hacia un razonamiento más complejo, el aprendizaje por refuerzo (RL) se vuelve crucial. Algoritmos como el Group Relative Policy Optimization (GRPO) facilitan esta transición, permitiendo que los modelos mejoren continuamente mediante retroalimentación iterativa.
En contraste con el ajuste fino supervisado tradicional, los bucles de entrenamiento en RL se dividen en dos fases distintas: una fase de generación con requisitos de latencia estrictos y una fase de entrenamiento que demanda alto rendimiento. Para que estas cargas de trabajo sean efectivas, se recurre a tipos de datos de baja precisión, como FP8, que mejoran el rendimiento durante el entrenamiento y la generación orientada a throughput.
FP8 para capas lineales en RL
Utilizamos la cuantización por bloques en FP8 introducida por el Informe Técnico DeepSeek-V3. La Tabla 1 detalla los formatos de tensor en las capas de proyección lineales.
| Tensor | Tipo de dato | Granularidad de cuantización | Factor de escalado | Tipo de escalado |
| Ponderaciones | FP8 (E4M3) | [128, 128] | FP32 | Por bloques |
| Activaciones de entrada | FP8 (E4M3) | [1, 128] | FP32 | Por bloques |
| Gradientes de salida | FP8 (E4M3) | [1, 128] | FP32 | Por bloques |
Con esta receta, las capas lineales pueden calcularse con matemáticas FP8, que ofrecen un rendimiento de 2x en comparación con BF16. Otros módulos, como la atención y la normalización, utilizan matemáticas BF16.
El desafío del desacuerdo numérico en RL
Las tuberías de RL suelen utilizar motores separados: vLLM para rollouts y NVIDIA Megatron Core para el entrenamiento. Cada uno utiliza núcleos CUDA personalizados de NVIDIA para maximizar el rendimiento, lo que introduce diferencias numéricas que se amplifican en menor precisión debido a la lógica adicional de cuantización y desacualización.
Cuantificamos esta diferencia numérica como un error de probabilidad multiplicativa por token:
(texttt{token-mult-prob-error} = frac{1}{n}sum_{i=1}^{n(texttt{tokens})}exp(left| texttt{log-train-fwk}_i – texttt{logprobs-inference-fwk}_i right|))
Una alineación perfecta obtiene un puntaje de 1, y los valores ‘aceptables’ suelen ser <1.03-1.05 sin técnicas adicionales.
FP8 de extremo a extremo en capas lineales reduce el desacuerdo numérico
Durante el desarrollo de la receta FP8, experimentamos con tres enfoques:
- Receta base: BF16 para generación y entrenamiento.
- Candidato de receta 1: FP8 aplicado exclusivamente durante la generación, mientras que el entrenamiento del modelo de política se realiza en BF16.
- Receta final: FP8 de extremo a extremo: utilizamos FP8 en ambos motores de generación y entrenamiento.
Observamos que, en comparación con la receta candidata 1, la receta final muestra un desacuerdo numérico consistentemente menor entre generación y entrenamiento. Cabe señalar que la receta base siempre produce el desacuerdo numérico más bajo. La Figura 1 muestra el error de probabilidad multiplicativa por token de las tres recetas.
Mitigando el desacuerdo numérico con muestreo de importancia
El muestreo de importancia se utiliza para corregir la discrepancia de distribución entre el modelo que genera los datos y el modelo que se entrena. Se aplica un peso por token multiplicado por la pérdida. Puedes consultar nuestra documentación de GRPO para más detalles teóricos sobre el muestreo de importancia.
Los experimentos muestran que:
- Para la receta candidata 1 (FP8 generación y BF16 entrenamiento), el muestreo de importancia puede reducir la brecha de precisión con BF16 RL, pero no la cierra por completo.
- Para la receta final (FP8 de extremo a extremo), el muestreo de importancia cierra totalmente la brecha con el entrenamiento BF16. La Figura 2 muestra la precisión de validación durante el entrenamiento para diferentes recetas.
