TASP：拓扑感知序列并行
机器学习
2025-10-10 v2 分布式、并行与集群计算
摘要
长上下文大语言模型 (LLM) 面临自注意力机制二次复杂度的限制。主流的序列并行 (SP) 方法如环注意力 (Ring Attention) 试图通过将查询在加速器之间分布到多个查询块上，并借助环组合并集通信原语使每个 Q 张量能够访问来自其他加速器的所有 KV 张量来解决此问题。然而，它在通信效率方面表现不佳，限制了实际应用。这种不效率源于其采用的环组合并集通信原语与现代加速器的 AlltoAll 拓扑之间的不匹配。环组合并集通信原语由环形数据传输的迭代组成，只能利用极小比例的 AlltoAll 拓扑。灵感来自完全有向图的哈密顿分解，我们识别出现代加速器拓扑可分解为多个正交环数据通路，这些通路可在不相互干扰的情况下并行传输数据。基于此，我们进一步观察到环组合并集通信原语在每个迭代中也可分解为相同数量的并行环形数据传输。基于上述洞见，我们提出 TASP，这是一种面向长上下文 LLM 的拓扑感知 SP 方法，通过拓扑分解和原语分解充分利用现代加速器的通信容量。在单节点和多节点的 NVIDIA H100 系统上以及单节点的 AMD MI300X 系统上的实验结果表明，TASP 在这些现代加速器拓扑上在通信效率方面优于环注意力，并相较于环注意力及其变体 Zigzag-Ring Attention 实现最高可达 3.58 倍的加速。代码已公开于 https://github.com/infinigence/HamiltonAttention。
引用
@article{arxiv.2509.26541,
title = {TASP: Topology-aware Sequence Parallelism},
author = {Yida Wang and Ke Hong and Xiuhong Li and Yuanchao Xu and Wenxun Wang and Guohao Dai and Yu Wang},
journal= {arXiv preprint arXiv:2509.26541},
year = {2025}
}