ZeCO：线性注意力的零通信开销序列并行
机器学习
2025-07-03 v2
摘要
线性注意力机制为大型语言模型 (LLM) 带来显著优势，具有线性计算复杂度，使其能够高效处理超长序列（例如 1M 长度的上下文）。然而，现有的序列并行 (SP) 方法对于分布式这些工作负载至关重要，却成为主要瓶颈 due to substantial communication overhead。本文提出了用于线性注意力模型的 ZeCO（零通信开销）序列并行，这是一种新的 SP 方法，旨�克服这些限制，实现长序列训练的端到端近线性可扩展性。例如，使用 ZeCO 在64个设备上训练 1M 序列长度的模型所需时间大约与在单个设备上训练 16k 序列所需时间相同。ZeCO 的核心是 All-Scan，这是一种新的集合通信原语。All-Scan 为每个 SP 排位提供正好需要的初始运算状态，同时保持最小的通信足迹，有效消除了通信开销。理论上，我们证明了 ZeCO 的最优性，表明它仅引入可忽略的时间和空间开销。实验上，我们比较了不同序列并行策略的通信成本，证明 All-Scan 在 SP 场景中实现了最快的通信。具体而言，在 256 个 GPU 上进行 8M 序列长度的训练，ZeCO 相对于当前最先进 (SOTA) SP 方法实现了 60% 的加速。我们认为 ZeCO 为在previously不可及的序列长度上高效训练下一代 LLM 建立了清晰的路径。
引用
@article{arxiv.2507.01004,
title = {ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention},
author = {Yuhong Chou and Zehao Liu and Ruijie Zhu and Xinyi Wan and Tianjian Li and Congying Chu and Qian Liu and Jibin Wu and Zejun Ma},
journal= {arXiv preprint arXiv:2507.01004},
year = {2025}
}