LASP-2: Heroverwegen van Sequentiële Parallelisme voor Lineaire Aandacht en de Hybride

LASP-2: Rethinking Sequence Parallelism for Linear Attention and Its Hybrid

February 11, 2025

Auteurs: Weigao Sun, Disen Lan, Yiran Zhong, Xiaoye Qu, Yu Cheng

cs.AI

Samenvatting

Lineaire sequentiële modelleringsbenaderingen, zoals lineaire aandacht, bieden voordelen zoals lineaire trainingstijd en constante geheugeninferentie over sequentielengtes. Bestaande sequentie parallelisme (SP) methoden zijn echter ofwel niet geoptimaliseerd voor de juiste-product-eerst eigenschap van lineaire aandacht of maken gebruik van een ring-stijl communicatiestrategie, wat resulteert in lagere berekeningsparallelisme, waardoor hun schaalbaarheid voor langere sequenties in gedistribueerde systemen beperkt wordt. In dit artikel introduceren we LASP-2, een nieuwe SP methode om zowel communicatie als berekeningsparallelisme te verbeteren bij het trainen van lineaire aandacht transformer modellen met zeer lange invoersequenties. In vergelijking met eerdere werk LASP, heroverweegt LASP-2 de minimale communicatievereiste voor SP op lineaire aandachtlagen, herorganiseert het volledige communicatie-berekeningsworkflow van LASP. Op deze manier is slechts één enkele AllGather collectieve communicatie nodig over tussenliggende geheugenstaten, waarvan de groottes onafhankelijk zijn van de sequentielengte, wat leidt tot aanzienlijke verbeteringen van zowel communicatie als berekeningsparallelisme, evenals hun overlap. Daarnaast breiden we LASP-2 uit naar LASP-2H door een vergelijkbare communicatieherontwerp toe te passen op standaard aandachtsmodules, wat een efficiënte SP-oplossing biedt voor hybride modellen die lineaire en standaard aandachtlagen combineren. Onze evaluatie van een Lineaire-Llama3 model, een variant van Llama3 met lineaire aandacht ter vervanging van standaard aandacht, toont de effectiviteit van LASP-2 en LASP-2H aan. Specifiek behaalt LASP-2 trainingsnelheidsverbeteringen van 15,2% ten opzichte van LASP en 36,6% ten opzichte van Ring Aandacht, met een sequentielengte van 2048K over 64 GPU's. De code is vrijgegeven als onderdeel van: https://github.com/OpenSparseLLMs/Linear-MoE.

English

Linear sequence modeling approaches, such as linear attention, provide advantages like linear-time training and constant-memory inference over sequence lengths. However, existing sequence parallelism (SP) methods are either not optimized for the right-product-first feature of linear attention or use a ring-style communication strategy, which results in lower computation parallelism, limits their scalability for longer sequences in distributed systems. In this paper, we introduce LASP-2, a new SP method to enhance both communication and computation parallelism when training linear attention transformer models with very-long input sequences. Compared to previous work LASP, LASP-2 rethinks the minimal communication requirement for SP on linear attention layers, reorganizes the whole communication-computation workflow of LASP. In this way, only one single AllGather collective communication is needed on intermediate memory states, whose sizes are independent of the sequence length, leading to significant improvements of both communication and computation parallelism, as well as their overlap. Additionally, we extend LASP-2 to LASP-2H by applying similar communication redesign to standard attention modules, offering an efficient SP solution for hybrid models that blend linear and standard attention layers. Our evaluation on a Linear-Llama3 model, a variant of Llama3 with linear attention replacing standard attention, demonstrates the effectiveness of LASP-2 and LASP-2H. Specifically, LASP-2 achieves training speed improvements of 15.2% over LASP and 36.6% over Ring Attention, with a sequence length of 2048K across 64 GPUs. The Code is released as a part of: https://github.com/OpenSparseLLMs/Linear-MoE.