BurstAttention : Un cadre d'attention distribué efficace pour des séquences extrêmement longues

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

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

cs.AI

Résumé

Les modules d'attention efficaces ont joué un rôle crucial dans le succès des grands modèles de langage (LLMs) basés sur les Transformers, mais les complexités quadratiques en temps et en mémoire de ces modules d'attention posent également un défi lors du traitement de longues séquences. Une solution potentielle au problème des longues séquences consiste à utiliser des clusters distribués pour paralléliser le calcul des modules d'attention sur plusieurs appareils (par exemple, des GPU). Cependant, l'adoption d'une approche distribuée introduit inévitablement des surcharges de mémoire supplémentaires pour stocker les résultats d'attention locaux et entraîne des coûts de communication supplémentaires pour agréger les résultats locaux en résultats globaux. Dans cet article, nous proposons un cadre d'attention distribué nommé « BurstAttention » pour optimiser l'accès à la mémoire et les opérations de communication à la fois au niveau du cluster global et des appareils locaux. Dans nos expériences, nous comparons BurstAttention à d'autres solutions d'attention distribuées compétitives pour le traitement de longues séquences. Les résultats expérimentaux sous différents paramètres de longueur démontrent que BurstAttention offre des avantages significatifs pour le traitement de longues séquences par rapport à ces bases de référence compétitives, réduisant de 40 % les surcharges de communication et obtenant un accélération de 2 X lors de l'entraînement de séquences de 32K sur 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.