LASP-2 : Repenser le parallélisme de séquence pour l'attention linéaire et son 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

Résumé

Les approches de modélisation de séquences linéaires, telles que l'attention linéaire, offrent des avantages tels qu'un entraînement en temps linéaire et une inférence en mémoire constante sur des longueurs de séquence. Cependant, les méthodes existantes de parallélisme de séquence (PS) ne sont pas optimisées pour la caractéristique de produit-d'abord correct de l'attention linéaire ou utilisent une stratégie de communication de type anneau, ce qui entraîne un parallélisme de calcul plus faible, limitant leur extensibilité pour des séquences plus longues dans des systèmes distribués. Dans cet article, nous présentons LASP-2, une nouvelle méthode de PS pour améliorer à la fois la communication et le parallélisme de calcul lors de l'entraînement de modèles de transformateur d'attention linéaire avec des séquences d'entrée très longues. Comparé au travail précédent LASP, LASP-2 repense l'exigence minimale de communication pour le PS sur les couches d'attention linéaire, réorganise l'ensemble du flux de travail communication-calcul de LASP. De cette manière, un seul AllGather collectif de communication est nécessaire sur les états de mémoire intermédiaires, dont les tailles sont indépendantes de la longueur de la séquence, entraînant des améliorations significatives à la fois du parallélisme de communication et de calcul, ainsi que de leur chevauchement. De plus, nous étendons LASP-2 à LASP-2H en appliquant une refonte de la communication similaire aux modules d'attention standard, offrant une solution de PS efficace pour les modèles hybrides qui combinent des couches d'attention linéaire et standard. Notre évaluation sur un modèle Linear-Llama3, une variante de Llama3 avec une attention linéaire remplaçant l'attention standard, démontre l'efficacité de LASP-2 et LASP-2H. Plus précisément, LASP-2 obtient des améliorations de vitesse d'entraînement de 15,2 % par rapport à LASP et de 36,6 % par rapport à Ring Attention, avec une longueur de séquence de 2048K sur 64 GPU. Le code est publié dans le cadre 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.