華為GTS LocMoE+：高可擴充性親和力 MoE 架構，低開銷實現主動路由

AIxiv專欄是本站發布學術、技術內容的欄位。過去數年，本站AIxiv專欄接收通報了2,000多篇內容，涵蓋全球各大專院校與企業的頂尖實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或聯絡報道。投稿信箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本論文共同一作為李婧博士，孫訓志傑和林大超推博士，主要成員來自GTS AI計算Lab，主要研究及落地領域包含LLM加速、AI訓練保障和圖計算。

MoE 因其在訓推流程中低銷高效的特點，近兩年在大語言模型領域大放異彩。身為 MoE 的靈魂，專家如何能發揮最大的學習潛能，相關的研究與討論層出不窮。先前，華為 GTS AI 運算 Lab 的研究團隊提出了 LocMoE ，包括新穎的路由網路結構、輔助降低通訊開銷的本地性 loss 等，引發了廣泛關注。

LocMoE 的上述設計有效緩解了部分MoE 經典結構在訓練中的瓶頸，如: 專家路由演算法可能無法有效區分token，以及通訊同步效率受限於節點內和節點間的傳輸頻寬差異等。此外，LocMoE 證明並求解出在能夠成功處理判別性 token 的前提下專家容量的下限。此下限是根據 token 被動派發給專家的場景下，判別性 token 存在於 token 批次中的機率分佈而得出。那麼，如果專家同樣具備擇優 token 的能力，判別性 token 被處理的機率將大幅提高，專家容量下限將能進一步壓縮。

在上述設想的基礎上，該團隊進一步提出了基於低開銷主動路由的 MoE 架構，將其命名為 LocMoE+。 LocMoE+ 繼承了 LocMoE 高判別性專家以及本地性通訊優勢，進一步改造路由策略，定義 token 與專家之間的親和力指標，並以此指標入手，更有效率完成 token 分派，進而提升訓練效率。

• 論文連結：https://arxiv.org/pdf/2406.00023

論文簡介主動路由，提升一定容量下處理判別性token 的機率，進而降低取樣雜訊並提升訓練效率。論文從 token 和其被分派的專家之間的關係入手，以低計算開銷的方案量化和定義了專家與 token 之間的親和性。據此，該論文實現了全局的自適應路由策略，並基於親和性分數以專家維度對 token 進行重排和選擇。同時，專家容量的下限被證明隨著 token 特徵分佈的穩定而逐漸減小，訓練開銷得以降低。

該論文率先結合兩種路由機制，根據學習性路由策略中token 傾向於路由至與其夾角更小的專家的發現，打破了主動路由現存方案開銷過大影響訓練效率的障礙，並與被動路由的本質保持統一。

值得一提的是，作者選用了和LocMoE 工作完全不同的硬體環境（伺服器型號，NPU卡型號，集群組網方案）、訓練框架和骨幹模型，以證明該系列工作的高擴展性和易於移植性。

自適應雙向路由分派機制

背景介紹1

背景介紹

1 router，直接把整個token 特徵進行分派；

（2）soft router，把token 特徵的加權組合分派。

本文延續考慮（1）因其較低的計算代價。而對於hard router 場景，可分為1）Token Choice Router (TCR)， 也就是讓每個token 去選top-k 專家；2）Expert Choice Router (ECR)，也就是讓每個專家去選擇top -C 合適的token。由於容量限制，每位專家接收的 token 數有上限 C， 因此在場景 1）中，會對每個專家接收的 token進行截斷：

前人工作指出MoE 訓練分為兩階段：Phase 1. 路由訓練，保障路由可以合理分派token，也就是能把不同領域或差異大的token 進行區分和分派給不同專家；Phase 2. 由於token 路由的作用，每個專家接收到相同領域或性質相近的token，每個專家在經歷一定訓練後就可以習得相關領域和性質的知識。總結來說，MoE 每步訓練「成功」的關鍵在於 token 分派的正確和合理性。

該文貢獻

（1）透過softmax 活化函數推演，專家與token 之間的餘弦相似度則能較為精確地測量親和性：2）從理論建模角度，分析TCR 和ECR 兩個常見場景下單次訓練成功率：

基於理論，作者指出

在模型訓練初期，路由分

• 在模型訓練初期，路由分派token 能力分差每次訓練TCR 要比ECR 以更高機率成功訓練，而且需要較大的專家容量保證選到合適的token。
• 在模型訓練後期，路由有一定能力正確分派 token 時，每次訓練 ECR 要比 TCR 以更高概率成功訓練，此時只需要較小的容量就可以選到合適的 token。

該理論也很符合直觀，路由沒有分派能力時，不如讓token 隨機選擇專家，而當路由有一定分派能力，也就是專家能選到合適token 的時候，使用ECR 更合適。因此作者推薦 TCR 向 ECR 方式的轉變，提出了全局層級的自適應路由切換策略，同時基於專家容量的需求估計，在訓練後期使用較小的專家容量。

該論文的實驗在昇騰910B3 NPU 的自建集群上進行，得益於中國一致性協議多卡場景下實現的高效能設備間資料通信，以及專為昇騰處理器設計的華為集合通訊庫（HCCL）在HCCS 等高速鏈路上實現高效能分散式訓練。實驗採用相容於昇騰 NPU 的 PyTorch for Ascend 框架和專為昇騰設備定制的加速庫 AscendSpeed 和訓練框架 ModelLink，專注於 LLM 並行策略和通訊遮掩優化。

訓練效率

實驗結果表明，在不影響模型訓練收斂性或有效性的情況下，每個專家至少需要處理的token數量較基線可以減少60% 以上。結合通訊優化，在 32 卡、64 卡和 256 卡的集群規模下，訓練效率平均提高 5.4% 至 46.6%。

顯存佔用

LocMoE+ 對於顯存佔用也存在一定增益，尤其是集群規模較小計算較密集的場景。使用 Ascend Insight 工具對顯存監控樣本進行分析，可以看出，LocMoE+ 顯存使用相比於基線下降了 4.57% 至 16.27%，相比於 LocMoE 下降了 2.86% 至 10.5%。隨著群集規模的增大，顯存佔用的差距隨之縮小。

效果評估

開源評測集 C-Eval 和 TeleQnA，以及自主建構的 ICT 領域評測集 GDAD 被用於評估 LocMoE+ 在通用知識和領域的能力。其中，GDAD 共涵蓋了 47 個子項，包含 18,060 個樣本，檢視模型在領域任務，領域能力認證考試和通用能力三大評估系統中的表現。

經過充分 SFT 後，LocMoE+ 相比於基線在領域任務能力的 16 個子能力上平均提升了約 20.1%，相比於 LocMoE 則提升了約3.5%。領域能力認證考試則分別提升了 16% 和 4.8%。在通用能力的 18 個子能力中，LocMoE+ 分別提升約 13.9% 和 4.8%。整體而言，LocMoE+ 在 GDAD、C-Eval 和 TeleQnA 上分別表現出 9.7% 到 14.1% 的效果提升。

以上是華為GTS LocMoE+：高可擴充性親和力 MoE 架構，低開銷實現主動路由的詳細內容。更多資訊請關注PHP中文網其他相關文章！