Reduce costos de checkpoints en entrenamiento de modelos con Python y NVIDIA nvCOMP

Los modelos de lenguaje (LLMs) requieren puntos de control periódicos. Estas instantáneas de los pesos del modelo, estados del optimizador y gradientes se almacenan para poder reanudar el entrenamiento tras interrupciones. A gran escala, estos puntos de control pueden ser enormes (782 GB para un modelo de 70B) y frecuentes (cada 15-30 minutos), lo que representa un gran gasto en el presupuesto de entrenamiento. Sin embargo, muchos equipos de IA se centran en la utilización de GPU, la eficiencia del entrenamiento y la calidad del modelo, sin considerar el coste del almacenamiento de puntos de control.

Esta omisión puede resultar costosa. Por ejemplo, el costo de un modelo de 405B en 128 GPUs NVIDIA Blackwell puede alcanzar los $200,000 al mes solo por el almacenamiento de puntos de control. Al implementar un paso de compresión sin pérdidas con aproximadamente 30 líneas de Python, es posible reducir estos costos en $56,000 mensuales. A continuación, explicaremos cómo se realizó este cálculo y cómo NVIDIA nvComp puede mejorar la eficiencia del almacenamiento de puntos de control.

Detalles de un solo punto de control

Las interrupciones de hardware en escalas de más de 1000 GPUs no son raras. Meta reportó 419 interrupciones inesperadas en 54 días de entrenamiento de Llama 3 en 16,384 GPUs NVIDIA H100. Por esta razón, la mayoría de los equipos implementan puntos de control cada 15-30 minutos, ya que son una infraestructura fundamental.

Este desglose sorprende a quienes lo ven por primera vez. El estado del optimizador, que incluye las estimaciones de momento y varianza, es cuatro veces más grande que los pesos del modelo, constituyendo la mayor parte de cada punto de control.

Al realizar un punto de control cada 30 minutos, que es la práctica estándar por tolerancia a fallas, se generan 48 puntos de control diarios. En un mes de entrenamiento continuo, esto resulta en:

782 GB × 48/día × 30 días = 1.13 PB escritos en almacenamiento mensual

Además, durante cada escritura de punto de control, todas las GPUs permanecen inactivas. El ciclo de entrenamiento se bloquea hasta que se completan las operaciones de escritura.

El costo de las GPUs inactivas durante estos periodos de espera puede ser significativo y puede sumar más de $2,200 al mes antes de considerar las tarifas de almacenamiento.

NVIDIA nvCOMP y la compresión acelerada por GPU

La idea central es simple: comprimir el punto de control antes de que salga de la memoria de la GPU, evitando movimientos de datos innecesarios. NVIDIA nvCOMP es una biblioteca de compresión sin pérdidas acelerada por GPU que logra esto. Al ofrecer una única biblioteca con soporte tanto para algoritmos estándar, como Zstandard (ZSTD), como para formatos optimizados específicos de GPU, como gANS, se abordan los cuellos de botella de datos directamente en el dispositivo.

Comprobamos las relaciones de compresión al ajustar dos arquitecturas de modelo (transformador denso y mezcla de expertos) y comprimir cada componente con nvCOMP en GPUs NVIDIA H200 y Blackwell. Las tasas de compresión dependen de los datos, no del hardware, y son idénticas en todas las GPUs.

ZSTD, un algoritmo de compresión de propósito general desarrollado por Meta, ofrece una buena relación entre tasas de compresión y velocidad. ANS, un sistema de codificación de entropía moderna, logra tasas de compresión similares a una velocidad mucho mayor.

A horizontal grouped bar chart comparing the compression ratios of ZSTD (green bars) and ANS (blue bars) using nvCOMP on a Blackwell GPU.

Las diferencias entre ZSTD y ANS radican en el equilibrio entre la tasa de compresión y la velocidad de compresión. ZSTD ofrece una ligera ventaja en la tasa de compresión, pero ANS puede alcanzar velocidades significativamente más altas.

Ilustración de un hombre mayor con auriculares y chaqueta