AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：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）hard 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 的自建叢集上進行，得益於華為專有的快取記憶體一致性協議高效能運算系統（HCCS）在多卡場景下實現的高效能裝置間資料通訊，以及專為昇騰處理器設計的華為集合通訊庫（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 個子項，包含 18060 條樣本，考察模型在領域任務，領域能力認證考試和通用能力三大評估體系中的表現。經過充分 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% 的效果提升。