arXiv:2507.01004cs.LG2025-07NeurIPS被引 2

ZeCO让长序列模型训练零通信开销,64设备跑100万序列仅需单卡1.6万序列时间。

ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention

  • 提出All-Scan通信原语,实现序列并行中无通信开销的初始状态分发
  • 在256张GPU上处理800万序列时,比当前最佳方法快60%
  • 适合超长序列训练的下一代大模型开发者使用

线性注意力机制通过线性计算复杂度,显著提升大语言模型对超长序列(如100万上下文)的处理效率。然而,现有序列并行(SP)方法因通信开销过大成为主要瓶颈。本文提出针对线性注意力模型的零通信开销序列并行方法ZeCO,实现长序列训练的端到端近线性可扩展性。例如,在64个设备上训练100万序列长度的模型,耗时与单设备训练1.6万序列相当。核心是新提出的集体通信原语All-Scan,使每个SP节点获取精确所需的初始操作状态,同时保持最小通信量,彻底消除通信开销。理论上证明了ZeCO的最优性,其时间和空间开销可忽略。实证显示,相比现有最优方案,256张GPU上处理800万序列长度时,ZeCO通信速度最快。我们相信ZeCO为训练下一代超长序列大模型提供了清晰路径。

原文摘要 · Abstract (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.

序列并行线性注意力零通信大模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。