混合專家模型 (MoE) 詳解

HuggingFace發表於2023-12-26
模型

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

讓我們開始吧！

簡短總結

混合專家模型 (MoEs):

  • 與稠密模型相比， 預訓練速度更快
  • 與具有相同引數數量的模型相比，具有更快的 推理速度
  • 需要 大量視訊記憶體，因為所有專家系統都需要載入到記憶體中
  • 在 微調方面存在諸多挑戰，但 近期的研究 表明，對混合專家模型進行 指令調優具有很大的潛力。

讓我們開始吧！

什麼是混合專家模型？

模型規模是提升模型效能的關鍵因素之一。在有限的計算資源預算下，用更少的訓練步數訓練一個更大的模型，往往比用更多的步數訓練一個較小的模型效果更佳。

混合專家模型 (MoE) 的一個顯著優勢是它們能夠在遠少於稠密模型所需的計算資源下進行有效的預訓練。這意味著在相同的計算預算條件下，您可以顯著擴大模型或資料集的規模。特別是在預訓練階段，與稠密模型相比，混合專家模型通常能夠更快地達到相同的質量水平。

那麼，究竟什麼是一個混合專家模型 (MoE) 呢？作為一種基於 Transformer 架構的模型，混合專家模型主要由兩個關鍵部分組成:

  • 稀疏 MoE 層: 這些層代替了傳統 Transformer 模型中的前饋網路 (FFN) 層。MoE 層包含若干“專家”(例如 8 個)，每個專家本身是一個獨立的神經網路。在實際應用中，這些專家通常是前饋網路 (FFN)，但它們也可以是更復雜的網路結構，甚至可以是 MoE 層本身，從而形成層級式的 MoE 結構。
  • 門控網路或路由: 這個部分用於決定哪些令牌 (token) 被髮送到哪個專家。例如，在下圖中，“More”這個令牌可能被髮送到第二個專家，而“Parameters”這個令牌被髮送到第一個專家。有時，一個令牌甚至可以被髮送到多個專家。令牌的路由方式是 MoE 使用中的一個關鍵點，因為路由器由學習的引數組成，並且與網路的其他部分一同進行預訓練。

Switch Transformers paper 論文中的 MoE layer

總結來說，在混合專家模型 (MoE) 中，我們將傳統 Transformer 模型中的每個前饋網路 (FFN) 層替換為 MoE 層，其中 MoE 層由兩個核心部分組成: 一個門控網路和若干數量的專家。

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

  • 訓練挑戰: 雖然 MoE 能夠實現更高效的計算預訓練，但它們在微調階段往往面臨泛化能力不足的問題，長期以來易於引發過擬合現象。
  • 推理挑戰: MoE 模型雖然可能擁有大量引數，但在推理過程中只使用其中的一部分，這使得它們的推理速度快於具有相同數量引數的稠密模型。然而，這種模型需要將所有引數載入到記憶體中，因此對記憶體的需求非常高。以 Mixtral 8x7B 這樣的 MoE 為例，需要足夠的 VRAM 來容納一個 47B 引數的稠密模型。之所以是 47B 而不是 8 x 7B = 56B，是因為在 MoE 模型中，只有 FFN 層被視為獨立的專家，而模型的其他引數是共享的。此外，假設每個令牌只使用兩個專家，那麼推理速度 (以 FLOPs 計算) 類似於使用 12B 模型 (而不是 14B 模型)，因為雖然它進行了 2x7B 的矩陣乘法計算，但某些層是共享的。

瞭解了 MoE 的基本概念後，讓我們進一步探索推動這類模型發展的研究。

混合專家模型簡史

混合專家模型 (MoE) 的理念起源於 1991 年的論文 Adaptive Mixture of Local Experts。這個概念與整合學習方法相似，旨在為由多個單獨網路組成的系統建立一個監管機制。在這種系統中，每個網路 (被稱為“專家”) 處理訓練樣本的不同子集，專注於輸入空間的特定區域。那麼，如何選擇哪個專家來處理特定的輸入呢？這就是門控網路發揮作用的地方，它決定了分配給每個專家的權重。在訓練過程中，這些專家和門控網路都同時接受訓練，以最佳化它們的效能和決策能力。

在 2010 至 2015 年間，兩個獨立的研究領域為混合專家模型 (MoE) 的後續發展做出了顯著貢獻:

  1. 元件專家: 在傳統的 MoE 設定中，整個系統由一個門控網路和多個專家組成。在支援向量機 (SVMs) 、高斯過程和其他方法的研究中，MoE 通常被視為整個模型的一部分。然而，Eigen、Ranzato 和 Ilya 的研究 探索了將 MoE 作為更深層網路的一個元件。這種方法允許將 MoE 嵌入到多層網路中的某一層，使得模型既大又高效。
  2. 條件計算: 傳統的神經網路透過每一層處理所有輸入資料。在這一時期，Yoshua Bengio 等研究人員開始探索基於輸入令牌動態啟用或停用網路元件的方法。

