BurstAttention: Un Framework Distribuito Efficace per l'Attenzione su Sequenze Estremamente Lunghe

BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences

March 14, 2024

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

cs.AI

Abstract

I moduli di attenzione efficaci hanno svolto un ruolo cruciale nel successo dei modelli linguistici di grandi dimensioni (LLM) basati su Transformer, ma le complessità quadratiche in termini di tempo e memoria di questi moduli di attenzione rappresentano anche una sfida durante l'elaborazione di sequenze lunghe. Una potenziale soluzione al problema delle sequenze lunghe è quella di utilizzare cluster distribuiti per parallelizzare il calcolo dei moduli di attenzione su più dispositivi (ad esempio, GPU). Tuttavia, l'adozione di un approccio distribuito introduce inevitabilmente sovraccarichi di memoria aggiuntivi per memorizzare i risultati locali dell'attenzione e comporta costi di comunicazione aggiuntivi per aggregare i risultati locali in risultati globali. In questo articolo, proponiamo un framework di attenzione distribuito denominato "BurstAttention" per ottimizzare l'accesso alla memoria e le operazioni di comunicazione sia a livello di cluster globale che di dispositivo locale. Nei nostri esperimenti, confrontiamo BurstAttention con altre soluzioni competitive di attenzione distribuita per l'elaborazione di sequenze lunghe. I risultati sperimentali in diverse configurazioni di lunghezza dimostrano che BurstAttention offre vantaggi significativi per l'elaborazione di sequenze lunghe rispetto a queste baseline competitive, riducendo del 40% i sovraccarichi di comunicazione e ottenendo un aumento di velocità di 2 X durante l'addestramento di sequenze di lunghezza 32K su 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.