BurstAttention: Um Framework de Atenção Distribuída Eficiente para Sequências Extremamente Longas

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

Autores: Sun Ao, Weilin Zhao, Xu Han, Cheng Yang, Zhiyuan Liu, Chuan Shi, Maosong Sun, Shengnan Wang, Teng Su

cs.AI

Resumo

Módulos de atenção eficazes têm desempenhado um papel crucial no sucesso de grandes modelos de linguagem baseados em Transformers (LLMs), mas as complexidades quadráticas de tempo e memória desses módulos de atenção também representam um desafio ao processar sequências longas. Uma solução potencial para o problema de sequências longas é utilizar clusters distribuídos para paralelizar o cálculo dos módulos de atenção em vários dispositivos (por exemplo, GPUs). No entanto, adotar uma abordagem distribuída inevitavelmente introduz sobrecargas adicionais de memória para armazenar resultados locais de atenção e incorre em custos extras de comunicação para agregar resultados locais em resultados globais. Neste artigo, propomos um framework de atenção distribuída chamado ``BurstAttention'' para otimizar o acesso à memória e as operações de comunicação tanto no nível global do cluster quanto no nível local do dispositivo. Em nossos experimentos, comparamos o BurstAttention com outras soluções competitivas de atenção distribuída para o processamento de sequências longas. Os resultados experimentais sob diferentes configurações de comprimento demonstram que o BurstAttention oferece vantagens significativas para o processamento de sequências longas em comparação com essas linhas de base competitivas, reduzindo 40% da sobrecarga de comunicação e alcançando um aumento de velocidade de 2 X durante o treinamento de sequências de 32K em 8 X A100.

English

Effective attention modules have played a crucial role in the success of Transformer-based large language models (LLMs), but the quadratic time and memory complexities of these attention modules also pose a challenge when processing long sequences. One potential solution for the long sequence problem is to utilize distributed clusters to parallelize the computation of attention modules across multiple devices (e.g., GPUs). However, adopting a distributed approach inevitably introduces extra memory overheads to store local attention results and incurs additional communication costs to aggregate local results into global ones. In this paper, we propose a distributed attention framework named ``BurstAttention'' to optimize memory access and communication operations at both the global cluster and local device levels. In our experiments, we compare BurstAttention with other competitive distributed attention solutions for long sequence processing. The experimental results under different length settings demonstrate that BurstAttention offers significant advantages for processing long sequences compared with these competitive baselines, reducing 40% communication overheads and achieving 2 X speedup during training 32K sequence length on 8 X A100.