被OpenAI、Mistral AI帶火的MoE是怎么回事？一文貫通專家混合架構(gòu)部署

作者：機(jī)器之心 2024-01-22 09:02:00

本文將介紹 MoE 的構(gòu)建模塊、訓(xùn)練方法以及在使用它們進(jìn)行推理時(shí)需要考慮的權(quán)衡因素。

專家混合 (MoE) 是 LLM 中常用的一種技術(shù)，旨在提高其效率和準(zhǔn)確性。這種方法的工作原理是將復(fù)雜的任務(wù)劃分為更小、更易于管理的子任務(wù)，每個(gè)子任務(wù)都由專門的迷你模型或「專家」處理。

早些時(shí)候，有人爆料 GPT-4 是采用了由 8 個(gè)專家模型組成的集成系統(tǒng)。近日，Mistral AI 發(fā)布的 Mixtral 8x7B 同樣采用這種架構(gòu)，實(shí)現(xiàn)了非常不錯(cuò)的性能（傳送門：一條磁力鏈接席卷 AI 圈，87GB 種子直接開源 8x7B MoE 模型）。

OpenAI 和 Mistral AI 的兩波推力，讓 MoE 一時(shí)間成為開放人工智能社區(qū)最熱門的話題。

本文將介紹 MoE 的構(gòu)建模塊、訓(xùn)練方法以及在使用它們進(jìn)行推理時(shí)需要考慮的權(quán)衡因素。

混合專家架構(gòu)簡稱 MoE，它的特點(diǎn)如下：

與稠密模型相比，預(yù)訓(xùn)練速度更快；
與參數(shù)數(shù)量相同的模型相比，推理速度更快；
因?yàn)橐阉袑＜夷Ｐ投技虞d在內(nèi)存中，所以需要大量顯存；
在微調(diào)方面面臨許多挑戰(zhàn)，但最近在 MoE 指令微調(diào)方面的工作很有希望能解決這些問題。

什么是混合專家（MoE）？

模型的規(guī)模是決定模型質(zhì)量的最重要因素之一。在預(yù)算固定的情況下，用較少的步驟訓(xùn)練較大的模型要優(yōu)于用較多的步驟訓(xùn)練較小的模型。

MoE 可以用較少的計(jì)算開銷對(duì)模型進(jìn)行預(yù)訓(xùn)練，這意味著可以用與稠密模型相同的計(jì)算開銷，大幅擴(kuò)大模型或數(shù)據(jù)集的規(guī)模。特別是，在預(yù)訓(xùn)練過程中，MoE 模型能更快地達(dá)到與稠密模型相同的性能。

那么，究竟什么是 MoE？從 Transformer 模型的角度來說，MoE 包含兩個(gè)主要元素：

使用稀疏的 MoE 層代替稠密的前饋網(wǎng)絡(luò)（FFN）層。MoE 層中有一定數(shù)量（如 8 個(gè)）的「專家」，每個(gè)專家都是一個(gè)神經(jīng)網(wǎng)絡(luò)。實(shí)際上，專家可以是 FFN，也可以是更復(fù)雜的網(wǎng)絡(luò)，甚至是 MoE 本身，這樣就會(huì)形成有多層 MoE 的 MoE。
使用門控網(wǎng)絡(luò)或者路由來決定將哪個(gè) token 發(fā)送給哪個(gè)專家。例如，在下圖中，「More」被發(fā)送給第二個(gè)專家，而「Parameters」被發(fā)送到第一個(gè)專家。如何將 token 通過路由發(fā)送給專家是使用 MoE 時(shí)需要計(jì)劃的重點(diǎn)之一，這是因?yàn)槁酚赏瑯佑蓪W(xué)習(xí)到的參數(shù)組成，并與網(wǎng)絡(luò)的其他部分同時(shí)進(jìn)行預(yù)訓(xùn)練。

