LASP-2: Repensando o Paralelismo de Sequência para Atenção Linear e Sua Híbrida

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

February 11, 2025

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

cs.AI

Resumo

Abordagens de modelagem de sequência linear, como a atenção linear, oferecem vantagens como treinamento em tempo linear e inferência com memória constante em relação ao comprimento das sequências. No entanto, os métodos existentes de paralelismo de sequência (SP) não são otimizados para a característica de produto-direito-primeiro da atenção linear ou utilizam uma estratégia de comunicação em anel, o que resulta em menor paralelismo de computação, limitando sua escalabilidade para sequências mais longas em sistemas distribuídos. Neste artigo, apresentamos o LASP-2, um novo método de SP para aprimorar tanto o paralelismo de comunicação quanto o de computação ao treinar modelos de transformadores com atenção linear em sequências de entrada muito longas. Em comparação com o trabalho anterior LASP, o LASP-2 repensa o requisito mínimo de comunicação para SP em camadas de atenção linear, reorganizando todo o fluxo de trabalho de comunicação-computação do LASP. Dessa forma, apenas uma única comunicação coletiva AllGather é necessária para estados de memória intermediários, cujos tamanhos são independentes do comprimento da sequência, resultando em melhorias significativas tanto no paralelismo de comunicação quanto no de computação, bem como em sua sobreposição. Além disso, estendemos o LASP-2 para o LASP-2H ao aplicar uma reformulação similar de comunicação a módulos de atenção padrão, oferecendo uma solução eficiente de SP para modelos híbridos que combinam camadas de atenção linear e padrão. Nossa avaliação em um modelo Linear-Llama3, uma variante do Llama3 com atenção linear substituindo a atenção padrão, demonstra a eficácia do LASP-2 e do LASP-2H. Especificamente, o LASP-2 alcança melhorias de velocidade de treinamento de 15,2% em relação ao LASP e 36,6% em relação ao Ring Attention, com um comprimento de sequência de 2048K em 64 GPUs. O código foi liberado como parte de: 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.