2023 年 12 月，首個開源 MoE 大模型 Mixtral 8×7B 發布，在多種基準測試中，其表現近乎超越了 GPT-3.5 和 LLaMA 2 70B，而推理開銷僅相當于 12B 左右的稠密模型。為進一步提升模型性能，稠密 LLM 常由于其參數規模急劇擴張而面臨嚴峻的訓練成本。MoE 的廣泛應用，使得在計算成本相對不變的條件下，模型容量能夠得到顯著擴展。此特性無疑使得 MoE 成為推動 LLM 發展的關鍵技術。

MoE 設計的初衷，是使模型的學習更加 “術業有專攻”，其有效性已得到業界肯定。然而現有 MoE 架構訓練中的弊端也逐漸凸顯，主要包括：專家負載失衡、專家內樣本混雜而專家間同質化現象嚴重、額外的通信開銷等等。

為了緩解現有 MoE 普遍存在的訓練效率與性能瓶頸，專精于高性能計算、LLM 訓練加速的華為 GTS AI 計算 Lab的研究團隊提出了名為 LocMoE 的全新 MoE 架構，從路由機制角度出發，以期降低稀疏 LLM 訓練成本的同時，提升其性能。

論文鏈接：https://arxiv.org/abs/2401.13920

論文簡介

在這項工作中，作者發現之前的 MoE 路由機制往往會導致路由至同一專家網絡的 token 差異較大，干擾專家網絡的收斂；而路由至不同專家的 token 相似性較高，造成專家間同質化現象嚴重，最終影響模型語義理解與生成的能力。作者通過理論闡明了專家路由與輸入數據特征之間的關聯，并在 NLP 領域首次證明專家網絡存在容量下限。在此理論基礎上，專家路由的門控權重經正交化處理后，明顯增強了專家網絡間的區分度，處理遠小于原先規模的 token，能夠在領域評測中達到相近的效果。同時該研究針對 MoE 架構中固有的 All-To-All 通信瓶頸，結合負載 / 通信優化，提出高效高能的 MoE 架構。

具體來說，作者提出了一種名為 LocMoE 的新穎 MoE 架構，將其嵌入到盤古大模型的骨干中以增強其能力。LocMoE 旨在增強路由機制的可解釋性，同時降低額外通信與計算開銷。首先，作者發現 token 總傾向于路由至與該 token 夾角最小的專家，當專家間門控權重向量滿足正交時，專家網絡間處理的 token 能盡可能避免同質化。

因此，本文采用 GAP 層提取 token 特征，將其作為路由的依據。GAP 層特性上滿足門控權重正交的條件，計算量相比 Dense 層也得到大幅下降。基于上述結構，作者通過理論求解出在不影響模型 loss 前提下，專家處理的 token 規模的下限，以降低專家網絡的負載。此外，作者結合 auxiliary loss，提出了 locality loss 對路由進行軟約束，降低跨機 All-To-All 通信開銷。最后，采用通信遮掩等工程優化，進一步提升稀疏大模型整體的訓練性能。

作者將 LocMoE 架構嵌入到盤古-Σ 38B 模型中，采用語義相似度較高的 ICT 領域數據進行訓練，檢驗其領域知識的學習能力。在十項下游任務中，LocMoE 的準確性普遍高于原生盤古-Σ，訓練性能每步提升 10%~20%。該 MoE 架構還具有較強的通用性和易于移植性，能夠低成本嵌入到其他硬件規格和其他 MoE 架構的 LLM 骨干中。

當前，LocMoE 已部署至華為 ICT 服務領域專業知識問答應用 AskO3 上，AskO3 已上線華為 O3 知識社區，獲得數萬工程師用戶群體廣泛好評。

創新點剖析

路由與數據特征的關系

針對現有 MoE 路由機制普遍缺乏可解釋性的問題，作者分析了 token 路由的本質，并設計了能夠將 token 有效區分的結構。對于某個 token，學習性的路由策略普遍選擇門控權重與該 token 乘積更大的專家進行分配：

那么，如果專家的門控權重滿足正交，能夠使得專家具備更高的判別性。同時，能夠得出 token 傾向于被路由至與其夾角更小的專家：

