Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

El artículo presenta "Deep Optimizer States", una técnica que mejora la eficiencia del entrenamiento de modelos transformadores al dividir dinámicamente el estado del optimizador entre CPU y GPU basándose en fluctuaciones de memoria, logrando iteraciones 2,5 veces más rápidas que los enfoques actuales.

Lo esencial

¡Claro que sí! Imagina que entrenar una Inteligencia Artificial gigante (como un modelo de lenguaje o LLM) es como intentar construir un rascacielos de cristal, pero tienes un problema grave: tu camión de transporte es demasiado pequeño.

Aquí te explico la idea central del paper "Deep Optimizer States" usando analogías sencillas:

1. El Problema: El Camión Pequeño y el Montón de Ladrillos

Imagina que tienes que construir un edificio enorme (el modelo de IA). Para hacerlo, necesitas millones de ladrillos (datos) y un equipo de obreros muy rápido (las tarjetas gráficas o GPUs).

• El problema: Las tarjetas gráficas son muy rápidas, pero tienen una "caja de herramientas" muy pequeña (memoria limitada).
• La solución actual: Como la caja es pequeña, los ingenieros decidieron guardar la mitad de los ladrillos y las herramientas en un almacén al lado (la memoria del CPU o el ordenador principal).
• El cuello de botella: El problema es que el camino entre el almacén y la caja de herramientas (el cable PCIe) es estrecho y lento. Además, los obreros del almacén (el CPU) son mucho más lentos que los obreros de la caja (la GPU).
  • Resultado: La GPU pasa la mayor parte del tiempo esperando a que le traigan los ladrillos desde el almacén. Es como tener un Ferrari atascado en un semáforo rojo eterno.

2. La Observación Brillante: "La Caja se Vacía de vez en cuando"

Los autores del paper se dieron cuenta de algo interesante mientras observaban el proceso de construcción:

• Cuando la GPU está "pensando" (haciendo cálculos), ocupa toda su caja.
• Pero, en ciertos momentos (cuando termina de calcular y antes de empezar de nuevo), la caja se vacía un poco. Hay espacio libre momentáneo.
• Al mismo tiempo, el camino de transporte (PCIe) a menudo está vacío porque los obreros no están moviendo cosas en ese preciso instante.

La analogía: Es como un restaurante. El chef (GPU) cocina rápido, pero a veces espera a que el camarero (CPU) le traiga los ingredientes. El paper propone que, cuando el chef tiene un segundo libre, ¡que él mismo vaya a la cocina a buscar los ingredientes más urgentes en lugar de esperar al camarero lento!

3. La Solución: "Deep Optimizer States" (El Sistema de Entrega Intermitente)

En lugar de tener una parte fija de herramientas en la caja y otra fija en el almacén (como hacían los sistemas anteriores), este nuevo sistema es dinámico y flexible.

Imagina que tienes un equipo de trabajo dividido en grupos pequeños (subgrupos).

• El sistema anterior (DeepSpeed TwinFlow): Decía: "Los grupos 1 y 2 siempre trabajan en la caja, los grupos 3 al 10 siempre trabajan en el almacén". Esto es rígido. Si la caja tiene espacio libre, no lo usa.
• El nuevo sistema (Deep Optimizer States): Dice: "¡Vamos a mezclarlo!".
  1. Mientras el CPU trabaja en el grupo 3, la GPU (que tiene espacio libre) ya está cargando el grupo 4 para trabajar en él inmediatamente.
  2. La GPU y el CPU trabajan al mismo tiempo en diferentes grupos, y el transporte de datos ocurre mientras ellos trabajan.
  3. Es como si el chef y el camarero estuvieran bailando una coreografía perfecta: el chef cocina el plato A mientras el camarero trae los ingredientes del plato B, y viceversa. Nadie espera a nadie.

4. El Truco Extra: Enviar los Paquetes en el Formato Correcto

Otro problema era que, para mover los ingredientes del almacén a la caja, a veces tenían que cambiarles la etiqueta (convertir datos de un formato a otro), lo cual tomaba mucho tiempo.

• La solución: El nuevo sistema decide mover los ingredientes ya etiquetados (en el formato correcto) directamente, evitando tener que parar y re-etiquetar todo en el camino. Es como pedir una pizza que ya viene cortada en lugar de tener que cortarla tú mismo al llegar a casa.

5. ¿Qué logran con esto?

