Spotlight: Synergizing Seed Exploration and Spot GPUs for DiT RL Post-Training

Resumen Técnico: SPOTLIGHT

Declaración del Problema

El post-entrenamiento mediante aprendizaje por refuerzo (RL) para los Diffusion Transformers (DiTs) es prohibitivamente costoso, requiriendo a menudo miles de GPUs de alto rendimiento. El flujo de trabajo estándar de estilo GRPO consta de dos fases: una fase de rollout (generación de muestras con semillas aleatorias) y una fase de entrenamiento (cálculo de gradientes y actualización de pesos).

Las estrategias de reducción de costes existentes enfrentan dos cuellos de botella principales cuando se aplican a los DiTs:

1. Sobrecarga de Exploración de Semillas: Métodos como DanceGRPO y TreeGRPO exploran múltiples semillas aleatorias para encontrar muestras de alto contraste (recompensas altas y bajas) para acelerar la convergencia. Sin embargo, esta fase de exploración se sitúa en el camino crítico, añadiendo tiempo de cómputo a cada iteración y anulando los posibles ahorros de costes.
2. Ociosidad y Fragmentación de GPUs Spot: Aunque las GPUs spot ofrecen costes un 69–77% menores, son interrumpibles y poseen un ancho de banda inter-nodo limitado.
   • Ociosidad: A diferencia de los Modelos de Lenguaje Extensos (LLMs) donde las latencias de rollout varían (permitiendo el pipelining con el entrenamiento), los rollouts de los DiTs completan su ejecución en un tiempo casi uniforme. Esto crea una barrera de sincronización donde las GPUs spot permanecen ociosas durante toda la fase de entrenamiento.
   • Fragmentación: Las interrupciones frecuentes de las GPUs spot rompen los grupos de Paralelismo de Secuencia (SP), que son esenciales para ejecutar modelos DiT de gran tamaño. La reconfiguración de estos grupos suele requerir minutos para recargar modelos e inicializar motores, desperdiciando recursos significativos.

Metodología: SPOTLIGHT

Los autores presentan SPOTLIGHT, el primer sistema diseñado para aprovechar las GPUs spot para el post-entrenamiento de RL de DiTs mediante la sinergia entre la exploración de semillas y los recursos spot. El sistema se basa en dos ideas clave y tres técnicas principales.

Ideas Clave

1. La Exploración Tolera Pesos Obsoletos: Los autores demuestran que el ranking relativo de las semillas aleatorias (qué semillas producen recompensas altas vs. bajas) se mantiene altamente consistente incluso utilizando los pesos del modelo de la iteración anterior en lugar de los pesos actualizados de la iteración actual. Esto permite que la exploración de semillas se ejecute en GPUs spot utilizando pesos obsoletos mientras el clúster reservado realiza la fase de entrenamiento, efectivamente descargando la exploración del camino crítico.
2. La Reconfiguración de SP Puede Reutilizar el Estado: El alto coste de reconfigurar los grupos SP proviene de la inicialización del programador de la CPU y la carga de modelos remotos. Al desacoplar el programador y utilizar la copia de pesos intra-nodo (vía NVLink) de pares co-localizados, el tiempo de reconfiguración de SP puede reducirse de minutos a sub-segundos.

Arquitectura del Sistema y Técnicas

SPOTLIGHT está implementado sobre la plataforma de RL de código abierto ROLL y emplea tres técnicas principales:

1. Planificador de Exploración Basado en Bandidos:

   • El sistema formula la planificación de la exploración como un problema de bandidos multibrazo.
   • Selecciona dinámicamente el número de secuencias y los pasos de eliminación de ruido efectivos (usando estrategias de caché como Teacache) para maximizar la varianza de la recompensa dentro del presupuesto de tiempo proporcionado por la fase de entrenamiento.
   • Utiliza una estrategia de Banda de Confianza Superior (UCB) para equilibrar la exploración de nuevas configuraciones y la explotación de configuraciones de alta varianza conocidas.

