混合專家模型 (MoE)核心組件和訓(xùn)練方法介紹

隨著 Mixtral 8x7B (announcement, model card) 的推出，一種稱為混合專家模型 (Mixed Expert Models，簡稱 MoEs) 的 Transformer 模型在開源人工智能社區(qū)引起了廣泛關(guān)注。在本篇博文中，我們將深入探討 MoEs 的核心組件、訓(xùn)練方法，以及在推理過程中需要考量的各種因素。

讓我們開始吧！

簡短總結(jié)

混合專家模型 (MoEs):

與稠密模型相比，預(yù)訓(xùn)練速度更快

與具有相同參數(shù)數(shù)量的模型相比，具有更快的推理速度

需要大量顯存，因?yàn)樗袑＜蚁到y(tǒng)都需要加載到內(nèi)存中

在微調(diào)方面存在諸多挑戰(zhàn)，但 近期的研究 表明，對混合專家模型進(jìn)行指令調(diào)優(yōu)具有很大的潛力。

讓我們開始吧！

什么是混合專家模型？

模型規(guī)模是提升模型性能的關(guān)鍵因素之一。在有限的計(jì)算資源預(yù)算下，用更少的訓(xùn)練步數(shù)訓(xùn)練一個(gè)更大的模型，往往比用更多的步數(shù)訓(xùn)練一個(gè)較小的模型效果更佳。

混合專家模型 (MoE) 的一個(gè)顯著優(yōu)勢是它們能夠在遠(yuǎn)少于稠密模型所需的計(jì)算資源下進(jìn)行有效的預(yù)訓(xùn)練。這意味著在相同的計(jì)算預(yù)算條件下，您可以顯著擴(kuò)大模型或數(shù)據(jù)集的規(guī)模。特別是在預(yù)訓(xùn)練階段，與稠密模型相比，混合專家模型通常能夠更快地達(dá)到相同的質(zhì)量水平。

那么，究竟什么是一個(gè)混合專家模型 (MoE) 呢？作為一種基于 Transformer 架構(gòu)的模型，混合專家模型主要由兩個(gè)關(guān)鍵部分組成:

稀疏 MoE 層: 這些層代替了傳統(tǒng) Transformer 模型中的前饋網(wǎng)絡(luò) (FFN) 層。MoE 層包含若干“專家”(例如 8 個(gè))，每個(gè)專家本身是一個(gè)獨(dú)立的神經(jīng)網(wǎng)絡(luò)。在實(shí)際應(yīng)用中，這些專家通常是前饋網(wǎng)絡(luò) (FFN)，但它們也可以是更復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)，甚至可以是 MoE 層本身，從而形成層級式的 MoE 結(jié)構(gòu)。

門控網(wǎng)絡(luò)或路由: 這個(gè)部分用于決定哪些令牌 (token) 被發(fā)送到哪個(gè)專家。例如，在下圖中，“More”這個(gè)令牌可能被發(fā)送到第二個(gè)專家，而“Parameters”這個(gè)令牌被發(fā)送到第一個(gè)專家。有時(shí)，一個(gè)令牌甚至可以被發(fā)送到多個(gè)專家。令牌的路由方式是 MoE 使用中的一個(gè)關(guān)鍵點(diǎn)，因?yàn)槁酚善饔蓪W(xué)習(xí)的參數(shù)組成，并且與網(wǎng)絡(luò)的其他部分一同進(jìn)行預(yù)訓(xùn)練。

Switch Transformers paper 論文中的 MoE layer

總結(jié)來說，在混合專家模型 (MoE) 中，我們將傳統(tǒng) Transformer 模型中的每個(gè)前饋網(wǎng)絡(luò) (FFN) 層替換為 MoE 層，其中 MoE 層由兩個(gè)核心部分組成: 一個(gè)門控網(wǎng)絡(luò)和若干數(shù)量的專家。

盡管混合專家模型 (MoE) 提供了若干顯著優(yōu)勢，例如更高效的預(yù)訓(xùn)練和與稠密模型相比更快的推理速度，但它們也伴隨著一些挑戰(zhàn):

訓(xùn)練挑戰(zhàn): 雖然 MoE 能夠?qū)崿F(xiàn)更高效的計(jì)算預(yù)訓(xùn)練，但它們在微調(diào)階段往往面臨泛化能力不足的問題，長期以來易于引發(fā)過擬合現(xiàn)象。

