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

要約

線形の注意などの線形シーケンスモデリングアプローチは、線形時間トレーニングやシーケンス長にわたる定数メモリ推論などの利点を提供します。
ただし、既存のシーケンス並列性（SP）メソッドは、線形注意の適切な製品ファースト機能に最適化されていないか、リングスタイルの通信戦略を使用するため、計算並列が低くなり、分散システムのより長いシーケンスのスケーラビリティが制限されます。
このホワイトペーパーでは、非常に長い入力シーケンスで線形注意トランスモデルをトレーニングするときに、通信と計算の両方の並列性を強化する新しいSPメソッドであるLASP-2を紹介します。
以前の作業LASPと比較して、LASP-2は、線形注意層上のSPの最小限の通信要件を再考し、LASPのコミュニケーションコンピューションワークフロー全体を再編成します。
このように、サイズはシーケンスの長さとは独立している中間メモリ状態では、1つのAllgather Collective Communicationが必要であり、通信と計算の並列性の両方の大幅な改善、およびそれらの重複をもたらします。
さらに、LASP-2をLASP-2Hに拡張し、同様の通信再設計を標準の注意モジュールに適用し、線形および標準の注意層をブレンドするハイブリッドモデルに効率的なSPソリューションを提供します。
標準的な注意を置き換える線形注意を払ったLLAMA3のバリアントである線形-LLAMA3モデルに関する私たちの評価は、LASP-2とLASP-2Hの有効性を示しています。
具体的には、LASP-2は、LASPよりも15.2％のトレーニング速度の改善、リングの注意よりも36.6％のトレーニング速度の改善を達成し、64 GPUにわたって2048Kのシーケンス長があります。
このコードは、https：//github.com/opensparsellms/linear-moeの一部としてリリースされます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google