大语言模型并行策略(六):流水线并行
当模型参数量过大时,即便使用张量并行,也可能无法将模型参数切分到单个 GPU 上。而因为大语言模型通常由多个层堆叠而成,每个层的计算都只依赖于前一层的输出,因此可以将模型的不同层分布到不同的 GPU 上进行计算,每个 GPU 只负责计算自己负责的层。每个 GPU 就像是流水线中的一个工序,前一个 GPU 计算完成后,将结果传递给下一个 GPU 进行计算。这种将模型按层切分为多个阶段,并将不同阶段分布到不同 GPU 上进行计算的并行策略,称为流水线并行(Pipeline Parallelism,PP)。
流水线并行的原理
大语言模型是由多个层堆叠而成的,每一层的计算都仅依赖于前一层的输出。流水线并行将从模型层的维度切分模型,将模型的不同层分布到不同的 GPU 上进行计算。每个 GPU 只负责计算自己负责的层,多个 GPU 之间通过网络传递数据,形成一个流水线的计算模式。以 4 个 GPU 处理一个 4 层的模型为例,可以将模型切分为 4 份,每个 GPU 负责计算 1 层的模型,如下所示:
流水线气泡
前向计算和反向传播都需要经过 个阶段才能完成,如果不做任何优化,则任意时刻只有一个 GPU 在工作,其他 个 GPU 都处于空闲状态,流水线中的空闲时间被称为流水线气泡(Pipeline Bubble)。
为了提高流水线并行的效率,必须要减少流水线中的气泡时间。流水线气泡的产生主要是因为前向计算和反向传播需要依赖前一层的输出,而前一层的输出只有在前一层计算完成后才能得到,因此在前向计算和反向传播过程中,后续的 GPU 需要等待前一层的计算完成才能开始计算,从而产生了流水线气泡。
设前向计算和反向传播的时间分别为 和 ,流水线中有 个阶段,则 GPU 的利用率为:
减少流水线气泡的策略
有很多策略可以减少流水线并行中的气泡,让 GPU 减少等待的时间,从而提高流水线并行的效率,下面介绍几种常用的策略。
GPipe
最直接的方式是增大任务量,GPipe 将整个 batch 划分为多个 micro-batch,先执行所有 micro-batch 的前向传播,再执行所有 micro-batch 的反向传播。如下图所示:
设每个 micro-batch 的前向计算和反向传播的时间分别为 和 ,流水线中有 个阶段,整个 batch 划分为 个 micro-batch,则 GPU 的利用率为:
下面是使用 8 个流水线阶段下,不同 micro-batch 数量下的 GPU 利用率:
| micro-batch 数量 | GPU 利用率 |
|---|---|
| 1 | 12.5% |
| 2 | 22.2% |
| 4 | 36.4% |
| 8 | 53.3% |
| 16 | 69.6% |
| 32 | 82.1% |
| 64 | 90.1% |
| 128 | 94.4% |
使用了 micro-batch 后,流水线中的气泡率降低了,但是多批数据处理完毕后,需要停下来等待反向传播的梯度传递回来,这里还存在流水线气泡。
另外,micro-batch 的数量不能无限增大,因为每个 micro-batch 产生的激活值都需要保存下来,等到反向传播时再使用。如果 micro-batch 的数量过多,激活值的显存占用量也会过大,可能会超过单个 GPU 的显存容量。
另外,在训练过程中,通常 batch size 是固定的,即用于计算梯度的训练样本数量是固定的。如果增加 micro batch 的数量,则每个 micro batch 的样本数量就会减少,这会导致单次矩阵乘法规模变小,GPU 本身的算力利用率可能下降。此时虽然流水线利用率上升了,但是单个 GPU 的 MFU(Matrix Multiply Utilization)下降了,整体的训练速度可能并没有提升。
1F1B
使用了 GPipe 后,流水线气泡率降低了,但还是需要等待前向计算完成后,才能开始反向计算。为了让反向传播尽快开始,最后一层可以在处理完第一个 micro-batch 后,立刻开始反向计算,这样就可以尽快向后传递梯度,让前面的层尽快开始反向计算,从而减少流水线气泡的时间。这种策略中,最后一层在处理完第一个 micro-batch 后,立刻开始反向计算,这就形成了一个交错的计算模式,称为 1F1B(1 Forward 1 Backward)。