這些研究的融合促進了在自然語言處理 (NLP) 領域對混合專家模型的探索。特別是在 2017 年，Shazeer 等人 (團隊包括 Geoffrey Hinton 和 Jeff Dean，後者有時被戲稱為 “谷歌的 Chuck Norris”) 將這一概念應用於 137B 的 LSTM (當時被廣泛應用於 NLP 的架構，由 Schmidhuber 提出)。透過引入稀疏性，這項工作在保持極高規模的同時實現了快速的推理速度。這項工作主要集中在翻譯領域，但面臨著如高通訊成本和訓練不穩定性等多種挑戰。

Outrageously Large Neural Network 論文中的 MoE layer

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

什麼是稀疏性?

稀疏性的概念採用了條件計算的思想。在傳統的稠密模型中，所有的引數都會對所有輸入資料進行處理。相比之下，稀疏性允許我們僅針對整個系統的某些特定部分執行計算。這意味著並非所有引數都會在處理每個輸入時被啟用或使用，而是根據輸入的特定特徵或需求，只有部分引數集合被呼叫和執行。

讓我們深入分析 Shazeer 對混合專家模型 (MoE) 在翻譯應用中的貢獻。條件計算的概念 (即僅在每個樣本的基礎上啟用網路的不同部分) 使得在不增加額外計算負擔的情況下擴充套件模型規模成為可能。這一策略在每個 MoE 層中實現了數以千計甚至更多的專家的有效利用。

這種稀疏性設定確實帶來了一些挑戰。例如，在混合專家模型 (MoE) 中，儘管較大的批次大小通常有利於提高效能，但當資料透過啟用的專家時，實際的批次大小可能會減少。比如，假設我們的輸入批次包含 10 個令牌， 可能會有五個令牌被路由到同一個專家，而剩下的五個令牌分別被路由到不同的專家。這導致了批次大小的不均勻分配和資源利用效率不高的問題。在接下來的部分中，將會討論讓 MoE 高效執行的其他挑戰以及相應的解決方案。

那我們應該如何解決這個問題呢？一個可學習的門控網路 (G) 決定將輸入的哪一部分傳送給哪些專家 (E):

\[y = \sum_{i=1}^{n} G(x)_i E_i(x) \]

在這種設定下，雖然所有專家都會對所有輸入進行運算，但透過門控網路的輸出進行加權乘法操作。但是，如果 G (門控網路的輸出) 為 0 會發生什麼呢？如果是這種情況，就沒有必要計算相應的專家操作，因此我們可以節省計算資源。那麼一個典型的門控函式是什麼呢？一個典型的門控函式通常是一個帶有 softmax 函式的簡單的網路。這個網路將學習將輸入傳送給哪個專家。

\[G_\sigma(x) = \text{Softmax}(x \cdot W_g) \]

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

  1. 新增一些噪聲

\[H(x)_i = (x \cdot W_{\text{g}})_i + \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i) \]

  1. 選擇保留前 K 個值

\[\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise.} \end{cases} \]

  1. 應用 Softmax 函式

\[G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k)) \]

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

我們為什麼要新增噪聲呢？這是為了專家間的負載均衡！

混合專家模型中令牌的負載均衡

正如之前討論的，如果所有的令牌都被髮送到只有少數幾個受歡迎的專家，那麼訓練效率將會降低。在通常的混合專家模型 (MoE) 訓練中，門控網路往往傾向於主要啟用相同的幾個專家。這種情況可能會自我加強，因為受歡迎的專家訓練得更快，因此它們更容易被選擇。為了緩解這個問題，引入了一個 輔助損失，旨在鼓勵給予所有專家相同的重要性。這個損失確保所有專家接收到大致相等數量的訓練樣本，從而平衡了專家之間的選擇。接下來的部分還將探討專家容量的概念，它引入了一個關於專家可以處理多少令牌的閾值。在 transformers 庫中，可以透過 aux_loss 引數來控制輔助損失。

MoEs and Transformers

Transformer 類模型明確表明，增加引數數量可以提高效能，因此谷歌使用 GShard 嘗試將 Transformer 模型的引數量擴充套件到超過 6000 億並不令人驚訝。

