ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

TL;DR

ZeCO introduces a novel sequence parallelism method with zero communication overhead for linear attention models, enabling efficient training of ultra-long sequences on large-scale hardware with near-linear scalability.

Contribution

ZeCO presents a new SP method utilizing All-Scan to eliminate communication overhead, achieving near-linear scalability for long sequence training in linear attention models.

Findings

ZeCO achieves 60% speedup over SOTA on 256 GPUs with 8M sequences.

All-Scan provides minimal communication, enabling efficient parallelism.

ZeCO allows training of 1M sequence length models across 64 devices with comparable time to 16k sequences on a single device.

Abstract

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.