拆分注意力与前馈层，实现训练加速1.8倍
DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism
- 将注意力和前馈层分到不同GPU组，用单向多对多通信减少瓶颈
- 在16节点8xH800集群上，训练效率提升最高达1.8倍
- 适合大规模MoE模型训练，尤其适用于高带宽受限场景
混合专家（MoE）架构使万亿参数大语言模型成为可能，其稀疏激活特性依赖专家并行（EP）策略。但随着模型规模扩大，需跨GPU节点分布专家，导致严重的全对全通信瓶颈，且受节点间网络带宽限制加剧。现有方法尝试重叠注意力与前馈网络（FFN）的计算与通信，但仍因两层计算-通信比不均衡而存在残余网络等待。本文提出DisagMoE，一种解耦式MoE训练系统，通过将注意力与FFN层分配至独立的GPU组，引入单向多对多通信的多阶段流水线，并采用计算-通信屋顶模型动态平衡各组的计算与网络资源分配。该方案基于Megatron-LM实现，在多个MoE模型上验证，16节点8xH800集群下训练效率最高提升1.8倍。
原文摘要 · Abstract (English)
Mixture-of-experts (MoE) architectures enable trillion-parameter LLMs with sparsely activated experts. Expert parallelism (EP) is a widely adopted MoE training strategy, but it suffers from severe all-to-all communication bottlenecks, which is exaggerated by the limited inter-node network bandwidth as the growing model size requires distributing experts across GPU nodes. Prior work focused on overlapping these all-to-all communications with feed-forward network (FFN) and self-attention computations, which often leaves residual network-bound stalls due to inherent imbalance in attention and FFN layers' computation-communication ratios. We present DisagMoE, a disaggregated MoE training system that jointly optimizes model placement and scheduling for maximal efficiency. DisagMoE separates attention and FFN layers into disjoint GPU groups, introduces a multi-stage pipeline with uni-directional, many-to-many communications, and employs a computation-communication roofline model to balance GPU and network bandwidth allocation among the attention and FFN groups. DisagMoE is implemented on Megatron-LM, and evaluation shows that DisagMoE improves training efficiency across multiple MoE models with up to 1.8x speedup on 16-node 8xH800 clusters.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。