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

Аннотация

Линейные механизмы внимания предоставляют значительные преимущества для крупных языковых моделей (LLMs), обеспечивая линейную вычислительную сложность, что позволяет эффективно обрабатывать сверхдлинные последовательности (например, контекст длиной 1 млн). Однако существующие методы параллелизма последовательностей (Sequence Parallelism, SP), необходимые для распределения этих задач между устройствами, становятся основным узким местом из-за значительных накладных расходов на коммуникацию. В данной работе мы представляем ZeCO (Zero Communication Overhead) — новый метод SP для моделей с линейным вниманием, разработанный для преодоления этих ограничений и достижения сквозной почти линейной масштабируемости при обучении на длинных последовательностях. Например, обучение модели с длиной последовательности 1 млн на 64 устройствах с использованием ZeCO занимает примерно столько же времени, сколько обучение с последовательностью длиной 16 тыс. на одном устройстве. В основе ZeCO лежит All-Scan — новый примитив коллективной коммуникации. All-Scan предоставляет каждому рангу SP именно то начальное состояние оператора, которое ему требуется, при этом минимизируя объем коммуникации, что эффективно устраняет накладные расходы на передачу данных. Теоретически мы доказываем оптимальность ZeCO, показывая, что он вносит лишь незначительные накладные расходы по времени и памяти. Эмпирически мы сравниваем затраты на коммуникацию различных стратегий SP и демонстрируем, что All-Scan обеспечивает самую быструю коммуникацию в сценариях SP. В частности, на 256 GPU с длиной последовательности 8 млн ZeCO достигает ускорения на 60% по сравнению с текущим передовым методом SP. Мы считаем, что ZeCO открывает четкий путь к эффективному обучению LLMs следующего поколения на ранее недоступных длинах последовательностей.

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.