Performance Optimization and Comparative Analysis of Generative AI Models on Advanced Accelerators

Resumo Técnico: Otimização de Desempenho e Análise Comparativa de Modelos de IA Generativa em Aceleradores Avançados

Declaração do Problema
Apesar das impressionantes capacidades dos modelos de IA Generativa, incluindo Grandes Modelos de Linguagem (LLMs) e modelos de difusão, sua implantação enfrenta obstáculos significativos. Esses desafios incluem requisitos substanciais de memória, latência de inferência estendida, alta demanda computacional e custos elevados de hardware. Esses problemas são agravados quando se avaliam modelos em plataformas heterogêneas, onde variações em formatos numéricos, larguras de banda de memória e stacks de software interagem de forma complexa com as arquiteturas de modelos e as características de carga de trabalho. O artigo aborda a necessidade de um estudo sistemático para otimizar o desempenho e comparar esses modelos em diversas tarefas de downstream em sistemas modernos de Computação de Alto Desempenho (HPC) e aceleradores avançados.

Metodologia
Os autores realizaram uma avaliação abrangente em três sistemas de HPC distintos: Perlmutter (NERSC, apresentando GPUs NVIDIA A100), Expanse (SDSC, apresentando GPUs NVIDIA V100, A100 e H100) e Voyager (SDSC, apresentando aceleradores Intel Gaudi). O estudo focou em três áreas primárias:

1. Quantização Pós-Treinamento (PTQ) de Precisão Mista: Os autores introduziram um novo framework de PTQ de precisão mista sensível à sensibilidade, projetado para otimizar LLMs (especificamente o TinyLlama-1.1B) nos aceleradores NVIDIA A100 e Intel Gaudi HPU.

   • Análise de Sensibilidade: O framework emprega uma abordagem de duas fases. Ele calcula pontuações de sensibilidade ao nível da camada transformer analisando as magnitudes dos pesos (excluindo vieses e parâmetros de normalização) e aplicando um limiar de poda unificado baseado em um alvo de esparsidade (30%).
   • Estratégia de Quantização: Camadas sensíveis recebem maior precisão, enquanto camadas robustas recebem menor precisão. O estudo avaliou quantização "fake" (simulada em FP32) e quantização "real" usando FP16/INT8 para A100 e BF16/FP8 (via Intel Neural Compressor) para Gaudi.
   • Datasets: A avaliação foi realizada no WikiText-2 (modelagem de linguagem), HellaSwag (raciocínio de senso comum) e BoolQ (compreensão de leitura).

2. Desempenho de Fine-Tuning no Gaudi: O estudo avaliou o ajuste fino (fine-tuning) de três LLMs de código aberto (Llama3.1-8B, Gemma2-27B e Llama3.3-70B) no Voyager usando a biblioteca Optimum-Habana.

   • Paralelização: Os autores compararam estratégias de Distributed Data Parallel (DDP), DeepSpeed (ZeRO3) e Fully Sharded Data Parallel (FSDP).
   • Otimização: O impacto do Flash Attention (FA) e as diferenças geracionais entre Gaudi1, Gaudi2 e Gaudi3 foram analisados.

3. Treinamento de Modelo de Difusão: Um modelo de Super-Resolução, utilizado para melhorar imagens astrofísicas de poeira galáctica, foi treinado em nós Gaudi e comparado contra GPUs NVIDIA (V100, A100, H100) no Expanse. O modelo foi paralelizado usando PyTorch DDP e executado em modo SynapseAI Lazy.

Resultos Principais

• Desempenho de Quantização: A quantização real no A100 e Gaudi2 alcançou taxas de compressão entre 2x e 6x com perda mínima de precisão. Por exemplo, na tarefa WikiText-2, o Gaudi2 alcançou uma taxa de compressão de 5.93x com uma perplexidade de 12.39 (comparado a 12.02 para FP32), enquanto o A100 alcançou 3.77x de compressão com uma perplexidade de 12.02. As melhorias de throughput foram significativas, com o Gaudi2 mostrando até 34.081 tokens/seg em cenários de quantização simulada.
• Eficiência de Fine-Tuning:
  • Ganhos Geracionais: As placas Gaudi2 demonstraram uma redução de tempo de execução de 2.5x a 3x em comparação com o Gaudi1, atribuída ao aumento de 8 para 24 Tensor Processing Cores (TPCs). O Gaudi3 ofereceu um ganho de desempenho adicional de 1.8x a 2x sobre o Gaudi2.
  • Flash Attention: Habilitar o Flash Attention no Gaudi1 e Gaudi2 resultou em uma melhoria de desempenho de 10% a 20% ao minimizar a movimentação de dados entre tipos de memória.
  • Overhead de Sharding: Para o modelo Gemma2 de 27B parâmetros, o DDP superou decisivamente os métodos de sharding de modelo (DeepSpeed e FSDP). Os autores descobriram que, embora o sharding reduza a pegada de memória, ele introduz um overhead significativo que degrada o desempenho. Consequentemente, o sharding deve ser empregado apenas quando um modelo não cabe na HBM de uma única placa.
• Escalabilidade de Modelo de Difusão:
  • GPUs NVIDIA: Exibiram escalonamento quase perfeito com um aumento de velocidade de ~2x do V100 para o A100 e outro de ~2x do A100 para o H100.
  • Intel Gaudi: Mostrou um aumento de velocidade de 4x do Gaudi1 para o Gaudi2. No entanto, essa tendência não continuou linearmente do Gaudi2 para o Gaudi3, indicando um platô na eficiência de escalonamento para esta carga de trabalho específica.

Significância e Alegações
O artigo afirma fornecer uma análise comparativa sistemática do desempenho de IA Generativa através de ambientes de HPC heterogêneos, oferecendo insights práticos para implantação. Suas principais contribuições incluem:

1. Um novo framework de quantização de precisão mista sensível à sensibilidade que equilibra a precisão do modelo e a compressão entre diferentes arquiteturas de hardware.
2. Evidência empírica demonstrando que, para modelos que cabem na memória de uma única placa, o DDP padrão é superior às estratégias complexas de sharding como DeepSpeed ou FSDP, que introduzem overhead desnecessário.
3. Benchmarking detalhado da evolução dos aceleradores Intel Gaudi, destacando ganhos de desempenho específicos (e limitações) em atualizações geracionais e a eficácia do Flash Attention neste ecossistema.
4. Validação das capacidades de treinamento de modelos de difusão em aceleradores Gaudi, fornecendo uma base para aplicações de computação científica, como o aprimoramento de imagens astrofísicas.

O trabalho enfatiza que, embora os aceleradores avançados ofereçam ganhos substanciais de desempenho, a implantação ideal requer consideração cuidadosa do tamanho do modelo, restrições de memória e os trade-offs específicos entre quantização, estratégias de paralelização e gerações de hardware.