vLLM V0 a V1: priorizar la corrección del backend antes de ajustar el objetivo RL

Resumen y objetivo de la migración

PipelineRL usa vLLM como motor de inferencia para generar rollouts: el motor muestrea tokens y devuelve logprobs, y el entrenador utiliza esos logprobs para calcular ratios de política, KL, clip rate, entropía y recompensa. Cualquier discrepancia en cómo se computan esos logprobs puede alterar la dinámica de entrenamiento. Esa brecha entre entrenamiento e inferencia fue el problema central al migrar de vLLM V0 a V1.

El objetivo de la migración fue deliberadamente estrecho: verificar que V1 devolviera los logprobs en la forma que el entrenador esperaba, ejecutar la misma carga de trabajo contra la referencia V0 y evaluar cambios en el objetivo solo después de restaurar la paridad del backend. La referencia usó vLLM 0.8.5; las pruebas V1 se hicieron con vLLM 0.18.1.

Síntomas iniciales: dónde apareció la diferencia

Los primeros síntomas visibles se manifestaron en métricas del entrenador durante un experimento con GSPO —el objetivo empleado en la evaluación—: indicadores como clamp_log_ratio_new_old_indicator, kl_new_old, entropía y recompensa mostraron desviaciones tempranas. En la práctica, el indicador más sencillo de leer fue el clip rate: su comportamiento mostró claramente un desfase entre la política del rollout y la política evaluada por el entrenador.

Además, las curvas de logprobs calculadas por el entrenador se separaron de la referencia V0 al inicio del entrenamiento, confirmando que el problema no era gradual ni marginal.

Capas de diagnóstico: semántica, camino de inferencia y objetivo

Para organizar la investigación se definieron tres posibles capas causales:

• Semántica: el backend devuelve logprobs con un significado distinto al que espera el entrenador.
• Camino de inferencia: cambios en defaults de runtime, cachés, scheduling o manejo de requests hacen que los mismos prompts sigan una ruta de ejecución diferente.
• Objetivo: el objetivo de RL necesita ajuste para la latencia o staleness que persiste.

El equipo descartó la tercera categoría hasta que las dos primeras quedaran aclaradas —es decir, primero se trató el problema como comportamiento del backend.

Correcciones en el backend: logprob semantics

El primer problema fue semántico. V1, por defecto, devolvía logprobs derivados de las salidas crudas del modelo, es decir, antes de cualquier post-procesamiento de logits (temperature scaling, penalties, top-k/top-p). PipelineRL, en cambio, esperaba logprobs de la distribución procesada que utiliza el muestreador.

La solución fue explícita: activar la opción de logprobs procesados en la configuración (logprobs-mode=processed_logprobs). Con eso se eliminó el sesgo de media evidente en los logprobs del rollout y la razón de política promedio se centró cerca de 1.0, lo que confirmó la corrección del sesgo medio.

Sin embargo, aunque la media quedó corregida, permanecieron diferencias en clip rate, KL, entropía y comportamiento de entrenamiento, lo que apuntó al siguiente grupo de causas.

Defaults de runtime y reproducibilidad

La primera ejecución con V1 mezcló la nueva versión del motor con defaults de runtime propios de V1: cacheo de prefijos (prefix caching) activado por defecto, scheduling asíncrono y overrides ad-hoc para atención en cascada. Estas opciones no habían sido fijadas explícitamente y, por tanto, aplicaron comportamientos distintos frente a la referencia V0.

Para lograr paridad se hizo explícita la configuración de V1, por ejemplo marcando:

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

El caso del prefix caching merece énfasis: en condiciones de inferencia con estado de modelo fijo, es una optimización que preserva la corrección. Pero en un escenario de RL online, donde el actor maneja prefijos repetidos, solicitudes concurrentes y actualizaciones de pesos en vuelo, la vida útil y la reutilización del caché en V1 difería de V0. Un hit en la caché podía reutilizar estado calculado antes de una actualización de pesos si la política de caché ignoraba el límite entre versiones de pesos. Desactivar el cacheo de prefijos eliminó una fuente de variación específica de V1.

Actualizaciones de pesos en vuelo (inflight weight updates)

Otra diferencia importante fue cómo se sincronizaban las actualizaciones de pesos. Una solución conservadora habría sido hacer V1 más estricta que V0: drenar peticiones y limpiar cachés en cada actualización. Pero el objetivo era verificar si V1 podía reproducir el comportamiento real de V0, no forzar un nuevo protocolo.

El comportamiento efectivo de V0 se parecía más a una pausa en un límite del engine, cargar nuevos pesos y reanudar sin invalidar explícitamente el estado cacheado. En V1, la analogía técnica usada fue await engine.pause_generation, ajustada para replicar ese patrón de sincronización. Alinear este flujo fue crucial para evitar que requests reutilizaran estados inconsistentes respecto a la versión de pesos esperada por el entrenador.

Ablaciones del lm_head en fp32

Finalmente, otro factor relevante fue la proyección final usada para convertir logits en distribuciones: el lm_head. La experimentación mostró que usar una proyección final en fp32 (en lugar de una versión cuantizada o en menor precisión) era necesaria para reproducir la referencia. Este ajuste forma parte de la secuencia de correcciones que, en conjunto, llevaron a que V1 replicara la trayectoria de V0 en métricas como clip rate, KL, entropía y recompensa.

Por qué arreglamos el backend antes que el objetivo RL

Cambiar el objetivo de entrenamiento para compensar una discrepancia de backend es una solución peligrosa: enmascara problemas de inferencia que pueden reaparecer o empeorar al escalar o al cambiar infra. La decisión fue clara: primero restaurar el significado y la ruta de ejecución de los logprobs, igualar defaults de runtime y alinear la gestión de actualizaciones de pesos; sólo después, si quedó una diferencia residual justificable, considerar ajustes al objetivo.

Este orden permitió validar que la divergencia observada no provenía del algoritmo de optimización (GSPO en este caso) sino de diferencias en cómo se generaban y entregaban los logprobs.

Lecciones para equipos que implementan RL online

• Verifiquen la semántica de los logprobs: confirme si la inferencia devuelve logprobs crudos o procesados y haga la configuración explícita.
• Fijen los defaults de runtime durante migraciones entre versiones de motor: scheduling, cachés y overrides pueden introducir variación sutil pero crítica.
• En entornos con actualizaciones de pesos en vuelo, documenten y reproduzcan el comportamiento de sincronización del engine de referencia antes de cambiar la lógica del objetivo.
• Traten la corrección del backend como prioridad: sólo una vez que la inferencia y el entrenador están alineados tiene sentido evaluar cambios en el objetivo de RL.

Conclusión

La migración de vLLM V0 a V1 mostró que diferencias aparentemente pequeñas en la semántica de logprobs, defaults de runtime y manejo de actualizaciones pueden provocar desviaciones rápidas en métricas clave de entrenamiento RL. Al abordar estos puntos —logprobs procesados, defaults explícitos (sin prefix caching ni scheduling asíncrono), sincronización de pesos y uso de lm_head en fp32— el equipo consiguió que V1 reprodujera la trayectoria de V0 sin modificar el objetivo RL. Para equipos en América Latina que integran inferencia y entrenamiento online, la recomendación práctica es validar la paridad de backend antes de ajustar políticas u objetivos de entrenamiento.