LASP-2: Riconsiderazione della Parallelismo Sequenziale per l'Attenzione Lineare e il suo Ibrido

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

February 11, 2025

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

cs.AI

Abstract

Gli approcci di modellazione lineare delle sequenze, come l'attenzione lineare, forniscono vantaggi come addestramento in tempo lineare e inferenza in memoria costante su lunghezze di sequenza. Tuttavia, i metodi esistenti di parallelismo delle sequenze (SP) non sono ottimizzati per la caratteristica di prodotto-prima giusta dell'attenzione lineare o utilizzano una strategia di comunicazione a anello, che comporta un minor parallelismo di calcolo, limitando la scalabilità per sequenze più lunghe nei sistemi distribuiti. In questo articolo, presentiamo LASP-2, un nuovo metodo SP per migliorare sia il parallelismo di comunicazione che di calcolo durante l'addestramento dei modelli transformer con attenzione lineare su sequenze di input molto lunghe. Rispetto al lavoro precedente LASP, LASP-2 ripensa il requisito minimo di comunicazione per SP sui livelli di attenzione lineare, riorganizza l'intero flusso di lavoro di comunicazione-calcolo di LASP. In questo modo, è necessaria solo una singola comunicazione collettiva AllGather su stati di memoria intermedi, le cui dimensioni sono indipendenti dalla lunghezza della sequenza, portando a miglioramenti significativi sia del parallelismo di comunicazione che di calcolo, nonché della loro sovrapposizione. Inoltre, estendiamo LASP-2 a LASP-2H applicando una ridisegnazione della comunicazione simile ai moduli di attenzione standard, offrendo una soluzione SP efficiente per modelli ibridi che combinano livelli di attenzione lineare e standard. La nostra valutazione su un modello Linear-Llama3, una variante di Llama3 con attenzione lineare al posto dell'attenzione standard, dimostra l'efficacia di LASP-2 e LASP-2H. In particolare, LASP-2 ottiene miglioramenti della velocità di addestramento del 15,2% rispetto a LASP e del 36,6% rispetto a Ring Attention, con una lunghezza di sequenza di 2048K su 64 GPU. Il codice è rilasciato come parte di: 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.