基于NVIDIA Megatron Core的MOE LLM實現和訓練優化

本文將分享阿里云人工智能平臺 PAI 團隊與 NVIDIA Megatron-Core 團隊在 MoE (Mixture of Experts) 大語言模型（LLM）實現與訓練優化上的創新工作。分享內容將按以下脈絡展開：

首先簡短回顧 MoE 技術的發展歷程，提煉核心概念及其在實踐應用中亟待解決的關鍵挑戰。

接著詳述雙方合作研發的 MoE 框架所具有的獨特能力和卓越性能，展示其在提升模型訓練效率、資源利用以及模型表現等方面取得的驗證結果。

最后，扼要介紹阿里云基于此合作成果所搭建的平臺工具及推薦的最佳實踐方案，賦能開發者高效運用 MoE 技術，促進大規模模型訓練的深入探索與廣泛應用。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIA Megatron-Core 的 MoE LLM 實現和訓練優化》

MoE 是一種模型結構，它將稠密模型結構拆分為多個子結構，每個子結構成為一個專家，通過在訓練和推理過程中動態選擇一組專家進行計算，實現了模型參數的稀疏更新。

簡而言之，MoE 模型將整體模型拆分為多個專業子模塊（專家），每次僅激活和更新少數與輸入相關的專家子結構，而選擇哪些專家參與計算則是通過一個路由機制決定。在 GPT 等超大規模語言模型中，采用MoE 技術能夠顯著減少訓練和推理時的計算負擔，因為不是所有參數都需要在每次操作時都更新，這極大提高了訓練效率，并且在推理階段僅使用部分活躍的網絡參數，極大地削減了計算資源需求。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

2023 年末，Mistral AI 推出開源 MoE 架構大語言模型 Mixtral 8x7B，憑借 46.7B 參數量，在多項下游任務榜單的效果勝過當時的最佳開源稠密模型 Llama-2 70B。右側圖表揭示，在同等推理資源條件下，MoE 模型性能顯著優于稠密模型。這一成就引發業界對 MoE 模型的強烈關注，進而推動阿里云與 NVIDIA Megatron-Core 團隊，共同深化在大模型領域的 MoE 技術合作與應用。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

為了描述使用 MoE 結構來實現工業級應用需要解決的問題，本文首先回顧 MoE 的發展歷史。90 年代初提出的多專家模型概念為 MoE 打下了理論基礎，即通過集合多個專家模型協同解決任務，并采用門控路由進行專家選擇。

隨著 2017 年深度學習突飛猛進，參數量對模型性能的重要性日益顯現，傳統稠密結構深度學習框架面臨容量瓶頸。谷歌率先將 MoE 與深度學習結合，首次在 RNN 中驗證了 MoE 的可行性。至 2020 年，Transformer 結構在語言模型領域展現出卓越效果和優秀的可擴展性。此后，Google Gshard 項目將 MoE 融入 Transformer，通過深度設計與實驗顯著提升效果和性能，并在 Switch Transformers 中將參數總量推向萬億級別，奠定了 MoE 架構在大模型領域的基礎。

后續研究如 ST-MoE、Tutel MoE、FasterMOE 和 MegaBlocks 等，針對速度和穩定性進行深入探索。2023 年，GPT-4 的驚艷表現引發了業內對其是否采用 MoE 結構的熱議。同年，Mixtral 的實現似乎驗證了這一猜想，表明 MoE 在大模型架構中的應用正逐步走向成熟和廣泛認可。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

后續我們將簡要剖析幾個標志性工作，首先是將 MoE 與深度學習融為一體的 Sparsely-Gated MoE。該研究將 RNN 網絡劃分為多個專家子網絡，并采用 Softmax Gating 機制來控制 token 的路由分配。然而，原始 Softmax Gating 邏輯可能導致負載不均衡，隨著訓練推進，部分專家網絡可能過于活躍，而其他專家的參數訓練不足。

為解決此問題，研究者在計算 Softmax 之前，對門控矩陣參數和輸入特征進行噪聲 (noise) 注入，以實現更為均衡的專家選擇。此外，在門控網絡損失函數 (loss function) 中增設輔助損失，引導模型實現更佳的負載分配。同時，路由策略上嘗試每次挑選 top-2 至 top-4 的專家，這些優化舉措最終助力模型參數容量首次突破千億門檻。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

