繼Phi家族之后，微軟又開源了新的混合專家大模型——GRIN MoE。

與Phi-3.5同樣的個(gè)頭（16 * 3.8B），卻采用了截然不同的訓(xùn)練方法。

這個(gè)「不走尋常路」如果寫個(gè)太長(zhǎng)不看版，那就是兩句話：

1. 使用新一代SparseMixer來精確估計(jì)專家路由的梯度，解決傳統(tǒng)方案中利用門控梯度代替路由梯度的問題。

2. 專家并行不要了，訓(xùn)練中改用數(shù)據(jù)、pipeline和張量并行，避免了傳統(tǒng)方法丟棄token的問題。

論文地址：https://arxiv.org/abs/2409.12136

當(dāng)然了，上面兩句話是小編說的，多少有點(diǎn)糙，文中細(xì)節(jié)，還請(qǐng)諸君繼續(xù)閱讀~

這年頭，新來一個(gè)LLM，當(dāng)然要先刷分了——

參數(shù)要少，效果要好，所以要在左上角：

GRIN作為MoE架構(gòu)，總參數(shù)量約42B，推理時(shí)激活的參數(shù)為6.6B，打同級(jí)別（7B）的非MoE模型是手拿把攥，甚至比14B的Phi-3還要略勝一籌。

在上面的這份成績(jī)單中，GRIN MoE表現(xiàn)優(yōu)異，尤其是在編碼和數(shù)學(xué)測(cè)試中。

比如，在衡量數(shù)學(xué)問題解決能力的GSM-8K中，GRIN MoE得分為90.4，而在編碼任務(wù)基準(zhǔn)HumanEval上拿到了74.4分。

在MMLU（大規(guī)模多任務(wù)語言理解）基準(zhǔn)測(cè)試中GRIN得分為79.4，超過了同為MoE架構(gòu)的Mixtral（70.5分），以及自家的Phi-3.5（78.9分）。

如果對(duì)比流行的商用模型，GPT-3.5表示感受到時(shí)代的力量，默默退出群聊。

開放權(quán)重：https://huggingface.co/microsoft/GRIN-MoE

demo：https://github.com/microsoft/GRIN-MoE

MoE全新訓(xùn)練路徑

GRIN MoE由常規(guī)的Transformer塊構(gòu)成，采用分組查詢注意力（GQA）和滑動(dòng)窗口注意力來提高計(jì)算效率。

采用RoPE進(jìn)行位置編碼，以便在預(yù)訓(xùn)練后實(shí)現(xiàn)長(zhǎng)上下文能力。

在MoE架構(gòu)中，模型通過路由網(wǎng)絡(luò)為每個(gè)輸入token挑選適合的專家模塊。對(duì)于有n個(gè)專家的網(wǎng)絡(luò)，一個(gè)用于推理的MoE模塊的輸出為：

其中z = Router（x，r），本文中Router采用線性網(wǎng)絡(luò)，Gating是門控函數(shù)（通常為softmax），Expert是FNN層。

MoE通過TopK函數(shù)進(jìn)行專家分配，這個(gè)專家路由的過程是不可微的，所以反向傳播的時(shí)候沒法求導(dǎo)。

對(duì)此，傳統(tǒng)的MoE訓(xùn)練將TopK視為常數(shù)，僅通過Gating來反向傳播計(jì)算路由權(quán)重梯度，相當(dāng)于用門控的梯度代替了路由的梯度。

這多少有點(diǎn)糙。

不可導(dǎo)怎么辦

恰好，本文一作之前有一篇工作（SparseMixer）：

論文地址：https://arxiv.org/pdf/2310.00811

受到直通梯度估計(jì)器的啟發(fā)，作者擴(kuò)展了前作，提出了SparseMixer-v2。

作者首先將TopK函數(shù)替換為模型訓(xùn)練中離散變量的隨機(jī)采樣，然后應(yīng)用heun’s third order method來近似專家路由梯度，并構(gòu)建一個(gè)改進(jìn)的反向傳播，為專家路由給出數(shù)學(xué)上合理的梯度估計(jì)。

前作中，SparseMixer的有效性在神經(jīng)機(jī)器翻譯任務(wù)和ELECTRA語言模型訓(xùn)練中得到了證明。

而在GRIN MoE的開發(fā)過程中，SparseMixer-v2終于有機(jī)會(huì)大規(guī)模應(yīng)用于自回歸語言模型訓(xùn)練。

作者用2.5T token訓(xùn)練了兩個(gè)16×0.9B MoE。其中一個(gè)遵循GRIN MoE中使用的相同方案，另一個(gè)用傳統(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)練采用專家并行，簡(jiǎn)單理解就是把不同的專家分配到不同的顯卡上。

一個(gè)明顯的問題是負(fù)載不均衡，有的專家會(huì)分到更多的token，有的專家卻很閑。

之前的做法是設(shè)定一個(gè)閾值，比如1000個(gè)token分給4個(gè)專家，每人應(yīng)該是250，這時(shí)候每張卡就最多只算250個(gè)token，超過后直接丟棄（送到下一層）。

