大型模型的模型与流水线并行

数据并行通过分散数据有效地扩展了训练，但它依赖一个主要假设：模型本身可以放入单个加速器的内存中。对于先进模型，特别是在自然语言处理和计算机视觉方面，这个假设不再成立。当模型的参数、梯度和优化器状态超出单个GPU的可用显存，就需要一种不同的方法。将模型本身分散到多个设备上的方案能够解决这个问题。

层间模型并行 (张量并行)

处理单个无法放入一个GPU的庞大层最直接的方法是拆分层的计算，这种方法通常称为张量并行。考虑一个大型线性层，由方程 $Y = XA$ 定义，这里 $X$ 是输入激活，$A$ 是权重矩阵。如果矩阵 $A$ 对单个设备来说太大，我们可以将其划分到多个GPU上。

例如，将线性层的权重矩阵 $A$ 按列拆分到两个GPU上 ($A = [A_1, A_2]$)，我们就可以并行计算输出的部分。

$$Y = X[A_1, A_2] = [XA_1, XA_2]$$

GPU 1 计算 $Y_1 = XA_1$，GPU 2 计算 $Y_2 = XA_2$。输入 $X$ 被广播到两个GPU。并行计算完成后，结果 $Y_1$ 和 $Y_2$ 被聚合以形成最终输出 $Y$。这个过程需要较多的通信：输入必须发送到所有设备，并且部分结果必须组合。

一张张量并行图，其中单个层的权重矩阵被拆分到两个GPU上。输入被广播，并且部分结果通过All-Gather过程组合。

反向传播遵循类似的反向方式。关于输出的梯度 $\nabla_Y$ 被拆分并发送到各自的GPU。每个GPU计算其局部权重梯度，并且关于输入的梯度在使用Reduce-Scatter过程组合后，才传递给上一层。

这种细粒度并行带来了高频通信，使其对互连带宽高度敏感。它在具有NVLink等高速连接的多GPU服务器内部最为有效。

流水线并行

对于非常深的模型，另一种方法是在层之间而不是之内划分模型。这被称为流水线并行。在这种方法中，连续的层块（称为阶段）被放置在不同的加速器上。

例如，在一个运行在4个GPU上的32层模型中：

• GPU 0 负责处理第1-8层。
• GPU 1 负责处理第9-16层。
• GPU 2 负责处理第17-24层。
• GPU 3 负责处理第25-32层。

一个数据批次被送入GPU 0。一旦GPU 0完成其第1-8层的前向传播，它就将产生的激活传递给GPU 1，GPU 1随后开始其计算。这个过程沿着流水线顺序继续。

朴素流水线并行中的顺序流。每个GPU阶段都等待前一个完成，造成较多的空闲时间。

解决流水线气泡问题

上面所示的简单顺序方法效率低下。当GPU 0处理一个批次时，GPU 1、2和3完全处于空闲状态。这个空闲时间，被称为“流水线气泡”，严重影响了整体硬件利用率。

为了减轻这种情况，输入数据批次被拆分成更小的微批次。一旦GPU 0完成处理第一个微批次并将其传递给GPU 1，它就可以立即开始处理第二个微批次。这使得GPU可以并行处理不同的微批次，创建了一个真正的流水线。

4个微批次的4阶段流水线中GPU随时间变化的利用率。“Fwd”表示前向传播，“Bwd”表示反向传播。初始启动和最终清空阶段（即“气泡”）呈现为灰色区域，但中间阶段体现出高且重叠的利用率。

这个调度方案，通常被称为GPipe，提高了利用率但不尽理想。在处理完整批次的开始（启动阶段）和结束（下降阶段）时，气泡仍然存在。更先进的调度器，例如DeepSpeed中使用的，可以通过改变反向传播顺序来进一步改进此问题，以填充更多气泡。

选择合适的方法

模型并行和流水线并行可以同时使用；它们解决了同一个问题的不同方面，并经常一起使用。选择取决于模型架构和硬件限制。

• 模型并行

  • 何时使用： 当单个层太大，无法放入单个GPU内存时。这在大型语言模型中的注意力头和MLP块中很常见。
  • 通信方式： 单个层的前向/反向传播中需要高频、低延迟的通信（例如，all-reduce）。它最适合单个节点内紧密连接的GPU（例如，通过NVLink）。

• 流水线并行

  • 何时使用： 对于可以逻辑上划分为顺序阶段的非常深的模型。它很适合通过较慢的网络互连（如InfiniBand或以太网）在多个节点之间进行扩展。
  • 通信方式： 较低频率的通信，包括每个微批次一次性在阶段之间传递较大的激活张量。

在许多生产情况下，混合方法是最有效的解决方案。例如，一个大型模型可能使用流水线并行跨不同节点拆分成多个阶段。在每个节点内部，分配给单个阶段的GPU可能使用张量并行来管理该阶段层的内存，并使用数据并行来更快地处理微批次。这种组合使得模型能够扩展到大型加速器集群中的超大模型尺寸。