LASP-2: Neubewertung der Sequenzparallelität für lineare Aufmerksamkeit und deren Hybrid

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

February 11, 2025

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

cs.AI

Zusammenfassung

Ansätze zur linearen Sequenzmodellierung, wie lineare Aufmerksamkeit, bieten Vorteile wie ein Training in linearer Zeit und eine konstante Speichernutzung bei Sequenzlängen. Allerdings sind bestehende Methoden zur Sequenzparallelität (SP) entweder nicht für das Merkmal des richtigen Produkt-zuerst bei linearer Aufmerksamkeit optimiert oder verwenden eine Ring-Kommunikationsstrategie, die zu einer geringeren Rechenparallelität führt und ihre Skalierbarkeit für längere Sequenzen in verteilten Systemen einschränkt. In diesem Paper stellen wir LASP-2 vor, eine neue SP-Methode zur Verbesserung sowohl der Kommunikations- als auch der Rechenparallelität beim Training von Transformer-Modellen mit linearer Aufmerksamkeit und sehr langen Eingabesequenzen. Im Vergleich zur vorherigen Arbeit LASP überdenkt LASP-2 die minimale Kommunikationsanforderung für SP bei linearen Aufmerksamkeitsschichten, reorganisiert den gesamten Kommunikations-Rechen-Workflow von LASP. Auf diese Weise ist nur eine einzige AllGather-Kollektivkommunikation auf Zwischenspeicherzustände erforderlich, deren Größen unabhängig von der Sequenzlänge sind, was zu signifikanten Verbesserungen sowohl bei der Kommunikations- als auch der Rechenparallelität sowie deren Überlappung führt. Darüber hinaus erweitern wir LASP-2 zu LASP-2H, indem wir eine ähnliche Kommunikationsneugestaltung auf Standard-Aufmerksamkeitsmodule anwenden, um eine effiziente SP-Lösung für Hybridmodelle anzubieten, die lineare und Standard-Aufmerksamkeitsschichten kombinieren. Unsere Evaluation an einem Linear-Llama3-Modell, einer Variante von Llama3 mit linearer Aufmerksamkeit anstelle von Standard-Aufmerksamkeit, zeigt die Wirksamkeit von LASP-2 und LASP-2H. Insbesondere erzielt LASP-2 Schulungsgeschwindigkeitsverbesserungen von 15,2% gegenüber LASP und 36,6% gegenüber Ring-Aufmerksamkeit bei einer Sequenzlänge von 2048K auf 64 GPUs. Der Code ist als Teil von: https://github.com/OpenSparseLLMs/Linear-MoE veröffentlicht.

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.