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

摘要

超大神经网络在实际应用中展现了空前的性能。然而，由于存储限制，我们必须使用模型并行来存储大型模型，否则单个设备将无法容纳这些模型。包括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。

介绍

大模型对硬件和软件都带来了巨大挑战。为了应对这一挑战，研究人员除了建造超级计算机之外还开发了新的算法。诸如中间激活存档(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。总的来说，我们的贡献如下：

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

背景

在这里，我们用以下记号：

批处理大小(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相似的矩阵分割方式。事实上，其不可能完全消除矩阵复制造成的存储冗余，即不可能每个张量在每个维度上都没有复制。

我们的方法

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中，在每次通讯前需要复制本地的子矩阵。现在，我们从缓冲工作区里分出一块来存放复制的本地子矩阵。正向传播缓冲区用于存放正向传播的中间激活，反向传播缓冲区用于存放中间激活的梯度，参数梯度缓冲区用于存放参数的梯度，连接缓冲区用来存放层间的中间激活和梯度，以防止在重置缓冲区时丢失后续计算中需要的矩阵。

实验结果

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

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 内存效率

总结

在这篇工作中，我们提出了一种全新的模型并行范式。该范式成功地消除了存储的瓶颈且提升了通讯效率。从而为更大语言模型的开发铺平了道路。在近一段时间的工作中，混合专家模型(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/)。

往期推荐

1、MLIR编译技术应对硬件设计挑战？—初探CIRCT项目

2、图神经网络加速器深度调研（中）

3、图神经网络加速器深度调研（下）

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们

继续阅读

阅读原文

关键词

参数

模型并行

大小

方法

矩阵乘法