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 by partitioning attention along the sequence dimension across device.
Through our experiments, we compare BurstAttention with other competitive long-sequence distributed attention solutions.
The experimental results under different lengths demonstrate that BurstAttention offers significant advantages for processing long sequences compared with these competitive baselines, especially tensor parallelism (Megatron-V3)  with FlashAttention, reducing 40% communication overheads and achieving 2speedup during training 128K sequence length on 8A100.