Infraestrutura e Chips
FlashAttention: atenção rápida pensando na memória da GPU
O gargalo da atenção não é só a conta, é o vaivém de dados na memória. O FlashAttention reorganiza o cálculo para ler e escrever muito menos.
@raulelgreco1 min de leitura
A atenção do Transformer cria uma matriz que cresce com o quadrado do tamanho do texto. Em GPUs, o problema principal nem é a quantidade de contas, mas o tráfego entre a memória principal da placa (HBM) e a memória rápida e pequena dos núcleos (SRAM).
O FlashAttention, de Dao e colegas (2022), calcula a atenção em blocos que cabem na SRAM e evita gravar a matriz completa na memória principal. O resultado é exato: não é uma aproximação.
Ganhos relatados
- Treino de modelos do tipo GPT-2 até cerca de 3 vezes mais rápido.
- Uso de memória que cresce de forma linear, não quadrática, com o tamanho do texto.
- Viabilidade de contextos bem mais longos no mesmo hardware.
Versões posteriores melhoraram ainda mais o aproveitamento das GPUs modernas. Hoje, o FlashAttention ou técnicas inspiradas nele estão embutidos na maioria das bibliotecas de treino e inferência.
Foto: Winston Chen / Unsplash.