推理挑戰(zhàn): MoE 模型雖然可能擁有大量參數(shù)，但在推理過程中只使用其中的一部分，這使得它們的推理速度快于具有相同數(shù)量參數(shù)的稠密模型。然而，這種模型需要將所有參數(shù)加載到內(nèi)存中，因此對內(nèi)存的需求非常高。以 Mixtral 8x7B 這樣的 MoE 為例，需要足夠的 VRAM 來容納一個(gè) 47B 參數(shù)的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因?yàn)樵?MoE 模型中，只有 FFN 層被視為獨(dú)立的專家，而模型的其他參數(shù)是共享的。此外，假設(shè)每個(gè)令牌只使用兩個(gè)專家，那么推理速度 (以 FLOPs 計(jì)算) 類似于使用 12B 模型 (而不是 14B 模型)，因?yàn)殡m然它進(jìn)行了 2x7B 的矩陣乘法計(jì)算，但某些層是共享的。

了解了 MoE 的基本概念后，讓我們進(jìn)一步探索推動(dòng)這類模型發(fā)展的研究。

混合專家模型簡史

混合專家模型 (MoE) 的理念起源于 1991 年的論文 Adaptive Mixture of Local Experts。這個(gè)概念與集成學(xué)習(xí)方法相似，旨在為由多個(gè)單獨(dú)網(wǎng)絡(luò)組成的系統(tǒng)建立一個(gè)監(jiān)管機(jī)制。在這種系統(tǒng)中，每個(gè)網(wǎng)絡(luò) (被稱為“專家”) 處理訓(xùn)練樣本的不同子集，專注于輸入空間的特定區(qū)域。那么，如何選擇哪個(gè)專家來處理特定的輸入呢？這就是門控網(wǎng)絡(luò)發(fā)揮作用的地方，它決定了分配給每個(gè)專家的權(quán)重。在訓(xùn)練過程中，這些專家和門控網(wǎng)絡(luò)都同時(shí)接受訓(xùn)練，以優(yōu)化它們的性能和決策能力。

在 2010 至 2015 年間，兩個(gè)獨(dú)立的研究領(lǐng)域?yàn)榛旌蠈＜夷Ｐ?(MoE) 的后續(xù)發(fā)展做出了顯著貢獻(xiàn):

組件專家: 在傳統(tǒng)的 MoE 設(shè)置中，整個(gè)系統(tǒng)由一個(gè)門控網(wǎng)絡(luò)和多個(gè)專家組成。在支持向量機(jī) (SVMs) 、高斯過程和其他方法的研究中，MoE 通常被視為整個(gè)模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了將 MoE 作為更深層網(wǎng)絡(luò)的一個(gè)組件。這種方法允許將 MoE 嵌入到多層網(wǎng)絡(luò)中的某一層，使得模型既大又高效。

條件計(jì)算: 傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)通過每一層處理所有輸入數(shù)據(jù)。在這一時(shí)期，Yoshua Bengio 等研究人員開始探索基于輸入令牌動(dòng)態(tài)激活或停用網(wǎng)絡(luò)組件的方法。

這些研究的融合促進(jìn)了在自然語言處理 (NLP) 領(lǐng)域?qū)旌蠈＜夷Ｐ偷奶剿?。特別是在 2017 年，Shazeer 等人 (團(tuán)隊(duì)包括 Geoffrey Hinton 和 Jeff Dean，后者有時(shí)被戲稱為 “谷歌的 Chuck Norris”) 將這一概念應(yīng)用于 137B 的 LSTM (當(dāng)時(shí)被廣泛應(yīng)用于 NLP 的架構(gòu)，由 Schmidhuber 提出)。通過引入稀疏性，這項(xiàng)工作在保持極高規(guī)模的同時(shí)實(shí)現(xiàn)了快速的推理速度。這項(xiàng)工作主要集中在翻譯領(lǐng)域，但面臨著如高通信成本和訓(xùn)練不穩(wěn)定性等多種挑戰(zhàn)。

Outrageously Large Neural Network 論文中的 MoE layer

混合專家模型 (MoE) 的引入使得訓(xùn)練具有數(shù)千億甚至萬億參數(shù)的模型成為可能，如開源的 1.6 萬億參數(shù)的 Switch Transformers 等。這種技術(shù)不僅在自然語言處理 (NLP) 領(lǐng)域得到了廣泛應(yīng)用，也開始在計(jì)算機(jī)視覺領(lǐng)域進(jìn)行探索。然而，本篇博客文章將主要聚焦于自然語言處理領(lǐng)域的應(yīng)用和探討。

什么是稀疏性?

稀疏性的概念采用了條件計(jì)算的思想。在傳統(tǒng)的稠密模型中，所有的參數(shù)都會(huì)對所有輸入數(shù)據(jù)進(jìn)行處理。相比之下，稀疏性允許我們僅針對整個(gè)系統(tǒng)的某些特定部分執(zhí)行計(jì)算。這意味著并非所有參數(shù)都會(huì)在處理每個(gè)輸入時(shí)被激活或使用，而是根據(jù)輸入的特定特征或需求，只有部分參數(shù)集合被調(diào)用和運(yùn)行。

讓我們深入分析 Shazeer 對混合專家模型 (MoE) 在翻譯應(yīng)用中的貢獻(xiàn)。條件計(jì)算的概念 (即僅在每個(gè)樣本的基礎(chǔ)上激活網(wǎng)絡(luò)的不同部分) 使得在不增加額外計(jì)算負(fù)擔(dān)的情況下擴(kuò)展模型規(guī)模成為可能。這一策略在每個(gè) MoE 層中實(shí)現(xiàn)了數(shù)以千計(jì)甚至更多的專家的有效利用。

這種稀疏性設(shè)置確實(shí)帶來了一些挑戰(zhàn)。例如，在混合專家模型 (MoE) 中，盡管較大的批量大小通常有利于提高性能，但當(dāng)數(shù)據(jù)通過激活的專家時(shí)，實(shí)際的批量大小可能會(huì)減少。比如，假設(shè)我們的輸入批量包含 10 個(gè)令牌，可能會(huì)有五個(gè)令牌被路由到同一個(gè)專家，而剩下的五個(gè)令牌分別被路由到不同的專家。這導(dǎo)致了批量大小的不均勻分配和資源利用效率不高的問題。在接下來的部分中，將會(huì)討論讓 MoE 高效運(yùn)行的其他挑戰(zhàn)以及相應(yīng)的解決方案。

那我們應(yīng)該如何解決這個(gè)問題呢？一個(gè)可學(xué)習(xí)的門控網(wǎng)絡(luò) (G) 決定將輸入的哪一部分發(fā)送給哪些專家 (E):

在這種設(shè)置下，雖然所有專家都會(huì)對所有輸入進(jìn)行運(yùn)算，但通過門控網(wǎng)絡(luò)的輸出進(jìn)行加權(quán)乘法操作。但是，如果 G (門控網(wǎng)絡(luò)的輸出) 為 0 會(huì)發(fā)生什么呢？如果是這種情況，就沒有必要計(jì)算相應(yīng)的專家操作，因此我們可以節(jié)省計(jì)算資源。那么一個(gè)典型的門控函數(shù)是什么呢？一個(gè)典型的門控函數(shù)通常是一個(gè)帶有 softmax 函數(shù)的簡單的網(wǎng)絡(luò)。這個(gè)網(wǎng)絡(luò)將學(xué)習(xí)將輸入發(fā)送給哪個(gè)專家。

Shazeer 等人的工作還探索了其他的門控機(jī)制，其中包括帶噪聲的 TopK 門控 (Noisy Top-K Gating)。這種門控方法引入了一些可調(diào)整的噪聲，然后保留前 k 個(gè)值。具體來說:

添加一些噪聲

選擇保留前 K 個(gè)值

應(yīng)用 Softmax 函數(shù)

這種稀疏性引入了一些有趣的特性。通過使用較低的 k 值 (例如 1 或 2)，我們可以比激活多個(gè)專家時(shí)更快地進(jìn)行訓(xùn)練和推理。為什么不僅選擇最頂尖的專家呢？最初的假設(shè)是，需要將輸入路由到不止一個(gè)專家，以便門控學(xué)會(huì)如何進(jìn)行有效的路由選擇，因此至少需要選擇兩個(gè)專家。Switch Transformers 就這點(diǎn)進(jìn)行了更多的研究。

我們?yōu)槭裁匆砑釉肼暷?？這是為了專家間的負(fù)載均衡！

混合專家模型中令牌的負(fù)載均衡

正如之前討論的，如果所有的令牌都被發(fā)送到只有少數(shù)幾個(gè)受歡迎的專家，那么訓(xùn)練效率將會(huì)降低。在通常的混合專家模型 (MoE) 訓(xùn)練中，門控網(wǎng)絡(luò)往往傾向于主要激活相同的幾個(gè)專家。這種情況可能會(huì)自我加強(qiáng)，因?yàn)槭軞g迎的專家訓(xùn)練得更快，因此它們更容易被選擇。為了緩解這個(gè)問題，引入了一個(gè)輔助損失，旨在鼓勵(lì)給予所有專家相同的重要性。這個(gè)損失確保所有專家接收到大致相等數(shù)量的訓(xùn)練樣本，從而平衡了專家之間的選擇。接下來的部分還將探討專家容量的概念，它引入了一個(gè)關(guān)于專家可以處理多少令牌的閾值。在 transformers 庫中，可以通過 aux_loss 參數(shù)來控制輔助損失。

