2023 年 12 月，首個(gè)開(kāi)源 MoE 大模型 Mixtral 8×7B 發(fā)布，在多種基準(zhǔn)測(cè)試中，其表現(xiàn)近乎超越了 GPT-3.5 和 LLaMA 2 70B，而推理開(kāi)銷(xiāo)僅相當(dāng)于 12B 左右的稠密模型。為進(jìn)一步提升模型性能，稠密 LLM 常由于其參數(shù)規(guī)模急劇擴(kuò)張而面臨嚴(yán)峻的訓(xùn)練成本。MoE 的廣泛應(yīng)用，使得在計(jì)算成本相對(duì)不變的條件下，模型容量能夠得到顯著擴(kuò)展。此特性無(wú)疑使得 MoE 成為推動(dòng) LLM 發(fā)展的關(guān)鍵技術(shù)。

MoE 設(shè)計(jì)的初衷，是使模型的學(xué)習(xí)更加 “術(shù)業(yè)有專(zhuān)攻”，其有效性已得到業(yè)界肯定。然而現(xiàn)有 MoE 架構(gòu)訓(xùn)練中的弊端也逐漸凸顯，主要包括：專(zhuān)家負(fù)載失衡、專(zhuān)家內(nèi)樣本混雜而專(zhuān)家間同質(zhì)化現(xiàn)象嚴(yán)重、額外的通信開(kāi)銷(xiāo)等等。

為了緩解現(xiàn)有 MoE 普遍存在的訓(xùn)練效率與性能瓶頸，專(zhuān)精于高性能計(jì)算、LLM 訓(xùn)練加速的華為 GTS AI 計(jì)算 Lab的研究團(tuán)隊(duì)提出了名為 LocMoE 的全新 MoE 架構(gòu)，從路由機(jī)制角度出發(fā)，以期降低稀疏 LLM 訓(xùn)練成本的同時(shí)，提升其性能。

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

論文簡(jiǎn)介

在這項(xiàng)工作中，作者發(fā)現(xiàn)之前的 MoE 路由機(jī)制往往會(huì)導(dǎo)致路由至同一專(zhuān)家網(wǎng)絡(luò)的 token 差異較大，干擾專(zhuān)家網(wǎng)絡(luò)的收斂；而路由至不同專(zhuān)家的 token 相似性較高，造成專(zhuān)家間同質(zhì)化現(xiàn)象嚴(yán)重，最終影響模型語(yǔ)義理解與生成的能力。作者通過(guò)理論闡明了專(zhuān)家路由與輸入數(shù)據(jù)特征之間的關(guān)聯(lián)，并在 NLP 領(lǐng)域首次證明專(zhuān)家網(wǎng)絡(luò)存在容量下限。在此理論基礎(chǔ)上，專(zhuān)家路由的門(mén)控權(quán)重經(jīng)正交化處理后，明顯增強(qiáng)了專(zhuān)家網(wǎng)絡(luò)間的區(qū)分度，處理遠(yuǎn)小于原先規(guī)模的 token，能夠在領(lǐng)域評(píng)測(cè)中達(dá)到相近的效果。同時(shí)該研究針對(duì) MoE 架構(gòu)中固有的 All-To-All 通信瓶頸，結(jié)合負(fù)載 / 通信優(yōu)化，提出高效高能的 MoE 架構(gòu)。

具體來(lái)說(shuō)，作者提出了一種名為 LocMoE 的新穎 MoE 架構(gòu)，將其嵌入到盤(pán)古大模型的骨干中以增強(qiáng)其能力。LocMoE 旨在增強(qiáng)路由機(jī)制的可解釋性，同時(shí)降低額外通信與計(jì)算開(kāi)銷(xiāo)。首先，作者發(fā)現(xiàn) token 總傾向于路由至與該 token 夾角最小的專(zhuān)家，當(dāng)專(zhuān)家間門(mén)控權(quán)重向量滿(mǎn)足正交時(shí)，專(zhuān)家網(wǎng)絡(luò)間處理的 token 能盡可能避免同質(zhì)化。

