智源聯(lián)合清華發(fā)布首個支持PyTorch框架的高性能MoE系統(tǒng)

北京智源人工智能研究院（以下簡稱“智源研究院”）和清華大學(xué)聯(lián)合發(fā)布首個支持 PyTorch 框架的高性能 MoE 系統(tǒng)：FastMoE 。

FastMoE 系統(tǒng)具有易用性強(qiáng)、靈活性好、訓(xùn)練速度快的優(yōu)勢，打破行業(yè)限制，可在不同規(guī)模的計算機(jī)或集群上支持研究者探索不同的 MoE 模型在不同領(lǐng)域的應(yīng)用。相比直接使用 PyTorch 實現(xiàn)的版本，提速 47 倍。FastMoE 是智源研究院于 2020 年發(fā)起的新型超大規(guī)模預(yù)訓(xùn)練模型研發(fā)項目“悟道”的最新成果，由“悟道文匯”（面向認(rèn)知的超大規(guī)模新型預(yù)訓(xùn)練模型）和“悟道文溯”（超大規(guī)模蛋白質(zhì)序列預(yù)訓(xùn)練模型）兩個研究小組聯(lián)合完成。

MoE 是什么？萬億模型的核心技術(shù)，推動預(yù)訓(xùn)練模型跨越式發(fā)展，卻令 GPU 與 PyTorch 用戶望而卻步。

MoE（Mixture of Experts）是一個在神經(jīng)網(wǎng)絡(luò)中引入若干專家網(wǎng)絡(luò)（Expert Network）的技術(shù)，也是 Google 最近發(fā)布的 1.5 萬億參數(shù)預(yù)訓(xùn)練模型 Switch Transformer 的核心技術(shù)。它對于預(yù)訓(xùn)練模型經(jīng)從億級參數(shù)到萬億級參數(shù)的跨越，起了重要推動作用。然而由于其對 Google 分布式訓(xùn)練框架 mesh-tensorflow 和 Google 定制硬件 TPU 的依賴，給學(xué)術(shù)界和開源社區(qū)的使用與研究帶來了不便。

MoE 設(shè)計：顯著增加模型參數(shù)量

在 ICLR 2017 上，Google 研究者提出了 MoE（Mixture of Experts）層。該層包含一個門網(wǎng)絡(luò)（Gating Network）和 n 個專家網(wǎng)絡(luò)（Expert Network）。對于每一個輸入，動態(tài)地由門網(wǎng)絡(luò)選擇 k 個專家網(wǎng)絡(luò)進(jìn)行激活。在圖 1 的例子中，門網(wǎng)絡(luò)決定激活第 2 個專家網(wǎng)絡(luò)和第 n-1 個專家網(wǎng)絡(luò)。

圖 1：MoE 層的設(shè)計（圖片來源 https://arxiv.org/pdf/1701.06538.pdfFigure 1）

在具體設(shè)計中，每個輸入 x 激活的專家網(wǎng)絡(luò)數(shù)量 k 往往是一個非常小的數(shù)字。比如在 MoE 論文的一些實驗中，作者采用了 n=512，k=2 的設(shè)定，也就是每次只會從 512 個專家網(wǎng)絡(luò)中挑選兩個來激活。在模型運(yùn)算量（FLOPs）基本不變的情況下，可以顯著增加模型的參數(shù)量。

GShard 和 Switch Transformer，達(dá)到驚人的 1.5 萬億參數(shù)量級

在 ICLR 2021 上，Google 的進(jìn)一步將 MoE 應(yīng)用到了基于 Transformer 的神經(jīng)機(jī)器翻譯的任務(wù)上。GShard 將 Transformer 中的 Feedforward Network（FFN）層替換成了 MoE 層，并且將 MoE 層和數(shù)據(jù)并行巧妙地結(jié)合起來。在數(shù)據(jù)并行訓(xùn)練時，模型在訓(xùn)練集群中已經(jīng)被復(fù)制了若干份。GShard 通過將每路數(shù)據(jù)并行的 FFN 看成 MoE 中的一個專家來實現(xiàn) MoE 層，這樣的設(shè)計通過在多路數(shù)據(jù)并行中引入 All-to-All 通信來實現(xiàn) MoE 的功能。在論文中，Google 使用 2048 個 TPU v3 cores 花 4 天時間訓(xùn)練了一個 6 千億參數(shù)的模型。

圖 2：GShard 的設(shè)計（圖片來源 https://arxiv.org/pdf/2006.16668.pdfFigure 3）

在 2021 年 1 月，Google 進(jìn)一步發(fā)布了萬億規(guī)模的基于 MoE 的大規(guī)模預(yù)訓(xùn)練模型 Switch Transformer。Switch Transformer 用 MoE 改進(jìn)了 Google 已有的 T5 預(yù)訓(xùn)練模型，其中最大的模型 Switch-C 已經(jīng)達(dá)到了 1.5 萬億參數(shù)。