MoEs and Transformers

Transformer 類模型明確表明，增加參數(shù)數(shù)量可以提高性能，因此谷歌使用 GShard 嘗試將 Transformer 模型的參數(shù)量擴(kuò)展到超過 6000 億并不令人驚訝。

GShard 將在編碼器和解碼器中的每個(gè)前饋網(wǎng)絡(luò) (FFN) 層中的替換為使用 Top-2 門控的混合專家模型 (MoE) 層。下圖展示了編碼器部分的結(jié)構(gòu)。這種架構(gòu)對于大規(guī)模計(jì)算非常有效: 當(dāng)擴(kuò)展到多個(gè)設(shè)備時(shí)，MoE 層在不同設(shè)備間共享，而其他所有層則在每個(gè)設(shè)備上復(fù)制。我們將在 “讓 MoE 起飛”部分對這一點(diǎn)進(jìn)行更詳細(xì)的討論。

GShard 論文中的 MoE Transformer Encoder

為了保持負(fù)載平衡和訓(xùn)練效率，GShard 的作者除了引入了上一節(jié)中討論的類似輔助損失外，還引入了一些關(guān)鍵變化:

隨機(jī)路由: 在 Top-2 設(shè)置中，我們始終選擇排名最高的專家，但第二個(gè)專家是根據(jù)其權(quán)重比例隨機(jī)選擇的。

專家容量: 我們可以設(shè)定一個(gè)閾值，定義一個(gè)專家能處理多少令牌。如果兩個(gè)專家的容量都達(dá)到上限，令牌就會(huì)溢出，并通過殘差連接傳遞到下一層，或在某些情況下被完全丟棄。專家容量是 MoE 中最重要的概念之一。為什么需要專家容量呢？因?yàn)樗袕埩康男螤钤诰幾g時(shí)是靜態(tài)確定的，我們無法提前知道多少令牌會(huì)分配給每個(gè)專家，因此需要一個(gè)固定的容量因子。

GShard 的工作對適用于 MoE 的并行計(jì)算模式也做出了重要貢獻(xiàn)，但這些內(nèi)容的討論超出了這篇博客的范圍。

注意: 在推理過程中，只有部分專家被激活。同時(shí)，有些計(jì)算過程是共享的，例如自注意力 (self-attention) 機(jī)制，它適用于所有令牌。這就解釋了為什么我們可以使用相當(dāng)于 12B 稠密模型的計(jì)算資源來運(yùn)行一個(gè)包含 8 個(gè)專家的 47B 模型。如果我們采用 Top-2 門控，模型會(huì)使用高達(dá) 14B 的參數(shù)。但是，由于自注意力操作 (專家間共享) 的存在，實(shí)際上模型運(yùn)行時(shí)使用的參數(shù)數(shù)量是 12B。

Switch Transformers

盡管混合專家模型 (MoE) 顯示出了很大的潛力，但它們在訓(xùn)練和微調(diào)過程中存在穩(wěn)定性問題。Switch Transformers 是一項(xiàng)非常激動(dòng)人心的工作，它深入研究了這些話題。作者甚至在 Hugging Face 上發(fā)布了一個(gè) 1.6 萬億參數(shù)的 MoE，擁有 2048 個(gè)專家，你可以使用 transformers 庫來運(yùn)行它。Switch Transformers 實(shí)現(xiàn)了與 T5-XXL 相比 4 倍的預(yù)訓(xùn)練速度提升。

Switch Transformer 論文中的 Switch Transformer Layer

就像在 GShard 中一樣，作者用混合專家模型 (MoE) 層替換了前饋網(wǎng)絡(luò) (FFN) 層。Switch Transformers 提出了一個(gè) Switch Transformer 層，它接收兩個(gè)輸入 (兩個(gè)不同的令牌) 并擁有四個(gè)專家。

與最初使用至少兩個(gè)專家的想法相反，Switch Transformers 采用了簡化的單專家策略。這種方法的效果包括:

減少門控網(wǎng)絡(luò) (路由) 計(jì)算負(fù)擔(dān)

每個(gè)專家的批量大小至少可以減半

降低通信成本

保持模型質(zhì)量

Switch Transformers 也對專家容量這個(gè)概念進(jìn)行了研究。

上述建議的容量是將批次中的令牌數(shù)量均勻分配到各個(gè)專家。如果我們使用大于 1 的容量因子，我們?yōu)榱钆品峙洳煌耆胶鈺r(shí)提供了一個(gè)緩沖。增加容量因子會(huì)導(dǎo)致更高的設(shè)備間通信成本，因此這是一個(gè)需要考慮的權(quán)衡。特別值得注意的是，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表現(xiàn)出色。

