AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs

A continuación se presenta un resumen técnico detallado del artículo "AGoQ: Cuantificación de Activaciones y Gradientes para el Entrenamiento Distribuido Eficiente en Memoria de LLMs".

1. Planteamiento del Problema

El entrenamiento de Modelos de Lenguaje Grandes (LLMs) requiere una enorme memoria de GPU, consumida principalmente por los parámetros del modelo, los estados del optimizador, los gradientes y las activaciones.

• Cuello de Botella de Memoria: Las activaciones suelen ocupar la mayor parte de la memoria durante el entrenamiento, escalando linealmente con la longitud de la secuencia y el tamaño del lote.
• Limitaciones de los Métodos Existentes:
  • Recomputación/Descarga de Activaciones: Reduce la memoria pero introduce una sobrecarga computacional significativa (recomputación) o latencia de E/S (descarga).
  • Cuantificación Estándar: Aunque la cuantificación de 8 bits (INT8/FP8) es común, aplicarla a activaciones de 4 bits o gradientes de 8 bits a menudo conduce a una pérdida severa de precisión, convergencia lenta o inestabilidad en el entrenamiento debido a desbordamientos numéricos durante la acumulación y las operaciones All-Reduce.
  • Comunicación de Gradientes: Los sistemas actuales suelen almacenar los gradientes en alta precisión (FP32) y solo los cuantifican para la comunicación, fallando en reducir la huella de memoria de los gradientes almacenados.

2. Metodología: Marco AGoQ

Los autores proponen AGoQ, un marco holístico que integra dos técnicas novedosas para lograr un almacenamiento de activaciones cercano a 4 bits y un almacenamiento/comunicación de gradientes de 8 bits sin sacrificar la convergencia del modelo.

A. Cuantificación de Activaciones Consciente de la Capa (LAAQ)

En lugar de cuantificar uniformemente todas las capas a 4 bits (lo que causa pérdida de precisión), LAAQ asigna dinámicamente anchos de bits basándose en los tipos de capas y las etapas del pipeline.

• Análisis Teórico del Error: Los autores derivan límites de error para diferentes operaciones (RMSNorm, SiLU, Atención, GEMM).
  • Hallazgos: Cuantificar las activaciones del módulo de Atención (Q, K, V) introduce errores de gradiente significativamente mayores en comparación con otros módulos. Por lo tanto, las activaciones de Atención no se cuantifican (se mantienen en mayor precisión).
  • Estrategia: Para otras capas (RMSNorm, SiLU, FFN), el sistema utiliza cuantificación por bloques de 4 bits.
  • Recomputación vs. Caché: El análisis muestra que, para operaciones como RMSNorm y SiLU, almacenar solo la entrada cuantificada y recomputar los valores intermedios durante el paso hacia atrás produce un límite de error más ajustado que almacenar los valores intermedios en forma cuantificada. Por lo tanto, AGoQ almacena solo entradas de 4 bits y recomputa los intermedios.
• Compensación Dinámica del Ancho de Bits (DBCA-PP):
  • En la Paralelización de Pipeline (PP), diferentes etapas del pipeline mantienen diferentes números de lotes de activaciones (uso de memoria desigual).
  • AGoQ explota esto asignando anchos de bits más altos (por ejemplo, 5, 6 u 8 bits) a las etapas con menos lotes activos y anchos de bits más bajos (4 bits) a las etapas con más lotes. Esto equilibra el uso total de memoria mientras mantiene un promedio de ~4 bits por elemento, compensando la pérdida de precisión donde hay memoria disponible.

B. Cuantificación de Gradientes con Preservación de Precisión (QuanGrad)

Esta técnica reduce tanto el uso de memoria como el tiempo de comunicación al utilizar gradientes de 8 bits (FP8) durante todo el proceso.

• Acumulación Local: Para prevenir desbordamientos al acumular gradientes de múltiples mini-lotes, el sistema descuantifica el gradiente principal FP8 a FP16/BF16, suma el nuevo gradiente local y luego re-cuantifica la suma de vuelta a FP8.
• All-Reduce con Preservación de Precisión: El All-Reduce estándar en FP8 causa desbordamiento durante el paso de reducción. AGoQ descompone la operación All-Reduce en:
  1. All-to-All: Enviar gradientes FP8 comprimidos a todos los dispositivos.
  2. Reducción Local: Cada dispositivo descuantifica los datos recibidos a FP32, realiza la suma local y re-cuantifica a FP8.
  3. All-Gather: Recopilar los resultados reducidos.
  • Esto asegura que la acumulación ocurra en alta precisión (FP32) mientras la comunicación permanece en baja precisión (FP8), previniendo desbordamientos y manteniendo la convergencia.

