LASP-2：線形注意とそのハイブリッドのためのシーケンス並列性の再考

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

February 11, 2025

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

cs.AI

要旨

線形シーケンスモデリング手法、例えば線形アテンションは、シーケンスの長さに関わらず線形時間のトレーニングと一定のメモリ推論の利点を提供します。しかし、既存のシーケンス並列処理（SP）手法は、線形アテンションの右積優先機能に最適化されていないか、リング形式の通信戦略を使用しており、計算並列処理が低下し、分散システムにおける長いシーケンスのスケーラビリティが制限されています。本論文では、非常に長い入力シーケンスで線形アテンショントランスフォーマーモデルをトレーニングする際に、通信と計算の並列処理の両方を向上させる新しいSP手法であるLASP-2を紹介します。以前の研究LASPと比較して、LASP-2は線形アテンションレイヤーにおけるSPの最小通信要件を再考し、LASPの通信-計算ワークフロー全体を再構築します。この方法により、シーケンス長に依存しない中間メモリ状態での1つのAllGather集団通信のみが必要となり、通信と計算の並列処理、およびそれらのオーバーラップが大幅に改善されます。さらに、LASP-2をLASP-2Hに拡張し、標準アテンションモジュールに同様の通信再設計を適用することで、線形と標準アテンションレイヤーを組み合わせたハイブリッドモデルに効率的なSPソリューションを提供します。Llama3の変種であるLinear-Llama3モデルでの評価は、LASP-2およびLASP-2Hの効果を示しています。具体的には、64台のGPUで2048Kのシーケンス長で、LASP-2はLASPに比べてトレーニング速度が15.2％向上し、Ring Attentionに比べて36.6％向上します。コードは以下の一部として公開されています: 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.