中文

HierMoE:通过分层令牌去重与专家交换加速 MoE 训练

分布式、并行与集群计算 2025-08-14 v1 机器学习

摘要

稀疏激活的混合专家 (MoE) Transformer 因其稀疏性已成为大型语言模型 (LLMs) 的常见架构,这种稀疏性在轻松扩展模型规模的同时,需要更少的计算需求。在 MoE 模型中,每个 MoE 层需要动态选择令牌以激活特定的专家进行计算,而被激活的专家可能与令牌不在同一设备或 GPU 上。然而,这导致了所有 GPU 之间显著的通信和负载不平衡,阻碍了 GPU 集群内分布式系统的可扩展性。为此,我们引入了 HierMoE,通过两种拓扑感知技术来加速 MoE 模型的训练:1) 令牌去重以减少通信流量,2) 专家交换以平衡所有 GPU 之间的工作负载。为了使上述两种方法更具通用性,我们构建了理论模型,旨在不同模型配置和硬件环境下实现最佳的令牌去重和专家交换策略。我们在 Megatron-LM 之上实现了 HierMoE 原型系统,并在一个 32-GPU 集群上使用 DeepSeek-V3 和 Qwen3-30B-A3B 模型进行了实验。实验结果表明,与最先进的 MoE 训练系统 Tutel-2DH、SmartMoE 和 Megatron-LM 相比,我们的 HierMoE 实现了 1.55×1.55\times3.32×3.32\times 的通信加速,以及 1.18×1.18\times1.27×1.27\times 的端到端训练加速。

关键词

引用

@article{arxiv.2508.09591,
  title  = {HierMoE: Accelerating MoE Training with Hierarchical Token Deduplication and Expert Swap},
  author = {Wenxiang Lin and Xinglin Pan and Lin Zhang and Shaohuai Shi and Xuan Wang and Xiaowen Chu},
  journal= {arXiv preprint arXiv:2508.09591},
  year   = {2025}
}