ZeCO: Paralelismo de Secuencias con Cero Sobrecarga de Comunicación para Atención Lineal

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

Resumen

Los mecanismos de atención lineal ofrecen ventajas significativas para los Modelos de Lenguaje a Gran Escala (LLMs) al proporcionar una complejidad computacional lineal, permitiendo el procesamiento eficiente de secuencias ultra largas (por ejemplo, 1 millón de tokens de contexto). Sin embargo, los métodos existentes de Paralelismo de Secuencias (SP, por sus siglas en inglés), esenciales para distribuir estas cargas de trabajo entre dispositivos, se convierten en el principal cuello de botella debido a un sobrecarga sustancial de comunicación. En este artículo, presentamos ZeCO (Zero Communication Overhead) para paralelismo de secuencias en modelos de atención lineal, un nuevo método SP diseñado para superar estas limitaciones y lograr una escalabilidad casi lineal de extremo a extremo en el entrenamiento de secuencias largas. Por ejemplo, entrenar un modelo con una secuencia de 1 millón de tokens en 64 dispositivos utilizando ZeCO toma aproximadamente el mismo tiempo que entrenar con una secuencia de 16k tokens en un solo dispositivo. En el núcleo de ZeCO se encuentra All-Scan, una nueva primitiva de comunicación colectiva. All-Scan proporciona a cada rango de SP exactamente el estado inicial del operador que requiere, manteniendo una huella de comunicación mínima, eliminando efectivamente la sobrecarga de comunicación. Teóricamente, demostramos la optimalidad de ZeCO, mostrando que introduce solo un tiempo y espacio adicional insignificante. Empíricamente, comparamos los costos de comunicación de diferentes estrategias de paralelismo de secuencias y demostramos que All-Scan logra la comunicación más rápida en escenarios de SP. Específicamente, en 256 GPUs con una secuencia de 8 millones de tokens, ZeCO logra una aceleración del 60\% en comparación con el método SP más avanzado actual (SOTA). Creemos que ZeCO establece un camino claro hacia el entrenamiento eficiente de la próxima generación de LLMs en longitudes de secuencia previamente intratables.

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.