北京智源人工智能研究院和清华大学联合发布首个支持PyTorch框架的高性能MoE系统:FastMoE。开源地址:https://github.com/laekov/fastmoe。FastMoE系统具有易用性强、灵活性好、训练速度快的优势,打破行业限制,可在不同规模的计算机或集群上支持研究者探索不同的MoE模型在不同领域的应用。相比直接使用PyTorch实现的版本,提速47倍。
FastMoE是智源研究院于2020年发起的新型超大规模预训练模型研发项目“悟道”的最新成果,由“悟道文汇”(面向认知的超大规模新型预训练模型)和“悟道文溯”(超大规模蛋白质序列预训练模型)两个研究小组联合完成。
MoE(Mixture of Experts)是一个在神经网络中引入若干专家网络(Expert Network)的技术,也是Google最近发布的1.5万亿参数预训练模型Switch Transformer的核心技术。它对于预训练模型经从亿级参数到万亿级参数的跨越,起了重要作用。
然而由于其对Google分布式训练框架mesh-tensorflow和Google定制硬件TPU的依赖,给学术界和开源社区的使用与研究带来了不便。
在ICLR 2017上,Google研究者提出了MoE(Mixture of Experts)层。该层包含一个门网络(Gating Network)和n个专家网络(Expert Network)。对于每一个输入,动态地由门网络选择k个专家网络进行激活。
在ICLR 2021上,Google的进一步将MoE应用到了基于Transformer的神经机器翻译的任务上。GShard将Transformer中的Feedforward Network(FFN)层替换成了MoE层,并且将MoE层和数据并行巧妙地结合起来。在2021年1月,Google进一步发布了万亿规模的基于MoE的大规模预训练模型Switch Transformer。
除了在自然语言处理中大放异彩之外,MoE还在推荐系统中找到了一席之地。在KDD 2018中,Google的研究人员提出了MMoE(Multi-gate Mixture-of-Experts),并将其应用到了Google的推荐系统的多任务分类问题中,取得了十分好的效果。
FastMoE是首个基于当前最流行的PyTorch框架的MoE开源系统,使得普通的用户可以使用常见的GPU资源来尝试和研究自己的MoE模型。与朴素版本相比,实现了47倍的提速优化,更加简单、灵活、高效。
FastMoE系统既可以作为PyTorch网络中的一个模块使用,也可用于“改造”现有网络中某个层:将其复制多份,并引入Gate,变为MoE层。
除了传统的两层MLP网络,FastMoE也支持将任意神经网络模块作为专家网络,而进行这样的操作仅需通过修改MoE层构造函数中的一个参数即可实现。FastMoE中包含了一些专门优化的CUDA代码。在单块GPU上,相对于一个朴素的PyTorch实现,FastMoE的算子更加充分地利用了GPU大规模并行计算的能力,从而实现多达47倍的加速。
FastMoE还支持在多个worker间以模型并行的方式进行扩展,即不同的worker上放置不同的experts,输入数据在计算前将被传输到所需的worker上,计算后会被传回原来的worker以进行后续计算。通过这种并行方式,模型规模可以以线性扩展,从而支持研究者探索更大规模的模型。
智源研究院是在科技部和北京市委市政府的指导和支持下成立的新型研发机构,旨在聚焦原始创新和核心技术,建立自由探索与目标导向相结合的科研体制,支持科学家勇闯人工智能科技前沿“无人区”。