IT之家 11 月 27 日消息，據(jù)微軟亞洲研究院發(fā)布，作為目前唯一被證明能夠將深度學習模型擴展到萬億以上參數(shù)的方法，MoE 能讓模型學習更多信息，并為計算機視覺、語音識別、自然語言處理和機器翻譯系統(tǒng)等提供支持。近期，微軟亞洲研究院發(fā)布了一個高性能 MoE 庫 ——Tutel，并在 8x 80GB NVIDIA A100 和 8x 200Gbps 的 InfiniBand 網(wǎng)絡的 Azure NDm A100 v4 節(jié)點上進行了實驗。讓我們來看一看，這個用于促進大規(guī)模 DNN 模型開發(fā)的高性能 MoE 庫有哪些優(yōu)勢？其表現(xiàn)又如何？
混合專家（Mixture-of-Experts，簡稱 MoE）是一種深度學習模型架構，其計算成本與參數(shù)的數(shù)量呈次線性關系，因此更容易擴展。MoE 是目前唯一被證明能夠將深度學習模型擴展到萬億以上參數(shù)的方法，它能讓模型學習更多信息，并為計算機視覺、語音識別、自然語言處理和機器翻譯系統(tǒng)等提供支持，從而以全新的方式為人類社會提供幫助。
近日，微軟亞洲研究院發(fā)布了一個用于促進大規(guī)模 DNN 模型開發(fā)的高性能 MoE 庫 ——Tutel，并針對已普遍使用的新 Azure NDm A100 v4 系列進行了高度優(yōu)化。借助 Tutel 多樣化和靈活的 MoE 算法支持，AI 領域的開發(fā)人員可以更輕松、高效地執(zhí)行 MoE。與最先進的 MoE 實現(xiàn)方式，如 fairseq（Meta 的 Facebook AI 研究院基于 PyTorch 的 Sequence to Sequence 工具包）相比，對于單個 MoE 層，Tutel 在具有 8 個 GPU 的單個 NDm A100 v4 節(jié)點上實現(xiàn)了 8.49 倍的加速，在具有 512 個 A100 GPU 的 64 個 NDm A100 v4 節(jié)點上實現(xiàn)了 2.75 倍的加速。在端到端性能方面，得益于 all-to-all 通信優(yōu)化，Tutel 在 Meta（原 Facebook 公司）的 1.1 萬億參數(shù)的 MoE 語言模型中使用 64 個 NDm A100 v4 節(jié)點實現(xiàn)了 40% 以上的加速。
Tutel 具有良好的兼容性和豐富的功能，確保了其在 Azure NDm A100 v4 群集上運行時可以發(fā)揮出色的性能。目前，Tutel 已開源，并已集成到 fairseq 中。
Tutel GitHub 鏈接：https://github.com/microsoft/tutel
Tutel MoE 的三大優(yōu)勢
作為 fairseq、FastMoE 等其他高水平 MoE 解決方案的補充，Tutel 主要專注于優(yōu)化面向 MoE 的計算和 all-to-all 通信，以及其他多樣化和靈活的 MoE 算法支持。Tutel 具有簡潔的接口，可以輕松集成到其他 MoE 解決方案中。當然，開發(fā)人員也可以從頭開始，利用 Tutel 的接口將獨立的 MoE 層合并到他們自己的 DNN 模型中，直接從高度優(yōu)化的、最先進的 MoE 功能中受益。
與現(xiàn)有的 MoE 解決方案相比，Tutel 具有以下三個主要優(yōu)勢：
優(yōu)化面向 MoE 的計算。由于缺乏高效的實現(xiàn)方法，目前基于 MoE 的 DNN 模型依賴于深度學習框架（如 PyTorch、TensorFlow 等）提供的多個現(xiàn)成 DNN 運算符的拼接來組成 MoE 計算。由于需要冗余計算，這種做法會導致顯著的性能開銷。Tutel 設計并實現(xiàn)了多個高度優(yōu)化的 GPU 內核，為面向 MoE 的計算提供了運算符。例如，Tute l 將調度“輸出選通（gating output）”的時間復雜度從 O (N^3) 降低到 O (N^2)，顯著提高了數(shù)據(jù)調度的效率。Tutel 還實現(xiàn)了快速 cumsum-minus-one 運算符（fast cumsum-minus-one operator），與 fairseq 實現(xiàn)方式相比，達到了 24 倍的加速。此外，Tutel 還利用 NVRTC（CUDA C++ 的運行時編譯庫）進一步實時優(yōu)化了定制的 MoE 內核。
圖 1 對比了 Tutel 與 faireseq 在 Azure NDm A100 v4 平臺上的運行結果，如前所述，使用 Tutel 的單個 MoE 層在 8 個 A100 GPU 上實現(xiàn)了 8.49 倍的加速，在 512 個 A100 GPU 上實現(xiàn)了 2.75 倍的加速。
圖 1：與 fairseq 相比，對于單個 MoE 層，Tutel 在具有 8 個 GPU 的單個 NDm A100 v4 節(jié)點上實現(xiàn)了 8.49 倍的加速，在具有 512 個 A100 GPU 的 64 個 NDm A100 v4 節(jié)點上實現(xiàn)了 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 群集上實現(xiàn)底層 all-to-all 通信優(yōu)化。Tutel 針對 Azure NDm A100 v4 群集上的大規(guī)模 MoE 訓練，優(yōu)化了 all-to-all 聚合通信（collective communication），其中包括 CPU-GPU 綁定和自適應路由（AR）調整。在非一致存儲訪問結構（NUMA）系統(tǒng)上，尤其是在 NDv4 VM 上，正確的 CPU-GPU 綁定對于 all-to-all 性能非常關鍵。但是，現(xiàn)有的機器學習框架（TensorFlow、PyTorch 等）并未提供高效的 all-to-all 通信庫，導致大規(guī)模分布式訓練的性能下降。
Tutel 可以自動優(yōu)化綁定，并為用戶微調提供簡潔的接口。此外，Tutel 在 NDm A100 v4 集群上使用了多路徑技術，即 AR。對于 MoE 中的 all-to-all 通信，每個 GPU 通信的總數(shù)據(jù)流量規(guī)模并不會發(fā)生變化，但每個 GPU 之間的數(shù)據(jù)規(guī)模會隨著 GPU 數(shù)量的增加而變小。而更小的數(shù)據(jù)規(guī)模會在 all-to-all 通信中產(chǎn)生更大的開銷，導致 MoE 訓練性能下降。借助 Azure NDm A100 v4 集群提供的 AR 技術，Tutel 提高了小消息組的通信效率，并在 NDv4 系統(tǒng)上提供了高性能的 all-to-all 通信。得益于 CPU-GPU 綁定和 AR 調整，Tutel 使用 512 個 A100 GPU，對通常用于 MoE 訓練的每次交換的數(shù)百兆字節(jié)的消息，實現(xiàn)了 2.56 倍到 5.93 倍的 all-to-all 加速，如圖 2 所示。
圖 2：在應用 Tutel 前后，具有 64 個 NDm A100 v4 節(jié)點（512 個 A100 GPU）的不同消息大小的 all-to-all 帶寬。Tutel 使用 512 個 A100 GPU，對大小為數(shù)百兆字節(jié)的消息實現(xiàn)了 2.56 倍到 5.93 倍的 all-to-all 加速。
多樣、靈活的 MoE 算法支持。Tutel 為最先進的 MoE 算法提供了多樣化且靈活的支持，包括：
為 Top-K gating 算法設置任意 K 值（大多數(shù)實現(xiàn)方法僅支持 Top-1 和 Top-2 ）。
不同的探索策略，包括批量優(yōu)先路由、輸入信息丟失、輸入抖動。
不同的精度級別，包括半精度（FP16）、全精度（FP32）、混合精度等（下一個版本中將支持 BF16）。
不同的設備類型，包括 NVIDIA CUDA 和 AMD ROCm 設備等。
Tutel 將積極整合來自開源社區(qū)的各種新興 MoE 算法。
Tutel 與 Meta 的 MoE 語言模型集成
此前，Meta 就開源了自己的 MoE 語言模型，并利用 fairseq 實現(xiàn)了 MoE。微軟亞洲研究院與 Meta 合作將 Tutel 集成到了 fairseq 工具包中。Meta 也一直用 Tutel 在 Azure NDm A100 v4 上訓練其大型語言模型，該模型中基于注意力的神經(jīng)架構類似于 GPT-3。
微軟亞洲研究院的研究員們使用 Meta 的語言模型評估了 Tutel 的端到端性能。該模型有 32 個注意力層，每個層都有 32 個 128 維注意力頭（32x128-dimension heads）。每 2 層中包含 1 個 MoE 層，而每個 GPU 都配有一名專家（expert）。表 1 總結了模型的詳細設置參數(shù)，圖 3 則顯示了 Tutel 的加速效果。由于 all-to-all 通信成為瓶頸，所以隨著 GPU 數(shù)量的增加，Tutel 帶來的改進從 8 個 A100 GPU 的 131% 提升至 512 個 A100 GPU 的 40%。研究員們將在下一個版本中對此做進一步優(yōu)化。
表 1：512 個 A100 (80G) GPU 的 MoE 語言模型配置
圖 3：在端到端性能方面，Tutel 在 Meta 的 1.1 萬億參數(shù)的 MoE 語言模型中使用 64 個 NDm A100 v4 節(jié)點實現(xiàn)了 40% 以上的加速。
MoE 技術前景廣闊，它可以基于許多領域的技術，實現(xiàn)整體訓練，例如系統(tǒng)路由、海量節(jié)點的網(wǎng)絡平衡，甚至可以從基于 GPU 的加速中受益。微軟亞洲研究院展示了一個高效的 MoE 實現(xiàn)方式 ——Tutel，它比 fairseq 框架有著顯著的進步。
現(xiàn)在，Tutel 也已集成到了 DeepSpeed 框架中。相信 Tutel 和相關的集成將使更多 Azure 服務受益，特別是對于那些希望高效擴展自有大型模型的客戶。今天的 MoE 仍處于早期階段，需要付出更多努力才能充分發(fā)揮 MoE 的潛力，因此研究員們也將繼續(xù)優(yōu)化 Tutel，希望未來可以給大家?guī)砀嗔钊伺d奮的研究和應用成果。
廣告聲明：文內含有的對外跳轉鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。