Mejoras en vLLM: Enfoque en la corrección antes de ajustes en RL

PipelineRL utiliza vLLM como motor de inferencia para la generación de rollout. Este motor muestrea tokens y devuelve logprobs de tokens; el entrenador utiliza esos logprobs para calcular ratios de política, KL, tasa de clip, entropía y recompensa. Cualquier discrepancia en cómo se calculan esos logprobs puede alterar la dinámica de entrenamiento. Este es el desajuste entre entrenamiento e inferencia que necesitamos eliminar durante la migración de vLLM V0 a V1.

Resumen: vLLM V1 coincide con nuestra referencia de vLLM V0 después de corregir cuatro aspectos: logprobs de rollout procesados, valores predeterminados de ejecución específicos de V1, el camino de actualización de peso inflight y el lm_head en fp32 utilizado para la proyección final. Corregimos el comportamiento del backend antes de cambiar el objetivo de RL.

Objetivo de Migración

vLLM V1 es una reescritura sustancial del motor V0. Nuestro objetivo de migración fue, por tanto, deliberadamente restringido:

1. verificar que V1 devolviera logprobs de rollout en la forma que el entrenador esperaba
2. repetir la misma carga de trabajo contra la referencia V0
3. evaluar cambios a nivel de objetivo solo después de restaurar la paridad del backend

Los primeros síntomas visibles aparecieron en:

• clamp_log_ratio_new_old_indicator
• kl_new_old
• entropía
• recompensa

Estas métricas provienen de una ejecución de entrenamiento GSPO, el objetivo utilizado para este experimento. La misma clase de desajuste puede aparecer en PPO, GRPO o cualquier sistema de RL en línea que trate los logprobs del lado del rollout como parte del objetivo de optimización.

Modos de Fallo

Separamos las posibles causas en tres capas:

1. Desajuste semántico: el backend devuelve logprobs con un significado diferente al que espera el entrenador.
2. Desajuste en la ruta de inferencia: el backend utiliza diferentes valores predeterminados de ejecución para la caché, programación o manejo de solicitudes, por lo que los mismos prompts siguen un camino de ejecución diferente.
3. Desajuste de objetivo: el objetivo de RL necesita corrección por la cantidad de desactualización o desajuste del backend que permanece.

Inicialmente sospechamos de la tercera categoría demasiado pronto. El diagnóstico útil provino de tratar las dos primeras como problemas de comportamiento del backend y descartarlas primero.

Correciones del Backend V1

Semántica de Logprob

El primer problema fue semántico. vLLM V1 devuelve logprobs de las salidas del modelo en bruto por defecto, antes del procesamiento de logits como escalado de temperatura, penalizaciones y filtrado top-k/top-p. PipelineRL esperaba logprobs de la distribución procesada utilizada por el muestreador.

La configuración requerida fue:

• logprobs-mode=processed_logprobs

Esto eliminó el evidente desplazamiento medio en los logprobs de rollout. Sin embargo, las curvas de entrenamiento aún mostraban una brecha en relación con la referencia conocida, por lo que el siguiente problema debía estar en el camino de inferencia.

La gráfica de la relación de política muestra esto directamente. Una vez que se activa processed_logprobs para V1, la relación media de política se mantiene extremadamente cerca de 1.0 en todas las tres ejecuciones. Esto establece la corrección de sesgo medio. El desajuste restante aparece en la tasa de clip, KL, entropía y el comportamiento de entrenamiento posterior.

Valores Predeterminados de Ejecución

La ejecución temprana de V1 mezcló la versión del motor con los valores predeterminados de ejecución de V1:

• caché de prefijos, no configurado en la ejecución temprana, por lo que se aplicó el valor predeterminado de vLLM 0.18.1
• programación asíncrona, no configurada en la ejecución temprana, por lo que se aplicó el valor predeterminado de vLLM 0.18.1
• una anulación ad-hoc disable-cascade-attn que se configuró a través de pasaje de kwarg en el momento de lanzamiento y se encuentra fuera de la receta de paridad en la configuración comprometida

Para la ejecución de paridad, hicimos estas elecciones explícitas:

vllm_config: use_v1: true vllm_kwargs: logprobs-mode: processed_logprobs enable-prefix-caching: false async-scheduling: false 

El uso de caché de prefijos merece una nota aparte. Normalmente, es una optimización de inferencia que preserva la corrección para un estado de modelo fijo. En este entorno de RL en línea, fue una diferencia exclusiva de V1 en la vida útil de la caché y reutilización en relación con la ruta de referencia V0. El actor también estaba manejando prefijos repetidos, solicitudes concurrentes, programación asíncrona y actualizaciones de peso inflight.

Un acierto de caché de prefijo puede reutilizar el estado calculado antes de una actualización de peso cuando la política de caché ignora el límite de actualización de peso. Deshabilitar la caché de prefijo eliminó un grado de libertad exclusivo de V1 de la comparación de paridad.