而谷歌的 Gshard 項目進一步將 MoE 應用于 Transformer 架構中，其主要創新包括：

設計了 MoE Transformer 的基本結構和并行模式，每間隔一層前向層采用 MoE 層替代，專家分布在不同設備上獨立計算，其他層的參數共享。

針對專家負載均衡問題，提出了“專家容量（expert capacity）”概念，限制每個專家處理的 token 數量，并采用殘差連接策略繞過已滿負荷的專家。

引入輔助損失與隨機 top-k 路由策略，以優化專家選擇過程。

在通信效率上，創新提出“local group dispatching”方案，通過門控網絡預篩選后再按專家 ID 分組傳輸數據，有效提升通信效率。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

實驗結果顯示，Gshard 成功將 Transformer 參數量從 150B 提升至 600B，而計算開銷僅增加 50%，首次驗證了 MoE 在模型擴展性上的顯著優勢。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

之后的 Switch Transformer 工作對 MoE Transformer 的訓練流程進行了深度優化，以實現模型參數規模的極限拓展。首先，路由策略選擇了更為激進的 top-1 方案，即每次僅選擇一個專家進行參數更新。此舉不僅能顯著降低路由的計算負擔，與 top-2 策略相比計算量近乎減半，同時每個專家處理的 batch size 也因此減少，進一步減輕了通信成本。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

其次，針對先前存在的專家負載均衡問題，工作引入了輔助損失函數。該函數中，N 代表專家總數，f 表示第 i 個專家分配到的 token 比例，P 代表路由器分配給第i個專家的概率。優化目標旨在促使 token 均勻分配至每個專家。通過最小化損失函數，使得 f 和 P 趨向于 N 分之一，此時損失函數中的調節參數 α 通常取值 0.01 左右，能夠實現較好的負載均衡效果。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

在分布式訓練方面，除了先前提及的數據并行和專家并行外，該工作對非專家結構部分實施了模型并行策略，來進一步降低單卡顯存需求，使得模型參數量得以顯著提升。

針對路由負載均衡引起的模型訓練隨機性與不穩定問題，該工作提出了一系列穩定訓練的策略。其中包括在本地計算時采用較高的 FP32 精度，而在通信階段使用較低的 BF16 精度的隨機精度策略；采用具有較小縮放系數的正態分布初始化權重；引入專家 dropout 即專家內部的 dropout 技術以減少過擬合現象。

通過這些優化措施，MoE transformer 的訓練收斂速度相較于稠密模型在迭代次數和時間上均有數倍提升。其中，最大的 Switch-C 模型參數量高達 1.5 萬億。后續的 ST-MoE 工作則更深入地探究了如何進一步改善 MoE 模型的訓練穩定性和提升模型性能。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

在預訓練穩定性方面，為解決專家路由隨機性和數值精度造成的 roundoff 誤差，即路由網絡計算 logits 值過大導致的訓練不收斂問題，工作提出了 z-loss 函數。該函數中，B 表示一個 batch 的 token 數量，N 代表專家數量，X 則對應輸入路由網絡的 logits 的維度。通過引入 z-loss 有效抑制了由精度等因素導致的logits 值過高現象，從而增強了模型訓練的穩定性，z-loss 同樣適用于稠密 LLM 訓練的穩定性提升。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

此外，該工作還深入探究了影響 MoE 訓練穩定性和效果均衡的多種因素，并通過細致實驗歸納出一系列最佳實踐。

譬如，向路由網絡引入適量噪聲 (noise)，例如 dropout 操作，能夠增強訓練穩定性，但可能模型效果受損。在調優 (Fine-tuning) 階段，ST-MoE 的研究進一步發現，增大 dropout 參數對稠密 MoE 模型和非稠密 MoE 模型的影響存在顯著差異，對于稠密模型影響更大。同時，與稠密模型不同，小型 batch size 和較大學習率 (learning rate) 對 MoE 模型訓練效果的提升更為有利。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

Mixtral 模型在實現上借鑒了 Megablocks 框架的獨特思路，該框架的主要特點是將 MoE 層的計算表述為塊稀疏運算。面對前面提及的動態路由問題，即不同專家處理的 token 數量各異，導致用戶在計算過程中需抉擇使用 drop token 或 padding，前者關乎模型性能，后者則會增加訓練成本。

Megablocks 創新性地將多個矩陣乘的操作統一定義為一個大型塊稀疏矩陣，即將多位專家的計算視作一個固定尺寸的大矩陣，其中的計算任務則細分為多個小矩陣塊的 GEMM（General Matrix Multiply）操作。

為優化這種塊稀疏矩陣的處理，Megablocks 框架利用了 Block Compressed Sparse Row (BCSR) 數據結構，對矩陣的行和列訪問速度以及轉換操作進行了優化。通過這種方法實現了路由矩陣的高效操作。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

上述 MoE 發展史詳細說明了其開發需要解決的問題，下文將簡介阿里云為解決這些問題，與 NVIDIA Megatron-Core 團隊的技術合作及產出。Megatron-Core 是一個成熟且輕量化的大規模 LLM 訓練框架，集成了訓練大規模 LLM 的核心技術，比如多元化的模型并行支持、算子優化、通信優化、顯存優化，以及低精度訓練（如 FP8）等先進技術。

Megatron-Core 沿襲了 Megatron-LM 的優秀能力，并在代碼質量、穩定性、功能完備度及測試覆蓋范圍等維度全面提升。尤為關鍵的是，該框架設計上更注重解耦和模塊化，開發者在做二次開發或探索新模型架構時享有高度靈活性。因此我們選擇了與 Megatron-Core 團隊開展合作。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

接下來探討 Megatron-Core 對 MoE 架構的主要支持特性。在并行化方面，Megatron-Core MoE 不僅支持專家并行，還支持 3D 并行，包括數據并行、張量并行、流水并行以及序列并行等。對于超大規模 MoE 模型，它能夠靈活地將專家并行與其他并行策略有機結合。

在 token 分發機制上，Megatron-Core MoE 采用了 dropless MoE 操作，即不丟棄任何 token。在路由和負載均衡優化層面，它支持多種路由策略，如通用的 top-k，并在負載均衡算法上支持 Sinkhorn 算法、z-loss 以及 load balancing loss 等多種方案。

此外，為解決多個專家接收變長輸入問題，Megatron-Core MoE 引入了 GroupedGMM 技術，并優化效率較低的操作，將其替換為優化的 CUDA kernel。

同時，在模型遷移與適配上，Megatron-Core MoE 提供了豐富的模型 checkpoint 轉換功能，允許用戶導入HuggingFace 模型，并自由調整 TP（tensor parallelism）、PP（pipeline parallelism）和 EP（expert parallelism）等結構，隨后利用 Megatron-Core 高效啟動模型訓練任務。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

阿里云人工智能平臺 PAI 團隊基于 Megatron-Core，開發了一套簡易、高效的大模型訓練工具。依托阿里云靈駿計算服務，團隊推出了 PAI-Megatron-Patch 工具庫，可實現從十余種主流開源大模型的模型格式到Megatron-LM 和 Megatron-Core 的無縫轉換。用戶通過 PAI DSW 和 DLC 等產品，能輕松啟動 Megatron-LM 及 Megatron-Core 進行大規模稠密模型和 MoE 模型訓練。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

針對開源大模型多采用 HuggingFace 格式，與 Megatron 訓練框架存在一定差異的問題，PAI-Megatron-patch 的主要功能在于提供從 HuggingFace 模型到 Megatron 框架的權重轉換服務。該轉換實質上是對模型命名空間進行映射，涵蓋了 layernorm 層、attention 層以及 MLP 層等核心組件的定義轉換。PAI-Megatron-patch 內置簡潔的權重轉換腳本，使得用戶能夠便捷地執行模型格式轉換操作，極大地簡化了遷移流程。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

在進行 MoE 模型轉換時，需確保轉換后的文件能直接用于分布式訓練，為此我們也做了大量驗證工作。以Mixtral 模型為例，在擁有 8 個專家和 16 張卡的場景下，采取 TP=4 和 EP=4 的切分策略。每個 TP 分區內的文件夾中包含四個文件，存儲全部八個專家的 FFN 權重，即每個文件承載兩位專家的 FFN 權重信息。

經過轉換，模型在零樣本（Zero-shot）損失精度方面的表現如途中下表所示，數據顯示轉換前后模型的精度差異非常微小，確保了模型轉換的有效性和準確性。

我們的工作不僅關注模型的轉換環節，同樣嚴謹地驗證了整個訓練流程的穩健性。仍以 Mixtral 8x7B 模型為例，我們在三個訓練階段——從頭預訓練、基于 checkpoint 的續訓及指令微調 (Finetune)，均進行了細致測試。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

在從頭預訓練階段，我們設定的參數包括：global_size=256，LR (learning rate) =1e-4，seq_len=2048，TP=4。經過 24 小時至 2.4K 個訓練步驟后，損失 loss 成功收斂至約 1.9。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

在繼續預訓練階段，設定 global_size = 256，LR=5e-5，seq_len=2048，TP=4。在 18 小時達到 2,000 個訓練步驟后，loss 亦表現出正常的收斂行為。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

針對 Finetune 階段，我們做了三組測試:

橙色曲線代表 finetune base 模型，設定 global_size=128，LR=1e-5，seq_len=2048，TP=4。

藍色曲線代表 finetune instruct 模型，其參數配置與橙色曲線一致。

紅色曲線代表僅計算 answer loss 的 finetune instruct 模型，其 global_size=32，LR 依舊為1e-5，seq_len 調整為 128，TP 依舊為 4。

所有模型在訓練過程中均呈現出正常收斂態勢，其中 finetune base 模型導致 loss 下降幅度較大，這一現象符合預期。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

為了進一步驗證 Finetune 的效果，我們選取了一個代碼生成任務作為評估基準。我們利用大約 80K 條公開的代碼訓練樣本，測試集在 HumanEvol 平臺上進行評估，此處的參數配置沿襲了上述 SFT 設定。經過 2,500 步訓練后，模型在 HumanEvol 上的性能指標從最初的 45.73% 顯著提升至 53.05%，有力證明了訓練流程的合理性和有效性。此外，在速度對比方面，該方法優于同等資源條件下運行的 Megablocks。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

我們通過阿里云人工智能平臺 PAI 提供了一套官方最佳實踐指南，指南是對我們從數據處理、訓練、推理、評估、直至在線服務部署的完整 AI 開發流程的測試成果總結：

https://mp.weixin.qq.com/s?__biz=Mzg4MzgxNDk2OA==&mid=2247492431&idx=1&sn=1e135a9e61ac65e88a156969d66ab5be&chksm=cf430418f8348d0e1ed97db178536e73d3ae7093e428c6e3e25bc0dc84a5754f49e3c20ce98e&cur_album_id=2918365856378880004&scene=189#wechat_redirect

該指南涵蓋了從數據準備階段起，指導用戶如何從 OSS （對象存儲服務）和 NAS （網絡附加存儲）讀取原始數據，并在 PAI DSW 環境中執行高效的數據預處理操作。預處理后的數據能夠方便地回存至 NAS 或 OSS，以為后續的模型訓練做好準備。

這種方式的模型訓練支持在 PAI DLC 進行大規模分布式訓練，同時也兼容 DSW 提供的單機訓練環境，確保用戶可根據實際需求靈活選擇訓練方式。訓練完成后，模型的 checkpoint 可以直接導出至 OSS 或 NAS 存儲系統。

完成訓練的 checkpoint 可先進行離線推理和模型性能評估。一旦推理和評估驗證無誤，開發者只需一鍵即可部署至 EAS 提供的模型在線服務。依托此服務的接口，開發者可以輕松構建各類 APP 和業務場景。

圖片來源于 GTC 2024 大會 China AI Day 線上專場的演講

《基于 NVIDIAMegatron-Core 的 MoE LLM 實現和訓練優化》

這份最佳實踐指南提供了詳細的步驟說明，確保用戶能清晰掌握每一步驟的操作方法。

未來，阿里云人工智能平臺 PAI 團隊將繼續深化與 NVIDIA Megatron-Core 團隊的合作，致力于在密集型和稀疏型模型的訓練表現和效率上取得更大突破，為推進 AGI（通用人工智能）技術的發展貢獻力量。我們熱忱歡迎全球開發者共同參與到開源社區項目以及阿里云的建設之中，攜手共進，共創智能未來。

審核編輯：劉清