El equipo de DeepSpeed ha presentado ZenFlow, un nuevo motor de offloading diseñado para resolver un gran problema en el entrenamiento de modelos de lenguaje grandes (LLM): las pausas en la GPU causadas por la CPU. Aunque descargar optimizadores y gradientes a la memoria de la CPU ayuda a reducir la presión sobre la memoria de la GPU, los métodos tradicionales como ZeRO-Offload y ZeRO-Infinity suelen dejar las costosas GPUs esperando durante gran parte del paso de entrenamiento debido a la lentitud de las actualizaciones en CPU y las transferencias PCIe. Por ejemplo, al afinar Llama 2-7B en 4 GPUs A100 con offloading completo, el tiempo por paso puede aumentar de 0.5 segundos a más de 7 segundos, un retraso de 14 veces. ZenFlow elimina esas pausas separando los cálculos de la GPU y la CPU mediante un pipeline que prioriza las partes más importantes, logrando hasta 5 veces más rapidez en todo el proceso y reduciendo las esperas en GPU en más del 85%.
¿Cómo funciona ZenFlow?
- Actualizaciones de gradientes según importancia: ZenFlow da prioridad a un pequeño grupo de los gradientes más relevantes para actualizarlos inmediatamente en la GPU, mientras que los menos importantes se acumulan de manera asíncrona en la CPU. Esto reduce casi a la mitad la cantidad de gradientes que se transfieren cada paso y disminuye a la mitad el uso del ancho de banda PCIe en comparación con ZeRO-Offload.
- Acumulación asíncrona con límite: Los gradientes menos críticos se agrupan y actualizan en la CPU en segundo plano, escondiendo ese trabajo detrás de la carga que soporta la GPU. Esto asegura que las GPUs se mantengan ocupadas sin interrupciones, aprovechando al máximo el hardware disponible.
- Selección ligera de gradientes: En lugar de recopilar todos los gradientes completos, ZenFlow utiliza un método simple basado en la norma por columna para identificar cuáles son más importantes, lo que reduce la comunicación hasta 4,000 veces sin afectar la precisión. Esto facilita escalar el entrenamiento en múltiples GPUs.
- Sin cambios en el código y configuración sencilla: ZenFlow está integrado dentro de DeepSpeed y solo requiere ajustes básicos en el archivo JSON de configuración, como definir qué porcentaje de gradientes priorizar y activar estrategias adaptativas que se ajustan automáticamente durante el entrenamiento.
- Optimización automática: El motor ajusta de forma dinámica los intervalos de actualización durante el entrenamiento, evitando la necesidad de configuraciones manuales y garantizando la máxima eficiencia según cómo evoluciona el proceso.
Resultados destacados
- Hasta 5 veces más velocidad total, acelerando la convergencia y reduciendo costos.
- Más del 85% menos pausas en la GPU, para un uso más eficiente del hardware.
- Consumo de ancho de banda PCIe casi a la mitad, aliviando la presión en el clúster.
- Sin pérdida de precisión en benchmarks como GLUE.
- Escalabilidad eficiente para configuraciones con múltiples GPUs.
- Ajustes automáticos que eliminan la necesidad de retoques manuales.
Uso práctico
ZenFlow se integra fácilmente como una extensión de ZeRO-Offload en DeepSpeed sin requerir modificaciones en el código, solo actualizando la configuración JSON. Por ejemplo, en el repositorio de DeepSpeedExamples hay un caso práctico para afinar modelos en el benchmark GLUE usando ZenFlow, que se puede ejecutar con un simple script y sirve como punto de partida para experimentar con esta tecnología.
Un ejemplo básico de configuración JSON sería:
```json "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "zenflow": { "topk_ratio": 0.05, "select_strategy": "auto", "select_interval": "auto", "update_interval": 4, "full_warm_up_rounds": 0, "overlap_step": true } } ```
Para comenzar, se recomienda consultar el ejemplo de fine-tuning con ZenFlow y el tutorial oficial, que guían paso a paso en la configuración y uso.
En resumen, ZenFlow representa un avance importante para quienes entrenan o afinan grandes modelos de lenguaje con recursos limitados en GPU. Al eliminar las esperas provocadas por la CPU, permite un rendimiento mucho más fluido y rentable, sin sacrificar la calidad del modelo. Su integración sencilla con DeepSpeed, junto con la optimización automática y documentación accesible, facilitan su adopción para equipos técnicos que buscan maximizar el uso del hardware, tanto en la nube como en instalaciones propias.
ZenFlow redefine la forma de hacer offloading en entrenamiento de LLMs, ofreciendo un proceso sin pausas y de alta velocidad con mínima complejidad en la configuración, siendo una propuesta imperdible para quienes trabajan con inteligencia artificial a gran escala.