Gracias a esta "coreografía" inteligente entre la memoria rápida (GPU) y la memoria lenta (CPU):

• El Ferrari ya no espera: La GPU está trabajando casi todo el tiempo.
• Velocidad: Logran entrenar los modelos 2.5 veces más rápido que los métodos actuales.
• Eficiencia: Pueden entrenar modelos gigantes en computadoras que antes parecían demasiado pequeñas para ello.

En resumen

El paper "Deep Optimizer States" es como inventar un sistema de logística inteligente para una fábrica. En lugar de tener una línea de montaje donde los trabajadores esperan a que les entreguen materiales, crean un sistema donde los trabajadores y los transportistas se coordinan perfectamente, aprovechando cada segundo de espacio libre y cada segundo de tiempo de transporte para que la producción nunca se detenga.

¡Es como pasar de un tráfico caótico a una autopista con carriles dinámicos donde nadie se queda atascado!

Resumen técnico

1. El Problema: La "Pared de Memoria" en el Entrenamiento de LLMs

El entrenamiento de Modelos de Lenguaje Grandes (LLMs) y Transformers ha enfrentado un desafío crítico conocido como la "pared de memoria". A medida que los modelos crecen (de cientos de miles de millones a billones de parámetros), los requisitos de memoria para almacenar no solo los parámetros del modelo, sino también los estados del optimizador (momentum, varianza, parámetros en FP32), superan la capacidad de la memoria de las GPUs, incluso utilizando paralelismo 3D (pipeline, tensor y de datos).

Para mitigar esto, enfoques actuales como DeepSpeed Offload y ZeRO-Offload mueven los estados del optimizador a la memoria del host (CPU). Sin embargo, esto introduce dos cuellos de botella severos:

1. Ancho de banda limitado de PCIe: La transferencia de datos entre CPU y GPU está restringida por enlaces PCIe (típicamente 25-50 GB/s), mucho más lentos que la memoria interna de la GPU.
2. Baja capacidad computacional de la CPU: Las actualizaciones de parámetros en la CPU son órdenes de magnitud más lentas que en la GPU. Por ejemplo, en el entorno de pruebas del artículo, las GPUs actualizan ~100 mil millones de parámetros/segundo, mientras que las CPUs solo logran ~8 mil millones/segundo.

Los enfoques existentes (como TwinFlow o ZeRO-Offload++) utilizan una estrategia estática: asignan una fracción fija de los estados del optimizador a la GPU y el resto a la CPU. Esto resulta ineficiente porque:

• La memoria de la GPU queda subutilizada durante las fases de forward y backward (cuando se liberan las activaciones).
• Los enlaces PCIe se subutilizan durante las fases de actualización.
• Existe una falta de superposición (overlap) entre el movimiento de datos y el cómputo, lo que genera tiempos de espera inactivos.

2. Metodología: Deep Optimizer States

Los autores proponen Deep Optimizer States, una técnica de middleware que introduce un desplazamiento entrelazado (interleaved offloading) dinámico de los estados del optimizador. En lugar de una asignación estática, el sistema decide dinámicamente qué subgrupos del optimizador se actualizan en la GPU y cuáles en la CPU en cada iteración.

Principios Clave de Diseño:

1. Actualizaciones Entrelazadas (Interleaved Updates):

   • Se aprovecha la observación de que el uso de memoria de la GPU fluctúa drásticamente durante la fase de actualización (cuando las activaciones se liberan).
   • El optimizador se divide en subgrupos (basado en la arquitectura ZeRO-3).
   • El sistema programa dinámicamente una fracción de estos subgrupos para su actualización en la GPU, mientras el resto se actualiza en la CPU, maximizando el uso de recursos.

2. Superposición de Movimiento y Ejecución:

   • Se utiliza un diseño centrado en la superposición (overlap-centric). Mientras la CPU calcula la actualización de un subgrupo, la GPU puede estar actualizando otro, y simultáneamente se realizan transferencias asíncronas (H2D y D2H) de otros subgrupos.
   • Se evitan bloqueos: las transferencias de datos no detienen el cómputo de otros componentes.