GShard 將在編碼器和解碼器中的每個前饋網路 (FFN) 層中的替換為使用 Top-2 門控的混合專家模型 (MoE) 層。下圖展示了編碼器部分的結構。這種架構對於大規模計算非常有效: 當擴充套件到多個裝置時，MoE 層在不同裝置間共享，而其他所有層則在每個裝置上覆制。我們將在 “讓 MoE 起飛”部分對這一點進行更詳細的討論。

GShard 論文中的 MoE Transformer Encoder

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

  • 隨機路由: 在 Top-2 設定中，我們始終選擇排名最高的專家，但第二個專家是根據其權重比例隨機選擇的。
  • 專家容量: 我們可以設定一個閾值，定義一個專家能處理多少令牌。如果兩個專家的容量都達到上限，令牌就會溢位，並透過殘差連線傳遞到下一層，或在某些情況下被完全丟棄。專家容量是 MoE 中最重要的概念之一。為什麼需要專家容量呢？因為所有張量的形狀在編譯時是靜態確定的，我們無法提前知道多少令牌會分配給每個專家，因此需要一個固定的容量因子。

GShard 的工作對適用於 MoE 的平行計算模式也做出了重要貢獻，但這些內容的討論超出了這篇部落格的範圍。

注意: 在推理過程中，只有部分專家被啟用。同時，有些計算過程是共享的，例如自注意力 (self-attention) 機制，它適用於所有令牌。這就解釋了為什麼我們可以使用相當於 12B 稠密模型的計算資源來執行一個包含 8 個專家的 47B 模型。如果我們採用 Top-2 門控，模型會使用高達 14B 的引數。但是，由於自注意力操作 (專家間共享) 的存在，實際上模型執行時使用的引數數量是 12B。

Switch Transformers

儘管混合專家模型 (MoE) 顯示出了很大的潛力，但它們在訓練和微調過程中存在穩定性問題。Switch Transformers 是一項非常激動人心的工作，它深入研究了這些話題。作者甚至在 Hugging Face 上釋出了一個 1.6 萬億引數的 MoE，擁有 2048 個專家，你可以使用 transformers 庫來執行它。Switch Transformers 實現了與 T5-XXL 相比 4 倍的預訓練速度提升。

Switch Transformer 論文中的 Switch Transformer Layer

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

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

  • 減少門控網路 (路由) 計算負擔
  • 每個專家的批次大小至少可以減半
  • 降低通訊成本
  • 保持模型質量

Switch Transformers 也對 專家容量 這個概念進行了研究。

\[\text{Expert Capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) \times \text{capacity factor} \]

上述建議的容量是將批次中的令牌數量均勻分配到各個專家。如果我們使用大於 1 的容量因子，我們為令牌分配不完全平衡時提供了一個緩衝。增加容量因子會導致更高的裝置間通訊成本，因此這是一個需要考慮的權衡。特別值得注意的是，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表現出色。

Switch Transformer 的作者還重新審視並簡化了前面章節中提到的負載均衡損失。在訓練期間，對於每個 Switch 層的輔助損失被新增到總模型損失中。這種損失鼓勵均勻路由，並可以使用超引數進行加權。

作者還嘗試了混合精度的方法，例如用 bfloat16 精度訓練專家，同時對其餘計算使用全精度進行。較低的精度可以減少處理器間的通訊成本、計算成本以及儲存張量的記憶體。然而，在最初的實驗中，當專家和門控網路都使用 bfloat16 精度訓練時，出現了不穩定的訓練現象。這種不穩定性特別是由路由計算引起的，因為路由涉及指數函式等操作，這些操作對精度要求較高。因此，為了保持計算的穩定性和精確性，保持更高的精度是重要的。為了減輕不穩定性，路由過程也使用了全精度。

使用混合精度不會降低模型質量並可實現更快的訓練

這個 Jupyter Notebook 展示瞭如何對 Switch Transformers 進行微調以進行摘要生成的詳細指南。然而，在開始微調 Switch Transformers 之前，強烈建議您先閱讀關於微調混合專家模型部分的內容。

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

用 Router z-loss 穩定模型訓練

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

ST-MoE 引入的 Router z-loss 在保持了模型效能的同時顯著提升了訓練的穩定性。這種損失機制透過懲罰門控網路輸入的較大 logits 來起作用，目的是促使數值的絕對大小保持較小，這樣可以有效減少計算中的舍入誤差。這一點對於那些依賴指數函式進行計算的門控網路尤其重要。為了深入瞭解這一機制，建議參考原始論文以獲得更全面的細節。

專家如何學習？

ST-MoE 的研究者們發現，編碼器中不同的專家傾向於專注於特定型別的令牌或淺層概念。例如，某些專家可能專門處理標點符號，而其他專家則專注於專有名詞等。與此相反，解碼器中的專家通常具有較低的專業化程度。此外，研究者們還對這一模型進行了多語言訓練。儘管人們可能會預期每個專家處理一種特定語言，但實際上並非如此。由於令牌路由和負載均衡的機制，沒有任何專家被特定配置以專門處理某一特定語言。

ST-MoE 論文中顯示了哪些令牌組被髮送給了哪個專家的表格

專家的數量對預訓練有何影響？

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

微調混合專家模型

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

稠密模型和稀疏模型在過擬合的動態表現上存在顯著差異。稀疏模型更易於出現過擬合現象，因此在處理這些模型時，嘗試更強的內部正則化措施是有益的，比如使用更高比例的 dropout。例如，我們可以為稠密層設定一個較低的 dropout 率，而為稀疏層設定一個更高的 dropout 率，以此來最佳化模型效能。

在微調過程中是否使用輔助損失是一個需要決策的問題。ST-MoE 的作者嘗試關閉輔助損失，發現即使高達 11% 的令牌被丟棄，模型的質量也沒有顯著受到影響。令牌丟棄可能是一種正則化形式，有助於防止過擬合。

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

在小任務 (左圖) 中，我們可以看到明顯的過擬合，因為稀疏模型在驗證集中的表現要差得多。在較大的任務 (右圖) 中，MoE 則表現良好。該圖來自 ST-MoE 論文

一種可行的微調策略是嘗試凍結所有非專家層的權重。實踐中，這會導致效能大幅下降，但這符合我們的預期，因為混合專家模型 (MoE) 層佔據了網路的主要部分。我們可以嘗試相反的方法: 僅凍結 MoE 層的引數。實驗結果顯示，這種方法幾乎與更新所有引數的效果相當。這種做法可以加速微調過程，並降低視訊記憶體需求。

透過僅凍結 MoE 層，我們可以在保持質量的同時加快訓練速度。該圖來自 ST-MoE 論文

在微調稀疏混合專家模型 (MoE) 時需要考慮的最後一個問題是，它們有特別的微調超引數設定——例如，稀疏模型往往更適合使用較小的批次大小和較高的學習率，這樣可以獲得更好的訓練效果。

降低學習率和調大批次可以提升稀疏模型微調質量。該圖來自 ST-MoE 論文

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

  • 單任務微調
  • 多工指令微調
  • 多工指令微調後接單任務微調

當研究者們對 MoE 和對應效能相當的 T5 模型進行微調時，他們發現 T5 的對應模型表現更為出色。然而，當研究者們對 Flan T5 (一種 T5 的指令最佳化版本) 的 MoE 版本進行微調時，MoE 的效能顯著提升。更值得注意的是，Flan-MoE 相比原始 MoE 的效能提升幅度超過了 Flan T5 相對於原始 T5 的提升，這意味著 MoE 模型可能從指令式微調中獲益更多，甚至超過了稠密模型。此外，MoE 在多工學習中表現更佳。與之前關閉 輔助損失 函式的做法相反，實際上這種損失函式可以幫助防止過擬合。

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

稀疏 VS 稠密，如何選擇?

稀疏混合專家模型 (MoE) 適用於擁有多臺機器且要求高吞吐量的場景。在固定的預訓練計算資源下，稀疏模型往往能夠實現更優的效果。相反，在視訊記憶體較少且吞吐量要求不高的場景，稠密模型則是更合適的選擇。

注意: 直接比較稀疏模型和稠密模型的引數數量是不恰當的，因為這兩類模型基於的概念和引數量的計算方法完全不同。

讓 MoE 起飛

最初的混合專家模型 (MoE) 設計採用了分支結構，這導致了計算效率低下。這種低效主要是因為 GPU 並不是為處理這種結構而設計的，而且由於裝置間需要傳遞資料，網路頻寬常常成為效能瓶頸。在接下來的討論中，我們會討論一些現有的研究成果，旨在使這些模型在預訓練和推理階段更加高效和實用。我們來看看如何最佳化 MoE 模型，讓 MoE 起飛。

平行計算

讓我們簡要回顧一下平行計算的幾種形式:

  • 資料並行: 相同的權重在所有節點上覆制，資料在節點之間分割。
  • 模型並行: 模型在節點之間分割，相同的資料在所有節點上覆制。
  • 模型和資料並行: 我們可以在節點之間同時分割模型和資料。注意，不同的節點處理不同批次的資料。
  • 專家並行: 專家被放置在不同的節點上。如果與資料並行結合，每個節點擁有不同的專家，資料在所有節點之間分割。