Switch Transformers 中的 MoE 層 (https://arxiv.org/abs/2101.03961)

簡而言之，在 MoE 中，一個(gè) MoE 層取代了 transformer 中的每個(gè) FFN 層，MoE 層由一個(gè)門控網(wǎng)絡(luò)和一定數(shù)量的專家網(wǎng)絡(luò)組成。

雖然與稠密模型相比，MoE 具有高效預(yù)訓(xùn)練和快速推理等優(yōu)點(diǎn)，但也面臨著一些挑戰(zhàn)：

訓(xùn)練：MoE 能夠大大提高預(yù)訓(xùn)練的計(jì)算效率，但在微調(diào)過程中難以實(shí)現(xiàn)泛化，從而導(dǎo)致過擬合。
推理：雖然 MoE 可能有很多參數(shù)，但在推理過程中只使用其中的一部分。與參數(shù)數(shù)量相同的稠密模型相比，推理速度要快得多。然而，所有參數(shù)都需要加載到 RAM 中，因此對(duì)內(nèi)存的要求很高。例如，給定一個(gè)像 Mixtral 8x7B 這樣的 MoE，需要有足夠的顯存來容納一個(gè) 47B 參數(shù)的稠密模型。為什么是 47B 參數(shù)而不是 8 x 7B = 56B？這是因?yàn)樵?MoE 模型中，只有 FFN 層被視為獨(dú)立的專家網(wǎng)絡(luò)，其余模型的參數(shù)都是共享的。同時(shí)，假設(shè)每個(gè) token 只傳入兩個(gè)專家網(wǎng)絡(luò)，那么推理速度（FLOPs）就像使用 12B 模型（而不是 14B 模型），因?yàn)樗M(jìn)行的是 2x7B 的矩陣乘法運(yùn)算，同時(shí)有些層是共享的（后文將詳細(xì)介紹）。

在對(duì) MoE 有了一個(gè)大致的介紹后，一起來了解一下 MoE 的發(fā)展軌跡。

MoE 簡史

MoE 起源于 1991 年的論文《Adaptive Mixture of Local Experts》（https://www.cs.toronto.edu/~hinton/absps/jjnh91.pdf）。該論文的理念與集合方法類似，都是為由不同網(wǎng)絡(luò)組成的系統(tǒng)提供監(jiān)督程序，每個(gè)網(wǎng)絡(luò)處理不同的訓(xùn)練集子集。每個(gè)獨(dú)立的網(wǎng)絡(luò)或者說專家擅長于輸入空間的不同區(qū)域。至于如何選擇專家這個(gè)問題，是由門控網(wǎng)絡(luò)來決定每個(gè)專家網(wǎng)絡(luò)的權(quán)重。在訓(xùn)練過程中，專家網(wǎng)絡(luò)和門控網(wǎng)絡(luò)都要接受訓(xùn)練。

2010-2015 年間，兩個(gè)研究領(lǐng)域的發(fā)展共同促成了 MoE 后來的進(jìn)步：

專家網(wǎng)絡(luò)組件化：在傳統(tǒng)的 MoE 中，整個(gè)系統(tǒng)由一個(gè)門控網(wǎng)絡(luò)和多個(gè)專家網(wǎng)絡(luò)組成。在 SVM、高斯過程和其他方法中，研究者們對(duì)作為整體模型的 MoE 進(jìn)行了探索。Eigen、Ranzato 和 Ilya 的研究，將 MoE 作為更深層網(wǎng)絡(luò)的組成部分進(jìn)行了探索。MoE 能夠作為多層網(wǎng)絡(luò)中的組成部分，使得模型既龐大又高效成為可能。
條件計(jì)算：傳統(tǒng)網(wǎng)絡(luò)每一層都會(huì)對(duì)所有輸入數(shù)據(jù)進(jìn)行處理。接著，Yoshua Bengio 研究了根據(jù)輸入的 token 動(dòng)態(tài)激活或停用組件網(wǎng)絡(luò)的方法。