图片来自于 ZERO BUBBLE PIPELINE PARALLELISM
1F1B 的 GPU 利用率和 GPipe 接近,都是 ,1F1B 的优势不在于减少气泡,它的优势在于可以让反向传播尽快开始,这可以减少同时驻留在内存中的激活值。GPipe 的激活值与 micro-batch 的数量成正比,而 1F1B 的激活值与流水线的阶段数成正比。
Zero bubble
1F1B 虽然可以减少流水线中的气泡,然后并未完全消除气泡。为了进一步减少流水线气泡,可以在更细的粒度上调度各级流水线上的计算任务。
上面的图片中可以看到,前向计算和反向传播耗时是不同的,这是因为反向传播阶段的计算量比前向计算阶段的计算量大。以全连接层为例,在前向计算时,只涉及输入矩阵 需要与权重矩阵 的矩阵乘法运行。而在反向传播时,需要计算权重矩阵 的梯度和输入矩阵 的梯度。因此,反向传播的计算量约为前向计算的两倍。

而 的梯度只有在更新参数时才会使用,而输入矩阵 的梯度则需要尽快传递给前一层进行反向计算。将前向计算中的矩阵乘法记为 F,将反向传播中 的计算记为 W,将 的计算记为 B。通过优先执行 B,可以让前一层尽快开始反向计算,从而减少流水线中因为等待依赖而产生的气泡。
下图是 zero bubble 的计算示意图,它在 1F1B 的基础上,优先执行 的计算,并将 的计算延后到最后执行,从而让前一层尽快开始反向计算,减少流水线气泡。

通过精细化地调整 F、B、W 的计算顺序,可以让流水线中的气泡完全消除,从而实现 zero bubble。下图对比了 ZB-H1 (handcrafted-1)、ZB-H2 (handcrafted-2) 两种不同的 zero bubble 策略:

显存使用量
使用流水线并行将模型切分为 个阶段后,每个 GPU 只需要保存 的模型参数。参数量降低后,梯度和优化器状态的显存占用量也会相应地降低至原来的 。
通常流水线并行会使用多个 micro-batch 来提高流水线的利用率,因此每个 GPU 需要保存多个 micro-batch 的激活值。如果 batch size 保持固定,使用流水线并行时的激活值总量降低至原来的 ,因为每个 GPU 只需要保存自己负责的层的激活值。使用 1F1B 或者 zero bubble 后,因为可以尽快释放部分 micro-batch 的激活值,因此每个 GPU 需要保存的激活值数量还会进一步减少。
通信量分析
流水线并行主要使用点对点通信,每个 GPU 只需要将自己的输出传递给下一个 GPU,或者将梯度传递给前一个 GPU。这里激活张量或者梯度的大小为:
其中:
- 是 micro-batch size
- 是序列长度
- 是隐藏层维度
- 是每个元素的字节数,例如 BF16 为 2 字节
设 ,使用 BF16,则激活张量的大小为:
因此,流水线并行中每个 GPU 需要向前传递 128 MB 的激活值,向后传递 128 MB 的梯度,每个 micro-batch 的通信量为 256 MB。
和张量并行相比,流水线并行中的通信频率比较低,每个流水线阶段可以包含多层 Transformer,而通信只发生在边界部分。另外流水线并行可以很好地实现计算和通信的重叠,因此只需要保证通信带宽足够大,对时延的要求比张量并行要低。最常见的部署策略是在一个机器的多张 GPU 部署张量并行,而在多台机器之间部署流水线并行。
总结
流水线并行将模型按层切分为多个阶段,并将不同阶段分布到不同 GPU 上进行计算。通过使用 GPipe、1F1B 或者 zero bubble 等策略,可以减少流水线中的气泡,提高 GPU 的利用率。流水线并行中的通信频率相对较低,并且可以实现计算和通信的重叠,所以对通信时延的要求较低,可以跨机器部署流水线并行的各个阶段。
流水线并行通过多个 micro-batch 来屏蔽网络时延的影响,但是在推理阶段通常需要快速完成一个 batch 的处理,流水线并行会大幅增加时延。因此流水线并行往往只使用在训练阶段,在推理阶段很少使用。