BurstAttention: An Efficient Distributed Attention Framework for
  Extremely Long Sequences
Authors: name: Sun Ao; name: Weilin Zhao; name: Xu Han; name: Cheng Yang; name: Zhiyuan Liu; name: Chuan Shi; name: Maosong Sun; name: Shengnan Wang; name: Teng Su
Published: 2024-03-14 12:51:58

Abstract: 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.