ZeCO: 零通信开销的线性注意力序列并行

线性注意力机制通过提供线性计算复杂度，为大型语言模型（LLM）带来了显著优势，从而能够高效处理超长序列（例如1M上下文）。然而，现有的序列并行（SP）方法对于跨设备分配这些工作负载至关重要，但由于大量的通信开销而成为主要瓶颈。在本文中，我们引入了 ZeCO（零通信开销）序列并行，这是一种专为线性注意力模型设计的新型 SP 方法，旨在克服这些限制，并为长序列训练实现端到端的近线性扩展性。例如，使用 ZeCO 在 64 个设备上训练一个具有 1M 序列长度的模型，所需时间与在单个设备上训练一个 16k 序列大致相同。ZeCO 的核心是 All-Scan，这是一种新的集合通信原语。All-Scan 为每个 SP 节点精确提供其所需的初始操作符状态，同时保持最小的通信占用，从而有效地消除了通信开销。在理论上，我们证明了 ZeCO 的最优性，表明它只引入了可忽略的时间和空间开销。在实践中，我们比较了不同序列并行策略的通信成本，并证明 All-Scan 在 SP 场景中实现了最快的通信。具体来说，在 256 个 GPU 上，处理 8M 序列长度时，ZeCO 比当前最先进（SOTA）的 SP 方法实现了 60% 的加速。我们相信 ZeCO 为在以前无法处理的序列长度上高效训练下一代 LLM 建立了一条清晰的路径。