BurstAttention: Ein effizientes verteiltes Aufmerksamkeits-Framework für extrem lange Sequenzen

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

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

cs.AI

Zusammenfassung

Effektive Aufmerksamkeitsmodule haben eine entscheidende Rolle im Erfolg von auf Transformatoren basierenden großen Sprachmodellen (LLMs) gespielt, aber die quadratischen Zeit- und Speicherkomplexitäten dieser Aufmerksamkeitsmodule stellen auch eine Herausforderung beim Verarbeiten langer Sequenzen dar. Eine mögliche Lösung für das Problem langer Sequenzen besteht darin, verteilte Cluster zu nutzen, um die Berechnung von Aufmerksamkeitsmodulen über mehrere Geräte (z. B. GPUs) parallelisiert durchzuführen. Die Annahme eines verteilten Ansatzes führt jedoch zwangsläufig zu zusätzlichen Speicherüberlastungen, um lokale Aufmerksamkeitsergebnisse zu speichern, und verursacht zusätzliche Kommunikationskosten, um lokale Ergebnisse zu globalen zusammenzuführen. In diesem Papier schlagen wir ein verteiltes Aufmerksamkeitsframework namens "BurstAttention" vor, um Speicherzugriff und Kommunikationsoperationen sowohl auf globaler Cluster- als auch auf lokaler Geräteebene zu optimieren. In unseren Experimenten vergleichen wir BurstAttention mit anderen wettbewerbsfähigen verteilten Aufmerksamkeitslösungen für die Verarbeitung langer Sequenzen. Die experimentellen Ergebnisse unter verschiedenen Längeneinstellungen zeigen, dass BurstAttention im Vergleich zu diesen wettbewerbsfähigen Baselines signifikante Vorteile bei der Verarbeitung langer Sequenzen bietet, wobei die Kommunikationsüberlastungen um 40 % reduziert werden und bei einer Trainingssequenzlänge von 32K auf 8 X A100 eine 2-fache Beschleunigung erreicht wird.

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.