针对长序列大模型通信瓶颈,提出拓扑感知的并行方法,提升训练效率。
TASP: Topology-aware Sequence Parallelism
- 基于图论分解加速器拓扑,实现多路并行通信
- 在H100和MI300X上相比环形通信提速最高达3.58倍
- 适合部署长上下文大模型的高性能分布式训练
长上下文大语言模型受限于自注意力机制的二次复杂度。主流序列并行方法Ring Attention通过将查询分块分布到多个加速器,并利用环形全部归约(Ring AllGather)通信使每个查询张量能访问其他加速器上的所有键值张量。然而,该方法通信效率低,限制了实际应用,根源在于其采用的环形全部归约通信原语与现代加速器的AlltoAll拓扑不匹配。环形全部归约由多轮环状数据传输构成,仅能利用AlltoAll拓扑中极小部分带宽。受完全有向图的哈密顿分解启发,我们发现现代加速器拓扑可分解为多个正交环路路径,可并发传输数据且互不干扰。基于此,我们进一步观察到环形全部归约也可分解为相同数量的并发环状数据传输。据此,我们提出TASP:一种面向长上下文大模型的拓扑感知序列并行方法,通过拓扑分解与原语分解,充分挖掘现代加速器的通信能力。在单节点与多节点NVIDIA H100系统及单节点AMD MI300X系统上的实验表明,TASP在这些现代加速器拓扑上均实现更高通信效率,相比Ring Attention及其变体Zigzag-Ring Attention最高提速3.58倍。代码已开源:https://github.com/infinigence/HamiltonAttention。
原文摘要 · Abstract (English)
Long-context large language models (LLMs) face constraints due to the quadratic complexity of the self-attention mechanism. The mainstream sequence parallelism (SP) method, Ring Attention, attempts to solve this by distributing the query into multiple query chunks across accelerators and enable each Q tensor to access all KV tensors from other accelerators via the Ring AllGather communication primitive. However, it exhibits low communication efficiency, restricting its practical applicability. This inefficiency stems from the mismatch between the Ring AllGather communication primitive it adopts and the AlltoAll topology of modern accelerators. A Ring AllGather primitive is composed of iterations of ring-styled data transfer, which can only utilize a very limited fraction of an AlltoAll topology. Inspired by the Hamiltonian decomposition of complete directed graphs, we identify that modern accelerator topology can be decomposed into multiple orthogonal ring datapaths which can concurrently transfer data without interference. Based on this, we further observe that the Ring AllGather primitive can also be decomposed into the same number of concurrent ring-styled data transfer at every iteration. Based on these insights, we propose TASP, a topology-aware SP method for long-context LLMs that fully utilizes the communication capacity of modern accelerators via topology decomposition and primitive decomposition. Experimental results on both single-node and multi-node NVIDIA H100 systems and a single-node AMD MI300X system demonstrate that TASP achieves higher communication efficiency than Ring Attention on these modern accelerator topologies and achieves up to 3.58 speedup than Ring Attention and its variant Zigzag-Ring Attention. The code is available at https://github.com/infinigence/HamiltonAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。