本文主要参考ColossalAI论文 Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training ColossalAI框架开源提供了本文介绍的所有并行训练: https://github.com/hpcaitech/ColossalAI
本文会介绍几种流行的并行方法,包括
下图给出了这些并行方法的示意图,非常直观好懂。
不过在介绍各种并行训练方法之前,我们首先对一些概念做一个声明,方便后面理解
模型训练过程中涉及到的参数主要包含两大类,model data 和 non-model data,具体表示如下:
pipeline parallelism是比较常见的模型并行算法,它是模型做层间划分,即inter-layer parallelism。以下图为例,如果模型原本有6层,你想在2个GPU之间运行pipeline,那么每个GPU只要按照先后顺序存3层模型即可。
已经有很多Pipeline相关的研究工作了,例如PipeDream,GPipe,和Chimera。它们的主要目的都是降低bubble time。这里不做过多介绍。
前面介绍的Pipeline Parallelism是对模型层间做划分,叫inter-layer parallelism。那么另一种方式则是对模型层内做划分,即intra-layer Parallelism,也叫Tensor Parallelism。
Megatron-LM [1] 是最早提出1D Tensor并行的工作。该工作主要是为了优化transformer训练效率,把线性层按照行或者列维度对权重进行划分。如图4所示,原本线性层为
,这里将
按列进行划分,将
按行进行划分。这样,每个GPU只需要存一半的权重即可,最后通过All-reduce操作来同步Y的结果。当GPU数量为
时,每个GPU只需要存
的权重即可,只不过每层输出需要用All-reduce来补全结果之后才能继续下一层的计算。
对于土豪公司,可以使用NVLink来连接GPU(如图5a),从而提供高带宽来降低通信开销。但是土豪终归是少数的,大部分公司和个人是没法承担这昂贵的硬件费用,因此比较常见的GPU连接方式是图5b,即节点内花点钱实现NVLink连接,节点之间通过PCIe连接。
1D Tensor并行对通信速度要求较高,不过1D在每层的输入和输出都有冗余的内存开销。以图4为例,我们可以看到虽然模型权重被划分了,但是每个GPU都有重复的输入
,另外All-reduce之后每个GPU也会有重复的输出
,所以后续一些工作尝试从这里做进一步改进,包括2D, 2.5D,和3D tensor并行。
2D Tensor Parallel [2] 基于SUMMA和Cannon矩阵相乘算法沿着两个不同的维度对 输入数据,模型权重,每层的输出 进行划分。给定
个GPU,tensor会被划分成
个chunk(使用torch.chunk
),每个GPU保存一个chunk。这
个GPU呈方形网络拓扑结构,即每行每列均为
个GPU。图6b展示了常见的4-GPU的节点划分示意图,假设tensor的维度大小是
,那么划分后每个GPU上存的chunk大小即为
。至此,每个GPU都只会保存部分的输入输出以及部分的权重。虽然相比于1D Tensor并行,2D额外增加了模型权重的通信,但是需要注意的是当GPU数量很多的时候,每个GPU上分配的模型权重就会小很多,而且因为使用的All-reduce通信方式,所以2D也还是要比1D更高效的。
2.5D Tensor Parallel [3] 是受2.5D矩阵乘法算法 [4] 启发进一步对2D Tensor并行的优化。具体来说2.5D增加了 depth 维度。当 depth=1 时等价于2D;当 depth>1 时,
同样假设有
个GPU,其中
,
类似于原来2D正方形拓扑结构的边长,而
则是新增加的维度 depth 。
可以由用户指定,
则会自动计算出来了。所以一般来说至少需要8个GPU才能运行2.5D算法,即
。
3D Tensor Parallel [5] 是基于3D矩阵乘法算法 [6] 实现的。假设有
个 GPU,tensor维度大小为
,那么每个chunk的大小即为
。当tensor维度小于3时,以全连接层为例,假设权重维度大小为
,那么可以对第一个维度划分两次,即每个chunk的维度大小为
。3D Tensor并行的通信开销复杂度是
,计算和内存开销都均摊在所有GPU上。
1D Tensor并行每一层的输出是不完整的,所以在传入下一层之前都需要做一次All-gather操作,从而使得每个GPU都有完整的输入,如图7a所示。
2D/2.5D/3D Tensor 并行算法因为在一开始就对输入进行了划分, 所以中间层不需要做通信,只需要在最后做一次通信即可。在扩展到大量设备(如GPU)时,通信开销可以降到很小。这3个改进的Tensor并行算法可以很好地和Pipeline并行方法兼容。
Tensor parallelism主要是为了解决由 model data (模型权重,梯度和优化器状态)导致的内存瓶颈,但是 non-model data也可能成为性能瓶颈。比如像AlphaFold和NAS任务中会存在很多中间特征值(也叫activations)。
以DARTS算法为例,它的模型参数量其实并不多,但是它有很多分支,所以activations会消耗大量GPU内存,这也是为什么很多NAS算法只能在CIFAR-10上搜索到合适的模型结构后,再做人工扩展,最后应用到ImageNet上做性能验证。
同样地,在使用Transformer训练语言模型时,由于Transformer层中的Self-attention机制的复杂度是
,其中
是序列长度。换言之,长序列数据将增加中间activation内存使用量,从而限制设备的训练能力。
Sequential Parallelism (SP) [7] 就为了解决non-model data导致的性能瓶颈而提出的。下图给出了SP在Transform并行训练上的应用,具体的原理可以查看原论文。
训练过程中GPU内存开销主要包含以下几个方面:
ZeRO针对模型状态的三部分都做了对应的内存改进方法:
下图给出了三种方法带来的内存开销收益
不管采用三种方法的哪一种,ZeRO简单理解就是给定
个设备,然后把一堆data等分到这些设备上,每个设备只存
的数据量,并且每次也只负责更新这
的数据。
因为对数据做了划分,ZeRO在每一层都需要有通信操作。我们考虑ZeRO在某一层的具体操作:
注意ZeRO对数据划分方式并没有什么具体的要求,可以是随意划分,因为最后反正会用all-gather使得所有设备商都有用完整的数据;当然,也可以使用前面提到的Tensor Parallelism的划分方式,这样一来可以有效降低通信开销,进一步提高效率。
关于ZeRO更详细的介绍可以查看原论文或者看看这篇博客 [8]
未完待续。。。
MARSGGBO♥原创 如有意合作或学术讨论欢迎私戳联系~ 邮箱:marsggbo@foxmail.com