這些工作促使研究者們?cè)?NLP 的背景下探索混合專家模型。具體來說，Shazeer 及 Geoffrey Hinton 、Jeff Dean，谷歌的 Chuck Norris 通過對(duì)引入稀疏網(wǎng)絡(luò)，將這一想法擴(kuò)展到了 137B LSTM（https://arxiv.org/abs/1701.06538），從而即使在大規(guī)模下也能保持極快的推理速度。這項(xiàng)工作的重點(diǎn)目標(biāo)是機(jī)器翻譯，同時(shí)也存在一些缺點(diǎn)，如通信成本高和訓(xùn)練不穩(wěn)定。

論文 Outrageously Large Neural Network 中的 MoE 層

MoE 可以訓(xùn)練數(shù)萬億級(jí)參數(shù)的模型，例如開源的 1.6T 參數(shù)的 Switch Transformer。計(jì)算機(jī)視覺領(lǐng)域也在探索 MoE，但這里還是先側(cè)重講解 NLP 領(lǐng)域。

什么是稀疏化？

稀疏化一詞來源于條件計(jì)算理念。在稠密模型中，所有參數(shù)都發(fā)揮作用，而稀疏化可以只運(yùn)行整個(gè)系統(tǒng)的某些部分。

前文提到 Shazeer 對(duì)機(jī)器翻譯中的 MoE 進(jìn)行了探索。條件計(jì)算（網(wǎng)絡(luò)中只有某些部分處于活動(dòng)狀態(tài)）使得在不增加計(jì)算量的情況下能夠擴(kuò)大模型的規(guī)模，因此，每層 MoE 都可以包含成千上萬的專家網(wǎng)絡(luò)。

但是這種設(shè)計(jì)帶來了一些挑戰(zhàn)。例如，雖然擴(kuò)大 batch size 通常更有利于提高模型性能，但 MOE 中的 batch size 會(huì)隨著數(shù)據(jù)在激活狀態(tài)的專家網(wǎng)絡(luò)中的流動(dòng)而縮小。例如，如果 batch size 為 10 個(gè) token，其中 5 個(gè) token 可能在一個(gè)專家網(wǎng)絡(luò)中結(jié)束，而另外 5 個(gè) token 可能在 5 個(gè)不同的專家網(wǎng)絡(luò)中結(jié)束，從而導(dǎo)致 batch size 大小不均和利用率不足的情況。

如何解決這個(gè)問題？方法之一是讓學(xué)習(xí)后的門控網(wǎng)絡(luò)（G）決定向哪些專家網(wǎng)絡(luò)（E）傳達(dá)輸入信息：

在這種情況下，全部的專家網(wǎng)絡(luò)會(huì)對(duì)所有輸入進(jìn)行運(yùn)算 — 用一種加權(quán)乘法的方式。但是，如果 G 為 0 時(shí)會(huì)怎樣呢？這種情況下，就不需要經(jīng)過相應(yīng)的專家網(wǎng)絡(luò)運(yùn)算，節(jié)省了計(jì)算開銷。那么典型的門控函數(shù)又是怎么樣的呢？在最傳統(tǒng)的設(shè)置中，只需使用一個(gè)帶有 softmax 函數(shù)的簡單網(wǎng)絡(luò)。該網(wǎng)絡(luò)將學(xué)習(xí)向哪個(gè)專家傳遞輸入數(shù)據(jù)。

Shazeer 的研究還探索了其他門控機(jī)制，如 Top-K 噪聲門控。這種門控方法會(huì)引入一些（可調(diào)整的）噪聲，然后保留前 K 個(gè)。也就是說：

1. 添加一些噪音

2. 只選擇前 k 個(gè)

3. 用 softmax 激活。

這種稀疏化操作帶來了一些有趣的特性。通過使用足夠小的 k（例如一個(gè)或兩個(gè)），研究者發(fā)現(xiàn)訓(xùn)練和推理速度比激活許多專家網(wǎng)絡(luò)的設(shè)置更快。那為什么不直接保留 top=1 時(shí)的專家網(wǎng)絡(luò)呢？研究者最初的猜想是，需要路由到一個(gè)以上的專家，才能讓門控網(wǎng)絡(luò)學(xué)習(xí)如何路由到不同的專家，因此至少要選擇兩個(gè)專家。Switch Transformer 的部分將重新討論了這一決策。

