編者按：作為目前唯一被證明能夠將深度學習模型擴充套件到萬億以上引數的方法，MoE 能讓模型學習更多資訊，併為計算機視覺、語音識別、自然語言處理和機器翻譯系統等提供支援。近期，微軟亞洲研究院釋出了一個高效能 MoE 庫——Tutel，並在 8x 80GB NVIDIA A100 和 8x 200Gbps 的 InfiniBand 網路的 Azure NDm A100 v4 節點上進行了實驗。讓我們來看一看，這個用於促進大規模 DNN 模型開發的高效能 MoE 庫有哪些優勢？其表現又如何？
混合專家（Mixture-of-Experts，簡稱 MoE）是一種深度學習模型架構，其計算成本與引數的數量呈次線性關係，因此更容易擴充套件。MoE 是目前唯一被證明能夠將深度學習模型擴充套件到萬億以上引數的方法，它能讓模型學習更多資訊，併為計算機視覺、語音識別、自然語言處理和機器翻譯系統等提供支援，從而以全新的方式為人類社會提供幫助。
近日，微軟亞洲研究院釋出了一個用於促進大規模 DNN 模型開發的高效能 MoE 庫——Tutel，並針對已普遍使用的新 Azure NDm A100 v4 系列進行了高度最佳化。藉助 Tutel 多樣化和靈活的 MoE 演算法支援，AI 領域的開發人員可以更輕鬆、高效地執行 MoE。與最先進的 MoE 實現方式，如 fairseq（Meta的Facebook AI 研究院基於 PyTorch 的 Sequence to Sequence 工具包）相比，對於單個 MoE 層，Tutel 在具有8個 GPU 的單個 NDm A100 v4 節點上實現了8.49倍的加速，在具有512個 A100 GPU 的64個 NDm A100 v4 節點上實現了2.75倍的加速。在端到端效能方面，得益於 all-to-all 通訊最佳化，Tutel 在 Meta（原 Facebook 公司）的1.1萬億引數的 MoE 語言模型中使用64個 NDm A100 v4 節點實現了40%以上的加速。
Tutel 具有良好的相容性和豐富的功能，確保了其在 Azure NDm A100 v4 群集上執行時可以發揮出色的效能。目前，Tutel 已開源，並已整合到 fairseq 中。
Tutel GitHub 連結：
https://github.com/microsoft/tutel
Tutel MoE的三大優勢
作為 fairseq、FastMoE 等其他高水平 MoE 解決方案的補充，Tutel 主要專注於最佳化面向 MoE 的計算和 all-to-all 通訊，以及其他多樣化和靈活的 MoE 演算法支援。Tutel 具有簡潔的介面，可以輕鬆整合到其他 MoE 解決方案中。當然，開發人員也可以從頭開始，利用 Tutel 的介面將獨立的 MoE 層合併到他們自己的 DNN 模型中，直接從高度最佳化的、最先進的 MoE 功能中受益。
與現有的 MoE 解決方案相比，Tutel 具有以下三個主要優勢：
最佳化面向 MoE 的計算。由於缺乏高效的實現方法，目前基於 MoE 的 DNN 模型依賴於深度學習框架（如 PyTorch、TensorFlow 等）提供的多個現成 DNN 運算子的拼接來組成 MoE 計算。由於需要冗餘計算，這種做法會導致顯著的效能開銷。Tutel 設計並實現了多個高度最佳化的 GPU 核心，為面向 MoE 的計算提供了運算子。例如，Tute l將排程“輸出選通（gating output）”的時間複雜度從 O(N^3)降低到 O(N^2)，顯著提高了資料排程的效率。Tutel 還實現了快速 cumsum-minus-one 運算子（fast cumsum-minus-one operator），與 fairseq 實現方式相比，達到了24倍的加速。此外，Tutel 還利用 NVRTC（CUDA C++ 的執行時編譯庫）進一步實時優化了定製的 MoE 核心。
圖1對比了 Tutel 與 faireseq 在 Azure NDm A100 v4 平臺上的執行結果，如前所述，使用 Tutel 的單個 MoE 層在8個 A100 GPU 上實現了8.49倍的加速，在512個 A100 GPU 上實現了2.75倍的加速。
圖1：與 fairseq 相比，對於單個 MoE 層，Tutel 在具有8個 GPU 的單個 NDm A100 v4 節點上實現了8.49倍的加速，在具有512個 A100 GPU 的64個 NDm A100 v4 節點上實現了2.75倍的加速。詳細設定為：batch_size = 32, sequence_length = 1,024, Top_K = 2, model_dim = 2,048, ands hidden_size = 2,048
在 Azure NDm A100 v4 群集上實現底層 all-to-all 通訊最佳化。Tutel 針對 Azure NDm A100 v4 群集上的大規模 MoE 訓練，優化了 all-to-all 聚合通訊（collective communication），其中包括 CPU-GPU 繫結和自適應路由（AR）調整。在非一致儲存訪問結構（NUMA）系統上，尤其是在 NDv4 VM上，正確的 CPU-GPU 繫結對於 all-to-all 效能非常關鍵。但是，現有的機器學習框架（TensorFlow、PyTorch 等）並未提供高效的 all-to-all 通訊庫，導致大規模分散式訓練的效能下降。
Tutel 可以自動最佳化繫結，併為使用者微調提供簡潔的介面。此外，Tutel 在 NDm A100 v4 叢集上使用了多路徑技術，即 AR。對於 MoE 中的 all-to-all 通訊，每個 GPU 通訊的總資料流量規模並不會發生變化，但每個 GPU 之間的資料規模會隨著 GPU 數量的增加而變小。而更小的資料規模會在 all-to-all 通訊中產生更大的開銷，導致 MoE 訓練效能下降。藉助 Azure NDm A100 v4 叢集提供的 AR 技術，Tutel 提高了小訊息組的通訊效率，並在 NDv4 系統上提供了高效能的 all-to-all 通訊。得益於 CPU-GPU 繫結和AR調整，Tutel 使用512個 A100 GPU，對通常用於 MoE 訓練的每次交換的數百兆位元組的訊息，實現了2.56倍到5.93倍的 all-to-all 加速，如圖2所示。
圖2：在應用 Tutel 前後，具有64個 NDm A100 v4 節點（512個 A100 GPU）的不同訊息大小的 all-to-all 頻寬。Tutel 使用512個 A100 GPU，對大小為數百兆位元組的訊息實現了2.56倍到5.93倍的 all-to-all 加速。
多樣、靈活的 MoE 演算法支援。Tutel 為最先進的 MoE 演算法提供了多樣化且靈活的支援，包括：
- 為 Top-K gating 演算法設定任意K值（大多數實現方法僅支援 Top-1 和 Top-2 ）。
- 不同的探索策略，包括批次優先路由、輸入資訊丟失、輸入抖動。
- 不同的精度級別，包括半精度（FP16）、全精度（FP32）、混合精度等（下一個版本中將支援 BF16）。
- 不同的裝置型別，包括 NVIDIA CUDA 和 AMD ROCm 裝置等。
Tutel 將積極整合來自開源社群的各種新興 MoE 演算法。
Tutel與Meta的MoE語言模型整合
此前，Meta 就開源了自己的 MoE 語言模型，並利用 fairseq 實現了 MoE。微軟亞洲研究院與 Meta 合作將 Tutel 整合到了 fairseq 工具包中。Meta 也一直用 Tutel 在 Azure NDm A100 v4 上訓練其大型語言模型，該模型中基於注意力的神經架構類似於 GPT-3。
微軟亞洲研究院的研究員們使用 Meta 的語言模型評估了 Tutel 的端到端效能。該模型有32個注意力層，每個層都有32個128維注意力頭（32x128-dimension heads）。每2層中包含1個 MoE 層，而每個 GPU 都配有一名專家（expert）。表1總結了模型的詳細設定引數，圖3則顯示了Tutel的加速效果。由於 all-to-all 通訊成為瓶頸，所以隨著 GPU 數量的增加，Tutel 帶來的改進從8個 A100 GPU 的131%提升至512個 A100 GPU 的40%。研究員們將在下一個版本中對此做進一步最佳化。
表1：512個 A100 (80G) GPU 的 MoE 語言模型配置
圖3：在端到端效能方面，Tutel 在 Meta 的1.1萬億引數的 MoE 語言模型中使用64個 NDm A100 v4 節點實現40%以上的加速
Tutel將繼續進化，充分發揮MoE潛力
MoE 技術前景廣闊，它可以基於許多領域的技術，實現整體訓練，例如系統路由、海量節點的網路平衡，甚至可以從基於 GPU 的加速中受益。微軟亞洲研究院展示了一個高效的 MoE 實現方式——Tutel，它比 fairseq 框架有著顯著的進步。
現在，Tutel 也已整合到了 DeepSpeed 框架中。相信 Tutel 和相關的整合將使更多 Azure 服務受益，特別是對於那些希望高效擴充套件自有大型模型的客戶。今天的 MoE 仍處於早期階段，需要付出更多努力才能充分發揮 MoE 的潛力，因此研究員們也將繼續最佳化 Tutel，希望未來可以給大家帶來更多令人興奮的研究和應用成果。