3. Gestión Eficiente de Gradientes y Precisión:

   • Almacenamiento de Gradientes: Se utilizan los espacios de memoria liberados por las activaciones en la GPU para almacenar gradientes de los subgrupos programados para actualización en la GPU.
   • Conversión de Precisión en el Dispositivo: Para evitar el costoso movimiento de datos con conversión de precisión en tiempo real (FP16 a FP32 en la CPU), el sistema realiza la conversión en la GPU (a 1.2 TB/s) y transfiere los bloques FP32 a la memoria del host. Esto es mucho más rápido que transferir FP16 y convertir en la CPU.

4. Modelo de Rendimiento para la Decisión de Programación:

   • Se introduce un modelo matemático (Ecuación 1) para calcular el "stride de actualización" ($k$) óptimo. Este modelo determina cuántos subgrupos deben actualizarse en la CPU por cada uno actualizado en la GPU, basándose en:
     • Tasa de actualización de CPU ($U_c$) vs. GPU ($U_g$).
     • Tasa de transferencia PCIe ($B$).
     • Tasa de conversión de precisión ($D_c$).
   • El objetivo es equilibrar los tiempos de cómputo y transferencia para maximizar la utilización de todos los recursos.

3. Contribuciones Clave

1. Análisis de Características del Sistema: Demostraron que el tamaño del subgrupo no afecta el tiempo de iteración, pero que el uso de memoria de la GPU y los enlaces PCIe tienen fluctuaciones significativas que pueden ser explotadas.
2. Diseño de Middleware: Implementación de Deep Optimizer States integrado en DeepSpeed y Megatron-LM, orquestando el paralelismo de fondo y la interacción con componentes existentes.
3. Modelo de Rendimiento y Algoritmo: Desarrollo de un algoritmo que decide dinámicamente la frecuencia de offloading a la GPU para maximizar la superposición, evitando la configuración manual de ratios estáticos.
4. Optimizaciones de Transferencia: Uso de conversiones de precisión en la GPU y buffers pinned para evitar cuellos de botella en la memoria del host y en la CPU.

4. Resultados Experimentales

Los autores evaluaron su enfoque en un entorno con 4x GPUs H100 (80 GB cada una) y CPUs Intel Xeon, entrenando modelos de hasta 20 mil millones de parámetros.

• Velocidad de Iteración: Deep Optimizer States logra ser 2.5 veces más rápido en las iteraciones de entrenamiento en comparación con los enfoques más avanzados (DeepSpeed ZeRO-3 y TwinFlow).
• Rendimiento de Actualización: La tasa de actualización de parámetros aumenta un 70% en promedio. En configuraciones específicas, se logra hasta un 3x de aceleración en la fase de actualización.
• Eficiencia de Memoria: El enfoque permite entrenar modelos de 20B parámetros con el mismo tiempo que un modelo de 7B parámetros usando métodos actuales. Además, ofrece iteraciones un 40% más rápidas que TwinFlow incluso con un 45% menos de uso de memoria de GPU.
• Escalabilidad: El sistema escala bien con diferentes tamaños de modelo (7B a 20B), tamaños de microbatch y grados de paralelismo de datos.
• Validación del Modelo: Los experimentos confirmaron que el modelo de rendimiento calcula correctamente el ratio óptimo (ej. actualizar 1 de cada 2 subgrupos en la GPU, $k=2$), logrando el máximo rendimiento en TFLOPs (hasta 75 TFLOPs frente a 30 en la base).

5. Significado e Impacto

Este trabajo es significativo porque aborda la ineficiencia fundamental de los métodos actuales de offloading híbrido (CPU-GPU). En lugar de tratar el offloading como una solución estática y pasiva, propone una gestión dinámica y activa de los recursos.

• Viabilidad para Recursos Limitados: Permite el entrenamiento y ajuste fino (fine-tuning) de LLMs moderadamente grandes (hasta 20B parámetros) en nodos con recursos limitados, lo cual es crucial para usuarios que no tienen acceso a clusters masivos.
• Preparación para Futuras Arquitecturas: La técnica es particularmente relevante para las próximas generaciones de hardware (como los sistemas Grace Hopper con interconexiones CPU-GPU de 200 GB/s), donde la capacidad de superponer cómputo y transferencia será aún más crítica.
• Generalidad: Aunque se centra en Transformers, la arquitectura es genérica y aplicable a otros modelos grandes (visión, modelos científicos) y frameworks de entrenamiento.

En resumen, Deep Optimizer States transforma el cuello de botella del offloading de optimizadores de un problema de latencia estática a una oportunidad de optimización dinámica, logrando aceleraciones sustanciales sin requerir hardware adicional.