因此，本文采用 GAP 層提取 token 特征，將其作為路由的依據(jù)。GAP 層特性上滿(mǎn)足門(mén)控權(quán)重正交的條件，計(jì)算量相比 Dense 層也得到大幅下降。基于上述結(jié)構(gòu)，作者通過(guò)理論求解出在不影響模型 loss 前提下，專(zhuān)家處理的 token 規(guī)模的下限，以降低專(zhuān)家網(wǎng)絡(luò)的負(fù)載。此外，作者結(jié)合 auxiliary loss，提出了 locality loss 對(duì)路由進(jìn)行軟約束，降低跨機(jī) All-To-All 通信開(kāi)銷(xiāo)。最后，采用通信遮掩等工程優(yōu)化，進(jìn)一步提升稀疏大模型整體的訓(xùn)練性能。

作者將 LocMoE 架構(gòu)嵌入到盤(pán)古-Σ 38B 模型中，采用語(yǔ)義相似度較高的 ICT 領(lǐng)域數(shù)據(jù)進(jìn)行訓(xùn)練，檢驗(yàn)其領(lǐng)域知識(shí)的學(xué)習(xí)能力。在十項(xiàng)下游任務(wù)中，LocMoE 的準(zhǔn)確性普遍高于原生盤(pán)古-Σ，訓(xùn)練性能每步提升 10%~20%。該 MoE 架構(gòu)還具有較強(qiáng)的通用性和易于移植性，能夠低成本嵌入到其他硬件規(guī)格和其他 MoE 架構(gòu)的 LLM 骨干中。

當(dāng)前，LocMoE 已部署至華為 ICT 服務(wù)領(lǐng)域?qū)I(yè)知識(shí)問(wèn)答應(yīng)用 AskO3 上，AskO3 已上線華為 O3 知識(shí)社區(qū)，獲得數(shù)萬(wàn)工程師用戶(hù)群體廣泛好評(píng)。

創(chuàng)新點(diǎn)剖析

路由與數(shù)據(jù)特征的關(guān)系

針對(duì)現(xiàn)有 MoE 路由機(jī)制普遍缺乏可解釋性的問(wèn)題，作者分析了 token 路由的本質(zhì)，并設(shè)計(jì)了能夠?qū)?token 有效區(qū)分的結(jié)構(gòu)。對(duì)于某個(gè) token，學(xué)習(xí)性的路由策略普遍選擇門(mén)控權(quán)重與該 token 乘積更大的專(zhuān)家進(jìn)行分配：

那么，如果專(zhuān)家的門(mén)控權(quán)重滿(mǎn)足正交，能夠使得專(zhuān)家具備更高的判別性。同時(shí)，能夠得出 token 傾向于被路由至與其夾角更小的專(zhuān)家：

作者最終選取 GAP 作為提取 token 特征的結(jié)構(gòu)，其門(mén)控權(quán)重能夠滿(mǎn)足正交的條件：

上述路由機(jī)制的實(shí)質(zhì)描述了路由判別能力與專(zhuān)家 token 間最小夾角之間的關(guān)聯(lián)，如圖所示。

圖：LocMoE 路由機(jī)制示意圖

專(zhuān)家容量下界

在上述理論的基礎(chǔ)上，作者發(fā)現(xiàn)，專(zhuān)家容量存在下界，即，在輸入數(shù)據(jù)確定的情況下，專(zhuān)家處理遠(yuǎn)小于經(jīng)驗(yàn)值規(guī)模的 token，也能達(dá)到相當(dāng)?shù)男阅堋Ｔ搯?wèn)題可以轉(zhuǎn)化為，求解使得至少一個(gè)具有類(lèi)別判別性 token 被路由至某個(gè)專(zhuān)家的最小 token 規(guī)模，作為所有專(zhuān)家容量拉齊時(shí)的下界。同時(shí)能夠得出，合理的專(zhuān)家容量與 token 和門(mén)控權(quán)重間的最小夾角呈負(fù)相關(guān)，并隨著夾角的減小呈指數(shù)級(jí)增長(zhǎng)。經(jīng)實(shí)驗(yàn)證實(shí)，專(zhuān)家容量設(shè)為該下界時(shí)，未對(duì)模型收斂性和 loss 曲線產(chǎn)生影響。

本地性約束

LocMoE 在 MoE 層的 loss 包含兩部分：auxiliary loss 和 locality loss。auxiliary loss 首次在 sparsely-gated MoE 中提出，同時(shí)應(yīng)用于 SwitchTransformer，用以提升專(zhuān)家負(fù)載均衡性：

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

