动态负载均衡提升MoE训练效率，实现1.69倍加速
LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training
- 采用全分片专家并行架构，按设备数拆分专家参数
- 训练中动态重排专家分布，解决热点专家负载不均问题
- 适配大规模MoE模型训练，适合追求高效推理部署的研究者
专家并行对高效训练混合专家（MoE）模型至关重要，使不同设备可托管不同专家，并各自处理不同输入数据。然而，在专家并行训练中，动态路由导致专家间显著负载不均：少数过载专家成为整体迭代瓶颈。本文提出LAER-MoE，一种高效MoE训练框架。核心是新型并行范式——全分片专家并行（FSEP），将每个专家参数按设备数完全分片，并在训练中通过All-to-All通信以专家粒度恢复部分专家。该机制支持训练过程中灵活重排专家参数，提升负载均衡。我们还设计细粒度通信调度策略，最小化通信开销；并开发负载均衡规划器，制定专家重排与令牌路由策略。在A100集群上实验表明，本系统相比当前最优训练系统最高提速1.69倍。源码见https://github.com/PKU-DAIR/Hetu-Galvatron/tree/laer-moe。
原文摘要 · Abstract (English)
Expert parallelism is vital for effectively training Mixture-of-Experts (MoE) models, enabling different devices to host distinct experts, with each device processing different input data. However, during expert parallel training, dynamic routing results in significant load imbalance among experts: a handful of overloaded experts hinder overall iteration, emerging as a training bottleneck. In this paper, we introduce LAER-MoE, an efficient MoE training framework. The core of LAER-MoE is a novel parallel paradigm, Fully Sharded Expert Parallel (FSEP), which fully partitions each expert parameter by the number of devices and restores partial experts at expert granularity through All-to-All communication during training. This allows for flexible re-layout of expert parameters during training to enhance load balancing. In particular, we perform fine-grained scheduling of communication operations to minimize communication overhead. Additionally, we develop a load balancing planner to formulate re-layout strategies of experts and routing schemes for tokens during training. We perform experiments on an A100 cluster, and the results indicate that our system achieves up to 1.69x acceleration compared to the current state-of-the-art training systems. Source code available at https://github.com/PKU-DAIR/Hetu-Galvatron/tree/laer-moe.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。