DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866913113376292864 |
|---|---|
| author | Zeng, Zhichen Chang, Chi-Chih Wang, Jiayi Wang, Zezhou Zheng, Ningxin Zhong, Zheng Stuardo, Cesar A. Wang, Dongyang Abdelfattah, Mohamed S. Lin, Haibin Zhu, Banghua Li, Ang Jiang, Ziheng |
| author_facet | Zeng, Zhichen Chang, Chi-Chih Wang, Jiayi Wang, Zezhou Zheng, Ningxin Zhong, Zheng Stuardo, Cesar A. Wang, Dongyang Abdelfattah, Mohamed S. Lin, Haibin Zhu, Banghua Li, Ang Jiang, Ziheng |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_11005 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism Zeng, Zhichen Chang, Chi-Chih Wang, Jiayi Wang, Zezhou Zheng, Ningxin Zhong, Zheng Stuardo, Cesar A. Wang, Dongyang Abdelfattah, Mohamed S. Lin, Haibin Zhu, Banghua Li, Ang Jiang, Ziheng Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing 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. |
| title | DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2605.11005 |