大模型训练流水线并行(PP)性能评价指标与分析方法
在大模型训练与推理的性能分析与调试调优过程中,直接查看底层每个计算、通信或线程的空隙和执行时间往往过于细粒度,难以快速定位瓶颈。因此,我们需要一种能够从高粒度算法层面分析大模型训练及推理各阶段性能的手段。
核心问题在于:如何快速分析出大模型的训练与推理在各个阶段的耗时数据?特别是分布式并行策略(DP/TP/PP)下的通信耗时、PP Bubble Ratio(气泡率)等关键指标,而不是依赖细粒度的 Profiling 结果进行人工提取。
流水并行 PP 评价指标
在流水线并行(Pipeline Parallelism, PP)中,假设将模型切分为 $p$ 个 stage,总样本数为 $m$,每个 micro-batch 的前向执行时间为 $t_f$,后向执行时间为 $t_b$。在理想环境下,前向和后向的合计执行时间应满足理论最优值。
然而,在实际 Pipeline 中将网络模型切分成不同 stage 并部署在不同机器上时,会导致正向的 $p-1$ 和反向的 $p-1$ 时间内出现计算等待时间。这个中间空白区域被称为 Bubble(空泡)。
Bubble 产生的原因
空泡的产生是因为算子内并行和算子间并行是为了解决单设备内存不足的问题。模型经过拆分后,上一个 rank 的 stage 需要长期处于空闲状态,等待其他 rank 的 stage 计算完成才能开始计算。这极大降低了设备的平均使用率,这种现象被称为并行空泡(Parallelism Bubble)。
为了充分提升算力利用率,即充分发挥硬件算力,最佳做法是将大模型按照模型并行(优先采用张量并行 TP)进行切分,如果可以的话加上流水并行 PP,尽可能占满模型并行机器节点的内存。在此前提下,可以自由调整流水线并行(PP)和张量并行(TP)的比例。
假设数据并行(DP)参数设置为 $d=1$,节点数为 $n$,张量并行(TP)设置为 $t$,流水并行(PP)设置为 $p$,且满足 $n = t \times p$。
性能评价结论
实际验证分析发现,增加张量并行(TP)的大小,将会减少泡沫比率(Bubble Ratio),从而提高整体利用率,减少卡间不同步等待的时间。不过,张量并行(TP)的大小不能无限增加,因为张量并行会拆分模型降低整体的计算粒度。一般默认张量并行(TP)最大不跨节点,节点内可以增大。
因此可以得出一个简单的结论:张量并行(TP)的大小应该随着模型尺寸的增加而缓慢增加,但不超过单个节点。经过一系列实验验证,张量并行(TP)的最佳配置范围通常为:
$$ 4 \le t \le 8 $$
其中参数 $t$ 为张量并行的维度。在 GLM-130B 模型验证过程中,当 micro-batches ($m$) = 176,PP ($p$) = 8 时,Bubble Size 仅为 3.8%,足以证明流水线模型并行的效率。如果在 910B 架构上是 1*16 的机器,模型是千亿级别,TP 尽可能在 4~16 之间,同时 PP 应尽量减少跨服务器节点通信。
实际实验证明,当 $m \ge 4p$ 时,在整体计算过程中会加入重计算技术来降低内存消耗,此时 Bubble Ratio 的百分比会降低到一个可以忽略的比例。重计算过程主要是针对网络模型的正向计算,反向计算的时候会跟计算通信重叠,因此流水线并行(PP)引入的气泡不会严重消耗训练效率。
主流流水并行算法盘点
目前主要是 Google、微软、NVIDIA 三家厂商提出了多种不同的流水线并行算法。以下是对几种主流算法的特点及适用场景的分析。
1. Native (朴素流水线)
朴素的流水线并行,同一时刻只有一个设备进行计算,其余设备处于空闲状态,计算设备利用率通常较低。这是最基础的实现方式,适用于对延迟要求极高但对吞吐量要求不高的场景。
2. GPipe
GPipe 将朴素流水线并行的 batch 再进行切分,减小设备间空闲状态的时间,可以显著提升流水线并行设备利用率。它采用 F-then-B 调度方式,将原本的 mini-batch(数据并行切分后的 batch)划分成多个 micro-batch(再切分后的 batch)。每个 pipeline stage(流水线并行的计算单元)先整体进行前向计算,再进行反向计算。通过在同一时刻分别计算模型的不同部分,F-then-B 模式可以显著提升设备资源利用率。但也不难看出这种模式由于缓存了多个 micro-batch 的中间变量和梯度,显存的实际利用率并不高。
3. 1F1B (One Forward One Backward)
1F1B 流水线并行方式解决了显存占用过高的问题。在该模式下,pipeline stage 前向计算和反向计算交叉进行,可以及时释放不必要的中间变量。相比 GPipe,1F1B 显著减少了显存峰值占用,同时保持了较高的计算密度。
4. PipeDream
PipeDream 在单个 GPU 上进行短暂的运行时性能分析后,可以自动决定怎样分割这些 DNN 算子,如何平衡不同 stage 之间的计算负载,而同时尽可能减少目标平台上的通信量。PipeDream 将 DNN 的这些层划分为多个阶段——每个阶段(stage)由模型中的一组连续层组成。该阶段对本阶段中所有层执行向前和向后传递。PipeDream 将包含输入层的阶段称为输入阶段,将包含输出层的阶段称为输出阶段。其特点是允许异步更新,但在收敛性上可能面临挑战。
5. Virtual Pipe
Virtual Pipe 是 Megatron-LM 论文中最主要的一个创新点。传统的 pipeline 并行通常会在一个 Device 上放置几个 block,这是为了扩展效率考虑,在计算强度和通信强度中间取一个平衡。但 virtual pipeline 却反其道而行之,在 device 数量不变的情况下,分出更多的 pipeline stage,以更多的通信量,换取空泡比率降低,减小了 step e2e 用时。
性能分析实践建议
在大模型训练过程中,分析 PP 的性能占比和耗时应遵循以下步骤:
- 建立理论基准:首先有一个理论的评价指标,根据理论的计算公式(如 Bubble Ratio 理论值)设定预期。
- 选择并行策略:对比具体使用哪种流水并行策略(如 1F1B vs GPipe),结合模型大小和显存限制进行选择。
- Profiling 实测:利用工具(如 NVIDIA Nsight Systems, PyTorch Profiler)获取实际大模型训练的 Profiling 结果。
- 差异分析:重点观察流水并行 PP 策略下,理论 Bubble Size 跟实测 Bubble Size 之间的差异。如果差异过大,需检查通信带宽是否受限、Stage 负载均衡是否不均或是否存在同步等待。
通过上述方法,可以有效评估流水线并行的效率,并为后续的超参数调整和系统优化提供数据支持。