C. Fusión de Kernels

Para eliminar la sobrecarga de las operaciones de cuantificación/descuantificación, AGoQ fusiona estas operaciones elemento a elemento con los kernels GEMM (Multiplicación General de Matrices) adyacentes. Esto aprovecha los Tensor Cores para el trabajo pesado mientras mantiene la lógica de cuantificación en los núcleos CUDA dentro del mismo kernel, minimizando la latencia.

3. Contribuciones Clave

1. Estrategia de Cuantificación Consciente de la Capa: Un enfoque teórico y práctico para identificar qué capas pueden comprimirse de forma segura a 4 bits (excluyendo la Atención) y cuáles requieren estrategias de recomputación para minimizar el error.
2. Compensación Dinámica del Ancho de Bits: Un algoritmo novedoso para la Paralelización de Pipeline que ajusta dinámicamente la precisión de cuantificación basándose en la disponibilidad de memoria en cada etapa, maximizando la eficiencia de memoria sin aumentar el uso máximo.
3. Comunicación de Gradientes con Preservación de Precisión: Un método para realizar All-Reduce de 8 bits dividiendo la operación para permitir una acumulación de alta precisión, resolviendo los problemas de desbordamiento que anteriormente impedían el entrenamiento con gradientes de bajo ancho de bits.
4. Integración del Sistema: Integración completa en Megatron-LM, compatible con optimizadores Adam de 8 bits, soportando tanto GPUs de NVIDIA como NPUs de Huawei Ascend.

4. Resultados Experimentales

Los experimentos se realizaron en clústeres con hasta 64 GPUs (NVIDIA A6000/Pro 6000) utilizando modelos LLaMA (7B, 8B, 13B, 34B) y OLMo-1B.

• Reducción de Memoria:

  • AGoQ reduce el uso máximo de memoria hasta un 52% en comparación con el entrenamiento de precisión completa (Megatron-LM).
  • En comparación con el sistema de vanguardia COAT, AGoQ reduce la memoria de activaciones en un 30% adicional y la memoria de gradientes en un 75%.
  • Permite entrenar longitudes de secuencia más largas (por ejemplo, 80K tokens) que causan errores de Memoria Agotada (OOM) en sistemas de referencia.

• Velocidad de Entrenamiento:

  • Logra una aceleración de hasta 1.34× en el tiempo de entrenamiento de extremo a extremo en comparación con Megatron-LM (con o sin ZeRO) y DeepSpeed.
  • La aceleración es más pronunciada con longitudes de secuencia más largas debido a la reducción de la sobrecarga de comunicación y las restricciones de memoria.

• Convergencia y Precisión:

  • Convergencia: Las curvas de pérdida de entrenamiento en LLaMA2-7B y LLaMA3-8B coinciden estrechamente con la línea base de precisión completa. Por el contrario, los métodos ingenuos de All-Reduce FP8 mostraron una pérdida significativamente mayor.
  • Tareas de Descarga: La precisión en cero disparos (zero-shot) en conjuntos de datos (ARC, HellaSwag, PIQA, etc.) permanece comparable a la línea base, con algunas tareas mostrando ligeras mejoras.
  • Estabilidad: Las normas de los gradientes permanecen estables durante todo el entrenamiento, confirmando la efectividad de la acumulación con preservación de precisión.

5. Significado

• Habilitación del Entrenamiento de Contexto Largo: Al reducir drásticamente las huellas de memoria, AGoQ hace viable entrenar LLMs con ventanas de contexto muy largas (por ejemplo, 64K–80K tokens) en hardware existente sin requerir clústeres masivos.
• Cerrando la Brecha: Empuja exitosamente la cuantificación a 4 bits para activaciones y 8 bits para gradientes, un régimen previamente considerado demasiado inestable para el entrenamiento de LLMs, sin comprometer la calidad del modelo.
• Escalabilidad: El método es altamente escalable a través de diferentes estrategias de paralelismo (DP, TP, PP) y arquitecturas de hardware, ofreciendo un camino práctico hacia el entrenamiento de modelos de billones de parámetros con recursos limitados.

En resumen, AGoQ representa un avance significativo en los sistemas para el entrenamiento de LLMs, yendo más allá de la compresión simple hacia un marco de cuantificación adaptativo y fundamentado matemáticamente que equilibra la eficiencia de memoria, la velocidad de comunicación y la precisión del modelo.