\u200E
流水线并行技术与飞桨优化实现详解
发布日期:2022-05-10T12:25:15.000+0000 浏览量:179次

实践证明,更大的模型在不少场景取得了更好的效果。但随着参数规模的扩大,AI加速卡存储(如GPU显存)容量问题和卡的协同计算问题成为了训练超大模型的瓶颈。流水线并行从模型切分和调度执行两个角度解决了这些问题。本文将从流水线并行介绍、主流实现方式以及一般调优技巧三方面从浅到深讲解流水线并行,希望能帮助您更好地理解和使用流水线并行功能。



模型切分一般有两种方式:参数(Tensor)切分和图切分。在上一篇张量模型并行中介绍过,张量模型并行可以把shape较大的参数(Tensor)切分到多个卡,从而有效减小在每个卡上的参数量。但是,采用这种方式单机8卡(V100,32GB)最多训练10B左右的Dense模型,而且由于其通信和计算不能重叠的特点一般不适合多机之间使用。更大的模型需要从层(Layer)级别进行进一步的图切分以减少单卡存储的容量需求,同时隐藏卡之间的通信时间,更加充分利用GPU卡的算力。为此业界提出了流水线并行的方式解决上述切分和调度的问题。



流水线并行介绍




图(a)


广义上讲,流水线并行从图切分的角度,将模型进一步按层切分到不同的设备上,相邻设备间在计算时只需要传递邻接层的中间变量和梯度。如上图(a)所示,这里引用上一篇张量模型并行中矩阵乘的介绍,对FFN(Feedforward Network)采用张量模型并行,把fc层参数切分到设备内的不同卡上,在同一设备内进行Allreduce求和的全量通信参数为 2*M*K(M:矩阵行维度, K:矩阵列维度)。而采用流水线并行将FFN层作为切分点切分模型层,在device间只需要 send/receive (发送\接收)M*K 的参数量,因此相比张量模型并行,流水线并行通信参数量更少。如下图(b)所示,以Gpipe论文中的示例进行说明,朴素的流水线并行,同一时刻只有一个设备进行计算,其余设备处于空闲状态,计算设备利用率通常较低。


图(b)


与之相对应,将朴素流水线并行的 batch 再进行切分,减小设备间空闲状态的时间,可以显著提升流水线并行设备利用率。如图(c)F-then-B 调度方式所示,将原本的 mini-batch(数据并行切分后的batch)划分成多个 micro-batch(mini-batch再切分后的batch),每个 pipeline stage (流水线并行的计算单元)先整体进行前向计算,再进行反向计算。通过在同一时刻分别计算模型的不同部分,F-then-B 可以显著提升设备资源利用率。但也不难看出这种 F-then-B 模式由于缓存了多个 micro-batch 的中间变量和梯度,显存的实际利用率并不高。

 

由此产生出 1F1B (在流水线并行中,pipeline stage 前向计算和反向计算交叉进行的方式)流水线并行方式解决了这个问题。在 1F1B 模式下,前向计算和反向计算交叉进行,可以及时释放不必要的中间变量。我们以下图(c)1F1B中 stage4 的 F42(stage4的第2个 micro-batch 的前向计算)为例,F42 在计算前,F41 的反向 B41(stage4的第1个 micro-batch 的反向计算)已经计算结束,即可释放 F41 的中间变量,从而 F42 可以复用 F41 中间变量的显存。1F1B 方式相比 F-then-B 方式峰值显存可以节省37.5%,对比朴素流水线并行峰值显存明显下降,设备资源利用率显著提升。


图(c)


F-then-B和 1F1B 调度方式都采用增加 batch 数量,即将 mini-batch 切分为 micro-batch 增加设备资源利用率。下面我们从 micro-batch 角度从理论上推导下计算空闲,即 bubble 和 micro-batch 实际计算时间之间的关系。引用 NVIDIA 在其训练超大语言模型论文对于 bubble 时间的计算公式:

其中:

 tpb: 一个 mini-batch 中bubble的时间, ppipeline stage的数量

   :每个 micro-batch 的前向时间,   :每个 micro-batch 的反向时间

其中:

   : 一个mini-batch的计算时间, m: 一个 mini−batch 中 micro-batch 的数量



由上述两个公式可以得到:


由此不难看出,在pipeline stage数p一定的情况下,影响 bubble 占比的是一个 mini-batch 中 micro-batch的数量 m 。在保证 m >> p 的情况下,可以有效降低 bubble 占比。



流水线并行实现




飞桨分别实现了上面介绍过的 F-then-B 和 1F1B 两种流水线调度方式。虽然 1F1B 方式显存利用率高于 F-then-B 的方式,但是 1F1B 实现比较复杂,使用门槛比较高,需要用户明确自己模型在哪些位置需要做通信;而 F-then-B 的实现相对 1F1B 更加简单,通信不容易出现问题,因此建议用户调模型时选择 F-then-B 流水线调度方式,性能优化时再改为 1F1B 调度方式。在实现流水线并行策略时,飞桨主要针对调度执行环节做了一些优化设计。下面我们从通信调度、显存调度两方面分别进行介绍。




通信调度

