LASP-2: Repensando el Paralelismo de Secuencia para Atención Lineal y su Híbrido

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

Resumen

Los enfoques de modelado de secuencias lineales, como la atención lineal, ofrecen ventajas como el entrenamiento en tiempo lineal y la inferencia en memoria constante sobre longitudes de secuencia. Sin embargo, los métodos existentes de paralelismo de secuencias (SP) no están optimizados para la característica de producto primero correcto de la atención lineal o utilizan una estrategia de comunicación de estilo anular, lo que resulta en un menor paralelismo de cálculo, limitando su escalabilidad para secuencias más largas en sistemas distribuidos. En este documento, presentamos LASP-2, un nuevo método de SP para mejorar tanto el paralelismo de comunicación como de cálculo al entrenar modelos transformadores de atención lineal con secuencias de entrada muy largas. En comparación con el trabajo anterior LASP, LASP-2 reconsidera el requisito mínimo de comunicación para SP en capas de atención lineal, reorganiza todo el flujo de trabajo de comunicación-cálculo de LASP. De esta manera, solo se necesita una única comunicación colectiva AllGather en estados de memoria intermedios, cuyos tamaños son independientes de la longitud de la secuencia, lo que conduce a mejoras significativas tanto en la comunicación como en el paralelismo de cálculo, así como en su superposición. Además, extendemos LASP-2 a LASP-2H aplicando un rediseño de comunicación similar a los módulos de atención estándar, ofreciendo una solución de SP eficiente para modelos híbridos que combinan capas de atención lineal y estándar. Nuestra evaluación en un modelo Linear-Llama3, una variante de Llama3 con atención lineal reemplazando la atención estándar, demuestra la efectividad de LASP-2 y LASP-2H. Específicamente, LASP-2 logra mejoras en la velocidad de entrenamiento del 15.2% sobre LASP y del 36.6% sobre Ring Attention, con una longitud de secuencia de 2048K en 64 GPUs. El código se encuentra disponible en: 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.