擎天柱: 训练超大模型的高效二维方法

擎天柱: 训练超大模型的高效二维方法

  • 发布时间:2021-05-15 11:10
  • 访问量:

【概要描述】在这项工作中,我们提出了Optimus(擎天柱),一个利用二维矩阵分割的高效可扩展模型并行范式,它将有助于训练无穷大的语言模型。

擎天柱: 训练超大模型的高效二维方法

【概要描述】在这项工作中,我们提出了Optimus(擎天柱),一个利用二维矩阵分割的高效可扩展模型并行范式,它将有助于训练无穷大的语言模型。

  • 分类:研究院
  • 作者:许启帆,尤洋
  • 来源:
  • 发布时间:2021-05-15 11:10
  • 访问量:
详情

​摘要

超大神经网络在实际应用中展现了空前的性能。然而,由于存储限制,我们必须使用模型并行来存储大型模型,否则单个设备将无法容纳这些模型。包括Megatron(威震天)在内的先前的工作将整个模型的参数划分到多个设备,然而每个设备必须在正向和反向传播中存储冗余的中间激活。在这项工作中,我们提出了Optimus(擎天柱),一个利用二维矩阵分割的高效可扩展模型并行范式,它将有助于训练无穷大的语言模型。在Optimus中,中间激活也被分割并分布到诸多设备中,由此进一步减少了冗余。就等效效率而言,Optimus显著胜过了Megatron。在美国德州高级计算中心的64个GPU上,Optimus实现了1.48倍Megatron的训练速度,1.78倍的推理速度,以及8倍的最大批处理量。Optimus在扩展效率上大大超过了Megatron。原文请参见https://arxiv.org/abs/2104.05343。代码请参见http://github.com/xuqifan897/Optimus。

 

第1节:介绍

 

大模型对硬件和软件都带来了巨大挑战。为了应对这一挑战,研究人员除了建造超级计算机之外还开发了新的算法。诸如中间激活存档(activation checkpointing),混合精度训练(mixed precision training),数据转移(data offload)等方法从不同角度来辅助大模型的训练。除这些工作之外,模型并行是另外一个重要的范式。它把大模型的参数分布到诸多设备上。流水线并行(pipeline parallelism)是将整个模型按照层间结构划分。在这种情况下,每个处理器从前一个处理器获得输入,然后将输出结果传给下一个处理器。另外一种并行方式是层内并行,即将一个运算分布到多个处理器同时进行。Megatron就抓住了语言模型矩阵乘法的本质,将权重参数逐行或逐列地切分并分布到不同处理器上。这样计算就是分块矩阵乘法了。这种方法虽然将参数分布在了诸多处理器上,每个处理器却需要存储一整个中间激活。所以在高并行度下,中间激活会成为存储的瓶颈。

值得注意的是,Megatron本质上采用了一维矩阵划分,在这种情况下,每一步通讯都会涉及所有的处理器。然而,如果我们考虑诸如SUMMA和Cannon's algorithm的二维矩阵分割,就有可能加速大模型的推理和训练。在这篇工作中,我们提出了Optimus,一种新型的模型并行范式。它同时分布了参数和激活,实现了高超的通讯效率和前所未有的存储性能。Optimus在PyTorch上实现,并且可以和前述训练大模型的技巧协同使用。Optimus的等效率函数(isoefficiency function)可达到,显著优于(小于)Megatron的。我们相信新型超级计算机的网格拓扑可以进一步解放Optimus的卓越性能。相对于Megatron,Optimus可以帮助我们将模型训练扩展到无穷大。我们在美国德州高级计算中心的64个GPU上实验演示了Optimus相对Megatron 1.48倍的训练速度,1.78倍的推理速度和8倍的最大批处理量。证明了Optimus在强扩展效率和弱扩展效率上均超过Megatron。总的来说,我们的贡献如下:

  1. 基于二维矩阵划分的思想,我们设计和实现了一种用于超大模型的高效训练系统。
  2. 我们提出了诸多高性能方法,例如二维分割梯度计算,高存储效率激活管理,以及系统缓冲技巧。
  3. 我们的方法可以基于PyTorch轻松实现,无需涉及编译层。

 

第2节:背景

 

在这里,我们用以下记号:

  • 批处理大小(batch size): b

  • 序列长度(sequence length): s

  • 隐层大小(hidden size): h

  • 注意力头(attention heads)数量: n

  • 词汇量大小(vocabulary size): v

  • SUMMA维度: q

  • Transformer 层数: N

  • 处理器个数: p

我们有。此后我们用方括号表示张量的形状。

Megatron的模型并行包括多层感知机(MLP)和自注意力组件(self-attention),其结构如下图所示.

 

Figure1 Megatron中MLP和MHA的多处理器映射

 