首先,从模型切分角度来考虑,通过流水线并行切分后的模型仍然可以看作是一个完整的计算图;只不过需要用通信来串联起每个计算单元,即 pipeline stage。这里我们可以简单理解,由于在相邻 pipeline stage 间需要交互中间变量和梯度,不可避免需要对交互后的变量进行同步操作,保证交互的变量已经被发送或是接收。如果我们只是在单一计算流(计算进程)进行这一操作,那么增加的通信操作无疑会增加不必要的计算空闲。所以我们在实现时,对于不同 pipeline stage 做不同流的 send\receive(发送\接收)操作,如下图(d)所示。在计算流进行第二步 matmul 计算时,同时在通信流(通信进程)发送第一步 matmul 运算后的中间变量结果给下一个 pipeline stage。通过在时间维度复用计算和通信操作,可以有效减少计算空闲,显著提升流水线并行效率。


图(d)


正如上文介绍的,流水线并行由于采用多流通信机制。通常需要在每个 micro-batch 的反向阶段增加同步操作,保证当前 micro-batch 在前向阶段向下一个 pipeline stage 传输的变量已经被接收。这样的同步操作往往也增加了很多不必要的通信调度时间。如下图(e)所示,飞桨在实现多流通信时,为了通信和计算可以更好的 overlap;通过拓扑分析,在 micro-batch 的 receive 之后增加 nop op(nop的蓝色箭头),hold住同 micro-batch 前向 send 变量(nop的红色箭头),取消一次同步操作。经测试在同样规模的模型下,由于在调度侧减少同步op的数量,并通过垃圾回收机制控制变量的生命周期。调度时间和峰值显存都有明显下降。


图(e)





显存调度

之后介绍下显存调度。 由于流水线并行在同一时间维度需要同时进行多个 micro-batch 的计算。 如果只采用一套变量域,那么同一 pipeline stage 下的不同批次的同名变量将无法有效进行版本管理,进而无法及时释放不必要的变量。 由此在实现上我们需要两种变量域,即每个 micro-batch 的变量域和全局变量域。 每个 micro-batch 的变量域负责保存每一个 micro-step 前向计算中的中间变量,用于反向计算 micro-batch 的梯度; 全局变量域负责保存 mini-batch 需要用到其相对的各个 micro-batch 的梯度。 通过调度不同变量域,可以有效对变量进行版本管理。 及时释放多余变量,减少显存开销。

图(f)

进而还可以发现,由于 mini-batch 被切分为多个 micro-batch,如果对每个 micro-batch 都进行参数更新,无疑在计算和显存上会增加大量开销。为了进一步降低计算和显存开销,飞桨对同一 mini-batch 下的 micro-batch 先进行梯度累加,再统一更新参数。具体来说,就是使用 mirco-batch 数据进行训练,即通过“前向+反向”网络计算得到梯度,其间会有一部分显存/内存用于存放梯度,然后对每个 micro-batch 计算出的梯度进行累加,当累加的次数达到 acc_step (需要梯度聚合的step数)时,使用累加的梯度对模型进行参数更新,从而达到 mini-batch 数据训练的效果。


在宏观上看,梯度累加是将训练一个 mini-batch 的过程由原来的 “前向 + 反向 + 更新” 改变成 “(前向 + 反向 + 梯度累加)* k + 更新”,通过在最终更新前进行k 次梯度累加模拟出 mini-batch 的效果,从而减少计算和显存多余开销。



流水线并行调优技巧






模型切分

这里主要介绍下流水线并行模型切分的一般调优策略。使用流水线并行做模型切分时,需要注意模型切分后各个 pipeline stage 的计算量是否均衡,避免由于某一 pipeline stage 的计算量明显大于其他 pipeline stage,导致 pipeline stage 间的等待。


图(g)

以上图(g)为例,假设流水线并行最后一个 pipeline stage 的计算量明显大于其他pipeline stage。在时间维度表现为计算时间明显大于其他 pipeline stage。这样会造成其他 pipeline stage 的计算等待,降低设备利用率。所以在对模型进行切分时,需要综合考虑模型各个部分的计算量、集群通信资源和设备拓扑方式,保证每个 pipeline stage 的计算、通信相对一致。


如果想深入了解模型切分的一般规则,欢迎感兴趣的同学点击阅读原文链接查看更多示例代码:

https://github.com/PaddlePaddle/FleetX





总结



综上所述,流水线并行将模型按照模型层切分到多个设备上,同时在调度上优化设备间的调度逻辑,使得超大规模模型训练既在设备上可以减小参数规模,也可以提升设备利用率。如果想进一步了解飞桨分布式使用方法,可以参考飞桨分布式训练Fleet下的相关API,通过使用Fleet下的相关分布式接口,即可快速开始使用飞桨进行大规模分布式训练。



参考引用


[1]Brown, Tom B. et al. Language models are few-shot learners.In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877–1901. Curran Associates, Inc., 2020.

[2]Huang Y, Cheng Y, Bapna A, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism[J]. Advances in neural information processing systems, 2019, 32: 103-112.(https://arxiv.org/abs/1811.06965)

[3]Narayanan D ,  Shoeybi M ,  Casper J , et al. Efficient Large-Scale Language Model Training on GPU Clusters[J]. 2021.(https://arxiv.org/abs/2104.04473v3)

[4]Paddle代码:https://github.com/PaddlePaddle/Paddle




相关阅读


关注【飞桨PaddlePaddle】公众号

获取更多技术内容~