模型并行 | 从Megatron谈Pipeline

关于 Megatron 的解读与分析，在网上还是能找到很多相关技术博客的。这些文章最关注的往往是 tensor parallel (TP) 的技术实现。然而，尽管 TP 在 nvidia 的实验中表现了强大的效率，它需要极高的 intra-node 带宽作为支持。在没有 nvlink 的硬件条件下，比如萤火二号，TP反而会起到反效果（详见我们过去的实验）。另一方面，作为一项非常普通的训练方法，pipeline parallel (PP) 并没有受到大量关注。人们仅仅把 PP 作为一种将大模型装入显存的方法。PyTorch 提供的 torch.distributed.pipeline.sync.Pipe 也已经足够满足大部分用户的需求。

然而，我们认为，对于萤火二号来说，PP 是一种值得深入研究的方法，PyTorch 所提供的 PP 实现也存在较大的改进空间。在这篇文档中，我们从 Megatron 出发，对 PP 进行了更多的讨论。

Pipeline 中的显卡分配

在萤火二号平台上，我们有 $n$ 个计算节点。每个节点有 8 块显卡，每 4 块组成一个 numa。数据传输速度 intra-numa $\gg$ inter-numa / intra-node $\gg$ inter-node。这也是大部分计算集群的情况。

假设我们将模型切成 4 个 stage，每个 stage 占用一块显卡。那么， 一个非常符合直觉的显卡分配策略是，4 个 stage 分别分给每个 numa 中的 4 块显卡，构成一个 model parallel group。

然而，通过阅读 Megatron 的源码，我们会发现 Megatron并不是这样做的，而是采用了一个相对复杂的做法。而且，无论是从理论上还是实际上，Megatron 的做法才是最优的。

Megatron 的显卡分配方法

源码见 megatron/mpu/initialize.initialize_model_parallel。

直观来说，Megatron 的显卡分配策略是优先把 TP 分配到相邻的显卡上，把 PP 分配到 尽可能远 的显卡上。

一个官方的例子是，如果节点数为 2（一共 16 块显卡），TP = 2，PP = 4，那么一共会分出两个 group 作为 data parallel (DP)，每个 group 会被拆成两个 pipeline（对应 TP = 2）。第一个 group 的一个 pipeline 使用的显卡会是 [0, 4, 8, 12]，另一个则是 [1, 5, 9, 13]。在这种分配下，TP 的两个显卡总是相邻的（0 和 1，4 和 5，等等），但是相邻的 stage 会跨 numa （0 和 4）甚至 node（4 和 8）。

直觉上的一种分配方式可能是 [0, 1, 2, 3] + [4, 5, 6, 7] 的组合。这样做肯定是不好的，因为 TP 的显卡已经跨 numa 了（0 和 4，1 和 5，等等）。

另一种符合直觉的分配方式是 [0, 2, 4, 6] + [1, 3, 5, 7] 的组合。既满足了 TP 的显卡相邻，也满足了每个 stage 都不会跨 node。但是 Megatron 也没有这样做。

一个帮助进一步理解的例子是，如果节点数为 4（一共 32 块显卡），TP = 1（符合我们的需求），PP = 4，那么 Megatron 给出的显卡分配会是 [0, 8, 16, 24] 一组，[1, 9, 17, 25] 一组，以此类推，而不是 [0, 1, 2, 3] 一组。

下面用一个简化的例子解释为什么 Megatron 的反直觉方法更好。假设我们有两个节点，每个节点两块显卡，TP = 1，PP = 2，那么 DP = 2。

一种分配方法是：[0, 1] + [2, 3]，即在显卡 0 和 1 之间做 pipeline 的数据传输。这意味着 DDP 时要在显卡 0 和 2 之间对梯度做 AllReduce 。

另一种分配方法是：[0, 2] + [1, 3]，即在显卡 0 和 2 之间做 pipeline 的数据传输。这意味着 DDP 时只需要在 0 和 1 之间对梯度做 AllReduce 。

当相邻 stage 被分配在接近的显卡上时，尽管 forward 与 backward 在计算时的数据传输开销降低了，代价却是在 DDP 中对梯度进行 AllReduce 的开销显著增加了。反过来讲，把 pipeline 尽可能分散，就意味着 AllReduce 的对应显卡会相邻，相应的数据传输开销就会更低。

理论计算

下面做一些简单的计算。

假设不使用 TP（TP = 1），每个 group 分到的 batch size 为 $b$，序列长度为 $s$，hidden size 为 $h$，每个节点的显卡数量为 $g$，pipeline 的深度为 $p$。另外假设模型共有 $l$ 层 transformer layer。

当优先让 pipeline 位于同一节点内时，每个节点发送的数据约为模型参数量的 2 倍，近似为 $24lh^2$。

