Si ha estado ejecutando aprendizaje por refuerzo (RL) después del entrenamiento en un modelo de lenguaje para razonamiento matemático, generación de código o cualquier tarea verificable, es casi seguro que se habrá quedado mirando una barra de progreso mientras su clúster de GPU se quema durante la generación de implementación. Un equipo de investigadores de NVIDIA propone una solución precisa integrando la decodificación especulativa en el propio ciclo de entrenamiento de RL, y hacerlo de una manera que preserve la distribución de salida exacta del modelo objetivo.
El equipo de investigación integró la decodificación especulativa directamente en NeMo RL v0.6.0 con un backend vLLM, lo que ofrece una aceleración de implementación sin pérdidas en escalas de modelo 8B y 235B proyectadas. La última versión de NeMo RL v0.6.0 incluye oficialmente la decodificación especulativa como una característica compatible junto con el backend SGLang, el optimizador Muon y el entrenamiento de contexto largo de YaRN.
Por qué la generación de implementación es el cuello de botella
Para comprender el problema, es útil saber cómo se descompone un paso de entrenamiento de RL sincrónico. En NeMo RL, cada paso consta de cinco etapas: carga de datos, sincronización de peso y preparación del backend (preparar), generación de implementación (gen), recómputo de probabilidad de registro (logprob) y optimización de políticas (entrenamiento).
El equipo de investigación midió este desglose en Qwen3-8B bajo dos cargas de trabajo: RL-Think, que continúa entrenando un modelo capaz de razonar, y RL-Zero, que comienza desde un modelo base y aprende a razonar desde cero. En ambos casos, la generación de implementación representa entre el 65% y el 72% del tiempo total del paso. El recálculo de la probabilidad logarítmica y el entrenamiento en conjunto requieren sólo entre el 27% y el 33%. Esto hace que la generación sea la única etapa a la que vale la pena apuntar para la aceleración, y la que determina el límite para cualquier optimización del lado del lanzamiento.
Lo que realmente hace la decodificación especulativa
La decodificación especulativa es una técnica en la que un modelo borrador más pequeño y más rápido propone varios tokens a la vez, y el modelo objetivo más grande (el que realmente estás entrenando) los verifica mediante un procedimiento de muestreo de rechazo. La propiedad clave y por qué es importante para RL es que se garantiza matemáticamente que el procedimiento de rechazo producirá la misma distribución de salida que si el modelo objetivo hubiera generado esos tokens de forma autorregresiva. No hay discrepancias en la distribución, no se necesitan correcciones fuera de las políticas, ni cambios en la señal de entrenamiento.
Esto es importante porque en la capacitación posterior de RL, la recompensa de la capacitación depende de las muestras propias de la política. Métodos como la ejecución asincrónica, la repetición fuera de política o las implementaciones de baja precisión intercambian cierta fidelidad de entrenamiento por rendimiento. La decodificación especulativa no cambia nada: los lanzamientos son idénticos en distribución a lo que el modelo objetivo habría generado por sí solo, solo que se produjo más rápido.
El desafío de la integración del sistema
Agregar un modelo borrador a un backend de servicio es sencillo. Agregar uno a un ciclo de entrenamiento de RL no lo es. Cada vez que se actualiza la política, el motor de implementación debe recibir nuevas ponderaciones. El borrador del modelo debe permanecer alineado con la política en evolución. Las probabilidades de registro, las penalizaciones de KL y la pérdida de la política GRPO deben calcularse en función de la política de destino (verificador), no del borrador, o el objetivo de optimización se corrompe silenciosamente.
El equipo de investigación de NVIDIA maneja esto en NeMo RL con una arquitectura de dos rutas. La ruta general utiliza EAGLE-3, un marco de redacción que funciona con cualquier modelo previamente entrenado sin requerir soporte nativo de predicción de múltiples tokens (MTP). También hay disponible una ruta nativa para los modelos que se envían con cabezales MTP integrados. Cuando se habilita la adaptación del borrador en línea, los estados ocultos y las probabilidades de registro del paso directo del verificador MegatronLM se almacenan en caché y se reutilizan para supervisar el cabezal del borrador a través de una ruta separada por gradiente, de modo que el entrenamiento del borrador nunca interfiere con la señal del gradiente de la política.
Resultados medidos en escala 8B
En 32 GPU GB200 (8 nodos GB200 NVL72, 4 GPU por nodo), EAGLE-3 reduce la latencia de generación de 100 segundos a 56,6 segundos en RL-Zero, una aceleración de generación de 1,8 veces. En RL-Think, cae de 133,6 segundos a 87,0 segundos, una aceleración de 1,54 veces. Debido a que el recálculo y el entrenamiento de la probabilidad logarítmica no cambian, estas ganancias del lado generacional se traducen en aceleraciones generales de pasos de 1,41× en RL-Zero y 1,35× en RL-Think. La precisión de la validación en AIME-2024 evoluciona de manera idéntica bajo la decodificación autorregresiva y especulativa durante el entrenamiento, lo que confirma que la garantía sin pérdidas se mantiene en la práctica.
El equipo de investigación también prueba la redacción de n-gramas como base especulativa sin modelos. A pesar de lograr longitudes de aceptación de 2,47 en RL-Zero y 2,05 en RL-Think, el dibujo de n-gramas es más lento que la línea de base autorregresiva en ambas configuraciones: 0,7 × y 0,5 × respectivamente. Este es un hallazgo fundamental para los profesionales: una duración de aceptación positiva es necesaria pero no suficiente. Si los gastos generales de verificación son lo suficientemente altos, la especulación empeora las cosas.
Tres decisiones de configuración que determinan la aceleración obtenida
El equipo de investigación aísla tres opciones operativas que los profesionales deben acertar.
La inicialización del borrador importa más que la capacidad de redacción genérica. Un borrador de EAGLE-3 inicializado en el conjunto de datos posterior al entrenamiento de DAPO logra una aceleración de generación de 1,77 veces en RL-Zero, mientras que un borrador inicializado en los conjuntos de datos de uso general UltraChat y Magpie logra solo 1,51 veces con la misma longitud de borrador. El borrador debe estar alineado con la distribución de implementación real que se encontró durante la RL, no solo con una distribución amplia del chat.
La longitud del calado tiene un óptimo no obvio. Con una longitud de borrador k = 3, RL-Zero logra una aceleración de 1,77 × y RL-Think logra 1,53 ×. Aumentar a k=5 aumenta la longitud de aceptación pero reduce la aceleración a 1,44× en RL-Zero y 0,84× en RL-Think; este último ya es más lento que el autorregresivo. En k=7, RL-Zero cae aún más hasta 1,21× y RL-Think hasta 0,71×. El contraste importa: los lanzamientos de RL-Zero se generan a partir de un modelo base que comienza con resultados cortos, lo que hace que sea más fácil de predecir para el borrador incluso con k alto. Es más difícil especular sobre los rastros de razonamiento completamente desarrollados de RL-Think, por lo que la sobrecarga de borradores más largos borra el beneficio antes. Un mayor trabajo especulativo por paso puede borrar por completo el beneficio de una mayor aceptación, especialmente en regímenes de generación más duros.
Adaptación del borrador en línea: actualizar el borrador durante RL mediante implementaciones generadas por la política actual ayuda más cuando el borrador se inicializa débilmente. Para un borrador inicializado con DAPO, las configuraciones en línea y fuera de línea funcionan de manera casi idéntica (1,77 × frente a 1,78 × en RL-Zero). Para un borrador inicializado con UltraChat, la actualización en línea mejora la velocidad de 1,51× a 1,63× en RL-Zero.
La interacción con la ejecución asincrónica también se probó directamente a escala 8B, no solo en simulación. El equipo de investigación ejecutó RL-Think con un retraso de política 1 en una configuración no ubicada de 16 nodos, con 12 nodos dedicados a la generación y 4 a la capacitación. En el modo asincrónico, la mayor parte de la generación de implementación ya está oculta detrás del recálculo de la probabilidad de registro y las actualizaciones de políticas, por lo que la cantidad relevante es el tiempo de generación expuesto que permanece en la ruta crítica. La decodificación especulativa reduce el tiempo de generación expuesto de 10,4 segundos a 0,6 segundos por paso y reduce el tiempo de paso efectivo de 75,0 segundos a 60,5 segundos (1,24×). La ganancia es menor que en la RL sincrónica (lo que se esperaba, ya que la superposición asincrónica ya oculta gran parte del costo de implementación), pero confirma que los dos mecanismos son genuinamente complementarios en lugar de redundantes.
Ganancias proyectadas a escala 235B
Utilizando un simulador de rendimiento de GPU patentado calibrado para características de interconexión, memoria y computación a nivel de dispositivo, el equipo de investigación proyectó ganancias de decodificación especulativas a escalas más grandes. Para Qwen3-235B-A22B que ejecuta RL síncrono en GPU de 512 GB200, la longitud del borrador k=3 con una longitud de aceptación de 3 tokens produce una aceleración de implementación de 2,72 veces y una aceleración de extremo a extremo de 1,70 veces.
En el punto operativo simulado más favorable (Qwen3-235B-A22B en GPU de 2048 GB200 con RL asíncrono en un retraso de política 2), la velocidad de implementación alcanza aproximadamente 3,5 veces, lo que se traduce en una aceleración de entrenamiento de extremo a extremo proyectada de 2,5 veces. La decodificación especulativa y la ejecución asincrónica se describen como complementarias: la especulación reduce el costo de cada implementación individual, mientras que la superposición asincrónica oculta el tiempo de generación restante detrás del entrenamiento y el cálculo de probabilidad logarítmica.
Conclusiones clave
La generación de implementación es el cuello de botella dominante en el post-entrenamiento de RL, y representa entre el 65 % y el 72 % del tiempo total del paso en cargas de trabajo de RL sincrónicas, lo que la convierte en la única etapa donde la aceleración tiene un impacto significativo en la velocidad del entrenamiento de un extremo a otro. La decodificación especulativa a través de EAGLE-3 ofrece una aceleración de implementación sin pérdidas, logrando una velocidad de generación de 1,8 veces a una escala de 8B (1,41 veces una aceleración de paso general) sin cambiar la distribución de salida del modelo de destino, a diferencia de la ejecución asincrónica, la reproducción fuera de política o las implementaciones de baja precisión, que intercambian fidelidad de entrenamiento por rendimiento. La calidad de la inicialización del borrador importa más que la longitud del borrador, ya que los borradores en el dominio (capacitados por DAPO) superan a los borradores generales del dominio de chat por un margen significativo; longitudes de borrador más largas (k≥5) consistentemente resultan contraproducentes en cargas de trabajo de razonamiento más difíciles, lo que hace que k=3 sea el valor predeterminado confiable. Las proyecciones del simulador muestran que las ganancias aumentan significativamente, alcanzando una aceleración de implementación de ~3,5 veces y una aceleración de entrenamiento de extremo a extremo proyectada de ~2,5 veces a una escala de 235 B en GPU de 2048 GB200, y la técnica ya está disponible en NeMo RL v0.6.0 bajo Apache 2.0.
Consulte el documento completo y el repositorio de Nemo RL. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 130.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros