FastMoE 系統
https://github.com/laekov/fastmoe
簡介
FastMoE 是一個易用且高效的基於 PyTorch 的 MoE 模型訓練系統.
安裝
依賴
啓用了 CUDA 的 PyTorch 是必要的. 當前版本的 FastMoE 在 PyTorch v1.8.0 和 CUDA 10
的平臺上經過了測試. 本系統從設計上也支持更舊的 PyTorch 版本.
如果需要使能 FastMoE 模型並行特性, 那麼支持點對點通信的 NCCL 庫 (即不舊於
2.7.5 版本) 也是必需的.
安裝
FastMoE 包含一些定製的 PyTorch 算子, 包含一些 C 的組件. 用 python setup.py install
來簡單地安裝 FastMoE.
FastMoE 分佈式模型並行特性默認是不被啓用的. 如果它需要被啓用,
則需要在運行上述命令時加入環境變量 USE_NCCL=1.
注意, 由於 PyTorch 框架通常僅集成了 NCCL 的運行時組件, 額外的 NCCL
開發包需要被安裝在編譯環境中, 而且它的版本需要與 PyTorch 的版本相對應. 推薦使用
PyTorch 官方 Docker 鏡像,
因爲那裏的環境較爲乾淨. 如果您希望手工配置環境, 可以在 NCCL
全部版本的下載鏈接
下載合適版本的 NCCL 開發包.
使用
將一個 Transformer 模型 FMoE 化
Transformer 是當前最流行的可被 MoE 化的模型. FastMoE 可以一鍵將一個普通的
Transformer 模型變爲一個 MoE 的模型. 其使用方法如下.
例如在 Megatron-LM 中,
添加如下的代碼即可將 Transformer 中的每個 MLP 層變爲多個 MLP 層構成的 MoE 網絡.
model = ...
from fmoe.megatron import fmoefy
model = fmoefy(model, num_experts=<number of experts per worker>)
train(model, ...)
一個更詳細的在 Megatron-LM 中使用 fmoefy 函數的樣例參見此處.
將 FastMoE 作爲一個網絡模塊使用
一個使用 FastMoE 的 Transformer 模型見這個示例.
最簡單的使用方式是使用 FMoE 層來代替 MLP 層.
分佈式地使用 FastMoE
FastMoE 支持數據並行和模型並行.
數據並行.
在 FastMoE 的數據並行模式下,
門網絡(gate)和專家網絡都被複制地放置在各個運算單元上.
下圖展示了一個有三個專家的兩路數據並行MoE模型進行前向計算的方式.
對於數據並行, 額外的代碼是不需要的. FastMoE 與 PyTorch 的 DataParallel 和
DistributedDataParallel 模塊都可以無縫對接. 該方式唯一的問題是,
專家的數量受到單個計算單元(如GPU)的內存大小限制.
模型並行
在 FastMoE 的模型並行模式中, 門網絡依然是複製地被放置在每個計算單元上的,
但是專家網絡被獨立地分別放置在各個計算單元上. 因此, 通過引入額外的通信操作,
FastMoE 可以允許更多的專家網絡們同時被訓練,
而其數量限制與計算單元的數量是正相關的.
下圖展示了一個有六個專家網絡的模型被兩路模型並行地訓練.
注意專家1-3被放置在第一個計算單元上, 而專家4-6被放置在第二個計算單元上.
FastMoE 的模型並行模式需要專門的並行策略, 而 PyTorch 和 Megatron-LM
都不支持這樣的策略. 因此, 需要使用 fmoe.DistributedGroupedDataParallel
模塊來代替 PyTorch 的 DDP 模塊.