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

1. Declaração do Problema

O treinamento de Modelos de Linguagem de Grande Escala (LLMs) requer memória de GPU massiva, consumida principalmente por parâmetros do modelo, estados do otimizador, gradientes e ativações.

• Gargalo de Memória: As ativações normalmente ocupam a maior parte da memória durante o treinamento, escalando linearmente com o comprimento da sequência e o tamanho do lote.
• Limitações dos Métodos Existentes:
  • Recomputação/Descarga de Ativações: Reduz a memória, mas introduz sobrecarga computacional significativa (recomputação) ou latência de E/S (descarga).
  • Quantização Padrão: Embora a quantização de 8 bits (INT8/FP8) seja comum, aplicá-la a ativações de 4 bits ou gradientes de 8 bits frequentemente leva a perda severa de precisão, convergência lenta ou instabilidade no treinamento devido a estouro numérico durante a acumulação e operações All-Reduce.
  • Comunicação de Gradientes: Sistemas atuais frequentemente armazenam gradientes em alta precisão (FP32) e os quantizam apenas para comunicação, falhando em reduzir a pegada de memória dos gradientes armazenados.

2. Metodologia: Framework AGoQ

Os autores propõem o AGoQ, um framework holístico que integra duas técnicas inovadoras para alcançar armazenamento de ativações próximo a 4 bits e armazenamento/comunicação de gradientes de 8 bits sem sacrificar a convergência do modelo.

A. Quantização de Ativação Sensível à Camada (LAAQ)

Em vez de quantizar uniformemente todas as camadas para 4 bits (o que causa perda de precisão), o LAAQ atribui dinamicamente larguras de bits com base nos tipos de camada e estágios do pipeline.

• Análise Teórica de Erro: Os autores derivam limites de erro para diferentes operações (RMSNorm, SiLU, Attention, GEMM).
  • Descobertas: Quantizar as ativações do módulo Attention (Q, K, V) introduz erros de gradiente significativamente maiores em comparação com outros módulos. Portanto, as ativações de Attention não são quantizadas (mantidas em maior precisão).
  • Estratégia: Para outras camadas (RMSNorm, SiLU, FFN), o sistema usa quantização em blocos de 4 bits.
  • Recomputação vs. Cache: A análise mostra que, para operações como RMSNorm e SiLU, armazenar apenas a entrada quantizada e recomputar os valores intermediários durante a passagem de retropropagação produz um limite de erro mais apertado do que armazenar os valores intermediários em forma quantizada. Assim, o AGoQ armazena apenas entradas de 4 bits e recomputa os intermediários.
• Compensação Dinâmica de Largura de Bits (DBCA-PP):
  • No Parallelismo de Pipeline (PP), diferentes estágios do pipeline mantêm diferentes números de lotes de ativação (uso de memória desigual).
  • O AGoQ explora isso atribuindo larguras de bits mais altas (por exemplo, 5, 6 ou 8 bits) a estágios com menos lotes ativos e larguras de bits mais baixas (4 bits) a estágios com mais lotes. Isso equilibra o uso total de memória enquanto mantém uma média de ~4 bits por elemento, compensando a perda de precisão onde há memória disponível.

B. Quantização de Gradiente com Precisão Preservada (QuanGrad)

Esta técnica reduz tanto o uso de memória quanto o tempo de comunicação ao usar gradientes de 8 bits (FP8) durante todo o processo.

• Acumulação Local: Para prevenir estouro ao acumular gradientes de vários mini-lotes, o sistema desquantiza o gradiente principal FP8 para FP16/BF16, adiciona o novo gradiente local e, em seguida, requantiza a soma de volta para FP8.
• All-Reduce com Precisão Preservada: O All-Reduce padrão em FP8 causa estouro durante a etapa de redução. O AGoQ decompõe a operação All-Reduce em:
  1. All-to-All: Enviar gradientes FP8 comprimidos para todos os dispositivos.
  2. Redução Local: Cada dispositivo desquantiza os dados recebidos para FP32, realiza a soma local e requantiza para FP8.
  3. All-Gather: Coletar os resultados reduzidos.
  • Isso garante que a acumulação ocorra em alta precisão (FP32) enquanto a comunicação permanece em baixa precisão (FP8), prevenindo estouro e mantendo a convergência.

