專家混合 (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ì)算開銷對模型進(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 中，因此對內(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ì)介紹）。

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

這些工作促使研究者們在 NLP 的背景下探索混合專家模型。具體來說，Shazeer 及 Geoffrey Hinton 、Jeff Dean，谷歌的 Chuck Norris 通過對引入稀疏網(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ù)萬億級參數(shù)的模型，例如開源的 1.6T 參數(shù)的 Switch Transformer。計(jì)算機(jī)視覺領(lǐng)域也在探索 MoE，但這里還是先側(cè)重講解 NLP 領(lǐng)域。

什么是稀疏化？

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

前文提到 Shazeer 對機(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ì)對所有輸入進(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è)置對大規(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。這就是為什么對于一個(gè)由 8 個(gè)專家組成的 47B 模型時(shí)，GShard 可以用一個(gè) 12B 的稠密模型進(jìn)行計(jì)算。如果使用 top-2，則需要 14B 個(gè)參數(shù)。但考慮到注意力操作是共享的，實(shí)際使用的參數(shù)數(shù)量為 12B。

Switch Transformer

盡管 MoE 前景廣闊，但它們在訓(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ù)載平衡損耗。對于每個(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 對應(yīng)。GLaM （https://arxiv.org/abs/2112.06905）探討了如何用三分之一的計(jì)算開銷訓(xùn)練出與 GPT-3 質(zhì)量相匹配的模型，從而擴(kuò)大模型的規(guī)模。GLaM 作者的研究重點(diǎn)是純解碼器模型以及小樣本及零樣本的評估結(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ò)的較大的對數(shù)值，在不降低質(zhì)量的情況下顯著提高了訓(xùn)練的穩(wěn)定性。由于這種損失會(huì)鼓勵(lì)數(shù)值變小，因此舍入誤差會(huì)減少，這對門控中的指數(shù)函數(shù)會(huì)產(chǎn)生很大影響。

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

根據(jù) ST-MoE 作者的觀察，編碼專家網(wǎng)絡(luò)專注于組別 token 或淺層概念。例如，標(biāo)點(diǎn)符號專家、專有名詞專家等。另一方面，解碼專家網(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ù)量對預(yù)訓(xùn)練的影響

專家數(shù)量越多，采樣效率越高，速度越快，但收益也會(huì)遞減（尤其是在 256 或 512 的量級之后），推理需要更多的顯存。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)。對泛化問題的另一個(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 論文。

研究者們一直在努力對 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)作者對 Flan T5 和 MoE 進(jìn)行微調(diào)后，MoE 的表現(xiàn)明顯更好。不僅如此，F(xiàn)lan-MoE 相對于 MoE 的改進(jìn)幅度大于 Flan T5 相對于 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ù)算固定的情況下，稀疏模型將更為理想。對于顯存較少的低吞吐量場景，稠密模型會(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ù)。請注意，不同的內(nèi)核處理不同批次的數(shù)據(jù)。
• 專家并行：將專家放在不同的工作站上。如果與數(shù)據(jù)并行相結(jié)合，則每個(gè)內(nèi)核都有不同的專家，數(shù)據(jù)在所有內(nèi)核之間進(jìn)行分割

在專家并行模式下，專家被置于不同的工作站上，每個(gè)工作站采集不同批次的訓(xùn)練樣本。對于非 MoE 層，專家并行的行為與數(shù)據(jù)并行相同。對于 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)不平衡分配。

針對不同大小的專家和 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，推理速度也更快。

未來展望

以下是一些值得探索的有趣的領(lǐng)域：

• 將 Mixtral 提煉為稠密模型
• 探索專家模型的合并技術(shù)，分析其對推理時(shí)間的影響
• Mixtral 的量化技術(shù)