實(shí)驗(yàn)結(jié)果

作者分別在包含 64 張、128 張和 256 張昇騰 910A NPU 的集群上進(jìn)行了實(shí)驗(yàn)，主要與兩款經(jīng)典的 MoE 結(jié)構(gòu)：Hash (來(lái)自 Facebook) 和 Switch (來(lái)自 Google) 進(jìn)行比較。

訓(xùn)練效率

作者記錄了各實(shí)驗(yàn)組訓(xùn)練過(guò)程中計(jì)算、通信、遮掩以及閑置的耗時(shí)。其中，在 64N 和 128N 的情況下，LocMoE 的計(jì)算開(kāi)銷(xiāo)和通信開(kāi)銷(xiāo)都是最低的。盡管 256N 時(shí) LocMoE 的計(jì)算開(kāi)銷(xiāo)仍然最低，但部分設(shè)備不包含專(zhuān)家使得本地性通信轉(zhuǎn)換失效，說(shuō)明了 LocMoE 在計(jì)算及通信方面同時(shí)存在顯著增益的適用條件是專(zhuān)家數(shù)至少大于等于節(jié)點(diǎn)數(shù)。

圖：多種 MoE 結(jié)構(gòu)在不同集群配置下的訓(xùn)練效率

專(zhuān)家負(fù)載

為了驗(yàn)證約束項(xiàng)對(duì)于專(zhuān)家負(fù)載的影響，作者分析了路由至每個(gè)專(zhuān)家的 token 分布情況。為了達(dá)到負(fù)載均衡，通過(guò) RRE 模塊實(shí)現(xiàn)的 Hash 路由采用靜態(tài)路由表的硬約束確保分配的均衡性，LocMoE 和 Switch 則考慮到 token 的具體特征而進(jìn)行路由。作為學(xué)習(xí)型路由，在 auxiliary 和本地性約束項(xiàng)的作用下，LocMoE 專(zhuān)家的均衡性明顯優(yōu)于 Switch，表現(xiàn)出穩(wěn)定且較高的資源利用率。

圖：多種 MoE 結(jié)構(gòu)下專(zhuān)家負(fù)載情況

分配給專(zhuān)家的樣本相似性

對(duì)于支撐 LocMoE 提出的關(guān)鍵理論，作者采用實(shí)驗(yàn)對(duì)其進(jìn)行了驗(yàn)證。左圖表明路由到同一專(zhuān)家的 token 相似性更高，接近于 1。右圖則表明 token 與其路由至的專(zhuān)家對(duì)應(yīng)的門(mén)控權(quán)重相似度分布相較其他專(zhuān)家更靠右，從而證實(shí)了 token 傾向于路由至與其夾角最小的專(zhuān)家的理論前提，并標(biāo)記出專(zhuān)家容量下限求解的關(guān)鍵參數(shù) δ。

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

NLP 領(lǐng)域下游任務(wù)

盤(pán)古-Σ 已經(jīng)使用了來(lái)自金融、健康等超過(guò) 40 個(gè)領(lǐng)域的語(yǔ)料進(jìn)行預(yù)訓(xùn)練，證明其從多領(lǐng)域文本數(shù)據(jù)中學(xué)習(xí)知識(shí)的能力。在本項(xiàng)工作中，作者使用 ICT 服務(wù)的領(lǐng)域數(shù)據(jù)，包含無(wú)線網(wǎng)絡(luò)、光、運(yùn)營(yíng)商 IT 等產(chǎn)品線的技術(shù)報(bào)告和工具手冊(cè)等，評(píng)估 LocMoE 在專(zhuān)業(yè)領(lǐng)域知識(shí)的學(xué)習(xí)表現(xiàn)。根據(jù)概念間邏輯復(fù)雜程度分為 L1 至 L3，梳理出十類(lèi) NLP 領(lǐng)域任務(wù)的評(píng)測(cè)數(shù)據(jù)集。與原生盤(pán)古-Σ 相比，LocMoE 使得模型語(yǔ)義理解和表達(dá)能力都有一定程度的提高。

圖：與原生盤(pán)古-Σ 相比，NLP 領(lǐng)域下游任務(wù)表現(xiàn)