C. Fusão de Kernel

Para eliminar a sobrecarga das operações de quantização/desquantização, o AGoQ funde essas operações elemento a elemento com kernels GEMM (Multiplicação de Matriz Geral) adjacentes. Isso aproveita os Tensor Cores para o trabalho pesado, mantendo a lógica de quantização nos núcleos CUDA dentro do mesmo kernel, minimizando a latência.

3. Principais Contribuições

1. Estratégia de Quantização Sensível à Camada: Uma abordagem teórica e prática para identificar quais camadas podem ser com segurança comprimidas para 4 bits (excluindo Attention) e quais requerem estratégias de recomputação para minimizar o erro.
2. Compensação Dinâmica de Largura de Bits: Um algoritmo inovador para Parallelismo de Pipeline que ajusta dinamicamente a precisão da quantização com base na disponibilidade de memória em cada estágio, maximizando a eficiência de memória sem aumentar o uso de pico.
3. Comunicação de Gradiente com Precisão Preservada: Um método para realizar All-Reduce de 8 bits dividindo a operação para permitir acumulação de alta precisão, resolvendo os problemas de estouro que anteriormente impediam o treinamento de gradientes de baixo bit.
4. Integração de Sistema: Integração completa no Megatron-LM, compatível com otimizadores Adam de 8 bits, suportando tanto GPUs NVIDIA quanto NPUs Ascend da Huawei.

4. Resultados Experimentais

Os experimentos foram conduzidos em clusters com até 64 GPUs (NVIDIA A6000/Pro 6000) usando modelos LLaMA (7B, 8B, 13B, 34B) e OLMo-1B.

• Redução de Memória:

  • O AGoQ reduz o uso de memória de pico em até 52% em comparação com o treinamento de precisão total (Megatron-LM).
  • Em comparação com o sistema de última geração COAT, o AGoQ reduz a memória de ativação em 30% adicionais e a memória de gradiente em 75%.
  • Permite o treinamento de comprimentos de sequência mais longos (por exemplo, 80K tokens) que causam erros de Memória Insuficiente (OOM) em sistemas de base.

• Velocidade de Treinamento:

  • Alcança até 1,34× de aceleração no tempo de treinamento de ponta a ponta em comparação com Megatron-LM (com ou sem ZeRO) e DeepSpeed.
  • A aceleração é mais pronunciada com comprimentos de sequência mais longos devido à redução da sobrecarga de comunicação e restrições de memória.

• Convergência e Precisão:

  • Convergência: As curvas de perda de treinamento no LLaMA2-7B e LLaMA3-8B correspondem estreitamente à linha de base de precisão total. Em contraste, métodos ingênuos de All-Reduce FP8 mostraram perda significativamente maior.
  • Tarefas de Descendente: A precisão zero-shot em conjuntos de dados (ARC, HellaSwag, PIQA, etc.) permanece comparável à linha de base, com algumas tarefas mostrando ligeiras melhorias.
  • Estabilidade: As normas de gradiente permanecem estáveis durante todo o treinamento, confirmando a eficácia da acumulação com precisão preservada.

5. Significado

• Habilitando Treinamento de Longo Contexto: Ao reduzir drasticamente as pegadas de memória, o AGoQ torna viável treinar LLMs com janelas de contexto muito longas (por exemplo, 64K–80K tokens) em hardware existente sem exigir clusters massivos.
• Preenchendo a Lacuna: Ele empurra com sucesso a quantização para 4 bits para ativações e 8 bits para gradientes, um regime anteriormente considerado instável demais para treinamento de LLMs, sem comprometer a qualidade do modelo.
• Escalabilidade: O método é altamente escalável através de diferentes estratégias de paralelismo (DP, TP, PP) e arquiteturas de hardware, oferecendo um caminho prático para treinar modelos com trilhões de parâmetros em recursos limitados.

Em resumo, o AGoQ representa um avanço significativo nos sistemas para treinamento de LLMs, indo além da simples compressão para um framework de quantização adaptativo e fundamentado matematicamente que equilibra eficiência de memória, velocidade de comunicação e precisão do modelo.