2023 年 7 月，清華大學(xué)計(jì)算機(jī)系 PACMAN 實(shí)驗(yàn)室發(fā)布稀疏大模型訓(xùn)練系統(tǒng) SmartMoE，支持用戶一鍵實(shí)現(xiàn) MoE 模型分布式訓(xùn)練，通過自動(dòng)搜索復(fù)雜并行策略，達(dá)到開源 MoE 訓(xùn)練系統(tǒng)領(lǐng)先性能。同時(shí)，PACMAN 實(shí)驗(yàn)室在國(guó)際頂級(jí)系統(tǒng)會(huì)議 USENIX ATC’23 發(fā)表長(zhǎng)文，作者包括博士生翟明書、何家傲等，通訊作者為翟季冬教授。PACMAN 實(shí)驗(yàn)室在機(jī)器學(xué)習(xí)系統(tǒng)領(lǐng)域持續(xù)深入研究，SmartMoE 是繼 FastMoE, FasterMoE 和 “八卦爐” 后在大模型分布式訓(xùn)練系統(tǒng)上的又一次探索。欲了解更多相關(guān)成果可查看翟季冬教授首頁(yè)：https://pacman.cs.tsinghua.edu.cn/~zjd

Mixture-of-Experts (MoE) 是一種模型稀疏化技術(shù)，因其高效擴(kuò)展大模型參數(shù)量的特性而備受研究者關(guān)注。為了提高 MoE 模型的易用性、優(yōu)化 MoE 模型訓(xùn)練性能，PACMAN 實(shí)驗(yàn)室在 MoE 大模型訓(xùn)練系統(tǒng)上進(jìn)行了系統(tǒng)深入的研究。2021 年初，開源發(fā)布了 FastMoE 系統(tǒng)，它是第一個(gè)基于 PyTorch 的 MoE 分布式訓(xùn)練系統(tǒng)開源實(shí)現(xiàn)，在業(yè)界產(chǎn)生了較大的影響力。進(jìn)一步，為了解決專家并行的稀疏、動(dòng)態(tài)計(jì)算模式帶來(lái)的嚴(yán)重性能問題，F(xiàn)asterMoE 系統(tǒng)地分析、優(yōu)化了專家并行策略。FasterMoE 中設(shè)計(jì)的「影子專家」技術(shù)顯著緩解了負(fù)載不均問題、通信 - 計(jì)算協(xié)同調(diào)度算法有效隱藏了 all-to-all 通信的高延遲。FasterMoE 成果發(fā)表在 PPoPP’22 國(guó)際會(huì)議。

• 論文地址：https://www.usenix.org/system/files/atc23-zhai.pdf
• 項(xiàng)目地址：https://github.com/zms1999/SmartMoE

MoE 模型遇到的難題

不同于稠密模型直接通過增大模型尺寸實(shí)現(xiàn)擴(kuò)展，如圖一所示，MoE 技術(shù)通過將一個(gè)小模型轉(zhuǎn)變?yōu)槎鄠€(gè)稀疏激活的小模型實(shí)現(xiàn)參數(shù)擴(kuò)展。由于各個(gè)專家在訓(xùn)練時(shí)稀疏激活，MoE 模型得以在不增加每輪迭代計(jì)算量的前提下增加模型參數(shù)量；從而有望在相同訓(xùn)練時(shí)間內(nèi)獲得更強(qiáng)的模型能力。

圖一：通過 MoE 技術(shù)擴(kuò)展模型規(guī)模

為了實(shí)現(xiàn) MoE 大模型的分布式訓(xùn)練，業(yè)界提出了專家并行（Expert Parallelism）技術(shù)。如圖二所示，各個(gè)專家被分布式地存儲(chǔ)在不同節(jié)點(diǎn)上，在訓(xùn)練過程中通過 all-to-all 通信將訓(xùn)練數(shù)據(jù)發(fā)送至對(duì)應(yīng)專家所在節(jié)點(diǎn)。專家并行相較于數(shù)據(jù)并行（Data Parallelism）有更小的內(nèi)存開銷，因?yàn)閷＜覅?shù)無(wú)冗余存儲(chǔ)?？梢哉J(rèn)為專家并行是一種針對(duì) MoE 結(jié)構(gòu)的模型并行（Model Parallelism）。

圖二：專家并行示意圖

然而，使用樸素的專家并行技術(shù)訓(xùn)練 MoE 模型有嚴(yán)重的性能問題，其根因是 MoE 模型的稀疏激活模式。它會(huì)導(dǎo)致節(jié)點(diǎn)間產(chǎn)生大量不規(guī)則 all-to-all 通信增加延遲、計(jì)算量負(fù)載不均造成硬件利用率低。如圖三所示的真實(shí)模型訓(xùn)練過程中的專家選擇分布，可以觀察到專家間顯著的負(fù)載不均現(xiàn)象，且分布隨訓(xùn)練進(jìn)行動(dòng)態(tài)變化。

圖三：真實(shí)訓(xùn)練中的專家選擇分布