為什么要添加噪聲？這是為了負(fù)載平衡。

MoE 的 token 負(fù)載平衡

如前所述，如果所有的 token 都只發(fā)送給少數(shù)幾個(gè)受歡迎的專家網(wǎng)絡(luò)，訓(xùn)練效率將變得低下。通常 MoE 訓(xùn)練中，門控網(wǎng)絡(luò)會(huì)收斂到頻繁激活同樣的幾個(gè)專家網(wǎng)絡(luò)。這種情況會(huì)隨著訓(xùn)練的進(jìn)行趨于顯著，因?yàn)槭芮嗖A的專家網(wǎng)絡(luò)會(huì)更快地得到訓(xùn)練，從而更容易被選中。為了緩解這種情況，可以添加一個(gè)輔助損失，以鼓勵(lì)給予所有專家同等的重要性。該損失可確保所有專家獲得大致相同數(shù)量的訓(xùn)練樣本。后文將探討專家能力的概念，即專家可以處理 token 的數(shù)量閾值。在 transformer 中，輔助損失通過 aux_loss 參數(shù)顯示。

MoE 和 Transformer

Transformer 是一個(gè)增加參數(shù)數(shù)量可以提高性能的非常明顯的例子，因此，谷歌順理成章地在 GShard 上沿用這種思維，將其中 Transformers 的參數(shù)量擴(kuò)展到 6000 億以上。

GShard 在編碼器和解碼器中都使用了 top-2 門控技術(shù)，用 MoE 層取代了 FFN 層。下圖顯示了編碼器的部分情況。

GShard 中的 MoE Transformer 編碼器

這種設(shè)置對(duì)大規(guī)模計(jì)算非常有利：當(dāng)擴(kuò)展到多個(gè)設(shè)備時(shí)，MoE 層將在設(shè)備間共享，而所有其他層都將被復(fù)制。具體細(xì)節(jié)在下文的「高效訓(xùn)練 MoE」中進(jìn)一步討論。

為了保持均衡負(fù)載和規(guī)模效率，GShard 的作者除了采用與上一節(jié)討論的類似的輔助損耗外，還引入了一些改進(jìn)：

隨機(jī)路由：在 top-2 的設(shè)置中，研究者總是選擇排名第一的專家，同時(shí)第二位專家以與其權(quán)重成正比的概率被選中。
專家容量：還可以設(shè)定一個(gè)閾值，即一個(gè)專家可以處理多少 token。如果兩位專家的處理能力都已達(dá)到閾值，則 token 會(huì)被視為溢出，并通過剩余連接發(fā)送到下一層（或在其他項(xiàng)目中完全丟棄）。這一概念將成為 MoEs 最重要的概念之一。為什么需要設(shè)置專家容量呢？因?yàn)樗袕埩啃螤疃际窃诰幾g時(shí)靜態(tài)確定的，但無法提前知道每個(gè)專家將獲得多少 token 作為輸入，所以需要確定容量。

GShard 的貢獻(xiàn)在于為 MoEs 確定了并行計(jì)算模式。需要注意的是：推理時(shí)，只有部分專家網(wǎng)絡(luò)會(huì)被觸發(fā)。與此同時(shí)，還有一些需要共享數(shù)據(jù)的步驟，比如自注意力機(jī)制，適用于所有 token。這就是為什么對(duì)于一個(gè)由 8 個(gè)專家組成的 47B 模型時(shí)，GShard 可以用一個(gè) 12B 的稠密模型進(jìn)行計(jì)算。如果使用 top-2，則需要 14B 個(gè)參數(shù)。但考慮到注意力操作是共享的，實(shí)際使用的參數(shù)數(shù)量為 12B。

Switch Transformer

