智源x清华开源FastMoE，万亿AI模型基石

北京智源人工智能研究院（以下简称“智源研究院”）和清华大学联合发布首个支持PyTorch框架的高性能MoE系统：FastMoE 。

开源地址：

https://github.com/laekov/fastmoe

FastMoE系统具有易用性强、灵活性好、训练速度快的优势，打破行业限制，可在不同规模的计算机或集群上支持研究者探索不同的MoE模型在不同领域的应用。相比直接使用PyTorch实现的版本，提速47倍。FastMoE是智源研究院于2020年发起的新型超大规模预训练模型研发项目“悟道”的最新成果，由“悟道文汇”（面向认知的超大规模新型预训练模型）和“悟道文溯”（超大规模蛋白质序列预训练模型）两个研究小组联合完成。

MoE是什么？

万亿模型的核心技术

推动预训练模型跨越式发展

却令GPU与PyTorch用户望而却步

MoE（Mixture of Experts）是一个在神经网络中引入若干专家网络（Expert Network）的技术，也是Google最近发布的1.5万亿参数预训练模型Switch
Transformer的核心技术。它对于预训练模型经从亿级参数到万亿级参数的跨越，起了重要推动作用。然而由于其对Google分布式训练框架mesh-tensorflow和Google定制硬件TPU的依赖，给学术界和开源社区的使用与研究带来了不便。

1. MoE 设计：显著增加模型参数量

在ICLR 2017上，Google研究者提出了MoE（Mixture of Experts）层。该层包含一个门网络（Gating Network）和n个专家网络（Expert Network）。对于每一个输入，动态地由门网络选择k个专家网络进行激活。在图1的例子中，门网络决定激活第2个专家网络和第n-1个专家网络。

图1：MoE层的设计（图片来源https://arxiv.org/pdf/1701.06538.pdf Figure 1）

在具体设计中，每个输入x激活的专家网络数量k往往是一个非常小的数字。比如在MoE论文的一些实验中，作者采用了n=512，k=2的设定，也就是每次只会从512个专家网络中挑选两个来激活。在模型运算量（FLOPs）基本不变的情况下，可以显著增加模型的参数量。

2. GShard 和 Switch Transformer，达到惊人的1.5万亿参数量级

在ICLR 2021上，Google的进一步将MoE应用到了基于Transformer的神经机器翻译的任务上。GShard将Transformer中的Feedforward
Network（FFN）层替换成了MoE层，并且将MoE层和数据并行巧妙地结合起来。在数据并行训练时，模型在训练集群中已经被复制了若干份。GShard通过将每路数据并行的FFN看成MoE中的一个专家来实现MoE层，这样的设计通过在多路数据并行中引入All-to-All通信来实现MoE的功能。在论文中，Google使用2048个TPU v3 cores花4天时间训练了一个6千亿参数的模型。

图2：GShard的设计（图片来源https://arxiv.org/pdf/2006.16668.pdf Figure 3）

在2021年1月，Google进一步发布了万亿规模的基于MoE的大规模预训练模型Switch Transformer。Switch Transformer用MoE改进了Google已有的T5预训练模型，其中最大的模型Switch-C已经达到了1.5万亿参数。

3. MMoE：MoE的推荐系统应用

除了在自然语言处理中大放异彩之外，MoE还在推荐系统中找到了一席之地。在KDD 2018中，Google的研究人员提出了MMoE（Multi-gate Mixture-of-Experts），并将其应用到了Google的推荐系统的多任务分类问题中，取得了十分好的效果。随后，Google在RecSys
2019介绍了MMoE在YouTube视频推荐中的应用。类似的MMoE模型也被快手的研究员应用到了快手推荐系统的1.9万亿参数的大规模精排模型中。

FastMoE

首个支持PyTorch框架的MoE系统

简单，灵活，高性能

支持大规模并行训练

MoE潜力巨大，但因为绑定Google软硬件，无法直接应用于PyTorch框架。FastMoE是首个基于当前最流行的PyTorch框架的MoE开源系统，使得普通的用户可以使用常见的GPU资源来尝试和研究自己的MoE模型。与朴素版本相比，实现了47倍的提速优化，更加简单、灵活、高效。

特色一：简单易用，一行代码即可MoE

FastMoE系统既可以作为PyTorch网络中的一个模块使用，也可用于“改造”现有网络中某个层：将其复制多份，并引入Gate，变为MoE层。

例如，对于当前流行的 Megatron-LM 训练系统，仅需要对代码进行如下改动，就可以将 Transformer 模型中的前馈网络（Feed Forward Network）全部替换为 MoE 网络。