在MLP中有两个参数矩阵,第一个按列划分,第二个按行划分。每个处理器都持有一份相同的输入,并和同一处理器上的第一个子矩阵相乘,得到同样按列划分的中间结果。中间结果首先输入激活函数,然后再和第二个子矩阵相乘,得到了一个和输入相同大小的输出矩阵。不同处理器上输出矩阵经历一个归约(all-reduce)操作,得到最终的输出。

同样地,在self-attention中也有两个参数矩阵。第一个按列划分,第二个按行划分。相同的输入矩阵被每个处理器所持有,并和第一个子矩阵相乘,得到中间结果。由于attention heads被分布到不同的处理器,每个处理器就负责n/p个attention heads。中间结果被划分为n/p组Q,K和V矩阵,每一组计算得到。n/p个计算结果在一起重新排列,并和第二个子矩阵相乘,再归约得到最终结果。同样地,每个处理器都需要存储一整个的输入和输出。

在这种范式下,通讯仅发生在正向和反向传播的归约操作上。

Mesh-Tensorflow是谷歌提出的分布式机器学习的框架。其引入了张量分割的简约表达。它首先将处理器排布为一个多维网格。对于每个网格维度,一个张量要么在这个维度上分割,要么在这个维度上复制。虽然Mesh-Tensorflow允许多维分割,但其实质上还是采用了和Megatron相似的矩阵分割方式。事实上,其不可能完全消除矩阵复制造成的存储冗余,即不可能每个张量在每个维度上都没有复制。

 

第3节:我们的方法

 

SUMMA (Scalable Universal Matrix Multiplication Algorithm),或可扩展通用矩阵乘法,包含了四种矩阵乘法。在这里,我们仅讨论其中三个,即,。事实上,这三个乘法相互为对方的微分计算:

 

接下来,我们假设处理器组成一个的网格,同时每个张量都被平均划分为的子矩阵,每个都分布在相应的处理器上。表示矩阵A的第i行第j列的子矩阵。我们注意到这种划分方式和谷歌TPU的构型相符。

为例,SUMMA将矩阵乘法视为一组外积(outer product)的和。如果我们将A矩阵按照列划分:,将B矩阵按照行划分:。这里我们用上标来表示列矩阵,用下标来表示行矩阵以示区分。接下来,我们只看外积。我们可以推出,第j行第k列的处理器所持有的子矩阵为。也就是说,第j行的每个处理器都需要,第k列的每个处理器都需要。下图演示了其细节:

 

Figure2 SUMMA算法示意图 i=2

 

我们以16个处理器为例,且仅展示i=2的情况。最开始,每个处理器持有其自己的A和B的子矩阵。接下来,第二列的所有处理器将自己的A子矩阵广播到同一行的每个处理器;第二行的所有处理器将自己的B子矩阵广播到同一列的每个处理器。然后,每个处理器在本地计算其持有的A和B的子矩阵相乘,并与i=1时的结果相加。

Optimus的MLP和self-attention由下图所示:

 

Figure3 Optimus中MLP和MHA的多处理器映射

 

输入和输出矩阵形状为[b,s,h]。在这里,我们仅对b和h维度进行划分,s维度保持不变。在图中,s维度被省略。MLP为两个连续的矩阵乘法,中间有一个非线性变换。由此,其可以直接表达为SUMMA形式。在self-attention中,第一个矩阵乘法可以表达为SUMMA形式。在这之后,每个处理器都拥有大小为[b/q,s, 3h/q]的子矩阵。这时,我们将其划分为n/q组Q,K和V,形状均为[b/q,s, h/n]。这样,我们就可以在本地计算。我们将n/q个计算结果重新排布为[b/q,s, h/q]的形状,再按SUMMA的方法计算第二个矩阵乘法。

我们注意到,Transformer由N个结构相同的Transformer layer组成。这也就揭示了存储的复用性,即每一层的计算可以使用相同一块存储空间。由此,除了中间激活存档(activation checkpointing)外,我们使用了缓冲工作区(workspace buffer),正向传播缓冲区(forward buffer),反向传播缓冲区(backward buffer),参数梯度缓冲区(parameter gradient buffer)和连接缓冲区(conjunction buffer)。在SUMMA中,在每次通讯前需要复制本地的子矩阵。现在,我们从缓冲工作区里分出一块来存放复制的本地子矩阵。正向传播缓冲区用于存放正向传播的中间激活,反向传播缓冲区用于存放中间激活的梯度,参数梯度缓冲区用于存放参数的梯度,连接缓冲区用来存放层间的中间激活和梯度,以防止在重置缓冲区时丢失后续计算中需要的矩阵。

 

第4节:实验结果

 

我们在美国德州高级计算中心的超级计算机上做了实验,得到并比较了弱扩展效率,强扩展效率和存储性能。结果如下图所示:

 

Figure4 weak scaling和strong scaling

 

