4
点赞
0
评论
1
转载
我要入驻

ASPLOS 2026 | LAER-MoE: 通过负载自适应专家重布局方法实现高效MoE训练

 

ASPLOS(ACM International Conference on Architectural Support for Programming Languages and Operating Systems)是计算机科学领域顶级的国际学术会议之一,专注于计算机体系结构、编程语言与操作系统等领域。作为中国计算机学会(CCF)推荐的A类会议,ASPLOS是计算机系统领域最具影响力的会议之一,近年来论文接收率维持在15%-20%左右,每年接收论文约100-150篇。PKU-DAIR实验室的论文《LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training》被ASPLOS 2026录用。

 

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

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

代码链接https://github.com/PKU-DAIR/Hetu-Galvatron/tree/laer-moe

 

本工作的代码在 ASPLOS 2026 的 Artifact Evaluation 中获得了三项徽章(badge):Artifact Available 表示作者提供的代码与材料已置于可公开访问的归档仓库并配有永久链接;Artifact Functional 表示 artifact 文档完整、可运行且通过评审方的验证;Results Reproduced 表示论文的主要实验结果已由评审委员会在作者提供的 artifact 上独立复现。

 

一、背景

混合专家模型(Mixture-of-Experts, MoE)通过仅激活部分专家来处理输入token,能够在保持计算量不变的同时显著增加模型参数量,已成为大模型训练的重要架构。专家并行(Expert Parallelism, EP)是训练MoE模型的关键技术,将不同专家分布在多个设备上。然而,在专家并行训练中,动态路由导致专家之间出现显著的负载不均衡:少数过载的专家会阻碍整体迭代,成为训练瓶颈。

目前,主流的解决思路主要分为两类:一方面,算法层面引入辅助损失(auxiliary loss)或丢弃token,但这会影响模型收敛和效果;另一方面,系统层面通过专家复制、专家重定位等方式调整专家布局,但会引入高昂的通信和迁移开销。这种"算法约束+系统重布局"的组合在实现上各有取舍,却难以在动态变化的路由分布下既优化专家布局又最小化重布局开销。因此,如何在二者之间取得平衡,是当前亟待解决的问题。

图1:专家并行中的负载不均衡与通信开销问题

二、方 法

我们提出LAER-MoE,一个高效的MoE训练框架。其核心思想是通过并行策略与负载均衡规划器的协同设计,实现动态且高效的专家负载均衡。核心方法包括:

1. 完全分片专家并行(Fully Sharded Expert Parallel, FSEP):FSEP是一种新颖的并行范式,将每个专家参数按设备数量完全切分(类似于FSDP),并在前向/反向计算时通过All-to-All通信按需恢复完整的专家参数。该范式在保持与FSDP相同内存高效性的同时,允许在训练过程中灵活重布局专家参数(即决定哪个设备恢复哪个专家),从而为负载均衡提供优化空间。我们进一步通过细粒度的通信调度(预取、梯度同步延迟)有效掩盖FSEP引入的通信开销。

 

图2:FSEP示意图

2. 负载均衡规划器:我们设计了智能规划器动态制定专家重布局策略与token路由方案,主要包含两个组件:token调度器采用基于贪心的轻量级路由算法,优先最小化节点间通信以快速确定token去向;专家布局调优器基于优先队列确定专家副本数量,并用贪心算法求解专家位置,以最小化计算与通信的总成本。

图3:使用FSEP进行负载均衡的示例

 

3. 系统框架和整体工作流程:token调度器在训练过程中实时根据当前专家路由决定设备路由策略;专家布局调优器则基于训练过程中收集的历史路由信息,由CPU端规划器异步生成下一轮专家布局策略,实现负载均衡优化与训练执行的无缝流水线并行。

图4:系统架构图

三、实验

在32张NVIDIA A100集群上,我们使用Mixtral-8x7B、Mixtral-8x22B和Qwen-8x7B模型架构,在wikitext和C4数据集上进行了广泛的实验。结果显示,相较于业界主流的Megatron-LM和FSDP+EP,LAER-MoE在吞吐量上实现了高达1.69倍的加速。实验表明,LAER-MoE在多种模型配置下均展现出优越的性能。

图5:端到端实验对比

四、总结

本文提出了LAER-MoE系统,通过完全分片专家并行(FSEP)范式和智能负载均衡规划器,有效解决了MoE训练中的负载不均衡问题。FSEP在消除重布局开销的同时提供了极大的优化自由度,配合动态规划器,在不牺牲模型质量的前提下显著提升了训练效率。实验证明,LAER-MoE在多种模型配置下均展现出优越的性能,为大规模MoE模型的训练提供了高效的系统支持。

 

 

实验室简介

 

北京大学数据与智能实验室(Data And Intelligence Research Lab at Peking Univeristy,PKU-DAIR实验室)由北京大学计算机学院崔斌教授领导,长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究,在理论和技术创新以及系统研发上取得多项成果,已在国际顶级学术会议和期刊发表学术论文200余篇,发布多个开源项目。课题组同学曾数十次获得包括CCF优博、ACM中国优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作,与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索,解决实际问题,进行科研成果的转化落地。


北京大学数据与智能实验室,PKU-DAIR,Peking University Data And Intelligence Research Lab,负责人为北京大学计算机学院崔斌教授。
返回顶部