特色二：灵活性，支持多种扩展方式

除了传统的两层MLP网络，FastMoE也支持将任意神经网络模块作为专家网络，而进行这样的操作仅需通过修改MoE层构造函数中的一个参数即可实现。

此外，专家选择模块Gate也有较高的研究价值。FastMoE系统目前仅提供了基于单层全连接网络的基础版本，但是通过给定接口，研究者可以方便地使用自己编写的深度神经网络模块作为Gate，从而探索出更好的专家选择方案。

特色三：运行高效，专有性能优化

FastMoE中包含了一些专门优化的CUDA代码。在单块GPU上，相对于一个朴素的PyTorch实现，FastMoE的算子更加充分地利用了GPU大规模并行计算的能力，从而实现多达47倍的加速，从而使得模型研究者可以在更短的时间内验证他们的想法。

FastMoE支持在同一个worker上运行多个experts，从而减少模型研究者在探索更多experts数量时所需的硬件资源。当experts数量较多时，FastMoE针对传统的两层MLP全连接网络（即Transformer中的FFN网络）使用了更精细的并行策略，从而使得Transformer模型中MLP部分的运算速度相比朴素的实现较大的加速。

图3：单GPU多experts情况下，FastMoE相比普通PyTorch实现的加速比。性能的提升主要来自FastMoE 针对传统的两层 MLP 全连接网络（即Transformer中的FFN网络）使用了更精细的并行策略。

单GPU的FastMoE优化配合PyTorch的数据并行，已经可以支持少量专家的MoE分布式训练，这种训练模式被称为FastMoE的数据并行模式。图4展示了一个在2个workers（GPU）上对一个由3个experts构成的MoE网络进行前向计算的例子。

图4：FastMoE数据并行模式，每个worker放置多个experts，worker之间数据并行。top-2 gate指的是门网络会选择激活分数最高的2个专家网络。

FastMoE的数据并行模式已经可以支持许多应用，开发者在著名的Transformer-XL模型上进行了实验。具体来说，Transformer-XL模型中的每一个FFN层（两层的带ReLU激活函数的MLP，隐层大小为512->2048->512）都被一个64选2的专家网络替代（每个专家网络是两层的带ReLU激活函数的MLP，隐层大小为512->1024->512）。这样一来，改造后的FastMoE-Transformer-XL在模型计算量基本不变的情况下，可以获得原始Transformer-XL模型约20倍的参数。如图5所示，改造后的FastMoE-Transformer-XL收敛得比Transformer-XL更快。

图5：FastMoE-Transformer-XL （64个experts）在enwik8数据集上前100K步的Training Loss，其收敛速度显著快于Transformer-XL。

特色四：支持大规模并行训练

图6：FastMoE模型并行模式，每个worker放置多个experts，worker之间进行experts的模型并行。top-2 gate指的是门网络会选择激活分数最高的2个专家网络。

FastMoE还支持在多个worker间以模型并行的方式进行扩展（如图6所示），即不同的worker上放置不同的experts，输入数据在计算前将被传输到所需的worker上，计算后会被传回原来的worker以进行后续计算。通过这种并行方式，模型规模可以以线性扩展，从而支持研究者探索更大规模的模型。这种模式被称为FastMoE的模型并行模式。

值得一提的是，FastMoE已经和英伟达开发的超大规模预训练工具Megatron-LM进行了深度整合，从而使研究者对现有代码做尽量小的修改即可并行运行基于MoE的超大规模预训练模型。开发者在Megatron-LM的GPT模型上进行了测试。如图7所示，类似在Transformer-XL上观察到的现象，一个96个experts的GPT模型可以收敛得比GPT模型更快。

图7：FastMoE-GPT （96个experts）在GPT上前60K步的Training Loss，其收敛速度显著快于GPT。

智源研究院

新型人工智能研究机构

支持科学家勇闯 AI「无人区」

智源研究院是在科技部和北京市委市政府的指导和支持下成立的新型研发机构，旨在聚焦原始创新和核心技术，建立自由探索与目标导向相结合的科研体制，支持科学家勇闯人工智能科技前沿“无人区”。

FastMoE团队成员来自于智源研究院和清华大学计算机系KEG和PACMAN实验室，打通了算法、系统等不同背景的学术人才，由智源研究院学术副院长-清华大学计算机系唐杰教授、智源青年科学家-清华大学计算机系翟季冬副教授、智源青年科学家-循环智能创始人杨植麟博士领导，团队成员有清华大学计算机系博士研究生何家傲、裘捷中以及本科生曾奥涵。