BurstAttention: Een efficiënt gedistribueerd aandachtkader voor extreem lange sequenties

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

Samenvatting

Effectieve aandachtmodules hebben een cruciale rol gespeeld in het succes van Transformer-gebaseerde grote taalmodellen (LLM's), maar de kwadratische tijd- en geheugencomplexiteit van deze aandachtmodules vormen ook een uitdaging bij het verwerken van lange sequenties. Een mogelijke oplossing voor het probleem van lange sequenties is het gebruik van gedistribueerde clusters om de berekening van aandachtmodules te paralleliseren over meerdere apparaten (bijvoorbeeld GPU's). Het adopteren van een gedistribueerde aanpak introduceert echter onvermijdelijk extra geheugenoverhead om lokale aandachtresultaten op te slaan en brengt extra communicatiekosten met zich mee om lokale resultaten samen te voegen tot globale resultaten. In dit artikel stellen we een gedistribueerd aandachtframework genaamd ``BurstAttention'' voor om geheugentoegang en communicatieoperaties te optimaliseren op zowel het niveau van het globale cluster als het lokale apparaat. In onze experimenten vergelijken we BurstAttention met andere competitieve gedistribueerde aandachtoplossingen voor het verwerken van lange sequenties. De experimentele resultaten onder verschillende lengte-instellingen tonen aan dat BurstAttention aanzienlijke voordelen biedt voor het verwerken van lange sequenties in vergelijking met deze competitieve baselines, waarbij 40% communicatieoverhead wordt gereduceerd en een 2x versnelling wordt bereikt tijdens het trainen van sequenties met een lengte van 32K op 8x 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.