Switch Transformer 的作者還重新審視并簡化了前面章節(jié)中提到的負(fù)載均衡損失。在訓(xùn)練期間，對于每個(gè) Switch 層的輔助損失被添加到總模型損失中。這種損失鼓勵(lì)均勻路由，并可以使用超參數(shù)進(jìn)行加權(quán)。

作者還嘗試了混合精度的方法，例如用 bfloat16 精度訓(xùn)練專家，同時(shí)對其余計(jì)算使用全精度進(jìn)行。較低的精度可以減少處理器間的通信成本、計(jì)算成本以及存儲張量的內(nèi)存。然而，在最初的實(shí)驗(yàn)中，當(dāng)專家和門控網(wǎng)絡(luò)都使用 bfloat16 精度訓(xùn)練時(shí)，出現(xiàn)了不穩(wěn)定的訓(xùn)練現(xiàn)象。這種不穩(wěn)定性特別是由路由計(jì)算引起的，因?yàn)槁酚缮婕爸笖?shù)函數(shù)等操作，這些操作對精度要求較高。因此，為了保持計(jì)算的穩(wěn)定性和精確性，保持更高的精度是重要的。為了減輕不穩(wěn)定性，路由過程也使用了全精度。

使用混合精度不會(huì)降低模型質(zhì)量并可實(shí)現(xiàn)更快的訓(xùn)練

這個(gè) Jupyter Notebook 展示了如何對 Switch Transformers 進(jìn)行微調(diào)以進(jìn)行摘要生成的詳細(xì)指南。然而，在開始微調(diào) Switch Transformers 之前，強(qiáng)烈建議您先閱讀關(guān)于微調(diào)混合專家模型部分的內(nèi)容。

Switch Transformers 采用了編碼器 - 解碼器的架構(gòu)，實(shí)現(xiàn)了與 T5 類似的混合專家模型 (MoE) 版本。GLaM 這篇工作探索了如何使用僅為原來 1/3 的計(jì)算資源 (因?yàn)?MoE 模型在訓(xùn)練時(shí)需要的計(jì)算量較少，從而能夠顯著降低碳足跡) 來訓(xùn)練與 GPT-3 質(zhì)量相匹配的模型來提高這些模型的規(guī)模。作者專注于僅解碼器 (decoder-only) 的模型以及少樣本和單樣本評估，而不是微調(diào)。他們使用了 Top-2 路由和更大的容量因子。此外，他們探討了將容量因子作為一個(gè)動(dòng)態(tài)度量，根據(jù)訓(xùn)練和評估期間所使用的計(jì)算量進(jìn)行調(diào)整。

用 Router z-loss 穩(wěn)定模型訓(xùn)練

之前討論的平衡損失可能會(huì)導(dǎo)致穩(wěn)定性問題。我們可以使用許多方法來穩(wěn)定稀疏模型的訓(xùn)練，但這可能會(huì)犧牲模型質(zhì)量。例如，引入 dropout 可以提高穩(wěn)定性，但會(huì)導(dǎo)致模型質(zhì)量下降。另一方面，增加更多的乘法分量可以提高質(zhì)量，但會(huì)降低模型穩(wěn)定性。

ST-MoE 引入的 Router z-loss 在保持了模型性能的同時(shí)顯著提升了訓(xùn)練的穩(wěn)定性。這種損失機(jī)制通過懲罰門控網(wǎng)絡(luò)輸入的較大 logits 來起作用，目的是促使數(shù)值的絕對大小保持較小，這樣可以有效減少計(jì)算中的舍入誤差。這一點(diǎn)對于那些依賴指數(shù)函數(shù)進(jìn)行計(jì)算的門控網(wǎng)絡(luò)尤其重要。為了深入了解這一機(jī)制，建議參考原始論文以獲得更全面的細(xì)節(jié)。

專家如何學(xué)習(xí)？

ST-MoE 的研究者們發(fā)現(xiàn)，編碼器中不同的專家傾向于專注于特定類型的令牌或淺層概念。例如，某些專家可能專門處理標(biāo)點(diǎn)符號，而其他專家則專注于專有名詞等。與此相反，解碼器中的專家通常具有較低的專業(yè)化程度。此外，研究者們還對這一模型進(jìn)行了多語言訓(xùn)練。盡管人們可能會(huì)預(yù)期每個(gè)專家處理一種特定語言，但實(shí)際上并非如此。由于令牌路由和負(fù)載均衡的機(jī)制，沒有任何專家被特定配置以專門處理某一特定語言。

