ZeCO: 線形アテンションのためのゼロ通信オーバーヘッド並列処理

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

July 1, 2025

著者: Yuhong Chou, Zehao Liu, Ruijie Zhu, Xinyi Wan, Tianjian Li, Congying Chu, Qian Liu, Jibin Wu, Zejun Ma

cs.AI

要旨

線形アテンションメカニズムは、線形計算複雑性を提供し、超長系列（例：100万コンテキスト）の効率的な処理を可能にすることで、大規模言語モデル（LLM）に大きな利点をもたらします。しかし、これらのワークロードをデバイス間で分散するために不可欠な既存のシーケンス並列化（SP）手法は、多大な通信オーバーヘッドにより主要なボトルネックとなっています。本論文では、線形アテンションモデル向けのゼロ通信オーバーヘッド（ZeCO）シーケンス並列化を提案します。これは、これらの制限を克服し、長系列学習におけるエンドツーエンドのほぼ線形スケーラビリティを実現する新しいSP手法です。例えば、64デバイスで100万系列長のモデルをZeCOを用いて学習する場合、単一デバイスで16k系列を学習するのとほぼ同じ時間を要します。ZeCOの中核には、新しい集団通信プリミティブであるAll-Scanがあります。All-Scanは、各SPランクに必要な初期演算子状態を正確に提供しつつ、最小限の通信フットプリントを維持することで、通信オーバーヘッドを効果的に排除します。理論的には、ZeCOの最適性を証明し、わずかな時間と空間のオーバーヘッドしか導入しないことを示します。実証的には、異なるシーケンス並列化戦略の通信コストを比較し、All-ScanがSPシナリオで最速の通信を達成することを示します。具体的には、256GPUで800万系列長の場合、ZeCOは現状の最先端（SOTA）SP手法と比較して60％の高速化を実現します。我々は、ZeCOが以前は扱い難かった系列長での次世代LLMの効率的な学習に向けた明確な道筋を確立すると信じています。

English

Linear attention mechanisms deliver significant advantages for Large Language Models (LLMs) by providing linear computational complexity, enabling efficient processing of ultra-long sequences (e.g., 1M context). However, existing Sequence Parallelism (SP) methods, essential for distributing these workloads across devices, become the primary bottleneck due to substantial communication overhead. In this paper, we introduce ZeCO (Zero Communication Overhead) sequence parallelism for linear attention models, a new SP method designed to overcome these limitations and achieve end-to-end near-linear scalability for long sequence training. For example, training a model with a 1M sequence length across 64 devices using ZeCO takes roughly the same time as training with an 16k sequence on a single device. At the heart of ZeCO lies All-Scan, a new collective communication primitive. All-Scan provides each SP rank with precisely the initial operator state it requires while maintaining a minimal communication footprint, effectively eliminating communication overhead. Theoretically, we prove the optimaity of ZeCO, showing that it introduces only negligible time and space overhead. Empirically, we compare the communication costs of different sequence parallelism strategies and demonstrate that All-Scan achieves the fastest communication in SP scenarios. Specifically, on 256 GPUs with an 8M sequence length, ZeCO achieves a 60\% speedup compared to the current state-of-the-art (SOTA) SP method. We believe ZeCO establishes a clear path toward efficiently training next-generation LLMs on previously intractable sequence lengths.