Actualizaciones de Peso Inflight

La sincronización de pesos también tuvo que coincidir con el modelo de actualización de RL en línea. Una opción era hacer que V1 fuera más estricta que V0 drenando solicitudes y borrando cachés en cada actualización. Eso respondería a una pregunta separada. Primero necesitábamos verificar que V1 pudiera igualar el comportamiento existente de V0.

Lo que V0 efectivamente hacía era más parecido a:

• bloquear la ejecución en un límite de motor
• cargar los nuevos pesos
• reanudar sin una invalidación explícita del estado en caché

El análogo más cercano en V1 era:

await engine.pause_generation(mode="keep", clear_cache=False) await engine_client.collective_rpc_async( "receive_weight_update", args=(request.model_dump_json(),), ) await engine.resume_generation() 

Dos detalles son importantes:

• mode="keep" coincide más estrechamente con el antiguo modelo de actualización inflight que wait o abort
• clear_cache=False coincide con el comportamiento del envoltorio V0, que dejó el estado en caché intacto en la actualización

El retraso fue un diagnóstico útil en tiempo de ejecución. La ruta inicial de V1 lleva más retraso persistente más adelante en el entrenamiento que la ejecución V1 corregida.

La Brecha Restante: fp32 lm_head

Las correcciones del backend V1 anteriores eliminaron los problemas evidentes de migración, pero la paridad final aún requería igualar el camino numérico utilizado para calcular logits. El entrenador utilizó un lm_head en fp32 para la proyección final. El backend de rollout tuvo que igualar ese comportamiento.

Un problema relacionado aparece en el informe técnico MiniMax-M1: su ejecución de RL mostró un desajuste de probabilidad de token de entrenamiento/inferencia que rastrearon hasta la cabeza de salida de LM y corrigieron calculando la cabeza en fp32.

Esto es importante porque la actualización de RL consume logprobs de token directamente. Pequeños cambios en logits pueden hacerse visibles en ratios de política, KL y clipping. La precisión de la proyección final es, por lo tanto, parte de la superficie de corrección para RL en línea. El artículo ScaleRL incluye más tarde el cálculo de logits/cabeza en fp32 como parte de su receta de RL y lo abla como una elección de diseño útil para RL a gran escala.

Con el camino de lm_head en fp32 incluido, la recompensa proporciona una visión compacta del resultado final de paridad. En la Figura 6, la ejecución final de V1 sigue la referencia V0; el intento inicial de V1 produce una curva de recompensa claramente diferente.

Ablaciones

Los resultados negativos son importantes porque descartan explicaciones comunes.

• processed_logprobs solo: corrigió el error semántico de logprob; el desajuste de entrenamiento permaneció.
• Invarianza de Lote: el desajuste permaneció en una prueba separada, con mayor retraso, mayor tasa de clip y complicaciones de NCCL.
• Tratar la primera ejecución de V1 como una base justa: la primera ejecución de V1 tenía múltiples valores predeterminados únicos de V1 habilitados, por lo que fue una comparación de migración confusa.

Por qué Corregimos Primero la Corrección del Backend

Correcciones del lado del objetivo, como el muestreo de importancia truncado, la reponderación de la razón de importancia y métodos relacionados, son herramientas útiles. Si los rollouts están intencionadamente desactualizados, generados de manera asíncrona o producidos por un backend donde la equivalencia con la política del lado del entrenador no está disponible, entonces alguna forma de corrección suele ser lo correcto.

El primer problema aquí fue la corrección de inferencia. Después de pasar a V1, el backend de rollout devolvió logprobs y comportamiento en tiempo de ejecución que rompió la suposición del entrenador. Agregar una corrección del lado del objetivo en ese punto habría mezclado dos preguntas:

• ¿está el backend de inferencia produciendo los logprobs correctos?
• dada la logprobs correcta, ¿necesita el objetivo aún una corrección fuera de política o asíncrona?

Esas preguntas necesitan separarse. De lo contrario, una corrección del lado del objetivo puede compensar el comportamiento roto del backend de inferencia, lo que hace que la curva de entrenamiento sea más difícil de interpretar.

El objetivo actual aún puede mejorar. Después de restaurar la paridad de inferencia, la próxima mejora es la limpieza habitual de async/off-policy:

• mantener logprobs de política de comportamiento explícitos desde el tiempo de rollout
• recalcular logprobs de política antigua del lado del entrenador en el tiempo de optimización
• separar la corrección de desajuste del backend de la razón de actualización de política
• seguir diagnósticos como ESS para el término de corrección junto con métricas agregadas del entrenador

La lección principal de esta migración es más estrecha: corregir primero la corrección del backend, luego agregar correcciones para el desajuste que permanece.