3D 并行 (3D Parallelism)详解

3D 并行 (3D Parallelism) 实战演练：打造 LLM 训练的“超级工厂”

导读：在实际的千亿参数模型训练中，单一的并行策略往往捉襟见肘。我们通常需要混合使用 数据并行 (DP)、流水线并行 (PP) 和 张量并行 (TP) 这三种策略，构建一个高效的 3D 并行系统。本文将通过一个具体的“超级工厂”场景，带你一步步搭建这套复杂的系统。

我正在学习LLM，这是我用Gemini生成的内容，最下面我自己照着他举的例子，理解着画了一张图，大家可以一起看看，如果我画的不对，大家也可以批评指正。

1. 场景设定：训练一个“巨无霸”模型

假设我们要训练一个 Llama-3-70B（700 亿参数）级别的模型。

• 硬件资源：我们拥有一个计算集群，包含 4 台服务器 (Node)。
• 单机配置：每台服务器内部署了 8 张 NVIDIA A100 显卡。
• 总资源池：共计 $4\times 8=32$ 张显卡。

面临的挑战：

1. 模型太宽：单层权重矩阵过于庞大，单张显卡无法完成计算 -> 需要 TP。
2. 模型太深：层数极多（例如 80 层），总显存占用超过单卡上限 -> 需要 PP。
3. 训练太慢：需要快速处理海量训练数据以缩短训练周期 -> 需要 DP。

解决方案：我们将这 32 张卡组织成一个 3D 并行系统。设定并行度为：TP=4, PP=4, DP=2。

2. 第一维：张量并行 (TP) —— “车间小组”

核心逻辑：利用单机内部的高速互联，解决“单层算不动”的问题。

我们将视角聚焦在**第 1 台服务器（Node 1）**上。里面的 8 张卡不再是单独作战的个体。我们将每 4 张卡 捆绑成一个紧密协作的“TP 小组”。

• TP Group 1: 包含 GPU 0, 1, 2, 3

  • 这 4 张卡利用 NVLink 实现高速互联，这是 TP 高频通信的基础。
  • 它们合力存放并计算模型第 1 层（以及后续层）的权重矩阵。
  • 运作方式：正如 TP 原理所述，巨大的权重矩阵 $W$ 被切分为 4 份，每张卡负责一部分计算。每层计算结束时，必须立刻进行 All-Reduce 通信拼装结果。
  • 对外表现：在外界看来，这 4 张卡仿佛合体成了一张拥有超大显存和算力的“超级显卡”，能够处理单个 GPU 无法应对的巨大矩阵乘法。

• TP Group 2: 包含 GPU 4, 5, 6, 7

  • 这是同一个机器里的第二个 TP 小组，同样由 4 张卡合体而成。

现状总结：通过 TP，我们成功将 32 张独立的显卡，转化为了 8 个强力的“TP 小组”（每组 4 卡）。

3. 第二维：流水线并行 (PP) —— “接力流水线”

核心逻辑：通过跨机器的层数切分，解决“层数太多装不下”的问题。

现在的“TP 小组”虽然算力强大，但显存容量依然有限，无法一次性装下完整的 80 层模型。假设每个小组只能容纳 20 层。
因此，我们需要采用“接力赛”的方式，将 4 个分布在不同机器上的“TP 小组”串联起来。

我们将 80 层的模型切分为 4 个阶段（Stages）：

• Stage 1 (第 1-20 层)：部署在 Node 1 的 TP Group 1。
• Stage 2 (第 21-40 层)：部署在 Node 2 的 TP Group 1。
• Stage 3 (第 41-60 层)：部署在 Node 3 的 TP Group 1。
• Stage 4 (第 61-80 层)：部署在 Node 4 的 TP Group 1。

运作流程：

1. 数据首先进入 Node 1，由其 TP Group 1 完成前 20 层的计算，得到中间结果（Activations）。
2. 中间激活值通过网线跨机器传输给 Node 2。
3. Node 2 继续计算接下来的 20 层… 如此接力，直到 Node 4 完成最后计算并算出 Loss。

现状总结：这 4 个跨机器的 TP 小组，串联组成了一条完整的模型流水线 (Model Replica)。

• 资源消耗：这条流水线共占用了 $4(\text{TP})\times 4(\text{PP})=16$ 张显卡。

4. 第三维：数据并行 (DP) —— “复制生产线”

核心逻辑：通过复制完整模型实例，解决“数据吞吐量不足”的问题。

回顾总资源池，我们一共有 32 张卡。刚才搭建第一条完整的流水线只用了 16 张卡（即 Node 1-4 的前半部分资源）。
剩下的 16 张卡（即 Node 1-4 的后半部分，也就是各自的 TP Group 2）该如何利用呢？

答案是：利用 DP 思想，再建一条一模一样的流水线，进行并行训练。

• Pipeline A (副本 1)：使用 16 张卡（由每台机器的 TP Group 1 组成）。
• Pipeline B (副本 2)：使用剩下的 16 张卡（由每台机器的 TP Group 2 组成）。

运作流程：

1. 数据切分：假设一个大 Batch 有 100 条数据。Pipeline A 负责处理前 50 条，Pipeline B 负责处理后 50 条。
2. 并行狂奔：两条流水线同时开工，互不干扰地进行前向和反向计算。
3. 最后同步 (All-Reduce)：当两条流水线都完成一次反向传播并算出梯度后，处于相同逻辑位置的显卡（例如 Pipeline A 的 GPU 0 和 Pipeline B 的 GPU 4，它们都存储着完全相同的模型参数分片）会进行跨机器通信，将各自计算出的梯度进行平均，以保证模型权重更新的一致性。

5. 总结：一张全景图

为了更直观地理解这个复杂的 3D 结构，我们将其整理为一张架构图。
假设：GPU (i, j) 代表第 i 台机器的第 j 张卡。

3D 立方体计算公式验证

最后，我们用公式验证一下这个配置：
$$\text{Total GPUs}=\text{TP size}\times \text{PP size}\times \text{DP size}$$
$$32=4\times 4\times 2$$

• TP=4：4 张卡合力，解决显存墙和单层计算瓶颈。
• PP=4：4 个节点接力，解决超深网络层数瓶颈。
• DP=2：2 个副本并行跑，让整体训练速度翻倍。

导师点评：
通过这个具体的例子，3D 并行的概念就变得立体了：

• TP 是把层拆碎在单机内。
• PP 是把拆碎的层在机器间串联。
• DP 是把串联好的整体进行复制扩展。

这就是为什么训练 GPT-4、Llama-3 这种级别的模型，不是简单的“堆显卡”，而是需要极高水平的系统架构设计能力，才能让成千上万张卡在这三个维度上完美协作，发挥出最大效能。

为了更好的理解，我自己又画了一张图