BurstAttention: 極めて長いシーケンスのための効率的な分散型アテンションフレームワーク

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

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

cs.AI

要旨

効果的なアテンションモジュールは、Transformerベースの大規模言語モデル（LLMs）の成功において重要な役割を果たしてきたが、これらのアテンションモジュールの二次的な時間およびメモリ複雑性は、長いシーケンスを処理する際に課題を引き起こす。長いシーケンス問題に対する一つの潜在的な解決策は、分散クラスタを利用して、複数のデバイス（例えば、GPU）間でアテンションモジュールの計算を並列化することである。しかし、分散アプローチを採用すると、ローカルなアテンション結果を保存するための追加のメモリオーバーヘッドが発生し、ローカルな結果をグローバルな結果に集約するための追加の通信コストがかかる。本論文では、グローバルクラスタレベルとローカルデバイスレベルの両方でメモリアクセスと通信操作を最適化する「BurstAttention」という分散アテンションフレームワークを提案する。実験では、BurstAttentionを他の競合する分散アテンションソリューションと比較し、長いシーケンス処理における性能を評価する。異なる長さ設定での実験結果は、BurstAttentionがこれらの競合ベースラインと比較して長いシーケンス処理において顕著な利点を提供し、通信オーバーヘッドを40％削減し、8台のA100で32Kシーケンス長のトレーニング中に2倍の高速化を達成することを示している。

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.