ZeCO: Null-Kommunikations-Overhead-Sequenzparallelismus für lineare Aufmerksamkeit

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

July 1, 2025

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

cs.AI

Zusammenfassung

Lineare Aufmerksamkeitsmechanismen bieten erhebliche Vorteile für Large Language Models (LLMs), indem sie eine lineare Rechenkomplexität bereitstellen und die effiziente Verarbeitung ultra-langer Sequenzen (z.B. 1M Kontext) ermöglichen. Allerdings werden bestehende Sequence Parallelism (SP)-Methoden, die für die Verteilung dieser Arbeitslasten über Geräte unerlässlich sind, aufgrund des erheblichen Kommunikationsaufwands zum primären Engpass. In diesem Artikel stellen wir ZeCO (Zero Communication Overhead) Sequence Parallelism für lineare Aufmerksamkeitsmodelle vor, eine neue SP-Methode, die entwickelt wurde, um diese Einschränkungen zu überwinden und eine end-to-end nahezu lineare Skalierbarkeit für das Training langer Sequenzen zu erreichen. Zum Beispiel dauert das Training eines Modells mit einer Sequenzlänge von 1M über 64 Geräte mit ZeCO in etwa genauso lange wie das Training mit einer 16k-Sequenz auf einem einzelnen Gerät. Im Kern von ZeCO liegt All-Scan, ein neues kollektives Kommunikationsprimitiv. All-Scan stellt jedem SP-Rank genau den initialen Operatorzustand bereit, den er benötigt, während ein minimaler Kommunikationsaufwand aufrechterhalten wird, wodurch der Kommunikationsoverhead effektiv eliminiert wird. Theoretisch beweisen wir die Optimalität von ZeCO, indem wir zeigen, dass es nur vernachlässigbare Zeit- und Raumoverheads einführt. Empirisch vergleichen wir die Kommunikationskosten verschiedener Sequence Parallelism-Strategien und demonstrieren, dass All-Scan die schnellste Kommunikation in SP-Szenarien erreicht. Insbesondere erreicht ZeCO auf 256 GPUs mit einer 8M-Sequenzlänge eine 60\%ige Beschleunigung im Vergleich zur derzeit besten SP-Methode (SOTA). Wir glauben, dass ZeCO einen klaren Weg zur effizienten Ausbildung der nächsten Generation von LLMs auf bisher nicht handhabbaren Sequenzlängen ebnet.

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.