分布式并行中并行组

在分布式训练中，通常会采用 数据并行（DP）、张量并行（TP）、流水并行（PP） 结合的方式，以最大化计算效率和内存利用率。你的 tp_groups 和 dp_groups 定义了一个 混合并行策略，其中：

• TP (Tensor Parallelism, 张量并行)
• DP (Data Parallelism, 数据并行)
• PP (Pipeline Parallelism, 流水并行)
  我们来逐步解析 并行组 的含义。

1. 并行组的基本概念

在分布式训练中，每个 GPU（或者 NPU）被称为一个 rank，所有 GPU 组成一个 world。
假设 world_size=16，即共有 16 个 GPU（编号 g0 ~ g15），然后进行 划分并行组。

• TP (Tensor Parallelism, 张量并行)
  • 组内 的 GPU 共同处理 同一批数据，但分担计算（例如，分割矩阵乘法）。
  • 这里 tp_world_size=2，即 每组 2 个 GPU 负责 同一个模型层的部分计算。
  • tp_groups 结构（每组 2 个 GPU）：
    tp_groups = [ [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15] ]
    • 例如 g0, g1 组成一个 张量并行组，它们协同计算同一个 Transformer 层的一部分计算（如 FFN 分割）。
    • 共有 16 / 2 = 8 组，每组 2 个 GPU。
• DP (Data Parallelism, 数据并行)
  • 组内 GPU 计算相同的模型参数，但每个 GPU 处理 不同的数据 batch。
  • 这里 dp_world_size=8，即 每组 8 个 GPU 计算 相同的模型参数，但处理 不同数据。
  • dp_groups 结构：
    dp_groups = [ [g0, g2, g4, g6, g8, g10, g12, g14], # DP 组 1 [g1, g3, g5, g7, g9, g11, g13, g15] # DP 组 2 ]
    • 例如 g0, g2, g4, g6, g8, g10, g12, g14 是 一个数据并行组，它们每个都执行相同的计算，只是输入的数据不同。
    • 共有 2 组，每组 8 个 GPU。
• PP (Pipeline Parallelism, 流水并行)
  • 每个 GPU 组计算模型的不同部分（如不同的 Transformer 层）。
  • 这里 pp_world_size=1，说明 没有使用流水并行（即，每个 GPU 计算完整的 Transformer 层，而不是切分成多个 stage）。
  • 如果 pp_world_size=2，那么模型的前半部分由 前 8 个 GPU 计算，后半部分由 后 8 个 GPU 计算，并通过 pipeline 进行通信。

2. 结合 tp_groups 和 dp_groups 理解并行组

你的并行划分如下：

• 每 2 个 GPU 为一组 进行 张量并行（tp_groups）。
• 每 8 个 GPU 为一组 进行 数据并行（dp_groups）。
• 没有流水并行（pp_world_size=1，即模型不会拆分到不同的 GPU 进行流水计算）。
  示例
• g0, g1：在 tp_groups[0] 内，它们 共享参数，分割计算。
• g0, g2, g4, g6, g8, g10, g12, g14：在 dp_groups[0] 内，它们 独立计算，但同步梯度。
• g1, g3, g5, g7, g9, g11, g13, g15：在 dp_groups[1] 内，计算相同的内容，只是 batch 不同。

3. 为什么这样划分？

这种并行策略的优势： ✅ 张量并行 (TP) 降低单卡计算量：每个 GPU 只计算 FFN 或 Attention 的一部分，减少显存占用。
✅ 数据并行 (DP) 增加 batch size：每个 GPU 计算相同的模型参数，但处理不同的数据，提升吞吐量。
✅ 避免流水并行 (PP) 复杂度：因为 pp_world_size=1，省去了流水并行的同步和等待问题，简化实现。

4. 计算示例

假设：

• 模型参数 W 被 张量并行 (TP) 分割 成 2 份，每个 GPU 计算一部分。
• 16 张 GPU 分成 8 组，每组 2 个 GPU 共享 W 的一部分。
• 梯度同步 (AllReduce) 在数据并行 (DP) 内进行。

前向传播

1. 张量并行：
   • g0, g1 计算 同一个 Transformer 层，但分别计算不同的权重 W1, W2。
   • 计算完成后，它们通信合并结果。
2. 数据并行：
   • g0, g2, g4, g6, g8, g10, g12, g14 计算 相同的模型参数，但不同 batch 数据。
   • 计算完成后，它们不会通信（只是计算不同数据）。

反向传播

1. 数据并行 (DP) 需要梯度同步：
   • g0, g2, g4, g6, g8, g10, g12, g14 计算出梯度后，执行 AllReduce，确保所有 GPU 参数一致。
   • g1, g3, g5, g7, g9, g11, g13, g15 也执行 AllReduce。
2. 张量并行 (TP) 需要通信梯度：
   • 由于 g0, g1 计算了不同部分的 W，它们需要进行 梯度合并 (Reduce-Scatter)。
   • 例如：
     • g0 计算 ∇W1，g1 计算 ∇W2，然后 Reduce-Scatter 交换梯度。

5. 结论

• 张量并行 (TP)：用于 减少单个 GPU 计算量，并行计算 Transformer 层。
• 数据并行 (DP)：用于 增加 batch size，提高吞吐量。
• 流水并行 (PP)：未使用，但如果 pp_world_size=2，则模型的 前后部分会分配到不同 GPU 计算。
  这种 (TP=2, DP=8, PP=1) 组合适用于 大规模 Transformer 训练，在 计算效率、通信开销、吞吐量 之间取得了较好的平衡。

6 并行度

并行度可以从不同的维度分析，包括 张量并行（TP）、数据并行（DP）、流水并行（PP）。
在你的设定中：

• TP 并行度（张量并行度） = tp_world_size = 2
• DP 并行度（数据并行度） = dp_world_size = 8
• PP 并行度（流水并行度） = pp_world_size = 1

总并行度计算

并行度的核心是 并行计算任务的总 GPU 数量，也就是
总并行度=TP×DP×PP\text{总并行度} = \text{TP} \times \text{DP} \times \text{PP}总并行度=TP×DP×PP
在你的配置中：
总并行度=2×8×1=16\text{总并行度} = 2 \times 8 \times 1 = 16总并行度=2×8×1=16

解释

• 你的 16 张 GPU 被分配到了 数据并行、张量并行的组合，但 没有流水并行。
• 每个 DP 组有 8 个 GPU，这意味着 有 8 份独立的数据被同时训练，最终使用 AllReduce 同步梯度。
• 每个 TP 组有 2 个 GPU，它们计算 模型的一部分张量，通过 Reduce-Scatter + All-Gather 来合并结果。
  并行度的直观理解：
• 你可以理解为：整个模型 由 2 张 GPU 共同计算（TP），但数据被切分成 8 份（DP），总共 16 张 GPU 在运行。
• 总共有 8 份数据在同时计算，每份数据的计算任务被 2 张 GPU 共同完成。

并行度对性能的影响

• 模型参数被 2 份 TP 计算，降低单卡显存占用
• 8 份数据同时计算，增强吞吐量
• 没有 PP，避免了流水并行的同步问题（如 bubble）
  因此，这种 (TP=2, DP=8, PP=1) 的并行度 适用于大规模 Transformer 训练，可以在显存利用和计算效率之间取得平衡。