作者丨白强伟@知乎
来源丨https://zhuanlan.zhihu.com/p/617087561
编辑丨极市平台
流水线性并行和张量并行都是对模型本身进行划分，目的是利用有限的单卡显存训练更大的模型。简单来说，流水线并行水平划分模型，即按照层对模型进行划分；张量并行则是垂直划分模型。3D并行则是将流行线并行、张量并行和数据并行同时应用到模型训练中。
一、流水线并行
流水线并行的目标是训练更大的模型。本小节先介绍符合直觉的朴素层并行方法，并分析其局限性。然后，介绍流水线并行算法GPipe和PipeDream。
1. 朴素层并行
当一个模型大到单个GPU无法训练时，最直接的想法是对模型层进行划分，然后将划分后的部分放置在不同的GPU上。下面以一个4层的序列模型为例，介绍朴素层并行：
将其按层划分至两个GPU上：
GPU1负责计算：intermediate (input)) ;
GPU2负责计算：output (intermediate ;
整个朴素层并行前向传播和后向传播的过程如上图所示。GPU1执行前向传播, 并将激活 (activations)缓存下来。然后将 层的输出intermediate发送给GPU2, GPU2完成前向传播和loss计算后, 开始反向传播。当GPU2完成反向传播后, 会将 的梯度返还给GPU1。GPU1完成最终的反向传播。
根据上面的介绍，可以发现朴素层并行的缺点：
低GPU利用率。在任意时刻，有且仅有一个GPU在工作，其他GPU都是空闲的。
计算和通信没有重叠。在发送前向传播的中间结果(FWD)或者反向传播的中间结果(BWD)时，GPU也是空闲的。
高显存占用。GPU1需要保存整个minibatch的所有激活，直至最后完成参数更新。如果batch size很大，这将对显存带来巨大的挑战。
2. GPipe
2.1 GPipe的原理
GPipe通过将minibatch划分为更小且相等尺寸的microbatch来提高效率。具体来说，让每个microbatch独立的计算前后向传播，然后将每个mircobatch的梯度相加，就能得到整个batch的梯度。由于每个层仅在一个GPU上，对mircobatch的梯度求和仅需要在本地进行即可，不需要通信。
假设有4个GPU，并将模型按层划分为4个部分。朴素层并行的过程为
| Timestep | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|---|---|---|---|---|---|---|---|---|
| GPU3 | FWD | BWD | ||||||
| GPU2 | FWD | BWD | ||||||
| GPU1 | FWD | BWD | ||||||
| GPU0 | FWD | BWD |
可以看到，在某一时刻仅有1个GPU工作。并且每个timesep花费的时间也比较长，因为GPU需要跑完整个minibatch的前向传播。
GPipe将minibatch划分为4个microbatch，然后依次送入GPU0。GPU0前向传播后，再将结果送入GPU1，以此类推。整个过程如下表
| Timestep | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| GPU3 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | ||||||
| GPU2 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | ||||||
| GPU1 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 | ||||||
| GPU0 | F1 | F2 | F3 | F4 | B4 | B3 | B2 | B1 |
F1表示使用当前GPU上的层来对microbatch1进行前向传播。在GPipe的调度中，每个timestep上花费的时间要比朴素层并行更短，因为每个GPU仅需要处理microbatch。
2.2 GPipe的Bubbles问题
bubbles指的是流水线中没有进行任何有效工作的点。这是由于操作之间的依赖导致的。例如，在GPU3执行完F1之前，GPU4只能等待。整个流水线过程中的bubbles如下图所示。
bubbles浪费时间的比例依赖于pipeline的深度 和mincrobatch的数量 。假设单个GPU上完成前向传播或者后向传播的面积为 1 (也就是上图中的单个小方块面积为 1 )。上图中的总长度为 , 宽度为 , 总面积为 。其中, 彩色小方块占用的面积表示GPU执行的时间, 为其 。空白处面积的占比代表了浪费时间的比较, 其值为
因此，增大microbatch的数量m，可以降低bubbles的比例。
2.3 GPipe的显存需求
增大batch size就会线性增大需要被缓存激活的显存需求。在GPipe中，GPU需要在前向传播至反向传播这段时间内缓存激活(activations)。以GPU0为例，microbatch1的激活需要从timestep 0保存至timestep 13。
GPipe为了解决显存的问题，使用了gradient checkpointing。该技术不需要缓存所有的激活，而是在反向传播的过程中重新计算激活。这降低了对显存的需求，但是增加了计算代价。
假设所有层都大致相等。每个GPU缓存激活所需要的显存为
总层数数量
也就是与单个GPU上的层数以及batch size成正比。相反，使用gradient checkpointing仅需要缓存边界层(需要跨GPU发送结果的层)的输入。这可以降低每个GPU的显存峰值需求
(batchsize) 是缓存边界激活所需要的显存。当对给定的microbatch执行反向传播时, 需要重新计算该microbatch梯度所需要的激活。对于每个GPU上的 总层数数量 层需要 的显存空间。
3. PipeDream
GPipe需要等所有的microbatch前向传播完成后，才会开始反向传播。PipeDream则是当一个microbatch的前向传播完成后，立即进入反向传播阶段。 理论上，反向传播完成后就可以丢弃掉对应microbatch缓存的激活。由于PipeDream的反向传播完成的要比GPipe早，因此也会减少显存的需求。
下图是PipeDream的调度图，4个GPU和8个microbatchs。蓝色的方块表示前向传播，绿色表示反向传播，数字则是microbatch的id。
PipeDream在bubbles上与GPipe没有区别，但是由于PipeDream释放显存的时间更早，因此会降低对显存的需求。
4. 合并数据并行和流水线并行
数据并行和流水线并行是正交的，可以同时使用。
对于流水线并行。每个GPU需要与下个流水线阶段(前向传播)或者上个流水线阶段(反向传播)进行通信。
对于数据并行。每个GPU需要与分配了相同层的GPU进行通信。所有层的副本需要AllReduce对梯度进行平均。
这将在所有GPU上形成子组，并在子组中使用集合通信。任意给定的GPU都会有两部分的通信，一个是包含所有相同层的GPU(数据并行)，另一个与不同层的GPU(流水线并行)。下图是流水线并行度为2且数据并行度为2的示例图，水平方向是完整的一个模型，垂直方向是相同层的不同副本。
二、张量并行
Transformer中的主要部件是全连接层和注意力机制，其核心都是矩阵乘法。张量并行的核心就是将矩阵乘法进行拆分，从而降低模型对单卡的显存需求。
1. 1D张量并行
本小节以全链接层 为例, 介绍张量并行。其中, 和 是输入和输出向量, 是权重矩阵, 是非线性激活函数。总量来说张量并行可以分为列并行和行并行(以权重矩阵的分割方式命名), 上图展示了两种并行。
(1) 矩阵乘法角度
这里以矩阵乘法的方式辅助理解1D张量并行。
列并行
将矩阵行列划分为n份(不一定必须相等大小)可以表示为 ，那么矩阵乘 法表示为
显然，仅需要对权重进行划分。
行并行
对权重进行划分，那么必须对输入矩阵也进行划分。假设要将A水平划分为n份，则输入矩阵X必须垂直划分为n份，那么矩阵乘法表示为
(2) 激活函数与通信
显然，只观察上面的数据公式，无论是行并行还是列并行 ，都只需要在各个部分计算完后进行一次通常。只不过列并行将通信的结果进行拼接，而行并行则是对通信结果相加。
现在，我们将非线性激活GeLU加上，并模拟两层的全链接层。设X是输入，A和B则是两个全链接层的权重。
列并行
通过上面的公式可以看到。当我们将A和B提前划分好后，就可以独立进行计算，在计算出后再进行通信。也就是说，这个例子中虽然有两个全链接层，但是仅需要在得到最终结果前进行通信即可。
所以，多个全链接层堆叠时，仅需要在最终输出时进行一次通信即可(如上图所示) 。
行并行
由于 是非线性的，所以
因此，行并行每一个全链接层都需要进行通信来聚合最终的结果。
(3) 多头注意力并行
多头注意力并行不是1D张量并行，但是由于其是Megatron-LM中与1D张量并行同时提出的，所以这里也进行简单的介绍。
由于多头注意力的各个头之间本质上就是独立的，因此各个头完全可以并行运算。
注意：张量并行(TP)需要非常快的网络，因此不建议跨多个节点进行张量并行。实际中，若一个节点有4个GPU，最高的张量并行度为4。
2. 2D、2.5D张量并行
在1D张量并行后，又逐步提出了2D、2.5D和3D张量并行。这里对2D和2.5D张量并行进行简单介绍：
(1) 2D张量并行
1D张量并行并没有对激活(activations，也就是模型中间层的输出结果)进行划分，导致消耗大量的显存。
这里仍然以矩阵乘法 为例, 给定有 个处理器。这里假设 , 则将X和A都划分为 的块。
基于上面的矩阵乘法的变化，可以发现Y=XA可以分解为两个矩阵相加。具体来说，两个矩阵的结果仍然需要串行的计算。但是，单个矩阵中的4个子矩阵可以使用 的处理器来并行计算。
当 , 也就是第一步。对
和
进行广播, 所有 的处理器均拥有这 4 个子矩阵。然后分别执行 。经过这一步后就得到了第一个矩阵的结果。
当 。对
和
进行广播，各个处理器在分别计算。最终得到第二个矩阵的结果。
将两个矩阵的结果相加。
(2) 2.5D张量并行
仍然是矩阵乘法 , 并假设有 个处理器, 将 X划分为 行和 列。不妨设 , 将 和 分别划分为
那么有
其中， concat 表示两个矩阵的垂直拼接操作。
基于上面的推导, 可以发现被拼接的两个矩阵天然可以并行计算。即 和 可以并行计算。看到这里, 应该就可以发现这两个矩阵乘法就是上面的2D张量并行适用的形式。所以, 我们总计有 个处理器, 每 个处理器使用2D张量并行来处理对应的矩阵乘法。最后, 将两个2D张量并行的结果进行拼接即可。
三、3D并行
总的来说，3D并行是由数据并行(DP)、张量并行(TP)和流水线并行(PP)组成。前面已经分别介绍了TP和PP，ZeRO-DP是一种显存高效的数据并行策略，原理见文章：
白强伟：【深度学习】【分布式训练】DeepSpeed：AllReduce与ZeRO-DP142
https://zhuanlan.zhihu.com/p/610587671
下面介绍如何将三种并行技术结合在一起，形成3D并行技术。
1. 一个3D并行的例子
假设有两个节点Node1和Node2，每个节点有8个GPU，共计16个GPU。16个GPU的编号分别为Rank0、Rank1、...、Rank15。此外，假设用户设置流水线并行度为4，张量并行度为2。
流水线并行。流水线并行会将整个模型划分为4份，这里称为sub_model_1至sub_model_4。每连续的4张GPU负责一个sub_model。即上图右上角中，相同颜色的GPU负责相同的sub_model。
张量并行。张量并行会针对流水线并行中的sub_model来进行张量的拆分。即Rank0和Rank1负责一份sub_model_1，Rank2和Rank3负责另一份sub_model_1；Rank4和Rank5负责sub_model_2，Rank6和Rank7负责另一份sub_model_2；以此类推。上图右下角中，绿色线条表示单个张量并行组，每个张量并行组都共同负责某个具体的sub_model。
数据并行。数据并行的目的是要保证并行中的相同模型参数读取相同的数据。经过流水线并行和张量并行后，Rank0和Rank2负责相同的模型参数，所以Rank0和Rank2是同一个数据并行组。上图左上角中的红色线条表示数据并行组。
2. 3D并行分析
为什么3D并行需要按上面的方式划分GPU呢？ 首先，模型并行是三种策略中通信开销最大的，所以优先将模型并行组放置在一个节点中，以利用较大的节点内带宽。其次，流水线并行通信量最低，因此在不同节点之间调度流水线，这将不受通信带宽的限制。最后，若张量并行没有跨节点，则数据并行也不需要跨节点；否则数据并行组也需要跨节点。
流水线并行和张量并行减少了单个显卡的显存消耗，提高了显存效率。但是，模型划分的太多会增加通信开销，从而降低计算效率。ZeRO-DP不仅能够通过将优化器状态进行划分来改善显存效率，而且还不会显著的增加通信开销。
参考资料
https://siboehm.com/articles/22/pipeline-parallel-training
https://arxiv.org/pdf/1811.06965.pdf
https://huggingface.co/docs/transformers/v4.18.0/en/parallelism
https://www.colossalai.org/zh-Hans/docs/features/1D_tensor_parallel
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/?utm_source=wechat_session&utm_medium=social&utm_oi=738477582021832704
推荐阅读
欢迎大家加入DLer-计算机视觉技术交流群！
大家好，群里会第一时间发布计算机视觉方向的前沿论文解读和交流分享，主要方向有：图像分类、Transformer、目标检测、目标跟踪、点云与语义分割、GAN、超分辨率、人脸检测与识别、动作行为与时空运动、模型压缩和量化剪枝、迁移学习、人体姿态估计等内容。
进群请备注：研究方向+学校/公司+昵称（如图像分类+上交+小明）
👆 长按识别，邀请您进群！
扫一扫
666
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
