arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.11686cs.DCcs.LG
收藏

LAER-MoE:负载自适应专家重新布局用于高效混合专家训练

LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training

Peking University(北京大学);Shanghai Jiao Tong University(上海交通大学)

Xinyi Liu, Yujie Wang, Fangcheng Fu, Xuefeng Xiao, Huixia Li, Jiashi Li, Bin Cui

AI总结:

LAER-MoE通过负载自适应专家重新布局提升混合专家模型训练效率,实现1.69倍加速。

摘要:

专家并行性对于有效训练混合专家(MoE)模型至关重要,使不同设备可以托管不同的专家,每个设备处理不同的输入数据。然而,在专家并行训练过程中,动态路由导致专家之间显著的负载不平衡:少数过载的专家阻碍了整体迭代,成为训练瓶颈。在本文中,我们介绍了一种高效的MoE训练框架LAER-MoE。LAER-MoE的核心是一种新的并行范式,即完全分片专家并行(FSEP),通过设备数量将每个专家参数完全划分,并通过All-to-All通信在训练过程中在专家粒度上恢复部分专家。这允许在训练过程中灵活重新布局专家参数以提高负载平衡。特别是,我们进行了细粒度的通信操作调度以最小化通信开销。此外,我们开发了一个负载平衡计划器,以制定专家的重新布局策略和训练期间的令牌路由方案。我们在A100集群上进行了实验,结果表明,我们的系统相比当前最先进的训练系统实现了高达1.69倍的加速。源代码可在https://github.com/PKU-DAIR/Hetu-Galvatron/tree/laer-moe上获得。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