ASPLOS(ACM International Conference on Architectural Support for Programming Languages and Operating Systems)是計算機科學領(lǐng)域頂級的國際學術(shù)會議之一,專注于計算機體系結(jié)構(gòu)、編程語言與操作系統(tǒng)等領(lǐng)域。作為中國計算機學會(CCF)推薦的A類會議,ASPLOS是計算機系統(tǒng)領(lǐng)域最具影響力的會議之一,近年來論文接收率維持在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 表示論文的主要實驗結(jié)果已由評審委員會在作者提供的 artifact 上獨立復現(xiàn)。
一、背景
混合專家模型(Mixture-of-Experts, MoE)通過僅激活部分專家來處理輸入token,能夠在保持計算量不變的同時顯著增加模型參數(shù)量,已成為大模型訓練的重要架構(gòu)。專家并行(Expert Parallelism, EP)是訓練MoE模型的關(guān)鍵技術(shù),將不同專家分布在多個設(shè)備上。然而,在專家并行訓練中,動態(tài)路由導致專家之間出現(xiàn)顯著的負載不均衡:少數(shù)過載的專家會阻礙整體迭代,成為訓練瓶頸。
目前,主流的解決思路主要分為兩類:一方面,算法層面引入輔助損失(auxiliary loss)或丟棄token,但這會影響模型收斂和效果;另一方面,系統(tǒng)層面通過專家復制、專家重定位等方式調(diào)整專家布局,但會引入高昂的通信和遷移開銷。這種"算法約束+系統(tǒng)重布局"的組合在實現(xiàn)上各有取舍,卻難以在動態(tài)變化的路由分布下既優(yōu)化專家布局又最小化重布局開銷。因此,如何在二者之間取得平衡,是當前亟待解決的問題。

圖1:專家并行中的負載不均衡與通信開銷問題
二、方 法
我們提出LAER-MoE,一個高效的MoE訓練框架。其核心思想是通過并行策略與負載均衡規(guī)劃器的協(xié)同設(shè)計,實現(xiàn)動態(tài)且高效的專家負載均衡。核心方法包括:
1. 完全分片專家并行(Fully Sharded Expert Parallel, FSEP):FSEP是一種新穎的并行范式,將每個專家參數(shù)按設(shè)備數(shù)量完全切分(類似于FSDP),并在前向/反向計算時通過All-to-All通信按需恢復完整的專家參數(shù)。該范式在保持與FSDP相同內(nèi)存高效性的同時,允許在訓練過程中靈活重布局專家參數(shù)(即決定哪個設(shè)備恢復哪個專家),從而為負載均衡提供優(yōu)化空間。我們進一步通過細粒度的通信調(diào)度(預取、梯度同步延遲)有效掩蓋FSEP引入的通信開銷。

圖2:FSEP示意圖
2. 負載均衡規(guī)劃器:我們設(shè)計了智能規(guī)劃器動態(tài)制定專家重布局策略與token路由方案,主要包含兩個組件:token調(diào)度器采用基于貪心的輕量級路由算法,優(yōu)先最小化節(jié)點間通信以快速確定token去向;專家布局調(diào)優(yōu)器基于優(yōu)先隊列確定專家副本數(shù)量,并用貪心算法求解專家位置,以最小化計算與通信的總成本。

圖3:使用FSEP進行負載均衡的示例
3. 系統(tǒng)框架和整體工作流程:token調(diào)度器在訓練過程中實時根據(jù)當前專家路由決定設(shè)備路由策略;專家布局調(diào)優(yōu)器則基于訓練過程中收集的歷史路由信息,由CPU端規(guī)劃器異步生成下一輪專家布局策略,實現(xiàn)負載均衡優(yōu)化與訓練執(zhí)行的無縫流水線并行。

圖4:系統(tǒng)架構(gòu)圖
三、實驗
在32張NVIDIA A100集群上,我們使用Mixtral-8x7B、Mixtral-8x22B和Qwen-8x7B模型架構(gòu),在wikitext和C4數(shù)據(jù)集上進行了廣泛的實驗。結(jié)果顯示,相較于業(yè)界主流的Megatron-LM和FSDP+EP,LAER-MoE在吞吐量上實現(xiàn)了高達1.69倍的加速。實驗表明,LAER-MoE在多種模型配置下均展現(xiàn)出優(yōu)越的性能。

圖5:端到端實驗對比
四、總結(jié)
本文提出了LAER-MoE系統(tǒng),通過完全分片專家并行(FSEP)范式和智能負載均衡規(guī)劃器,有效解決了MoE訓練中的負載不均衡問題。FSEP在消除重布局開銷的同時提供了極大的優(yōu)化自由度,配合動態(tài)規(guī)劃器,在不犧牲模型質(zhì)量的前提下顯著提升了訓練效率。實驗證明,LAER-MoE在多種模型配置下均展現(xiàn)出優(yōu)越的性能,為大規(guī)模MoE模型的訓練提供了高效的系統(tǒng)支持。
實驗室簡介
北京大學數(shù)據(jù)與智能實驗室(Data And Intelligence Research Lab at Peking Univeristy,PKU-DAIR實驗室)由北京大學計算機學院崔斌教授領(lǐng)導,長期從事數(shù)據(jù)庫系統(tǒng)、大數(shù)據(jù)管理與分析、人工智能等領(lǐng)域的前沿研究,在理論和技術(shù)創(chuàng)新以及系統(tǒng)研發(fā)上取得多項成果,已在國際頂級學術(shù)會議和期刊發(fā)表學術(shù)論文200余篇,發(fā)布多個開源項目。課題組同學曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國優(yōu)博、北大優(yōu)博、微軟學者、蘋果獎學金、谷歌獎學金等榮譽。PKU-DAIR實驗室持續(xù)與工業(yè)界展開卓有成效的合作,與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業(yè)開展項目合作和前沿探索,解決實際問題,進行科研成果的轉(zhuǎn)化落地。

評論 0