盡管 MoE 前景廣闊，但它們?cè)谟?xùn)練和微調(diào)不穩(wěn)定性方面仍有弊端。Switch Transformer （https://arxiv.org/abs/2101.03961）的出現(xiàn)意義非凡，作者甚至在 Hugging Face 發(fā)布了一個(gè)擁有 2048 個(gè)專家的 1.6 萬億個(gè)參數(shù)的模型（https://huggingface.co/google/switch-c-2048）。與 T5-XXL 相比，Switch Transformer 的預(yù)訓(xùn)練速度提高了 4 倍。

Switch Transformer 論文中的 Switch Transformer 層

正如在 GShard 中作者用 MoE 層取代了 FFN 層一樣，Switch Transformer 論文提出了一個(gè) Switch Transformer 層，該層接收兩個(gè)輸入（兩個(gè)不同的 token），有四個(gè)專家網(wǎng)絡(luò)。

與至少使用兩個(gè)專家網(wǎng)絡(luò)的最初想法相反，Switch Transformers 采用了簡化的單一專家策略。這種方法的效果如下：

減少了路由的計(jì)算量
每個(gè)專家的 batch size 至少減半
通信成本降低
模型的質(zhì)量仍能得到保證

Switch Transformer 還探索了專家容量的概念。

上面推薦的容量計(jì)算方法是將 batch size 中的 token 數(shù)量平均分配給專家。如果使用大于 1 的容量系數(shù)，就能在 token 不完全平衡時(shí)提供緩沖。增加容量會(huì)導(dǎo)致設(shè)備間的通信費(fèi)用增加，因此需要注意 trade-off。尤其是在容量系數(shù)較低（1-1.25）的情況下，Switch Transformer 的表現(xiàn)尤為出色。

Switch Transformer 的作者還重新審視并簡化了章節(jié)中提到的負(fù)載平衡損耗。對(duì)于每個(gè)交換層，輔助損耗會(huì)在訓(xùn)練過程中添加到模型總損耗中。這種損耗會(huì)鼓勵(lì)模型傾向于統(tǒng)一路由，并可使用超參數(shù)加權(quán)。

Switch Transformer 的作者還嘗試了選擇性精度，例如使用 bfloat16 的參數(shù)精度來訓(xùn)練專家，而在其他計(jì)算中使用全精度。較低的精度可以降低處理器之間的通信成本、計(jì)算成本和用于存儲(chǔ)張量的內(nèi)存。在最初的實(shí)驗(yàn)中，專家和門控網(wǎng)絡(luò)都使用 bfloat16 進(jìn)行了訓(xùn)練，但訓(xùn)練結(jié)果并不穩(wěn)定。這主要是由于路由網(wǎng)絡(luò)參與計(jì)算造成：由于路由網(wǎng)絡(luò)具有指數(shù)函數(shù)，因此更高的精度非常重要。為了減少不穩(wěn)定性，路由也使用了全精度。

使用選擇性精度不會(huì)降低質(zhì)量，并能更快地訓(xùn)練模型。

微調(diào)部分中，Switch Transformer 使用了編碼器 - 解碼器設(shè)置，并將 T5 與 MoE 對(duì)應(yīng)。GLaM （https://arxiv.org/abs/2112.06905）探討了如何用三分之一的計(jì)算開銷訓(xùn)練出與 GPT-3 質(zhì)量相匹配的模型，從而擴(kuò)大模型的規(guī)模。GLaM 作者的研究重點(diǎn)是純解碼器模型以及小樣本及零樣本的評(píng)估結(jié)果，并不是微調(diào)。他們使用了 Top-2 路由和更大的容量系數(shù)。此外，他們還將容量因子作為一個(gè)指標(biāo)進(jìn)行了調(diào)研，在訓(xùn)練和驗(yàn)證過程中，可以根據(jù)想要使用的計(jì)算量進(jìn)行改變。

穩(wěn)定訓(xùn)練基于 Z-loss 損失函數(shù)的路由網(wǎng)絡(luò)

