一个易用且高效的基于 PyTorch 的 MoE 模型训练系统.

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 模块.