MMoE：MoE 的推薦系統(tǒng)應(yīng)用

除了在自然語言處理中大放異彩之外，MoE 還在推薦系統(tǒng)中找到了一席之地。在 KDD 2018 中，Google 的研究人員提出了 MMoE（Multi-gate Mixture-of-Experts），并將其應(yīng)用到了 Google 的推薦系統(tǒng)的多任務(wù)分類問題中，取得了十分好的效果。隨后，Google 在 RecSys 2019 介紹了 MMoE 在 YouTube 視頻推薦中的應(yīng)用。類似的 MMoE 模型也被快手的研究員應(yīng)用到了快手推薦系統(tǒng)的 1.9 萬億參數(shù)的大規(guī)模精排模型中。

FastMoE 是首個支持 PyTorch 框架的 MoE 系統(tǒng)，簡單，靈活，高性能，支持大規(guī)模并行訓(xùn)練

MoE 潛力巨大，但因為綁定 Google 軟硬件，無法直接應(yīng)用于 PyTorch 框架。FastMoE 是首個基于當(dāng)前最流行的 PyTorch 框架的 MoE 開源系統(tǒng)，使得普通的用戶可以使用常見的 GPU 資源來嘗試和研究自己的 MoE 模型。與樸素版本相比，實現(xiàn)了 47 倍的提速優(yōu)化，更加簡單、靈活、高效。

特色一：簡單易用，一行代碼即可 MoE

FastMoE 系統(tǒng)既可以作為 PyTorch 網(wǎng)絡(luò)中的一個模塊使用，也可用于“改造”現(xiàn)有網(wǎng)絡(luò)中某個層：將其復(fù)制多份，并引入 Gate，變?yōu)?MoE 層。

例如，對于當(dāng)前流行的 Megatron-LM 訓(xùn)練系統(tǒng)，僅需要對代碼進(jìn)行如下改動，就可以將 Transformer 模型中的前饋網(wǎng)絡(luò)（Feed Forward Network）全部替換為 MoE 網(wǎng)絡(luò)。

特色二：靈活性，支持多種擴(kuò)展方式

除了傳統(tǒng)的兩層 MLP 網(wǎng)絡(luò)，F(xiàn)astMoE 也支持將任意神經(jīng)網(wǎng)絡(luò)模塊作為專家網(wǎng)絡(luò)，而進(jìn)行這樣的操作僅需通過修改 MoE 層構(gòu)造函數(shù)中的一個參數(shù)即可實現(xiàn)。

此外，專家選擇模塊 Gate 也有較高的研究價值。FastMoE 系統(tǒng)目前僅提供了基于單層全連接網(wǎng)絡(luò)的基礎(chǔ)版本，但是通過給定接口，研究者可以方便地使用自己編寫的深度神經(jīng)網(wǎng)絡(luò)模塊作為 Gate，從而探索出更好的專家選擇方案。

特色三：運(yùn)行高效，專有性能優(yōu)化

FastMoE 中包含了一些專門優(yōu)化的 CUDA 代碼。在單塊 GPU 上，相對于一個樸素的 PyTorch 實現(xiàn)，F(xiàn)astMoE 的算子更加充分地利用了 GPU 大規(guī)模并行計算的能力，從而實現(xiàn)多達(dá) 47 倍的加速，從而使得模型研究者可以在更短的時間內(nèi)驗證他們的想法。

FastMoE 支持在同一個 worker 上運(yùn)行多個 experts，從而減少模型研究者在探索更多 experts 數(shù)量時所需的硬件資源。當(dāng) experts 數(shù)量較多時，F(xiàn)astMoE 針對傳統(tǒng)的兩層 MLP 全連接網(wǎng)絡(luò)（即 Transformer 中的 FFN 網(wǎng)絡(luò)）使用了更精細(xì)的并行策略，從而使得 Transformer 模型中 MLP 部分的運(yùn)算速度相比樸素的實現(xiàn)較大的加速。

圖 3：單 GPU 多 experts 情況下，F(xiàn)astMoE 相比普通 PyTorch 實現(xiàn)的加速比。性能的提升主要來自 FastMoE 針對傳統(tǒng)的兩層 MLP 全連接網(wǎng)絡(luò)（即 Transformer 中的 FFN 網(wǎng)絡(luò)）使用了更精細(xì)的并行策略。

單 GPU 的 FastMoE 優(yōu)化配合 PyTorch 的數(shù)據(jù)并行，已經(jīng)可以支持少量專家的 MoE 分布式訓(xùn)練，這種訓(xùn)練模式被稱為 FastMoE 的數(shù)據(jù)并行模式。圖 4 展示了一個在 2 個 workers（GPU）上對一個由 3 個 experts 構(gòu)成的 MoE 網(wǎng)絡(luò)進(jìn)行前向計算的例子。