上文討論的平衡損失可能會(huì)導(dǎo)致不穩(wěn)定問題。但是可以使用許多方法以犧牲質(zhì)量為代價(jià)來穩(wěn)定稀疏模型。例如，引入 dropout 可以提高穩(wěn)定性，但有損于模型的性能質(zhì)量。另一方面，添加更多的乘法分量可以提高模型的性能質(zhì)量，但會(huì)降低穩(wěn)定性。

ST-MoE（https://arxiv.org/abs/2202.08906）中引入的路由器 z 損失，通過懲罰進(jìn)入門控網(wǎng)絡(luò)的較大的對(duì)數(shù)值，在不降低質(zhì)量的情況下顯著提高了訓(xùn)練的穩(wěn)定性。由于這種損失會(huì)鼓勵(lì)數(shù)值變小，因此舍入誤差會(huì)減少，這對(duì)門控中的指數(shù)函數(shù)會(huì)產(chǎn)生很大影響。

專家網(wǎng)絡(luò)能學(xué)到什么？

根據(jù) ST-MoE 作者的觀察，編碼專家網(wǎng)絡(luò)專注于組別 token 或淺層概念。例如，標(biāo)點(diǎn)符號(hào)專家、專有名詞專家等。另一方面，解碼專家網(wǎng)絡(luò)的專業(yè)化程度較低。其作者還在多語言環(huán)境中進(jìn)行了訓(xùn)練。不同于想象之中，每個(gè)專家網(wǎng)絡(luò)都精通一種語言：由于 token 被路由分發(fā)和負(fù)載平衡的原因，沒有任何一種語言的專家網(wǎng)絡(luò)是專業(yè)的。

ST-MoE 論文中的表格，顯示不同的 token 組分別被發(fā)送給了哪個(gè)專家。

擴(kuò)大專家數(shù)量對(duì)預(yù)訓(xùn)練的影響

專家數(shù)量越多，采樣效率越高，速度越快，但收益也會(huì)遞減（尤其是在 256 或 512 的量級(jí)之后），推理需要更多的顯存。Switch Transformers 中研究的大規(guī)模特性在小規(guī)模中也是一致的，即使每層有 2、4 或 8 個(gè)專家也是如此。

微調(diào) MoE

Mixtral 支持 4.36.0 版本的 transformers 。使用 pip install "transformers==4.36.0 --upgrade 進(jìn)行更新。

稠密模型和稀疏模型的過擬合動(dòng)態(tài)截然不同。稀疏模型更容易出現(xiàn)過擬合，因此可以在專家本身內(nèi)部探索更高的正則化（例如，可以為稠密層設(shè)定一個(gè) dropout，為稀疏層設(shè)定另一個(gè)更高的 dropout）。

還有一個(gè)需要決策的問題是：是否使用輔助損失（auxiliary loss ）進(jìn)行微調(diào)。ST-MoE 的作者曾嘗試關(guān)閉輔助損耗，結(jié)果發(fā)現(xiàn)即使有高達(dá) 11% 的 token 被丟棄，質(zhì)量也沒有受到明顯影響。token dropping 可能是正則化的一種形式，有助于防止過擬合。

Switch Transformer 的作者觀察到，在固定的預(yù)訓(xùn)練困惑度下，稀疏模型在下游任務(wù)中的表現(xiàn)不如稠密模型，尤其是在推理任務(wù)繁重的任務(wù)中，如 SuperGLUE。另一方面，在 TriviaQA 等知識(shí)密集型語料集中，稀疏模型的表現(xiàn)卻好得出乎意料。作者還觀察到，專家較少有助于微調(diào)。對(duì)泛化問題的另一個(gè)觀察結(jié)果是，模型在較小的任務(wù)中表現(xiàn)較差，但在較大的任務(wù)中表現(xiàn)良好。

在小型任務(wù)（左圖）中，可以看到明顯的過擬合，稀疏模型在驗(yàn)證集中的表現(xiàn)要差得多。而在大型任務(wù)中（右圖），MoE 的表現(xiàn)很好。圖片來自 ST-MoE 論文。