ST-MoE 論文中顯示了哪些令牌組被發(fā)送給了哪個(gè)專家的表格

專家的數(shù)量對預(yù)訓(xùn)練有何影響？

增加更多專家可以提升處理樣本的效率和加速模型的運(yùn)算速度，但這些優(yōu)勢隨著專家數(shù)量的增加而遞減 (尤其是當(dāng)專家數(shù)量達(dá)到 256 或 512 之后更為明顯)。同時(shí)，這也意味著在推理過程中，需要更多的顯存來加載整個(gè)模型。值得注意的是，Switch Transformers 的研究表明，其在大規(guī)模模型中的特性在小規(guī)模模型下也同樣適用，即便是每層僅包含 2、4 或 8 個(gè)專家。

微調(diào)混合專家模型

4.36.0 版本的 transformers 庫支持 Mixtral 模型。你可以用以下命令進(jìn)行安裝: pip install "transformers==4.36.0 --upgrade

稠密模型和稀疏模型在過擬合的動(dòng)態(tài)表現(xiàn)上存在顯著差異。稀疏模型更易于出現(xiàn)過擬合現(xiàn)象，因此在處理這些模型時(shí)，嘗試更強(qiáng)的內(nèi)部正則化措施是有益的，比如使用更高比例的 dropout。例如，我們可以為稠密層設(shè)定一個(gè)較低的 dropout 率，而為稀疏層設(shè)置一個(gè)更高的 dropout 率，以此來優(yōu)化模型性能。

在微調(diào)過程中是否使用輔助損失是一個(gè)需要決策的問題。ST-MoE 的作者嘗試關(guān)閉輔助損失，發(fā)現(xiàn)即使高達(dá) 11% 的令牌被丟棄，模型的質(zhì)量也沒有顯著受到影響。令牌丟棄可能是一種正則化形式，有助于防止過擬合。

Switch Transformers 的作者觀察到，在相同的預(yù)訓(xùn)練困惑度下，稀疏模型在下游任務(wù)中的表現(xiàn)不如對應(yīng)的稠密模型，特別是在重理解任務(wù) (如 SuperGLUE) 上。另一方面，對于知識密集型任務(wù) (如 TriviaQA)，稀疏模型的表現(xiàn)異常出色。作者還觀察到，在微調(diào)過程中，較少的專家的數(shù)量有助于改善性能。另一個(gè)關(guān)于泛化問題確認(rèn)的發(fā)現(xiàn)是，模型在小型任務(wù)上表現(xiàn)較差，但在大型任務(wù)上表現(xiàn)良好。

在小任務(wù) (左圖) 中，我們可以看到明顯的過擬合，因?yàn)橄∈枘Ｐ驮隍?yàn)證集中的表現(xiàn)要差得多。在較大的任務(wù) (右圖) 中，MoE 則表現(xiàn)良好。該圖來自 ST-MoE 論文

一種可行的微調(diào)策略是嘗試凍結(jié)所有非專家層的權(quán)重。實(shí)踐中，這會(huì)導(dǎo)致性能大幅下降，但這符合我們的預(yù)期，因?yàn)榛旌蠈＜夷Ｐ?(MoE) 層占據(jù)了網(wǎng)絡(luò)的主要部分。我們可以嘗試相反的方法: 僅凍結(jié) MoE 層的參數(shù)。實(shí)驗(yàn)結(jié)果顯示，這種方法幾乎與更新所有參數(shù)的效果相當(dāng)。這種做法可以加速微調(diào)過程，并降低顯存需求。

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

在微調(diào)稀疏混合專家模型 (MoE) 時(shí)需要考慮的最后一個(gè)問題是，它們有特別的微調(diào)超參數(shù)設(shè)置——例如，稀疏模型往往更適合使用較小的批量大小和較高的學(xué)習(xí)率，這樣可以獲得更好的訓(xùn)練效果。

降低學(xué)習(xí)率和調(diào)大批量可以提升稀疏模型微調(diào)質(zhì)量。該圖來自 ST-MoE 論文

此時(shí)，您可能會(huì)對人們微調(diào) MoE 中遇到的這些挑戰(zhàn)而感到沮喪，但最近的一篇論文 《MoEs Meets Instruction Tuning》 (2023 年 7 月) 帶來了令人興奮的發(fā)現(xiàn)。這篇論文進(jìn)行了以下實(shí)驗(yàn):

單任務(wù)微調(diào)

多任務(wù)指令微調(diào)

多任務(wù)指令微調(diào)后接單任務(wù)微調(diào)

