專家模型不要專家并行！微軟開源MoE新路徑
共 3904字，需瀏覽 8分鐘
·
2024-11-11 13:05
新智元報道
新智元報道
【新智元導(dǎo)讀】近日，來自微軟的研究人員開源了使用全新方法訓(xùn)練的MoE大模型，不走尋常路，且編碼和數(shù)學(xué)表現(xiàn)出色。
繼Phi家族之后，微軟又開源了新的混合專家大模型——GRIN MoE。
與Phi-3.5同樣的個頭（16 * 3.8B），卻采用了截然不同的訓(xùn)練方法。
這個「不走尋常路」如果寫個太長不看版，那就是兩句話：
1. 使用新一代SparseMixer來精確估計專家路由的梯度，解決傳統(tǒng)方案中利用門控梯度代替路由梯度的問題。
2. 專家并行不要了，訓(xùn)練中改用數(shù)據(jù)、pipeline和張量并行，避免了傳統(tǒng)方法丟棄token的問題。
論文地址：https://arxiv.org/abs/2409.12136
當(dāng)然了，上面兩句話是小編說的，多少有點(diǎn)糙，文中細(xì)節(jié)，還請諸君繼續(xù)閱讀~
這年頭，新來一個LLM，當(dāng)然要先刷分了——
參數(shù)要少，效果要好，所以要在左上角：
GRIN作為MoE架構(gòu)，總參數(shù)量約42B，推理時激活的參數(shù)為6.6B，打同級別（7B）的非MoE模型是手拿把攥，甚至比14B的Phi-3還要略勝一籌。
在上面的這份成績單中，GRIN MoE表現(xiàn)優(yōu)異，尤其是在編碼和數(shù)學(xué)測試中。
比如，在衡量數(shù)學(xué)問題解決能力的GSM-8K中，GRIN MoE得分為90.4，而在編碼任務(wù)基準(zhǔn)HumanEval上拿到了74.4分。
在MMLU（大規(guī)模多任務(wù)語言理解）基準(zhǔn)測試中GRIN得分為79.4，超過了同為MoE架構(gòu)的Mixtral（70.5分），以及自家的Phi-3.5（78.9分）。
如果對比流行的商用模型，GPT-3.5表示感受到時代的力量，默默退出群聊。
demo：https://github.com/microsoft/GRIN-MoE
MoE全新訓(xùn)練路徑
采用RoPE進(jìn)行位置編碼，以便在預(yù)訓(xùn)練后實(shí)現(xiàn)長上下文能力。
在MoE架構(gòu)中，模型通過路由網(wǎng)絡(luò)為每個輸入token挑選適合的專家模塊。對于有n個專家的網(wǎng)絡(luò)，一個用于推理的MoE模塊的輸出為：
其中z = Router（x，r），本文中Router采用線性網(wǎng)絡(luò)，Gating是門控函數(shù)（通常為softmax），Expert是FNN層。
MoE通過TopK函數(shù)進(jìn)行專家分配，這個專家路由的過程是不可微的，所以反向傳播的時候沒法求導(dǎo)。
對此，傳統(tǒng)的MoE訓(xùn)練將TopK視為常數(shù)，僅通過Gating來反向傳播計算路由權(quán)重梯度，相當(dāng)于用門控的梯度代替了路由的梯度。
這多少有點(diǎn)糙。
不可導(dǎo)怎么辦
恰好，本文一作之前有一篇工作（SparseMixer）：
論文地址：https://arxiv.org/pdf/2310.00811
受到直通梯度估計器的啟發(fā)，作者擴(kuò)展了前作，提出了SparseMixer-v2。
作者首先將TopK函數(shù)替換為模型訓(xùn)練中離散變量的隨機(jī)采樣，然后應(yīng)用heun’s third order method來近似專家路由梯度，并構(gòu)建一個改進(jìn)的反向傳播，為專家路由給出數(shù)學(xué)上合理的梯度估計。
前作中，SparseMixer的有效性在神經(jīng)機(jī)器翻譯任務(wù)和ELECTRA語言模型訓(xùn)練中得到了證明。
而在GRIN MoE的開發(fā)過程中，SparseMixer-v2終于有機(jī)會大規(guī)模應(yīng)用于自回歸語言模型訓(xùn)練。
作者用2.5T token訓(xùn)練了兩個16×0.9B MoE。其中一個遵循GRIN MoE中使用的相同方案，另一個用傳統(tǒng)的GShard方法替換 SparseMixer-v2。
如上圖所示，將SparseMixer-v2的性能提升推廣到16×0.9B尺度的自回歸語言模型訓(xùn)練。
在前0.5T token上GShard表現(xiàn)更好，但SparseMixer-v2在訓(xùn)練后期取得了更強(qiáng)的性能。
專家模型不要專家并行
傳統(tǒng)的MoE訓(xùn)練采用專家并行，簡單理解就是把不同的專家分配到不同的顯卡上。
一個明顯的問題是負(fù)載不均衡，有的專家會分到更多的token，有的專家卻很閑。
之前的做法是設(shè)定一個閾值，比如1000個token分給4個專家，每人應(yīng)該是250，這時候每張卡就最多只算250個token，超過后直接丟棄（送到下一層）。
而在本文中，作者利用數(shù)據(jù)并行、pipeline并行和張量并行來訓(xùn)練GRIN MoE。
此外，對于沒有專家并行性的MoE計算，作者發(fā)現(xiàn)Megablocks包非常有用，它的grouped_GEMM內(nèi)核和包裝器的性能更好。
應(yīng)用這些新的工程化方法避免了專家并行，也就不用丟棄token了。
最終，與具有相同激活參數(shù)的密集模型相比，本文的方法實(shí)現(xiàn)了超過80%的訓(xùn)練效率提升。
上表中，作者將兩種不同大小的MoE模型與具有相同激活參數(shù)量的密集模型進(jìn)行了比較，使用相同的硬件測量了它們的訓(xùn)練吞吐量。
盡管MoE總的參數(shù)量是密集模型的六倍多，但在實(shí)驗(yàn)中達(dá)到了超過80%的相對吞吐量，證實(shí)了使用GRIN MoE方法的模型具有顯著的計算擴(kuò)展?jié)摿Α?/span>
（PS：密集模型的吞吐量是在與MoE模型相同的并行度設(shè)置下測量的，這里的比較是為了研究密集激活網(wǎng)絡(luò)（非MoE）和稀疏激活網(wǎng)絡(luò)（MoE）的GPU內(nèi)核效率）
此外，在擴(kuò)大模型大小時，密集模型和MoE模型顯示出相似的減速模式，比如6.6B密集模型的訓(xùn)練吞吐量大約比1.6B密集模型的訓(xùn)練吞吐量慢4.19倍（后者的參數(shù)少4倍）。同樣，42B MoE模型的訓(xùn)練吞吐量比10B MoE 模型的訓(xùn)練吞吐量慢約3.96倍（對應(yīng)參數(shù)少4.2倍）。
并行實(shí)驗(yàn)
在只使用pipeline并行的情況下，通過在GPU之間進(jìn)一步劃分不同層，可以將最大專家數(shù)量從16個擴(kuò)展到32個。但是，如果再增加專家數(shù)量，則會導(dǎo)致單個層的參數(shù)過多，一個GPU就放不下了。
所以下一個維度采用張量并行。
專家并行在前向和后向計算中有兩個all-to-all通信開銷，而張量并行在前向和后向計算中有兩個all-reduce通信開銷。
相比之下all-reduce操作的延遲更高一點(diǎn)，但可以通過精心排布前向和反向的計算來overlap掉一部分開銷。
如上圖所示，通過結(jié)合pipeline并行和張量并行，系統(tǒng)支持的最大專家數(shù)量擴(kuò)展到52個（總共132B參數(shù)）。
這個數(shù)量是因?yàn)閷?shí)驗(yàn)只用了64個GPU，最多能將模型劃分為64個階段，如果有更多的GPU，那么還能繼續(xù)向上擴(kuò)展。
不過作者也表示，使用更復(fù)雜的并行通常會導(dǎo)致計算吞吐量降低。
負(fù)載均衡
如前所述，本文沒有采用專家并行，但是負(fù)載不均衡的事實(shí)依然存在。
作者在這里通過調(diào)整負(fù)載均衡損失來調(diào)節(jié)全局的負(fù)載均衡。常見的負(fù)載均衡損失定義為：
其中α是超參數(shù)，n是專家數(shù)量，fi是調(diào)度給專家的token比例。
傳統(tǒng)方法在本地不同的GPU上計算fi，因此負(fù)載均衡損失將調(diào)節(jié)本地專家負(fù)載均衡并緩解token丟棄。
在本文中，作者通過計算全局的fi（比如數(shù)據(jù)并行過程中組內(nèi)的all-reduce）來修改負(fù)載均衡損失，調(diào)節(jié)專家負(fù)載以達(dá)到全局平衡。
盡管這種調(diào)整會產(chǎn)生額外的通信開銷，但類似于張量并行，這些通信也可以與計算overlap，從而在很大程度上減少額外的延遲。
最后，放一個測試結(jié)果來show一下GRIN MoE的數(shù)學(xué)推理能力：
作者注：我們對新發(fā)布的GAOKAO（即全國普通大學(xué)和學(xué)院入學(xué)統(tǒng)一考試）的數(shù)學(xué)問題進(jìn)行案例研究，這是中國一年一度的全國本科入學(xué)考試。
該考試以其嚴(yán)格的安全協(xié)議而聞名，是評估AI模型回答數(shù)學(xué)問題的能力的理想測試平臺。請注意，GRIN MoE的訓(xùn)練于太平洋標(biāo)準(zhǔn)時間6月3日結(jié)束，2024年GAOKAO于中國標(biāo)準(zhǔn)時間6月7日開始。