圖 4：FastMoE 數(shù)據(jù)并行模式，每個 worker 放置多個 experts，worker 之間數(shù)據(jù)并行。top-2 gate 指的是門網(wǎng)絡(luò)會選擇激活分?jǐn)?shù)最高的 2 個專家網(wǎng)絡(luò)。

FastMoE 的數(shù)據(jù)并行模式已經(jīng)可以支持許多應(yīng)用，開發(fā)者在著名的 Transformer-XL 模型上進(jìn)行了實驗。具體來說，Transformer-XL 模型中的每一個 FFN 層（兩層的帶 ReLU 激活函數(shù)的 MLP，隱層大小為 512->2048->512）都被一個 64 選 2 的專家網(wǎng)絡(luò)替代（每個專家網(wǎng)絡(luò)是兩層的帶 ReLU 激活函數(shù)的 MLP，隱層大小為 512->1024->512）。這樣一來，改造后的 FastMoE-Transformer-XL 在模型計算量基本不變的情況下，可以獲得原始 Transformer-XL 模型約 20 倍的參數(shù)。如圖 5 所示，改造后的 FastMoE-Transformer-XL 收斂得比 Transformer-XL 更快。

圖 5：FastMoE-Transformer-XL （64 個 experts）在 enwik8 數(shù)據(jù)集上前 100K 步的 Training Loss，其收斂速度顯著快于 Transformer-XL。

特色四：支持大規(guī)模并行訓(xùn)練

圖 6：FastMoE 模型并行模式，每個 worker 放置多個 experts，worker 之間進(jìn)行 experts 的模型并行。top-2 gate 指的是門網(wǎng)絡(luò)會選擇激活分?jǐn)?shù)最高的 2 個專家網(wǎng)絡(luò)。

FastMoE 還支持在多個 worker 間以模型并行的方式進(jìn)行擴(kuò)展（如圖 6 所示），即不同的 worker 上放置不同的 experts，輸入數(shù)據(jù)在計算前將被傳輸?shù)剿璧?worker 上，計算后會被傳回原來的 worker 以進(jìn)行后續(xù)計算。通過這種并行方式，模型規(guī)?？梢砸跃€性擴(kuò)展，從而支持研究者探索更大規(guī)模的模型。這種模式被稱為 FastMoE 的模型并行模式。

值得一提的是，F(xiàn)astMoE 已經(jīng)和英偉達(dá)開發(fā)的超大規(guī)模預(yù)訓(xùn)練工具 Megatron-LM 進(jìn)行了深度整合，從而使研究者對現(xiàn)有代碼做盡量小的修改即可并行運(yùn)行基于 MoE 的超大規(guī)模預(yù)訓(xùn)練模型。開發(fā)者在 Megatron-LM 的 GPT 模型上進(jìn)行了測試。如圖 7 所示，類似在 Transformer-XL 上觀察到的現(xiàn)象，一個 96 個 experts 的 GPT 模型可以收斂得比 GPT 模型更快。

圖 7：FastMoE-GPT （96 個 experts）在 GPT 上前 60K 步的 Training Loss，其收斂速度顯著快于 GPT。

智源研究院

新型人工智能研究機(jī)構(gòu)、支持科學(xué)家勇闖 AI「無人區(qū)」

智源研究院是在科技部和北京市委市政府的指導(dǎo)和支持下成立的新型研發(fā)機(jī)構(gòu)，旨在聚焦原始創(chuàng)新和核心技術(shù)，建立自由探索與目標(biāo)導(dǎo)向相結(jié)合的科研體制，支持科學(xué)家勇闖人工智能科技前沿“無人區(qū)”。

FastMoE 團(tuán)隊成員來自于智源研究院和清華大學(xué)計算機(jī)系 KEG 和 PACMAN 實驗室，打通了算法、系統(tǒng)等不同背景的學(xué)術(shù)人才，由智源研究院學(xué)術(shù)副院長 - 清華大學(xué)計算機(jī)系唐杰教授、智源青年科學(xué)家 - 清華大學(xué)計算機(jī)系翟季冬副教授、智源青年科學(xué)家 - 循環(huán)智能創(chuàng)始人楊植麟博士領(lǐng)導(dǎo)，團(tuán)隊成員有清華大學(xué)計算機(jī)系博士研究生何家傲、裘捷中以及本科生曾奧涵。

原文標(biāo)題：首個支持 PyTorch 框架的 MoE 系統(tǒng)來了！智源聯(lián)合清華開源FastMoE，萬億AI模型基石

文章出處：【微信公眾號：通信信號處理研究所】歡迎添加關(guān)注！文章轉(zhuǎn)載請注明出處。

責(zé)任編輯：haq