當(dāng)研究者們對 MoE 和對應(yīng)性能相當(dāng)?shù)?T5 模型進(jìn)行微調(diào)時(shí)，他們發(fā)現(xiàn) T5 的對應(yīng)模型表現(xiàn)更為出色。然而，當(dāng)研究者們對 Flan T5 (一種 T5 的指令優(yōu)化版本) 的 MoE 版本進(jìn)行微調(diào)時(shí)，MoE 的性能顯著提升。更值得注意的是，F(xiàn)lan-MoE 相比原始 MoE 的性能提升幅度超過了 Flan T5 相對于原始 T5 的提升，這意味著 MoE 模型可能從指令式微調(diào)中獲益更多，甚至超過了稠密模型。此外，MoE 在多任務(wù)學(xué)習(xí)中表現(xiàn)更佳。與之前關(guān)閉輔助損失函數(shù)的做法相反，實(shí)際上這種損失函數(shù)可以幫助防止過擬合。

與稠密模型相比，稀疏模型從指令微調(diào)中受益更多。該圖來自 MoEs Meets instructions Tuning 論文

稀疏 VS 稠密，如何選擇?

稀疏混合專家模型 (MoE) 適用于擁有多臺機(jī)器且要求高吞吐量的場景。在固定的預(yù)訓(xùn)練計(jì)算資源下，稀疏模型往往能夠?qū)崿F(xiàn)更優(yōu)的效果。相反，在顯存較少且吞吐量要求不高的場景，稠密模型則是更合適的選擇。

注意: 直接比較稀疏模型和稠密模型的參數(shù)數(shù)量是不恰當(dāng)?shù)模驗(yàn)檫@兩類模型基于的概念和參數(shù)量的計(jì)算方法完全不同。

讓 MoE 起飛

最初的混合專家模型 (MoE) 設(shè)計(jì)采用了分支結(jié)構(gòu)，這導(dǎo)致了計(jì)算效率低下。這種低效主要是因?yàn)?GPU 并不是為處理這種結(jié)構(gòu)而設(shè)計(jì)的，而且由于設(shè)備間需要傳遞數(shù)據(jù)，網(wǎng)絡(luò)帶寬常常成為性能瓶頸。在接下來的討論中，我們會(huì)討論一些現(xiàn)有的研究成果，旨在使這些模型在預(yù)訓(xùn)練和推理階段更加高效和實(shí)用。我們來看看如何優(yōu)化 MoE 模型，讓 MoE 起飛。

并行計(jì)算

讓我們簡要回顧一下并行計(jì)算的幾種形式:

數(shù)據(jù)并行: 相同的權(quán)重在所有節(jié)點(diǎn)上復(fù)制，數(shù)據(jù)在節(jié)點(diǎn)之間分割。

模型并行: 模型在節(jié)點(diǎn)之間分割，相同的數(shù)據(jù)在所有節(jié)點(diǎn)上復(fù)制。

模型和數(shù)據(jù)并行: 我們可以在節(jié)點(diǎn)之間同時(shí)分割模型和數(shù)據(jù)。注意，不同的節(jié)點(diǎn)處理不同批次的數(shù)據(jù)。

專家并行: 專家被放置在不同的節(jié)點(diǎn)上。如果與數(shù)據(jù)并行結(jié)合，每個(gè)節(jié)點(diǎn)擁有不同的專家，數(shù)據(jù)在所有節(jié)點(diǎn)之間分割。

在專家并行中，專家被放置在不同的節(jié)點(diǎn)上，每個(gè)節(jié)點(diǎn)處理不同批次的訓(xùn)練樣本。對于非 MoE 層，專家并行的行為與數(shù)據(jù)并行相同。對于 MoE 層，序列中的令牌被發(fā)送到擁有所需專家的節(jié)點(diǎn)。

Switch Transformers 論文中展示如何使用不同的并行技術(shù)在節(jié)點(diǎn)上分割數(shù)據(jù)和模型的插圖

容量因子和通信開銷

提高容量因子 (Capacity Factor, CF) 可以增強(qiáng)模型的性能，但這也意味著更高的通信成本和對保存激活值的顯存的需求。在設(shè)備通信帶寬有限的情況下，選擇較小的容量因子可能是更佳的策略。一個(gè)合理的初始設(shè)置是采用 Top-2 路由、1.25 的容量因子，同時(shí)每個(gè)節(jié)點(diǎn)配置一個(gè)專家。在評估性能時(shí)，應(yīng)根據(jù)需要調(diào)整容量因子，以在設(shè)備間的通信成本和計(jì)算成本之間找到一個(gè)平衡點(diǎn)。

部署技術(shù)

您可以在 Inference Endpoints 部署 mistralai/Mixtral-8x7B-Instruct-v0.1。

