ZeCO : Parallélisme de séquence sans surcharge de communication pour l'attention linéaire

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

July 1, 2025

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

cs.AI

Résumé

Les mécanismes d'attention linéaire offrent des avantages significatifs pour les grands modèles de langage (LLMs) en fournissant une complexité computationnelle linéaire, permettant un traitement efficace de séquences ultra-longues (par exemple, un contexte de 1 million). Cependant, les méthodes existantes de parallélisme de séquence (SP), essentielles pour répartir ces charges de travail sur plusieurs dispositifs, deviennent le principal goulot d'étranglement en raison d'un surcoût de communication substantiel. Dans cet article, nous présentons ZeCO (Zero Communication Overhead) pour les modèles d'attention linéaire, une nouvelle méthode de SP conçue pour surmonter ces limitations et atteindre une scalabilité quasi-linéaire de bout en bout pour l'entraînement de longues séquences. Par exemple, entraîner un modèle avec une séquence de 1 million sur 64 dispositifs en utilisant ZeCO prend à peu près le même temps qu'un entraînement avec une séquence de 16 000 sur un seul dispositif. Au cœur de ZeCO se trouve All-Scan, une nouvelle primitive de communication collective. All-Scan fournit à chaque rang SP exactement l'état initial de l'opérateur dont il a besoin tout en maintenant une empreinte de communication minimale, éliminant ainsi efficacement le surcoût de communication. Théoriquement, nous prouvons l'optimalité de ZeCO, montrant qu'il introduit seulement un surcoût temporel et spatial négligeable. Empiriquement, nous comparons les coûts de communication de différentes stratégies de parallélisme de séquence et démontrons qu'All-Scan réalise la communication la plus rapide dans les scénarios de SP. Plus précisément, sur 256 GPU avec une séquence de 8 millions, ZeCO atteint une accélération de 60 % par rapport à la méthode de SP actuellement la plus avancée (SOTA). Nous croyons que ZeCO établit une voie claire vers l'entraînement efficace des LLMs de nouvelle génération sur des longueurs de séquences auparavant inaccessibles.

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.