2. Paralelismo de Secuencia Elástico (SP):

   • Programador Persistente: El programador de la CPU se desacopla de los trabajadores de la GPU y se mantiene residente en el nodo a través de los cambios de SP, eliminando los costes de inicialización repetidos.
   • Carga de Pesos Intra-Nodo: Cuando un nuevo trabajador de GPU se une a un grupo SP en el mismo nodo, copia los pesos de un par co-localizado a través de NVLink en lugar de extraerlos de un nodo remoto, reduciendo drásticamente la latencia de reconfiguración.

3. Programador de Peticiones Basado en Pull Sensible a la Interrupción:

   • Los trabajadores de rollout extraen peticiones de una cola centralizada en lugar de tener tareas asignadas estáticamente, adaptándose a las velocidades heterogéneas de los trabajadores.
   • Compromiso de Estado: Los estados intermedios de los rollouts (latentes y metadatos) se comprometen a un Almacén de Tensores persistente en los nodos reservados. Si una GPU spot es interrumpida, el sistema puede recuperar el estado en vuelo y reanudar la petición en otro trabajador, minimizando la pérdida de progreso.

Resultados

El sistema fue evaluado en el post-entrenamiento de Qwen-Image (20B) utilizando los conjuntos de datos DeepSeek-OCR y Geneval en resoluciones de 512×512 y 1280×1280.

• Reducción de Costes: SPOTLIGHT reduce el coste total de entrenamiento de 1.4× a 6.4× en comparación con las líneas base de vanguardia (incluyendo sistemas que usan GPUs spot sin estas optimizaciones y sistemas de solo reserva).
• Velocidad de Convergencia: Alcanza el mismo puntaje de validación objetivo 4× más rápido que las líneas base. Esto se logra reduciendo el número de iteraciones requeridas (debido a la exploración efectiva de semillas) y superponiendo la exploración con el entrenamiento.
• Calidad: SPOTLIGHT logra puntuaciones de validación superiores en todos los conjuntos de datos en comparación con las líneas base, debido a la selección dinámica de semillas de alto contraste.
• Eficiencia:
  • La sobrecarga de exploración se mantiene baja (2–3% del tiempo de iteración) al ajustar la exploración dentro de la ventana de entrenamiento.
  • El SP elástico permite que el sistema rastree los cambios de capacidad spot en tiempo casi real, mientras que los enfoques ingenuos sufren retrasos de varios minutos durante la reconfiguración.
  • El sistema demuestra una escalabilidad casi lineal en el rendimiento (throughput) a medida que el número de GPUs spot aumenta de 8 a 64.

Significancia y Reivindicaciones

El artículo afirma que SPOTLIGHT es el primer sistema que logra cosechar GPUs spot para el post-entrenamiento de RL de DiTs. Su significancia radica en:

1. Romper la Dependencia del Camino Crítico: Al demostrar que la exploración de semillas puede tolerar pesos obsoletos, el sistema elimina el compromiso entre una convergencia más rápida (vía exploración) y la latencia de iteración.
2. Resolver la Ociosidad de las GPUs Spot: Resuelve el problema específico de ociosidad de los rollouts de DiT utilizando la capacidad ociosa de entrenamiento para la exploración, una sinergia pasada por alto por los sistemas previos centrados en LLMs.
3. Robustez ante Interrupciones: A través de un SP elástico y la persistencia de estado, hace que el entrenamiento de DiT sea resiliente a las interrupciones frecuentes inherentes a los mercados spot, convirtiendo un conjunto de recursos fragmentados en una infraestructura de entrenamiento viable y rentable.

Los autores concluyen que estas contribuciones permiten un sistema de post-entrenamiento de RL de DiT eficiente en costes que no sacrifica la calidad de la imagen, haciendo que el entrenamiento de modelos de difusión de alta calidad sea más accesible.