A atenção é um elemento fundamental da arquitetura de transformadores que sustenta os grandes modelos de linguagem (LLMs). No entanto, à medida que os LLMs se expandem e gerenciam sequências de entrada cada vez mais longas, o custo computacional da atenção se torna um gargalo significativo.
Para resolver esse problema, uma equipe colaborativa da Colfax Research, Meta, Nvidia, Georgia Tech, Princeton University e Together AI apresentou o FlashAttention-3. Essa técnica inovadora acelera substancialmente a computação de atenção em GPUs Nvidia Hopper (H100 e H800).
Entendendo o Desafio da Computação de Atenção em LLMs
O mecanismo de atenção do modelo transformer permite avaliar as relações entre vários tokens em uma sequência de entrada. Embora eficaz, esse mecanismo é intensivo em computação. Com o aumento do comprimento das sequências de entrada, o custo da computação de atenção cresce quadraticamente, levando a gargalos significativos na escalabilidade dos LLMs.
Além disso, as GPUs modernas são otimizadas principalmente para operações de multiplicação de matrizes (matmul), essenciais para o aprendizado profundo. Outras operações, como exponenciação, são muito mais lentas, agravando os problemas de desempenho. As computações de atenção, que incorporam multiplicações de matrizes junto com funções mais complexas, como a função softmax para normalizar os pesos de atenção, podem se tornar limitadas devido ao custo computacional elevado da última. Programar cargas de trabalho de maneira eficaz é fundamental para evitar conflitos operacionais e utilizar os recursos de memória de forma eficiente.
Aprimorando a Utilização de Recursos de Hardware com FlashAttention
Lançado em 2022, o FlashAttention abordou as ineficiências da computação de atenção, minimizando as transferências de memória entre a memória de alta largura de banda (HBM) e a memória estática de acesso aleatório (SRAM) nas GPUs. Ao processar os pesos de atenção em blocos menores ou “tiles”, o FlashAttention melhorou a eficiência, permitindo que os LLMs ampliassem suas janelas de contexto de milhares para potencialmente milhões de tokens.
Entretanto, com o avanço das capacidades de hardware, surgia a necessidade de mais otimizações. O FlashAttention-2, introduzido em 2023, otimizou o uso dos recursos da GPU, alcançando 70% do desempenho máximo nas GPUs Nvidia A100, mas utilizando apenas 35% das capacidades das H100.
Inovações do FlashAttention-3
O FlashAttention-3 aproveita novos recursos das GPUs Nvidia Hopper para melhorar o desempenho, incluindo um maior throughput para multiplicação de matrizes e transferências de dados mais rápidas entre segmentos de memória. Isso possibilita melhor eficiência em operações de baixa precisão.
As principais inovações do FlashAttention-3 incluem:
1. Agendamento Otimizado: As operações são organizadas para maximizar a sobreposição entre computação e movimentação de dados, reduzindo o tempo ocioso das GPUs.
2. Interleaving de Operações Sem Costura: Ao combinar operações de multiplicação de matrizes e softmax, o FlashAttention-3 minimiza possíveis gargalos.
3. Desempenho Aprimorado do Modelo Quantizado: Arranjos especiais nas operações garantem computações mais rápidas e precisas, mesmo ao usar representações de baixa bit, abordando a típica troca de precisão associada à quantização.
Pesquisas indicam que o FlashAttention-3 pode utilizar até 75% do desempenho máximo da GPU H100, proporcionando um aumento de 1,5 a 2 vezes em comparação com versões anteriores do FlashAttention para treinamento e implementação de LLMs.
Benefícios do FlashAttention-3
A computação de atenção mais rápida proporcionada pelo FlashAttention-3 tem implicações profundas para o desenvolvimento e aplicações de LLMs:
- Treinamento Acelerado: A eficiência aprimorada pode reduzir significativamente os tempos de treinamento, permitindo que pesquisadores explorem modelos e conjuntos de dados maiores.
- Janelas de Contexto Expandidas: Com a capacidade de processar sequências mais longas de forma eficiente, o FlashAttention-3 abre novas possibilidades para aplicações de LLM, como compreensão de documentos longos e aprendizado contextual em muitos exemplos.
- Eficiência de Custos: Maior utilização da GPU pode resultar em menos aceleradores necessários para as operações de LLM, reduzindo, assim, os custos de produção.
O FlashAttention-3 foi disponibilizado como código aberto sob uma licença permissiva, com planos para integração em bibliotecas populares de aprendizado profundo, como PyTorch e Hugging Face Transformers. Isso visa facilitar a adoção das inovações do FlashAttention-3 por pesquisadores e desenvolvedores. Conforme destacado por pesquisadores em um post no blog da Together AI, "Projetar algoritmos que aproveitam recursos de hardware pode gerar melhorias significativas na eficiência e desbloquear novas capacidades dos modelos." Eles aguardam novas otimizações para inferência de LLM e a aplicação de suas técnicas em diversas arquiteturas de hardware.
Classificação
