BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

• URL: http://arxiv.org/abs/2403.09347v4
• Date: Thu, 6 Jun 2024 05:43:52 GMT
• Title: BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences
• Authors: Ao Sun, Weilin Zhao, Xu Han, Cheng Yang, Zhiyuan Liu, Chuan Shi, Maosong Sun,
• Abstract summary: We propose a distributed attention framework named BurstAttention'' to optimize memory access and communication operations. The experimental results under different length settings demonstrate that BurstAttention offers significant advantages for processing long sequences.
• Score: 96.74779792715819
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 1.37 X speedup during training 128K sequence length on 32 X A100.

Related papers

• Star Attention: Efficient LLM Inference over Long Sequences [17.401430615714]
  We introduce Star Attention, a two-phase block-sparse approximation that improves computational efficiency by sharding attention across multiple hosts. Star Attention integrates seamlessly with most Transformer-based Large Language Models trained with global attention, reducing memory requirements and inference time by up to 11x while preserving 95-100% of accuracy.
  arXiv  Detail & Related papers  (2024-11-26T05:10:04Z)
• Squeezed Attention: Accelerating Long Context Length LLM Inference [64.11145320159126]
  We propose Squeezed Attention as a mechanism to accelerate LLM applications where a large portion of the input prompt is fixed. We use K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. We then compute exact attention using only these important keys from the fixed context, thereby reducing bandwidth and computational costs.
  arXiv  Detail & Related papers  (2024-11-14T18:54:19Z)
• Correlation-Aware Select and Merge Attention for Efficient Fine-Tuning and Context Length Extension [21.729875191721984]
  We introduce correlation-aware selection and merging mechanisms to facilitate efficient sparse attention. We also propose a novel data augmentation technique involving positional encodings to enhance generalization to unseen positions. Our method achieves 100% accuracy on the passkey task with a context length of 4M and maintains stable perplexity at a 1M context length.
  arXiv  Detail & Related papers  (2024-10-05T15:59:32Z)
• Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers [58.5711048151424]
  We introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome computational and memory obstacles. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods.
  arXiv  Detail & Related papers  (2024-06-24T15:55:59Z)
• Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences [60.489682735061415]
  We propose CHELA, which replaces state space models with short-long convolutions and implements linear attention in a divide-and-conquer manner. Our experiments on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.
  arXiv  Detail & Related papers  (2024-06-12T12:12:38Z)
• Bidirectional Long-Range Parser for Sequential Data Understanding [3.76054468268713]
  We introduce BLRP (Bidirectional Long-Range), a novel and versatile attention mechanism designed to increase performance and efficiency on long-sequence tasks. We show the benefits and versatility of our approach on vision and language domains by demonstrating competitive results against state-of-the-art methods.
  arXiv  Detail & Related papers  (2024-04-08T05:45:03Z)
• Bifurcated Attention: Accelerating Massively Parallel Decoding with Shared Prefixes in LLMs [39.16152482491236]
  Bifurcated attention is a method designed to enhance language model inference in shared-context batch decoding scenarios. Our approach addresses the challenge of redundant memory IO costs, a critical factor contributing to latency in high batch sizes and extended context lengths.
  arXiv  Detail & Related papers  (2024-03-13T16:30:57Z)
• UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation [93.88170217725805]
  We propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features. Our evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy.
  arXiv  Detail & Related papers  (2022-12-08T18:59:57Z)
• Cluster-Former: Clustering-based Sparse Transformer for Long-Range Dependency Encoding [90.77031668988661]
  Cluster-Former is a novel clustering-based sparse Transformer to perform attention across chunked sequences. The proposed framework is pivoted on two unique types of Transformer layer: Sliding-Window Layer and Cluster-Former Layer. Experiments show that Cluster-Former achieves state-of-the-art performance on several major QA benchmarks.
  arXiv  Detail & Related papers  (2020-09-13T22:09:30Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.