当优先让 pipeline 位于不同节点内时，对于每个节点，它需要发送的数据量为节点内不同 group 的数量（即为显卡数量） $g$ 乘以 transformer layer 的输出数据量 $2bsh$ （2 对应 forward + backward，考虑 activation checkpoint），即为 $2bshg$。

二者的比例为 $bsg / 12lh$。其中，$g$ 一般固定为 8，$s$ 一般固定为 1024 或 2048，$b$ 一般固定为 8 或 16。然而，随着模型规模的增加，$l$ 和 $h$ 是会增大的。这使得 $bsg/12lh$ 往往小于 0.1，也就是说 DDP 要传输的数据量是远大于 pipeline 的。因此，在分配显卡时，应当尽可能让保存相同权重的显卡尽量接近，尽管这样的代价是 pipeline 的显卡会相隔较远。

一点实验

我们修改了 Megatron 的代码使其能够支持优先将 pipeline 分配到相邻显卡上（megatron/mpu/initialize.initialize_model_parallel_intra_node），并与官方的分配策略进行了对比：

其中 intra-node 与 inter-node 表示 pipeline 是否跨节点。可以看到，在 intra-node 下，即使 send/recv 的时间开销都更低，但是在 all-reduce 一项上的开销是 inter-node 的 13.3 倍。作为参考，在对应的参数设置下，理论上 DDP 传输的数据量是 pipeline 的 12 倍，和实际情况比较接近。

一点总结

用户应该根据模型的规模来决定 pipeline 的显卡分配策略。但是在大多数情况下，让 pipeline 跨 node 都远比让 DDP 跨 node 划算。不幸的是，PyTorch 的 pipeline 实现并不支持跨节点。

Megatron 的跨节点 Pipeline 实现

我们已经说明，跨节点的 Pipeline 对于提高模型训练效率是有利的，而 PyTorch 目前并不支持这一特性。在这节中，我们简单介绍 Megatron 的 pipeline 实现思路。

PyTorch 的 Pipeline 实现

PyTorch 为用户封装好了 torch.distributed.pipeline.sync.Pipe。它会接收 torch.nn.sequential 对象作为参数并将其中的每个操作作为独立的 stage。

在这样的设置下，一个 pipeline 就对应一个完整的模型，各个 stage 都属于同一个进程，这个进程则拥有多块显卡。Pipe 要求用户事先将每个 stage 移动到相应的显卡上去，然后使用 RPC 作为后端来控制 pipeline 的进行。这样做很直观，但是问题在于，当我们希望 pipeline 跨节点的时候，pipeline 中的各个 stage 都不可能属于同一个进程。

Megatron 的 Pipeline 实现

Megatron 使用了更加灵活的方式来构建训练。在初始化时，对于每块显卡，Megatron 都启动一个独立的进程，并根据自身的 global rank 来构建模型的对应部分，或者说是一部分 transformer layer。

接下来，Megatron 根据不同的并行化方式构建相对应的 process group，包括：

• PIPELINE_MODEL_PARALLEL_GROUP：每个 group 包含 pipeline 中的所有 stage。
• TENSOR_MODEL_PARALLEL_GROUP：每个 group 包含对于同一个 transformer layer 的拆分。即每个 group 的显卡拼合在一起就是一个完整的 transformer layer。
• DATA_PARALLEL_GROUP：每个 group 包含模型的相同部分。即同一个 group 中的每块显卡都保存了相同的模型参数。
• MODEL_PARALLEL_GROUP：PIPELINE_MODEL_PARALLEL_GROUP 与 TENSOR_MODEL_PARALLEL_GROUP 的合并。

上一小节就是主要介绍了如何构建这些 GROUP。还是沿用之前的例子，假设显卡数为 16，TP = 2，PP = 4，那么：

• PIPELINE_MODEL_PARALLEL_GROUP = [0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15]
• TENSOR_MODEL_PARALLEL_GROUP = [0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15]
• DATA_PARALLEL_GROUP = [0, 2], [1, 3], [4, 6], [5, 7], [8, 10], [9, 11], [12, 14], [13, 15]
• MODEL_PARALLEL_GROUP = [0, 1, 4, 5, 8, 9, 12, 13], [2, 3, 6, 7, 10, 11, 14, 15]

这样，借助于 torch.distributed 中提供的操作，就可以比较方便地进行并行化的计算。

pipeline 中进行的操作都是 p2p 的。比如在 pipeline 运行一个 stage 时，可以直接用 ring_exchange （根据字面意思可以猜测用途，源码中有，但是实际上 PyTorch 没有提供这个 API）来进行一次轮转，或者在自己的 process group 中从上一个 rank 拿数据，计算后再发给下一个 rank。这些操作可以用 torch.distributed.P2POp 来简单实现。

类似地，在进行 TP 时，只需要在对应的 TENSOR_MODEL_PARALLEL_GROUP 里做 AllReduce 或者 AllGather；在做 DDP 时，只需要在 DATA_PARALLEL_GROUP 里做 AllReduce。