隨著學(xué)界對(duì)各并行策略的深入研究，使用各并行策略的復(fù)雜組合（稱為混合并行）進(jìn)行大模型訓(xùn)練成為必要模式?；旌喜⑿械牟呗哉{(diào)優(yōu)過程十分復(fù)雜，為了提高可用性，學(xué)界提出了自動(dòng)并行算法自動(dòng)搜索、調(diào)優(yōu)混合并行策略。然而，現(xiàn)有混合并行、自動(dòng)并行系統(tǒng)無(wú)法高效處理 MoE 大模型，他們?nèi)鄙賹?duì) MoE 模型訓(xùn)練稀疏激活、計(jì)算負(fù)載不均且動(dòng)態(tài)變化的特征的針對(duì)性設(shè)計(jì)。

SmartMoE 帶來(lái)解決方案

為了實(shí)現(xiàn) MoE 模型的高效訓(xùn)練，SmartMoE 系統(tǒng)對(duì) MoE 模型的分布式訓(xùn)練策略進(jìn)行了全面的支持。對(duì)于常用的四種并行策略（數(shù)據(jù)并行、流水線并行、模型并行和專家并行），SmartMoE 系統(tǒng)做出了全面的支持，允許用戶對(duì)它們?nèi)我饨M合；在論文投稿時(shí)（2023 年 1 月），尚未有其他系統(tǒng)能做到這一點(diǎn)（如圖四所示）。

為了處理 MoE 的動(dòng)態(tài)計(jì)算負(fù)載，SmartMoE 獨(dú)特設(shè)計(jì)了專家放置（Expert Placement）策略，在經(jīng)典并行策略組合的基礎(chǔ)上，實(shí)現(xiàn)了動(dòng)態(tài)負(fù)載均衡。如圖五所示，MoE 模型不同的計(jì)算負(fù)載（workload）會(huì)造成不同專家的過載；使用不同的專家放置順序，能在特定負(fù)載下實(shí)現(xiàn)節(jié)點(diǎn)間負(fù)載均衡。

圖四：開源分布式系統(tǒng)對(duì)各并行策略的支持情況對(duì)比

圖五：不同 MoE 訓(xùn)練負(fù)載需要不同專家放置策略

為了提高 MoE 模型復(fù)雜混合并行策略的易用性，SmartMoE 設(shè)計(jì)了一套輕量級(jí)且有效的兩階段自動(dòng)并行算法。現(xiàn)有自動(dòng)并行系統(tǒng)只能在訓(xùn)練開始前進(jìn)行策略搜索，無(wú)法根據(jù)負(fù)載情況動(dòng)態(tài)調(diào)整策略。簡(jiǎn)單的將現(xiàn)有自動(dòng)并行搜索算法在訓(xùn)練過程中周期性使用亦不可行，因?yàn)橛?xùn)練過程中的并行策略搜索和調(diào)整對(duì)延遲要求很高，現(xiàn)有算法的開銷過大。

SmartMoE 獨(dú)創(chuàng)性地將自動(dòng)并行搜索過程分為兩階段：

• 訓(xùn)練開始前，使用經(jīng)典算法搜索，獲得一個(gè)較小的候選策略集合
• 訓(xùn)練過程中，根據(jù)當(dāng)前負(fù)載，在候選策略集合中動(dòng)態(tài)調(diào)整，由于候選策略集合大小有限，此過程的開銷可以得到控制。

最終，SmartMoE 實(shí)現(xiàn)了輕量級(jí)且有效的自動(dòng)并行，達(dá)到了業(yè)界領(lǐng)先的性能。

在性能測(cè)試中，SmartMoE 在不同模型結(jié)構(gòu)、集群環(huán)境和規(guī)模下均有優(yōu)異的表現(xiàn)。例如，在 GPT-MoE 模型的訓(xùn)練性能測(cè)試中，相較于 FasterMoE，SmartMoE 有最高 1.88x 的加速比。值得注意的，在對(duì)每一輪迭代的性能觀察中發(fā)現(xiàn)，動(dòng)態(tài)的并行策略調(diào)整是必要的，且需要使用合適的調(diào)整頻率，如圖六所示。更多實(shí)驗(yàn)細(xì)節(jié)請(qǐng)參考論文原文。

圖六：MoE 模型在不同迭代的運(yùn)行時(shí)間?！眃yn.X” 表示每 X 輪進(jìn)行一次策略調(diào)整。

圖七：SmartMoE 在 GPT-MoE 模型端到端訓(xùn)練中的性能提升

結(jié)語(yǔ)

SmartMoE 現(xiàn)已開源，開發(fā)者維護(hù)活躍，且仍在持續(xù)優(yōu)化迭代，助力 MoE 大模型的發(fā)展。這是 PACMAN 實(shí)驗(yàn)室繼 FastMoE，[PPoPP’22] FasterMoE，[PPoPP’22] BaGuaLu 后在大模型分布式訓(xùn)練系統(tǒng)上的又一次探索。