北京智源人工智能研究院(以下简称“智源研究院”)和清华大学联合发布首个支持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实验室,打通了算法、系统等不同背景的学术人才,由智源研究院学术副院长-清华大学计算机系唐杰教授、智源青年科学家-清华大学计算机系翟季冬副教授、智源青年科学家-循环智能创始人杨植麟博士领导,团队成员有清华大学计算机系博士研究生何家傲、裘捷中以及本科生曾奥涵。

智源x清华开源FastMoE,万亿AI模型基石相关推荐

  1. 《智源社区周刊:预训练模型》第3期:智源x清华开源万亿AI模型基石FastMoE、英国机构发起世界最大图灵测试...

    超大规模预训练模型是当前人工智能领域研究的热点,为了帮助研究与工程人员了解这一领域的进展和资讯,智源社区整理了第3期<智源社区周刊:预训练模型>,从论文推荐.研究动态.热点讨论等几个维度推 ...

  2. 交互式多模型_26亿参数,智源、清华开源中文大规模预训练模型

    近日,北京智源人工智能研究院和清华大学研究团队联合发布了以中文为核心的大规模预训练语言模型 CPM-LM,参数规模达 26 亿,预训练中文数据规模 100 GB. 26亿参数,智源.清华开源中文大规模 ...

  3. 强强联合!智源x清华AIR,共启 AI 健康研究新篇章

    北京智源人工智能研究院 清华大学智能产业研究院(AIR) 两所极具代表性的创新型 AI 研究院,强强联手成立「清华(AIR)- 智源健康计算联合研究中心」 ,马维英任联合中心主任. 中心将致力于通过人 ...

  4. 2021北京智源大会圆满闭幕,现场8千人,参会7万人,发布「悟道2.0」全球最大万亿级模型...

    句点,是下一行诗篇的开始. 6月3日晚上10点整,2021北京智源大会正式闭幕,为北京乃至中国人工智能发展画上浓重一点. 定位于"内行的AI盛会",北京智源大会以国际性.权威性.专 ...

  5. 智源研究院:开源开放让AI快速发展,要建设大模型时代的Linux | 中国AIGC产业峰会...

    萧箫 整理自 AIGC峰会 量子位 | 公众号 QbitAI ChatGPT引爆了大模型,也彻底将大模型相关的AI产业生态带到了新的阶段-- 大模型的"涌现能力",让AI真正展现出 ...

  6. 中文版GPT-3来了?智源、清华发布清源 CPM——以中文为核心的大规模预训练模型

    2020-11-18 23:43:21 清源 CPM(Chinese Pretrained Models)是北京智源人工智能研究院和清华大学研究团队合作开展的大规模预训练模型开源计划,清源计划是以中文 ...

  7. 大模型技术实践路线图出炉,唐杰牵头,近百位智源、清华、北大、MILA等顶尖研究者联合推出...

    导读:随着以深度学习为代表的AI技术的快速发展,智能模型的训练应用模式逐渐由"大炼模型"向"炼大模型"转变.过去,人们根据不同的任务,利用特定的标注数据训练大量 ...

  8. 解密Elasticsearch技术,腾讯开源的万亿级分布式搜索分析引擎

    「免费学习 60+ 节公开课:投票页面,点击讲师头像」 作者 | johngqjiang,腾讯 TEG 云架构平台部研发工程师 来源 | 腾讯技术工程(ID:Tencent_TEG) [导读]Elas ...

  9. 免费开放!清华、智源、协和联合开发乳腺癌AI诊断工具,可预测分子亚型,准确率高达76%...

    文:贾伟 使用AI技术来检测乳腺癌,如今已有诸多研究,但大多是基于磁共振.乳腺X线影像进行预测.而超声是中国女性乳腺癌筛查的主要方式.如何利用好超声影像,进行乳腺癌筛查,并对乳腺癌做分子亚型分类,是中 ...

最新文章

  1. 张量功率谱CAMB参数调试
  2. java 关于日期的一些常用惯例
  3. 一站式学习Redis 从入门到高可用分布式实践(慕课)第十章 缓存设计与优化
  4. 新元素之video,audio,meter,datalist,keygen,output
  5. 为什么lp的最优解是一个概率_什么时候应该用概率图模型、消息传递替代传统组合优化求解器?未来工作?(持续更新)...
  6. 使用外置tomcat部署spring boot程序需添加servlet-api
  7. Spring cloud集成Rabbitmq
  8. Spring : Spring @Transactional事物创建
  9. 零基础如何学前端,如何规划?
  10. netsh winsock reset命令,作用是重置 Winsock 目录
  11. FFmpeg 转码压缩
  12. 各省市大学 全国各个省市重点民办本专科大学
  13. Labview实现AM调制
  14. 洛谷 P1725 琪露诺
  15. 如何手机桌面加计算机,教你自制手机、电脑壁纸|手机一键加字、加光晕特效,4步搞定!...
  16. Mac字体路径,解决matplotlib中文显示问题,超级好用~
  17. android 9.0 10.0 去掉音量键+Power键组合键的屏幕截图功能
  18. python kayb算法之从一组序列当中获取一组与目标值最接近的算法
  19. maven阿里镜像下载jar包报错
  20. 2610_lab2 命令行与快捷键

热门文章

  1. mysql 开仓函数_MySQL函数大全 及用法示例
  2. python分类流程_文本分类指南:你真的要错过 Python 吗?
  3. 支付宝支付成功之后的接口_支付宝放支付大招,花呗借呗之后又来一“呗”,快告诉家人吧!...
  4. Android怎么不显示手机模型,以编程方式获取Android手机模型
  5. 7-7 12-24小时制 (C语言)
  6. win10 64位SSDT函数索引动态查找
  7. redhat7操作系统安装
  8. 1084. [SCOI2005]最大子矩阵【网格DP】
  9. 通过反射获取及调用方法(Method)
  10. netstat常用命令