可以看到,尽管随着GPU数量的增加,Optimus和Megatron的弱扩展效率都在下降,但是Optimus相对Megatron的优势越来越显著。在64结点上,Optimus达到了1.48倍Megatron的训练速度和1.78倍Megatron的推理速度。在强扩展效率上,我们观察到了随着GPU数量的增加,Megatron的效率在下降,而Optimus的效率在上升。这是因为在SUMMA中,当矩阵大小不变时,处理器数量增加,通讯时间减少。

在测试弱扩展效率的相同条件下,我们测试了Optimus和Megatron的显存效率。即我们改变b的大小,直到显存耗尽(out-of-memory)。在下图中,表示在时,可以运行,但当时出现显存耗尽。随着GPU数量的增加,Megatron的批处理量在减少,而Optimus的批处理量在增加。在64GPU上,Optimus达到了Megatron 8倍的批处理量。

 

Figure5 内存效率

 

第5节:总结

 

在这篇工作中,我们提出了一种全新的模型并行范式。该范式成功地消除了存储的瓶颈且提升了通讯效率。从而为更大语言模型的开发铺平了道路。在近一段时间的工作中,混合专家模型(Mixtureof Experts,MOE)得到了越来越多的关注。许多强大的模型都是基于MOE实现。未来的工作可能会致力于解决MOE的通讯和存储问题。此外,算子融合也是一个新方向。该方法可以显著减少中间结果的存储,从而提高存储效率。比如在我们的实现中,注意力得分(attention score)会占据一个形状为[b,n, s, s]的张量。而中间激活的形状为[b,s, h]。当n=16,s=512,h=1024时,注意力得分所占据的存储是中间激活的8倍,但只占总计算量的一小部分。所以,利用算子融合,我们有希望解放这一部分存储。随着体量的增加,我们相信巨型模型将能够胜任更多令人惊艳的任务。

 

作者简介:

 

 

许启帆,加州大学洛杉矶分校一年级博士生,研究方向是高性能计算,医学AI,机器学习。许启帆本科毕业于中国科学技术大学,在读期间入选拔尖计划,于2019年入选加州大学洛杉矶分校CSST项目。2020-2021年在新加坡国立大学HPC-AI实验室实习期间完成这项工作。

 

尤洋,壁仞科技顾问,是新加坡国立大学计算机系的校长青年教授。他从加州大学伯克利分校计算机系获得了博士学位。他的导师是美国科学院院士,工程院院士,ACM/IEEE fellow,伯克利前计算机系主任以及EECS学院院长James Demmel教授。尤洋的研究兴趣包括高性能计算,并行算法,以及机器学习。他当前的研究重点是大规模深度学习训练算法的分布式优化。他曾创造ImageNet以及BERT训练速度的世界纪录,并被ScienceDaily,TheNextWeb,i-programmer等几十家媒体广泛报道。他设计的算法被广泛应用于谷歌,微软,英特尔,英伟达等科技巨头。尤洋近三年以第一作者身份在NIPS,ICLR,Supercomputing,IPDPS,ICS等国际重要会议或期刊上发表论文十余篇。他曾以第一作者身份获得了国际并行与分布式处理大会(IPDPS)的最佳论文(0.8%获奖率)和国际并行处理大会(ICPP)的最佳论文(0.3%获奖率)。尤洋曾获清华大学优秀毕业生,北京市优秀毕业生,国家奖学金,以及当时清华大学计算机系数额最高的西贝尔奖学金。他还在2017年获得美国计算机协会(ACM)官网上唯一颁给在读博士生的ACM-IEEECS George Michael Memorial HPC Fellowships。他也获得了颁发给伯克利优秀毕业生的LotfiA.Zadeh Prize。尤洋还被伯克利提名为ACM Doctoral Dissertation Award候选人(81名UC Berkeley EECS 2020博士毕业生中选2人)。2021年4月,尤洋入选福布斯亚洲30岁以下精英榜。更多信息请查看他的研究小组主页(https://ai.comp.nus.edu.sg/)。

近期文章

流体力学与物理导引神经网络

流体力学与物理导引神经网络

现代科学中众多复杂的关键应用依赖着对流体运动的精确预测。然而,由于流体力学方程数值求解的复杂性,基于传统数值方法的高分辨率长周期数值模拟运算量大且难以保证数值稳定性。本文将深入介绍PINN方法在流体计算中的应用并分析其与传统数值方法的差别。
2021-10-25
物理导引神经网络方法分析

物理导引神经网络方法分析

随着GPU能力的提升,支撑深度学习的软硬件生态得到了快速发展。通过深度学习来解决科学计算问题成了一种趋势,其中用深度学习来求解偏微分方程的方法也逐渐兴起。尤其引人注意的是一种称为物理导引神经网络的方法,其为科学计算领域注入了新的活力。
2021-10-18

联系我们

这是描述信息

招贤纳士

公众号二维码
这是描述信息

 版权所有   ©   上海壁仞科技有限公司   |    沪ICP备19047354号