laekov 在 3月1号下午 3点1分 开源了 FastMoE 系统的代码。合作者是 KEG 实验室 的 xptree 和 sengxian。

简单讲，MoE 模型就是把神经网络里的一层复制若干份，称它们为专家（experts）。在计算时，引入一个额外的门（gate）网络来对于每个输入决定它被哪个专家所计算。

举个例子，一个 nlp 模型里的动词和名词可能会被不同的 expert 处理。这样可能会得到更精准的模型。模型部分 laekov 不太懂，不做更多讨论。

FastMoE 这个名字是 laekov 起的，包含了 laekov 关注的两层问题：

  • 训练这样的模型怎么训得快（fast training）。
  • 如何使得这个系统像吃快餐一样（fastfood-like）开箱即用。

训练性能

通常，expert 网络是一个 MLP，即多层全连接，即矩阵乘。一个值得注意的问题是，gate 会使得一个 batch 被打散。若简单地用 pytorch 来实现，则会导致 GeMM 运算降级为 GeMV 运算，这对 GPU 这样的现代并行计算设备来说是十分不友好的。测试发现朴素的实现在 V100 上仅能达到 500GFLOPs, 即峰值性能的 3% 左右。

注意到 expert 的数量不会太多，故 batch size / expert，即平均每个 expert 处理的 batch 中的元素依然是一个较大的值。因此考虑对 batch 中的 samples 按照 expert 进行重排，每个 expert 依然以 batch 形式处理输入。这样会更充分地利用硬件资源。

另外考虑到当 expert 数量足够大（如 64），则可以挖掘 expert 之间的并行性，即使用不同的 cuda stream 来执行这些矩阵乘，也能得到不错的加速效果。

综合以上两点，laekov 在单 GPU 上实现了 reorder (在代码里被称为 scatter 和 gather）算子，并用 cublas 定制了并行的多个全连接层同时进行前向/反向计算的算子，从而达到了 mlp 层单轮迭代比 PyTorch 快至多 47x 的优化效果。

laekov 发现 scatter-expert-gather 模式也可以很容易地将多 GPU 之间的并行计算不同的 expert 的需求融合进来。即在本地 scatter 之后再进行一次全局的 scatter, gather 也类似。结合 nccl 最新推出的 p2p 和 all2all 通信支持，FastMoE 可以在多个 gpu 的显存里放不同的 expert，并让输入被任意 expert 计算。

然而这里的 all2all 通信模式对于树状的 pcie 连接和网络拓朴十分不友好。目前 laekov 发现仅在 nvlink 的节点内部有较高的性能。其它情况下通信是主要瓶颈。

系统设计

由于目前主流的框架是 pytorch, 故 FastMoE 的实现形式是 PyTorch 的一个 cuda 插件。又由于目前大家喜欢使用 Megatron 来训练自己的 transformer 模型，故 laekov 花了较多精力来支持 megatron. 目前可以做到在 megatron 代码中仅加入两行代码就能用上 fastmoe.

具体来讲，FastMoE 的 python 部分分为三层。

核心层是通用的 FMoE模块，它实现了 scatter-expert-gather 的计算模式。但是 expert 由用户定义，可以是任何神经网络模块，它的输入是目标为某个 expert 的整个 batch 的数据，输出是这个 batch 对应的输出。它也包含了可选+可定制的 Gate 网络。

模型层是 FMoETransformerMLP 模块，它利用前文提到的优化过的线性层 FMoELinear 来实现了一个计算比较快的 Transformer 中的两层 MLP 模块，这里支持了自定义激活函数 activation ，也支持自定义中间隐层的大小，但是总的网络结构在这一层是固定的了。

支持 Megatron 的应用层是 MegatronMLP 层和 fmoefy 函数。MegatronMLP 针对 Megatron 对于 mlp 层的接口进行了适配，而 fmoefy 函数实现了将 megatron 的 model 中的 mlp 层用 MegatronMLP 进行替换。此外，为了定制化同步参数，FastMoE 提供了一个额外的 DDP 模块来代替 Megatron 或 PyTorch 的 DDP，通过一些独特的 tag 来决定梯度同步的范围。

对于 Megatron 模型的使用者，可以用应用层的接口，直接魔改 megatron. 对于其它 transformer 玩家，可以使用模型层的接口，接入自己的 transormer 代码。对于更 general 的 moe 网络玩家，可以使用核心层的接口将任何网络变成 MoE 网络。

感想 & future work

laekov 从 2020 年 12 月开始研究 moe, 并向 xptree 提供了一些基础的 idea 和 sample, 由 xptree 在 transformer-xl 的基础上开发了初版可以跑的 moe 模型及代码。laekov 在一月整理了 prototype 的代码，初步形成了 fastmoe 的框架，在二月开发了 fastmoe 的主体。 sengxian 贡献了单元测试。laekov 和 xptree 和 sengxian 一起在 megatron 上进行了辛苦的 loss curve 的验证与对齐。xptree 承担了与智源宣传组的对接工作。nina 小姐姐在开源前夜绘制了 fastmoe 的 logo. laekov 在此对所有支持了此项目的老师和同学表示感谢。

laekov maybe 会被 challenge 说为什么辛辛苦苦奋战几个月写出来的东西不先拿去投 paper 而是先搞开源，会不会开源后被友商抢走，会不会博士毕不了业qwq。laekov 认为 FastMoE 目前更像是某种基础设施（轮子）。要致富，先修路。laekov 认为把这样类似于道路的基础设施开源，可以减少做科研的众人重复的劳动，从而有动力创造更多有价值的东西。

laekov 将会继续研究分布式 moe 训练系统及性能优化的路径。

Historical Comments

laekov at 2021-03-01T20:57:27

写完这个顺手修了一下自己博客的 github 第三方登录，看着自己 2015 年写的 php 代码，心情复杂。

hefaming at 2021-03-02T09:37:55

赞

hefaming at 2021-03-02T09:37:55

赞

mmmqaq at 2021-03-05T14:36:15

orz