ZeCO: Paralelismo de Sequência com Sobrecarga Zero de Comunicação para Atenção Linear

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

July 1, 2025

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

cs.AI

Resumo

Mecanismos de atenção linear oferecem vantagens significativas para Modelos de Linguagem de Grande Escala (LLMs) ao proporcionar complexidade computacional linear, permitindo o processamento eficiente de sequências ultra-longas (por exemplo, contexto de 1M). No entanto, os métodos existentes de Paralelismo de Sequência (SP), essenciais para distribuir essas cargas de trabalho entre dispositivos, tornam-se o principal gargalo devido à sobrecarga substancial de comunicação. Neste artigo, introduzimos o ZeCO (Zero Communication Overhead) para paralelismo de sequência em modelos de atenção linear, um novo método SP projetado para superar essas limitações e alcançar escalabilidade quase linear de ponta a ponta para o treinamento de sequências longas. Por exemplo, treinar um modelo com uma sequência de 1M em 64 dispositivos usando o ZeCO leva aproximadamente o mesmo tempo que treinar com uma sequência de 16k em um único dispositivo. No cerne do ZeCO está o All-Scan, um novo primitivo de comunicação coletiva. O All-Scan fornece a cada rank SP exatamente o estado inicial do operador que ele requer, mantendo uma pegada de comunicação mínima, eliminando efetivamente a sobrecarga de comunicação. Teoricamente, provamos a otimalidade do ZeCO, mostrando que ele introduz apenas sobrecargas de tempo e espaço insignificantes. Empiricamente, comparamos os custos de comunicação de diferentes estratégias de paralelismo de sequência e demonstramos que o All-Scan alcança a comunicação mais rápida em cenários de SP. Especificamente, em 256 GPUs com uma sequência de 8M, o ZeCO alcança um aumento de velocidade de 60\% em comparação com o método SP mais avançado atualmente (SOTA). Acreditamos que o ZeCO estabelece um caminho claro para o treinamento eficiente de LLMs de próxima geração em comprimentos de sequência anteriormente intratáveis.

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.