而在本文中，作者利用數(shù)據(jù)并行、pipeline并行和張量并行來訓(xùn)練GRIN MoE。

此外，對(duì)于沒有專家并行性的MoE計(jì)算，作者發(fā)現(xiàn)Megablocks包非常有用，它的grouped_GEMM內(nèi)核和包裝器的性能更好。

應(yīng)用這些新的工程化方法避免了專家并行，也就不用丟棄token了。

最終，與具有相同激活參數(shù)的密集模型相比，本文的方法實(shí)現(xiàn)了超過80%的訓(xùn)練效率提升。

上表中，作者將兩種不同大小的MoE模型與具有相同激活參數(shù)量的密集模型進(jìn)行了比較，使用相同的硬件測(cè)量了它們的訓(xùn)練吞吐量。

盡管MoE總的參數(shù)量是密集模型的六倍多，但在實(shí)驗(yàn)中達(dá)到了超過80%的相對(duì)吞吐量，證實(shí)了使用GRIN MoE方法的模型具有顯著的計(jì)算擴(kuò)展?jié)摿Α?/span>

（PS：密集模型的吞吐量是在與MoE模型相同的并行度設(shè)置下測(cè)量的，這里的比較是為了研究密集激活網(wǎng)絡(luò)（非MoE）和稀疏激活網(wǎng)絡(luò)（MoE）的GPU內(nèi)核效率）

此外，在擴(kuò)大模型大小時(shí)，密集模型和MoE模型顯示出相似的減速模式，比如6.6B密集模型的訓(xùn)練吞吐量大約比1.6B密集模型的訓(xùn)練吞吐量慢4.19倍（后者的參數(shù)少4倍）。同樣，42B MoE模型的訓(xùn)練吞吐量比10B MoE 模型的訓(xùn)練吞吐量慢約3.96倍（對(duì)應(yīng)參數(shù)少4.2倍）。

并行實(shí)驗(yàn)

在只使用pipeline并行的情況下，通過在GPU之間進(jìn)一步劃分不同層，可以將最大專家數(shù)量從16個(gè)擴(kuò)展到32個(gè)。但是，如果再增加專家數(shù)量，則會(huì)導(dǎo)致單個(gè)層的參數(shù)過多，一個(gè)GPU就放不下了。

所以下一個(gè)維度采用張量并行。

專家并行在前向和后向計(jì)算中有兩個(gè)all-to-all通信開銷，而張量并行在前向和后向計(jì)算中有兩個(gè)all-reduce通信開銷。

相比之下all-reduce操作的延遲更高一點(diǎn)，但可以通過精心排布前向和反向的計(jì)算來overlap掉一部分開銷。

如上圖所示，通過結(jié)合pipeline并行和張量并行，系統(tǒng)支持的最大專家數(shù)量擴(kuò)展到52個(gè)（總共132B參數(shù)）。

這個(gè)數(shù)量是因?yàn)閷?shí)驗(yàn)只用了64個(gè)GPU，最多能將模型劃分為64個(gè)階段，如果有更多的GPU，那么還能繼續(xù)向上擴(kuò)展。

不過作者也表示，使用更復(fù)雜的并行通常會(huì)導(dǎo)致計(jì)算吞吐量降低。

負(fù)載均衡

如前所述，本文沒有采用專家并行，但是負(fù)載不均衡的事實(shí)依然存在。

作者在這里通過調(diào)整負(fù)載均衡損失來調(diào)節(jié)全局的負(fù)載均衡。常見的負(fù)載均衡損失定義為：

其中α是超參數(shù)，n是專家數(shù)量，fi是調(diào)度給專家的token比例。

傳統(tǒng)方法在本地不同的GPU上計(jì)算fi，因此負(fù)載均衡損失將調(diào)節(jié)本地專家負(fù)載均衡并緩解token丟棄。

在本文中，作者通過計(jì)算全局的fi（比如數(shù)據(jù)并行過程中組內(nèi)的all-reduce）來修改負(fù)載均衡損失，調(diào)節(jié)專家負(fù)載以達(dá)到全局平衡。

盡管這種調(diào)整會(huì)產(chǎn)生額外的通信開銷，但類似于張量并行，這些通信也可以與計(jì)算overlap，從而在很大程度上減少額外的延遲。

最后，放一個(gè)測(cè)試結(jié)果來show一下GRIN MoE的數(shù)學(xué)推理能力：

作者注：我們對(duì)新發(fā)布的GAOKAO（即全國(guó)普通大學(xué)和學(xué)院入學(xué)統(tǒng)一考試）的數(shù)學(xué)問題進(jìn)行案例研究，這是中國(guó)一年一度的全國(guó)本科入學(xué)考試。

該考試以其嚴(yán)格的安全協(xié)議而聞名，是評(píng)估AI模型回答數(shù)學(xué)問題的能力的理想測(cè)試平臺(tái)。請(qǐng)注意，GRIN MoE的訓(xùn)練于太平洋標(biāo)準(zhǔn)時(shí)間6月3日結(jié)束，2024年GAOKAO于中國(guó)標(biāo)準(zhǔn)時(shí)間6月7日開始。