部署混合專家模型 (MoE) 的一個(gè)關(guān)鍵挑戰(zhàn)是其龐大的參數(shù)規(guī)模。對于本地使用情況，我們可能希望使用更小的模型。為了使模型更適合部署，下面是幾種有用的技術(shù):

預(yù)先蒸餾實(shí)驗(yàn): Switch Transformers 的研究者們進(jìn)行了預(yù)先蒸餾的實(shí)驗(yàn)。他們通過將 MoE 模型蒸餾回其對應(yīng)的稠密模型，成功保留了 30-40%的由稀疏性帶來的性能提升。預(yù)先蒸餾不僅加快了預(yù)訓(xùn)練速度，還使得在推理中使用更小型的模型成為可能。

任務(wù)級別路由: 最新的方法中，路由器被修改為將整個(gè)句子或任務(wù)直接路由到一個(gè)專家。這樣做可以提取出一個(gè)用于服務(wù)的子網(wǎng)絡(luò)，有助于簡化模型的結(jié)構(gòu)。

專家網(wǎng)絡(luò)聚合: 這項(xiàng)技術(shù)通過合并各個(gè)專家的權(quán)重，在推理時(shí)減少了所需的參數(shù)數(shù)量。這樣可以在不顯著犧牲性能的情況下降低模型的復(fù)雜度。

高效訓(xùn)練

FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同并行策略下的理論性能極限，并且探索了一系列創(chuàng)新技術(shù)，包括用于專家權(quán)重調(diào)整的方法、減少延遲的細(xì)粒度通信調(diào)度技術(shù)，以及一個(gè)基于最低延遲進(jìn)行專家選擇的拓?fù)涓兄T控機(jī)制。這些技術(shù)的結(jié)合使得 MoE 運(yùn)行速度提升高達(dá) 17 倍。

Megablocks (2022 年 11 月) 則專注于通過開發(fā)新的 GPU kernel 來處理 MoE 模型中的動(dòng)態(tài)性，以實(shí)現(xiàn)更高效的稀疏預(yù)訓(xùn)練。其核心優(yōu)勢在于，它不會(huì)丟棄任何令牌，并能高效地適應(yīng)現(xiàn)代硬件架構(gòu) (支持塊稀疏矩陣乘)，從而達(dá)到顯著的加速效果。Megablocks 的創(chuàng)新之處在于，它不像傳統(tǒng) MoE 那樣使用批量矩陣乘法 (這通常假設(shè)所有專家形狀相同且處理相同數(shù)量的令牌)，而是將 MoE 層表示為塊稀疏操作，可以靈活適應(yīng)不均衡的令牌分配。

針對不同規(guī)模的專家和令牌數(shù)量的塊稀疏矩陣乘法。該圖來自 MegaBlocks 論文

開源混合專家模型

目前，下面這些開源項(xiàng)目可以用于訓(xùn)練混合專家模型 (MoE):

Megablocks: https://github.com/stanford-futuredata/megablocks

Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm

OpenMoE: https://github.com/XueFuzhao/OpenMoE

對于開源的混合專家模型 (MoE)，你可以關(guān)注下面這些:

Switch Transformers (Google): 基于 T5 的 MoE 集合，專家數(shù)量從 8 名到 2048 名。最大的模型有 1.6 萬億個(gè)參數(shù)。

NLLB MoE (Meta): NLLB 翻譯模型的一個(gè) MoE 變體。

OpenMoE: 社區(qū)對基于 Llama 的模型的 MoE 嘗試。

Mixtral 8x7B (Mistral): 一個(gè)性能超越了 Llama 2 70B 的高質(zhì)量混合專家模型，并且具有更快的推理速度。此外，還發(fā)布了一個(gè)經(jīng)過指令微調(diào)的模型。有關(guān)更多信息，可以在 Mistral 的 公告博客文章 中了解。

一些有趣的研究方向

首先是嘗試將稀疏混合專家模型 (SMoE)蒸餾回到具有更少實(shí)際參數(shù)但相似等價(jià)參數(shù)量的稠密模型。

MoE 的量化也是一個(gè)有趣的研究領(lǐng)域。例如，QMoE (2023 年 10 月) 通過將 MoE 量化到每個(gè)參數(shù)不到 1 位，將 1.6 萬億參數(shù)的 Switch Transformer 所需的存儲從 3.2TB 壓縮到僅 160GB。

簡而言之，一些值得探索的有趣領(lǐng)域包括:

將 Mixtral 蒸餾成一個(gè)稠密模型。

探索合并專家模型的技術(shù)及其對推理時(shí)間的影響。

嘗試對 Mixtral 進(jìn)行極端量化的實(shí)驗(yàn)。