SkyRL y vLLM ajustan Mistral Large 4 con RL: AIME sube del 64% al 75%

El equipo del proyecto vLLM y el laboratorio NovaSky AI han anunciado que su framework SkyRL permite entrenar Mistral Large 4 con aprendizaje por refuerzo (RL) desde el primer día. Un clúster de 8 nodos Nvidia GB200 elevó, en 75 pasos de entrenamiento, el resultado del modelo en AIME 2024 de aproximadamente el 64% al 75%. AIME 2024 consta de 30 problemas de la olimpiada estadounidense de matemáticas, cuya respuesta es siempre un número entero. La métrica avg@8 significa que el modelo hace 8 intentos por cada problema y que la proporción final de aciertos se calcula como promedio. Así, el resultado depende menos de la suerte en un único intento.

Qué significa Day-0 y para qué sirve

Day-0 significa que el soporte está listo el día del lanzamiento del modelo. Normalmente, un modelo grande nuevo solo se puede ejecutar al principio, y el reentrenamiento llega más tarde, cuando la comunidad adapta el código a su arquitectura. En este caso, se afirma que la combinación de SkyRL y vLLM funcionó con Mistral Large 4 desde el primer momento.

El aprendizaje por refuerzo funciona así. El modelo recibe un problema y escribe varias variantes de solución. Esas variantes se llaman rollouts («ejecuciones»). Después se comprueba cada una: si la respuesta coincide con la de referencia. Por cada respuesta correcta, el modelo recibe una recompensa y los pesos se desplazan para que las cadenas de razonamiento acertadas se vuelvan más probables. En matemáticas, esta comprobación es barata e inequívoca, por lo que AIME y conjuntos de problemas similares se usan con frecuencia para RL.

Quién hace qué en esta combinación

La principal idea de ingeniería es el reparto del trabajo. Las dos etapas del ciclo exigen cosas distintas al hardware, y cada una la gestiona una herramienta diferente:

• vLLM genera los rollouts. Es un motor de inferencia rápida, es decir, de ejecución del modelo para producir respuestas. La generación ocupa la mayor parte del tiempo del ciclo de RL, así que hace falta un servidor rápido.
• FSDP2 o Megatron se encargan del entrenamiento. FSDP2 es la segunda versión del enfoque de PyTorch en el que los parámetros y los gradientes se fragmentan entre varias GPU. Megatron es el framework de Nvidia para entrenar modelos muy grandes. SkyRL permite elegir cualquiera de los dos.
• SkyRL, de NovaSky AI, une todo en un único proceso: toma los rollouts de vLLM, calcula las recompensas y pasa los datos al entrenamiento.

La ventaja práctica es que el ingeniero no tiene que escribir el código de enlace entre el motor de generación y el código de entrenamiento. Toma la combinación ya hecha y le conecta el modelo.

Por qué se entrenó LoRA sobre un checkpoint de 4 bits

La publicación menciona dos técnicas que abaratan el proceso.

La primera es LoRA. En lugar de actualizar todos los pesos de un modelo enorme, se entrenan pequeñas matrices adicionales que se acoplan sobre él. Ocupan una fracción reducida de memoria en comparación con el propio modelo, y Mistral Large 4 permanece congelado.

La segunda es un checkpoint en formato NVFP4. Es una representación numérica de 4 bits que Nvidia admite en la arquitectura Blackwell, a la que pertenece GB200. Los pesos en este formato ocupan varias veces menos que en 16 bits. Un modelo grande cabe en menos nodos y la generación es más rápida.

El resultado es una configuración modesta para los estándares de los modelos grandes: 8 nodos, 75 pasos y una mejora de unos 11 puntos porcentuales en un solo benchmark.

Cómo interpretar este resultado

Las cifras proceden del propio anuncio del proyecto vLLM y, por ahora, no hay una verificación independiente. La mejora en AIME 2024 demuestra que el ciclo de RL funciona con este modelo y este hardware. No demuestra que el modelo haya mejorado en general. En problemas matemáticos con verificación inequívoca, el RL suele producir un efecto notable, pero en tareas abiertas (escritura, escenarios con agentes, código en un repositorio grande) el resultado hay que medirlo por separado.

Para los ingenieros, la conclusión es esta. Reentrenar con refuerzo un modelo abierto grande es ahora más sencillo: una combinación ya preparada de componentes abiertos funciona desde el primer día del lanzamiento, y LoRA y NVFP4 reducen los requisitos de hardware. Para las empresas que tienen tareas propias con verificación automática de la respuesta, es una vía para adaptar el modelo a sus necesidades sin reentrenarlo por completo. Queda por comprobar con datos propios cuántos de esos 11 puntos se pueden reproducir.