通过分层去重与专家交换,显著提升MoE模型训练速度。
HierMoE: Accelerating MoE Training with Hierarchical Token Deduplication and Expert Swap
- 分层令牌去重减少通信量,专家交换平衡各GPU负载。
- 在32卡集群上,通信速度提升1.55至3.32倍,训练加速1.18至1.27倍。
- 适用于大规模MoE模型训练,适合追求高效分布式训练的研究者。
稀疏激活的混合专家(MoE)Transformer已成为大语言模型(LLM)的常见架构,因其稀疏性可在不增加计算负担的前提下轻松扩展模型规模。然而,在MoE模型中,每个MoE层需动态选择令牌以激活特定专家进行计算,而被激活的专家可能不在与令牌同一设备或GPU上,导致跨所有GPU产生大量通信和负载不均衡,阻碍了分布式系统在GPU集群中的可扩展性。为此,我们提出HierMoE,通过两种拓扑感知技术加速MoE模型训练:1)令牌去重以降低通信流量;2)专家交换以平衡各GPU的工作负载。为使上述方法更具通用性,我们构建理论模型,旨在不同模型配置和硬件环境下实现最优的令牌去重与专家交换策略。我们在Megatron-LM基础上实现了HierMoE原型系统,并在32卡集群上对DeepSeek-V3和Qwen3-30B-A3B模型进行了实验。结果表明,相较于最先进的MoE训练系统Tutel-2DH、SmartMoE和Megatron-LM,HierMoE实现了1.55至3.32倍的通信加速,以及1.18至1.27倍的端到端训练加速。
原文摘要 · Abstract (English)
The sparsely activated mixture-of-experts (MoE) transformer has become a common architecture for large language models (LLMs) due to its sparsity, which requires fewer computational demands while easily scaling the model size. In MoE models, each MoE layer requires to dynamically choose tokens to activate particular experts for computation while the activated experts may not be located in the same device or GPU as the token. However, this leads to substantial communication and load imbalances across all GPUs, which obstructs the scalability of distributed systems within a GPU cluster. To this end, we introduce HierMoE to accelerate the training of MoE models by two topology-aware techniques: 1) token deduplication to reduce the communication traffic, and 2) expert swap to balance the workloads among all GPUs. To enable the above two proposed approaches to be more general, we build theoretical models aimed at achieving the best token duplication and expert swap strategy under different model configurations and hardware environments. We implement our prototype HierMoE system atop Megatron-LM and conduct experiments on a 32-GPU cluster with DeepSeek-V3 and Qwen3-30B-A3B models. Experimental results show that our HierMoE achieves $1.55\times$ to $3.32\times$ faster communication and delivers $1.18\times$ to $1.27\times$ faster end-to-end training compared to state-of-the-art MoE training systems, Tutel-2DH, SmartMoE, and Megatron-LM.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。