ZeCO: Parallelismo Sequenziale a Sovraccarico di Comunicazione Zero per l'Attenzione Lineare

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

July 1, 2025

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

cs.AI

Abstract

I meccanismi di attenzione lineare offrono vantaggi significativi per i Large Language Models (LLM) garantendo una complessità computazionale lineare, consentendo un'elaborazione efficiente di sequenze ultra-lunghe (ad esempio, contesto di 1M). Tuttavia, i metodi esistenti di Parallelismo di Sequenza (SP), essenziali per distribuire questi carichi di lavoro tra i dispositivi, diventano il principale collo di bottiglia a causa di un sovraccarico di comunicazione sostanziale. In questo articolo, introduciamo ZeCO (Zero Communication Overhead) parallelismo di sequenza per modelli di attenzione lineare, un nuovo metodo SP progettato per superare queste limitazioni e raggiungere una scalabilità quasi lineare end-to-end per l'addestramento di sequenze lunghe. Ad esempio, addestrare un modello con una lunghezza di sequenza di 1M su 64 dispositivi utilizzando ZeCO richiede all'incirca lo stesso tempo dell'addestramento con una sequenza di 16k su un singolo dispositivo. Al centro di ZeCO si trova All-Scan, un nuovo primitivo di comunicazione collettiva. All-Scan fornisce a ciascun rango SP esattamente lo stato iniziale dell'operatore di cui ha bisogno, mantenendo un'impronta di comunicazione minima, eliminando efficacemente il sovraccarico di comunicazione. Teoricamente, dimostriamo l'ottimalità di ZeCO, mostrando che introduce solo un trascurabile sovraccarico di tempo e spazio. Empiricamente, confrontiamo i costi di comunicazione di diverse strategie di parallelismo di sequenza e dimostriamo che All-Scan raggiunge la comunicazione più veloce negli scenari SP. In particolare, su 256 GPU con una lunghezza di sequenza di 8M, ZeCO raggiunge un miglioramento del 60\% rispetto al metodo SP attualmente più avanzato (SOTA). Crediamo che ZeCO stabilisca un percorso chiaro verso l'addestramento efficiente delle prossime generazioni di LLM su lunghezze di sequenza precedentemente intrattabili.

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.