嘗試凍結(jié)所有非專家權(quán)重導(dǎo)致了性能的大幅下降，不過這在意料之中，因?yàn)?MoE 層占據(jù)了大部分網(wǎng)絡(luò)。嘗試相反的方法：只凍結(jié) MoE 層的參數(shù)后，結(jié)果發(fā)現(xiàn)效果幾乎與更新所有參數(shù)一樣好。這個(gè)發(fā)現(xiàn)有助于加快微調(diào)速度并減少內(nèi)存占用。

通過只凍結(jié) MoE 層，可以在保證質(zhì)量的同時(shí)加快訓(xùn)練速度。本圖來自 ST-MoE 論文。

在微調(diào)稀疏化 MoE 時(shí)需要考慮的最后一個(gè)問題是，不同的 MoE 有不同的微調(diào)超參數(shù) — 例如，稀疏模型往往更受益于較小的批大小和較高的學(xué)習(xí)率。

稀疏模型的微調(diào)質(zhì)量隨著學(xué)習(xí)率的增大和 batch size 的降低而提高。本圖來自 ST-MoE 論文。

研究者們一直在努力對(duì) MoE 進(jìn)行微調(diào)，這個(gè)過程是充滿曲折的。最近的一篇論文《MoEs Meets Instruction Tuning》行了這樣的實(shí)驗(yàn)：

單任務(wù)微調(diào)
多任務(wù)指令微調(diào)
多任務(wù)指令調(diào)整后進(jìn)行單任務(wù)微調(diào)

論文中，當(dāng)作者微調(diào) MoE 和 T5 后，T5 等效模擬輸出的效果更好。當(dāng)作者對(duì) Flan T5 和 MoE 進(jìn)行微調(diào)后，MoE 的表現(xiàn)明顯更好。不僅如此，F(xiàn)lan-MoE 相對(duì)于 MoE 的改進(jìn)幅度大于 Flan T5 相對(duì)于 T5 的改進(jìn)幅度，這表明 MoE 可能比密集模型更受益于指令調(diào)優(yōu)。任務(wù)變多會(huì)讓 MoE 受益更多。與之前建議關(guān)閉輔助損失函數(shù)的結(jié)論不同，損失函數(shù)實(shí)際上可以防止過擬合。

與稠密模型相比，稀疏模型更受益于 instruct-tuning。圖片來自《MoEs Meets Instruction Tuning》論文

何時(shí)使用稀疏 MoE，何時(shí)使用稠密 MoE？

專家模型適用于使用多臺(tái)機(jī)器的高吞吐量場景。在預(yù)訓(xùn)練運(yùn)算預(yù)算固定的情況下，稀疏模型將更為理想。對(duì)于顯存較少的低吞吐量場景，稠密模型會(huì)更好。

注意：不能直接比較稀疏模型和稠密模型的參數(shù)量，因?yàn)槎咚淼囊饬x明顯不同。

高效訓(xùn)練 MoE

最初的 MoE 工作將 MoE 層作為一個(gè)分支設(shè)置，導(dǎo)致模型計(jì)算速度緩慢，因?yàn)?GPU 并非為此而設(shè)計(jì)，而且由于設(shè)備需要向其他設(shè)備發(fā)送信息，網(wǎng)絡(luò)帶寬成為瓶頸。本節(jié)將討論現(xiàn)有的一些工作，以使這些模型的預(yù)訓(xùn)練和推理更加實(shí)用。

并行計(jì)算

并行計(jì)算的種類：

數(shù)據(jù)并行：相同的權(quán)重在所有內(nèi)核上復(fù)制，數(shù)據(jù)在內(nèi)核上分割。
模型并行：模型在不同內(nèi)核間分區(qū)，數(shù)據(jù)在不同內(nèi)核間復(fù)制。
模型和數(shù)據(jù)并行：在不同內(nèi)核間劃分模型和數(shù)據(jù)。請(qǐng)注意，不同的內(nèi)核處理不同批次的數(shù)據(jù)。
專家并行：將專家放在不同的工作站上。如果與數(shù)據(jù)并行相結(jié)合，則每個(gè)內(nèi)核都有不同的專家，數(shù)據(jù)在所有內(nèi)核之間進(jìn)行分割

