【前沿进展】训练参数规模万亿的预训练模型，对于超级计算机而言是不小的挑战。如何提升超算的计算效率，实现更大规模的参数训练，成为近年来研究者探索的课题。在近日举办的Big Model Meetup第二期活动中，清华大学长聘副教授，智源青年科学家翟季冬做特邀报告，介绍了在国产超级计算机上训练百万亿参数的超大规模预训练模型解决方案。智源社区对内容进行了整理。
活动回看地址：https://event.baai.ac.cn/activities/175
讲者简介：翟季冬，清华大学计算机系长聘副教授，博士生导师。现为清华大学计算机系高性能所副所长，ACM中国高性能计算专家委员会秘书长、北京智源青年科学家。主要研究方向包括高性能计算、性能评测和编译优化等。研究成果发表在相关领域顶级学术会议和期刊——SC、ICS、PPOPP、ASPLOS、MICRO、OSDI、ATC、IEEE TC、IEEE TPDS等。研究成果获ACM ICS 2021最佳学生论文奖、SC 2014 Best Paper Finalist、ICDCS 2020 Best Paper Honorable Mention奖。担任NPC 2018程序委员会主席、IEEE Cluster 2021领域主席、SC 2022领域副主席，SC、ICS、PPOPP等国际学术会议程序委员会委员。目前担任《IEEE Transactions on Computers》、《IEEE Transactions on Parallel and Distributed Systems》、《IEEE Transactions on Cloud Computing》等多个国际学术期刊编委。担任清华大学学生超算团队教练，指导的团队十一次获得世界冠军。获教育部科技进步一等奖、中国电子学会科学技术一等奖、中国计算机学会优秀博士学位论文奖、IEEE TPDS杰出编委奖、国家自然科学基金优秀青年科学基金（2017）、CCF-IEEE CS青年科学家奖。
背景
近几年，如果大家经常关注课题新闻，就能发现大规模预训练模型在很多领域都有较强的影响力，如谷歌最新的推荐系统、搜索引擎、国内阿里的淘宝等，它们的一些图像生成任务都采用了预训练模型。下图中右侧的牛油果形状的扶手椅的小图是在网上找的，这些图片本来是不存在的，都是通过预训练模型来生成的。
从计算的角度来考虑，预训练模型的核心主要是Transformer，该模型主要的计算集中在嵌入层、注意力层还有前馈网络。如果从底层的计算来看，这些问题的核心都是矩阵乘法，而在高性能计算领域，实际上对矩阵乘法已经开展做了大量的优化工作。
预训练模型的发展趋势
从下图可以看到，横坐标是模型的参数量，从最早的GPT大约是1.1亿个参数，然后到最新的GPT-3和Switch Transformer达到了上万亿的参数。至少从目前来看，探索更大参数量的模型仍具有很重要的意义。
从计算的角度来看，模型规模在变的越来越大、训练数据逐步增多，此时，单机就无法满足非常大规模模型的训练。
比如说最经典的串行训练，在训练过程中，每一轮迭代会产生梯度，然后去更新模型，这个过程不断重复迭代，最后会达到收敛的状态。当这个模型变得很大的时候，会有一些典型的并行策略来处理模型，最常见的是数据并行和模型并行。
数据并行是把输入数据进行拆分，然后模型在不同节点上有都会有一个完整的拷贝。数据并行过程中，节点两边各得到部分一个梯度，经过一次all-reduce的通信，交互之后更新模型，并最终产生一个全局的通信。
模型并行，就是当模型很大，可以把模型中间切一刀，然后分成两部分。由于切的方向不一样，比如对于一个矩阵乘，可以有不同的切法 ，可能会引入all-gather或all-reduce不同的通信，最后两边分别更新。
上面只是简单介绍了数据并行和模型并行，核心观点是当模型变得很大的时候，一定要以并行的方式把数据或者模型进行拆分。
并行训练必不可少
下面主要从计算的角度来介绍下MoE这一模型。看预训练模型，主要是看对一些输入产生的一些输出。模型增大有非常大的好处，模型增大的一个简单策略是把中间的模型变得非常大，原先的矩阵乘法可能是几万维的，现在可以把它变成几十万、上百万维的，这样就变成一个更大的矩阵乘。
实际上，对于输入，如中文翻译为英文，很有可能只用到了模型的一部分，此时在输入向量后面增加模型的一个网关，这样对于不同的输入可以选择不同的小的预训练模型，对于每次的输入向量只是选择一部分，最后再把选择的结果做一些加权求和来得到输出，这就是MoE的基本思想。
对于MoE而言，中间所有的小模型的总和是整个模型的参数量，它的量和稠密大模型相比也非常大。然而，对于每次运算，只是从离散的、稀疏的取一部分，计算量并没有增加很多，模型参数量却增加了很多。
MoE模型里面的每个小规模模型就叫做专家，MoE就是混合专家系统，MoE的优势就是扩展了模型参数，提高了模型规模，在相同的训练时间，效果会优于稠密大模型，谷歌最新的Switch Transformer就采用了MoE的架构。
MoE模型的并行训练机制
MoE大模型并行训练会引入两种划分，第一个是数据，对数据进行切分。第二个是专家，要对专家进行切分。可以看到，对数据进行拆分在计算运算过程中还需要之间的交互，这个交互引入了全局All-to-All通信，这就是MoE模型并行训练的方式。
大规模预训练模型在很大的系统上进行训练时会面临很多挑战，挑战主要分为四类，下面分别详细的介绍。
为了训练一个大的模型，可以有不同的并行策略，如数据并行、模型并行和MoE并行，另外还有流水线并行等各种各样的并行策略，这些并行策略从并行的角度把一个大的模型按照不同的方式并行之后，它的核心对底层计算和通信的需求会不一样。
简而言之，即很有可能部分并行策略的通信需求会比较小，会引入更多的计算，而有的并行策略计算比较小但会引入更多通信。那么，面对这样的问题该如何来选择最优策略呢？
举一个简单的示例，以16个节点为例，训练的时候包括两部分，一个是模型的参数，另一个是输入的数据，主要是这两部分。
如果是简单数据并行，模型参数并不划分，每台机器上都有一份完整复制的模型，数据并行的核心是把输入数据简单的切分成 16份。类似的，模型并行要把模型切成16份。
如果是MoE，则有16个专家，然后把它切成16个，每个节点一个专家。不过输入数据也要把它切成16份。当然，并行有多种组合方式，比如数据并行加模型并行的混合方式，对于模型并行，在4个节点上切分一份参数，对于数据并行来说，复制4份参数；对于数据并行，各组可以使用不同输入。
同理，也可以有数据和MoE的并行，混合到一起，这里就不详细讲了。既然有这么多的并行策略可以选择，那么如何选取高效的并行策略就变成了非常有挑战的一件事。
以万亿参数的模型为例，如果模型精度是32位，模型的参数就有4T，模型的梯度也有4T，优化器的更新参数有8T，计算中间值也有8T。如果用V100训练这个模型，并且把这些数据存储下来，就需要768块V100。
这些数据如何划分？这些参数是均匀划分还是部分划分？不同的划分方式对底层的计算和通信也会产生不同的影响，即高效存储相关数据来支持高效训练也非常有挑战。
大家知道采用更低的精度需要的内存量更小、性能更好，不过，低精度会带来模型准确度下降。如果用更高精度训练，如32位或者64位，需要的内存量会变大、计算的要求会更高。
以Multibox SSD模型为例，横坐标表示参数数值的指数分布。如果用FPfp16的范围来表示，紫色区域表示FPfp16可以精确表示的范围，蓝色的区域表示FPfp16开始失真，有很大的误差。红色区域表示FPfp16已经无法表示。
如果Multibox SSD这个模型大部分的参数采用是半精度，大部分参数也就是要么在失真的范围，要么在是无法表示这个范围。这样来看，如果直接替换成半精度，就无法进行训练。所以第三个挑战就是给定的一个模型，如何选择最优的混合精度策略，同样也是非常大的一个挑战。
正如前面提到的MoE模型，MoE模型的样本是通过网关对不同的输入选择不同的专家。这个过程和去医院看病类似，非常流行、非常受欢迎的专家会有很多患者都会挂他的号，如图中专家0，很多输入都会选择它，然而有一些专家比较冷门，选择它的计算就比较少。
此时，如果做非常大规模的并行训练，存在的问题就是负载会不均衡，某些节点会负载非常高，某些节点负载会非常低。如果存在这种情况，整个大系统的效率就发挥不上去。所以第四个挑战就是如何动态改善负载均衡，提高模型的训练效率。
国产超算助力MoE加速
此处简述一下新一代超级计算机，新一代系统采用的是神威处理器芯片，每个节点里包括六个核组，核组就是图片里的CG0、CG1等。它是对称结构，总共有六个部分，把每个核组放大，黑色的部分叫做主核，主核是一个通用的 处理器，就像通用的CPU一样，可以处理大部分工作。
真正用来加速的、用来计算的是从核阵列，从核阵列是一个8*8的mesh，也就是网格的结构，如果大家对NVIDIA的GPU熟悉，可以认为从核阵列是插在PCI-3E上的GPU，神威处理器是将主核和核组封装到一个处理器里面。还有就是存储控制器，这六个核组通过一个环形的网络把它们连接起来，每个节点上有两个网卡，所有的节点通过网卡再进行连接。
神威网络的架构是两层的网络，最下面的一层把256个节点组织成一个超节点，超节点可以理解成任何两个节点之间的延迟带宽超节点内部节点之间完全互连；，跨超节点的时候，通信需要走到上层网络，最终构成了两层胖树的结构。
不过，上层网络的带宽是底层超节点的1/8，这主要是出于成本的考虑，因为机器的规模非常大，如果全部用底层一样的网络成本会显著的增加。正如前面提到，数据并行和模型并行的需求是不一样的，当把模型和数据往一个异构的网络拓扑中去映射的时候，就该考虑如何去做高效的映射。
此外，大模型对内存计算和通信有非常大的需求，新一代的超级计算机有PB级的内存空间，有非常强的计算能力，有自主可控的高速网络，有灵活的MPI通信接口。因此，新一代的超级计算机提供了训练大规模模型的基础。
今年我们做了一个尝试，就是在国产超算进行大模型的加速，这个大模型是把MoE模型做了扩展，工作核心主要想验证当模型做到一定规模时，并行或者加速可以解决掉还面临哪些方面挑战，我们如何解决。
目前，训练的是174万亿的模型，达到了人脑的规模，针对前面的四个挑战，并行策略、数据存储、数据精度和负载均衡，均给出了一些解决方案。当然，这些策略并不一定是最优的，对于国产超算而言带来了挑战。
新一代超级计算机是一个两层的结构，上层是1/8的裁剪，无论是数据并行还是模型并行，使用单一的通信模式到了上层的时候通信必然会显著的下降，因为上层只有底层1/8的带宽，所以采用一种并行策略行不通，那么如何选择更合适的并行策略呢？
这个问题是高性能计算的基本问题核心，可以需要把应用程序的通信模式还有底层的网络拓扑去构建一个性能新的模型，比如都用一个图来表示，上层应用程序也用一个图来表示，然后核心是把这两个图做一个映射。
最后我们的解决方案是在超节点也就是256个节点里面做数据并行，在跨超节点做MoE并行，这个方案跟对称的方案，也就是超节点里面的MoE，节点间做数据并行相比，性能会提高1.6倍，当然，这个方案也不一定是最优。
第二个挑战，也就是前面提到的数据存储，也就是这些参数如何分散的去存储。优化器的参数会占用75%的空间，如果简单的采用数据并行，数据并行只是简单的把数据拆分，节点上有整个的模型参数，而模型参数是重复的存储和更新，是非常浪费的。
一个比较简单的策略是把这些参数也进行拆分，然后每个节点只存储一部分。这个部分做的比较好的是微软和NVIDIA的一个工作叫Zero，并在19年的高性能计算SC会议上成功发表发布，核心思想是把这些参数分布在不同节点进行更新，用通信来换取内存，这个方法本身并不一定会提高计算的性能，但是好处是节约内存。
新一代的国产系统既支持双精度也支持半精度，单精度是用双精度去模拟的，半精度的性能为单精度的4倍。我们发现有一个硬件特性就是如果用混合精度会有显著的性能提升，那么问题就变成了如何针对国产的超计算机设计合适的混合精度策略。
如果大家用 NVIDIA平台的话，应该对NVIDIA的这套系统比较了解，NVIDIA有一套系统叫APEX，这是针对混合系统开发的一个库，里面提供了4种策略，分别是不优化、只优化计算、除更新部分和全部优化。
既然存在这么几个策略，我们就就应当在国产系统上验证实验一下，后来发现这几个策略在国产系统上并不适用，那么我们此时需要找到一个办法，即国产系统上如何去做混合精度。
这里可分享一个我们的经验，称为分层的混合精度。经过对不同层进行小规模的一个验证，发现不同层对精度的要求是不一样的。FFN层用的是半精度，注意力层用的是半精度计算、单精度更新，其它层用的是单精度，这种情况下仍旧可以保证模型收敛。
混合精度训练有很多有意思的问题论题，比如在大的科学计算运算程序里面有很多迭代，在迭代的开始、中间和最后都可以选择不同的精度，以动态的混合精度达到计算和内存的最高的效率也是一个比较前沿开放的研究方向。
最后，如果采用更低精度的计算，需要采用loss scaling技术，这是一个比较常见的技术，不做缩放前，大部分的参数都没有进入到半精度可表示范围，乘以一个系数后，就可以保证很多参数能进来。所以我们的研究过程也在前面引入了缩放，在后面引入了反缩放，也就是动态的调整缩放系数，保证在不上溢出的情况下数值尽量大。
为了解决MoE模型负载不均衡的问题，我们提出了一个算法叫SWIPE，这个算法的思想很简单，当一个专家非常受欢迎，如果它超过某个threshold阈值时，就把它过多多余的输入分给提供给其它的专家，这个策略简单直接粗暴，但是验证的数据集结果显示仍旧可以收敛。
最终，把上面的方法都整合到神威的超算平台，然后促成了百万亿模型的参数训练。这些模型的数据主要是来自阿里巴巴的图文数据集M6，包括了百科、网页等多种数据。
结果显示，在1.93T参数规模的时候，计算性能算力是达到了1.18 EFLOPS，当模型增大到14.5T的时候，混合精度性能达到1.00 EFLOPS。最大的是174T，但是因为只有一维并行，这个模型太大了，没办法做混合数据并行，性能会有明显的下降，只达到0.230 EFLOPS，不过训练最终都有非常明显的收敛趋势。
我们的工作起名叫“八卦炉”，这是一次HPC和AI碰撞的得较好例子。在系统方面，发现可以结合这些优化策略做一个非常大的预训练模型系统加速，并且在这个过程中发现高性能计算遇到了很多挑战，比如网络等等，如果真要做一个面向预训练模型加速国产的硬件系统，网络如何来设计也就是非常大的挑战。
总的来说，该项研究的主要工作就是针对大型模型用国产的系统去做加速训练，并且在这个过程中给出提出了所遇到一些问题的解决方案。目前这些方案很多还在研究之中，希望有更多学者能一起参与讨论。
扫码参与智源社区 2021 AI 大调查，参与有奖活动！
内容中包含的图片若涉及版权问题，请及时与我们联系删除
评论
沙发等你来抢