作者最終選取 GAP 作為提取 token 特征的結構，其門控權重能夠滿足正交的條件：

上述路由機制的實質描述了路由判別能力與專家 token 間最小夾角之間的關聯，如圖所示。

圖：LocMoE 路由機制示意圖

專家容量下界

在上述理論的基礎上，作者發現，專家容量存在下界，即，在輸入數據確定的情況下，專家處理遠小于經驗值規模的 token，也能達到相當的性能。該問題可以轉化為，求解使得至少一個具有類別判別性 token 被路由至某個專家的最小 token 規模，作為所有專家容量拉齊時的下界。同時能夠得出，合理的專家容量與 token 和門控權重間的最小夾角呈負相關，并隨著夾角的減小呈指數級增長。經實驗證實，專家容量設為該下界時，未對模型收斂性和 loss 曲線產生影響。

本地性約束

LocMoE 在 MoE 層的 loss 包含兩部分：auxiliary loss 和 locality loss。auxiliary loss 首次在 sparsely-gated MoE 中提出，同時應用于 SwitchTransformer，用以提升專家負載均衡性：

然而，跨機 All-To-All 帶來的額外通信開銷仍無法避免。因此，作者添加了本地性約束，使得在專家負載均衡的前提下，token 更傾向于被分派給本地設備的專家，最終達到約束平衡。locality loss 采用當前 token 分布與完全本地化分布之間的差異即 KL 散度來量化，從而將部分機間通信轉為機內通信，充分利用機內互聯高帶寬。

實驗結果

作者分別在包含 64 張、128 張和 256 張昇騰 910A NPU 的集群上進行了實驗，主要與兩款經典的 MoE 結構：Hash (來自 Facebook) 和 Switch (來自 Google) 進行比較。

訓練效率

作者記錄了各實驗組訓練過程中計算、通信、遮掩以及閑置的耗時。其中，在 64N 和 128N 的情況下，LocMoE 的計算開銷和通信開銷都是最低的。盡管 256N 時 LocMoE 的計算開銷仍然最低，但部分設備不包含專家使得本地性通信轉換失效，說明了 LocMoE 在計算及通信方面同時存在顯著增益的適用條件是專家數至少大于等于節點數。

圖：多種 MoE 結構在不同集群配置下的訓練效率

專家負載

為了驗證約束項對于專家負載的影響，作者分析了路由至每個專家的 token 分布情況。為了達到負載均衡，通過 RRE 模塊實現的 Hash 路由采用靜態路由表的硬約束確保分配的均衡性，LocMoE 和 Switch 則考慮到 token 的具體特征而進行路由。作為學習型路由，在 auxiliary 和本地性約束項的作用下，LocMoE 專家的均衡性明顯優于 Switch，表現出穩定且較高的資源利用率。

圖：多種 MoE 結構下專家負載情況

分配給專家的樣本相似性

對于支撐 LocMoE 提出的關鍵理論，作者采用實驗對其進行了驗證。左圖表明路由到同一專家的 token 相似性更高，接近于 1。右圖則表明 token 與其路由至的專家對應的門控權重相似度分布相較其他專家更靠右，從而證實了 token 傾向于路由至與其夾角最小的專家的理論前提，并標記出專家容量下限求解的關鍵參數 δ。

圖：路由至同一 / 不同專家 token 相似性（左）；token 與其路由至的專家的相似性（右）

NLP 領域下游任務

盤古-Σ 已經使用了來自金融、健康等超過 40 個領域的語料進行預訓練，證明其從多領域文本數據中學習知識的能力。在本項工作中，作者使用 ICT 服務的領域數據，包含無線網絡、光、運營商 IT 等產品線的技術報告和工具手冊等，評估 LocMoE 在專業領域知識的學習表現。根據概念間邏輯復雜程度分為 L1 至 L3，梳理出十類 NLP 領域任務的評測數據集。與原生盤古-Σ 相比，LocMoE 使得模型語義理解和表達能力都有一定程度的提高。

圖：與原生盤古-Σ 相比，NLP 領域下游任務表現