在專家并行模式下，專家被置于不同的工作站上，每個(gè)工作站采集不同批次的訓(xùn)練樣本。對(duì)于非 MoE 層，專家并行的行為與數(shù)據(jù)并行相同。對(duì)于 MoE 層，序列中的 token 會(huì)被發(fā)送到所需的專家所在的工作站。

Switch Transformers 論文中的插圖，顯示了數(shù)據(jù)和模型是如何通過不同的并行技術(shù)在內(nèi)核上分割的。

容量系數(shù)和通信成本

提高容量因子 (CF) 可以提高模型質(zhì)量，但會(huì)增加通信成本和內(nèi)存開銷。如果全連接通信（all-to-all communications）速度較慢，使用較小的容量因子會(huì)更好。這里提供一個(gè)可以參考的配置：使用容量系數(shù)為 1.25 的 top-2 路由機(jī)制，每個(gè)內(nèi)核留有一名專家。在驗(yàn)證過程中，可以改變?nèi)萘肯禂?shù)以減少計(jì)算量。

服務(wù)器

mistralai/Mixtral-8x7B-Instruct-v0.1 可以被部署到推理終端上。

MoE 的一大缺點(diǎn)是參數(shù)較多。本地用例可能希望使用更小的模型。以下是幾種有助于本地部署的技術(shù)：

Switch Transformer 的作者進(jìn)行過蒸餾實(shí)驗(yàn)。將 MoE 蒸餾到稠密模型中，可以保留 30-40% 的稀疏化增益。因此，蒸餾可以帶來更快的預(yù)處理速度以及模型更小的好處。
創(chuàng)新的路由算法：將完整的句子或任務(wù)路由至專家網(wǎng)絡(luò)，使用用于提取的子網(wǎng)絡(luò)來提供服務(wù)。
MoE 聚合：這種技術(shù)可以合并專家權(quán)重，從而減少推理時(shí)使用的參數(shù)。

其他高效訓(xùn)練的方法

FasterMoE（2022 年 3 月提出）分析了高效分布式系統(tǒng)中 MoE 的性能，并分析了不同并行策略的理論極限，還分析了傾斜專家受歡迎程度的技術(shù)、減少延遲的細(xì)粒度通信調(diào)度，以及根據(jù)最低延遲挑選專家的拓?fù)涓兄T控，從而使速度提高了 17 倍。

Megablocks（https://arxiv.org/abs/2211.15841）推出了一款全新的 GPU 內(nèi)核，能處理 MoE 中存在的動(dòng)態(tài)問題，探索了高效的稀疏化預(yù)訓(xùn)練。論文中建議不要丟棄任何一個(gè) token，并實(shí)現(xiàn)了高效地 token 映射技術(shù)，從而顯著提高了速度。訣竅在于，傳統(tǒng)的 MoE 使用分批矩陣乘法，假定所有專家都具有相同的形狀和相同數(shù)量的 token。相比之下，Megablocks 將 MoE 層表示為塊稀疏運(yùn)算，可以適應(yīng)不平衡分配。

針對(duì)不同大小的專家和 token 數(shù)量的塊稀疏矩陣乘法。圖片摘自 MegaBlocks 論文。

開源的 MoE 算法

當(dāng)前的 MoE 開源項(xiàng)目：

Megablocks: https://github.com/stanford-futuredata/megablocks
Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
OpenMoE: https://github.com/XueFuzhao/OpenMoE

已發(fā)布的 MoE：

Switch Transformer (Google)：從 8 到 2048 個(gè)專家的基于 T5 的模擬引擎集合。最大的模型有 1.6 萬億個(gè)參數(shù)。
NLLB MoE (Meta)：NLLB 翻譯模型的 MoE 變體。
OpenMoE：基于 Llama 的 MoE 。
Mixtral 8x7B (Mistral)：一種高質(zhì)量的翻譯模型，其性能優(yōu)于 Llama 2 70B，推理速度也更快。