在專家並行中，專家被放置在不同的節點上，每個節點處理不同批次的訓練樣本。對於非 MoE 層，專家並行的行為與資料並行相同。對於 MoE 層，序列中的令牌被髮送到擁有所需專家的節點。

Switch Transformers 論文中展示如何使用不同的並行技術在節點上分割資料和模型的插圖

容量因子和通訊開銷

提高容量因子 (Capacity Factor, CF) 可以增強模型的效能，但這也意味著更高的通訊成本和對儲存啟用值的視訊記憶體的需求。在裝置通訊頻寬有限的情況下，選擇較小的容量因子可能是更佳的策略。一個合理的初始設定是採用 Top-2 路由、1.25 的容量因子，同時每個節點配置一個專家。在評估效能時，應根據需要調整容量因子，以在裝置間的通訊成本和計算成本之間找到一個平衡點。

部署技術

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

部署混合專家模型 (MoE) 的一個關鍵挑戰是其龐大的引數規模。對於本地使用情況，我們可能希望使用更小的模型。為了使模型更適合部署，下面是幾種有用的技術:

  • 預先蒸餾實驗: Switch Transformers 的研究者們進行了預先蒸餾的實驗。他們透過將 MoE 模型蒸餾回其對應的稠密模型，成功保留了 30-40%的由稀疏性帶來的效能提升。預先蒸餾不僅加快了預訓練速度，還使得在推理中使用更小型的模型成為可能。
  • 任務級別路由: 最新的方法中，路由器被修改為將整個句子或任務直接路由到一個專家。這樣做可以提取出一個用於服務的子網路，有助於簡化模型的結構。
  • 專家網路聚合: 這項技術透過合併各個專家的權重，在推理時減少了所需的引數數量。這樣可以在不顯著犧牲效能的情況下降低模型的複雜度。

高效訓練

FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同並行策略下的理論效能極限，並且探索了一系列創新技術，包括用於專家權重調整的方法、減少延遲的細粒度通訊排程技術，以及一個基於最低延遲進行專家選擇的拓撲感知門控機制。這些技術的結合使得 MoE 執行速度提升高達 17 倍。

Megablocks (2022 年 11 月) 則專注於透過開發新的 GPU kernel 來處理 MoE 模型中的動態性，以實現更高效的稀疏預訓練。其核心優勢在於，它不會丟棄任何令牌，並能高效地適應現代硬體架構 (支援塊稀疏矩陣乘)，從而達到顯著的加速效果。Megablocks 的創新之處在於，它不像傳統 MoE 那樣使用批次矩陣乘法 (這通常假設所有專家形狀相同且處理相同數量的令牌)，而是將 MoE 層表示為塊稀疏操作，可以靈活適應不均衡的令牌分配。

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

開源混合專家模型

目前，下面這些開源專案可以用於訓練混合專家模型 (MoE):

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

  • Switch Transformers (Google): 基於 T5 的 MoE 集合，專家數量從 8 名到 2048 名。最大的模型有 1.6 萬億個引數。
  • NLLB MoE (Meta): NLLB 翻譯模型的一個 MoE 變體。
  • OpenMoE: 社群對基於 Llama 的模型的 MoE 嘗試。
  • Mixtral 8x7B (Mistral): 一個效能超越了 Llama 2 70B 的高質量混合專家模型，並且具有更快的推理速度。此外，還發布了一個經過指令微調的模型。有關更多資訊，可以在 Mistral 的 公告部落格文章 中瞭解。

一些有趣的研究方向

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

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

簡而言之，一些值得探索的有趣領域包括:

  • 將 Mixtral 蒸餾成一個稠密模型。
  • 探索合並專家模型的技術及其對推理時間的影響。
  • 嘗試對 Mixtral 進行極端量化的實驗。

相關資源

Citation

@misc {sanseviero2023moe,
    author       = { Omar Sanseviero and
                     Lewis Tunstall and
                     Philipp Schmid and
                     Sourab Mangrulkar and
                     Younes Belkada and
                     Pedro Cuenca
                   },
    title        = { Mixture of Experts Explained },
    year         = 2023,
    url          = { https://huggingface.co/blog/moe },
    publisher    = { Hugging Face Blog }
}

Sanseviero, et al., "Mixture of Experts Explained", Hugging Face Blog, 2023.

英文原文: https://hf.co/blog/moe

原文作者: Omar Sanseviero, Lewis Tunstall, Philipp Schmid, Sourab Mangrulkar, Younes Belkada, Pedro Cuenca

譯者: xinyu66 (Xinyu Yang)

相關文章