数据并行、模型并行、流水线并行

分布式训练的三种主要方式。各有适用场景，大模型训练通常是混着用。

这篇把三种并行讲清楚，包括原理、优缺点、什么时候用。

数据并行（Data Parallelism）

原理

最简单的并行方式：每张卡放一份完整的模型，各自处理不同的数据。

输入数据：[batch 0, batch 1, batch 2, batch 3]

卡0：完整模型，处理 batch 0 → 梯度 0
卡1：完整模型，处理 batch 1 → 梯度 1
卡2：完整模型，处理 batch 2 → 梯度 2
卡3：完整模型，处理 batch 3 → 梯度 3

↓ 梯度同步（AllReduce）

所有卡得到平均梯度 → 更新参数

梯度同步

关键步骤是 AllReduce：把所有卡的梯度汇总求平均。

# 伪代码
for each GPU:
    gradient = backward(loss)

# AllReduce
avg_gradient = sum(all_gradients) / num_gpus

for each GPU:
    model.parameters -= lr * avg_gradient

AllReduce 要传输的数据量 = 模型参数量。7B 模型 FP16 就是 14GB。

通信模式

AllReduce 通常用 Ring AllReduce 实现：

GPU0 → GPU1 → GPU2 → GPU3 → GPU0
     ←      ←      ←      ←

每张卡把自己的一部分梯度传给下一张卡，同时接收上一张卡的梯度。环形传一圈就完成了。

优点

  • 实现简单，代码改动小
  • PyTorch 原生支持（DDP）
  • 效率高（通信可以和计算重叠）

缺点

  • 每张卡要放完整模型，显存受限
  • 模型太大就放不下

适用场景

  • 模型能放进单卡
  • 想通过多卡加速训练
  • 入门分布式训练

代码示例

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化分布式环境
dist.init_process_group(backend='nccl')

# 包装模型
model = DDP(model, device_ids=[local_rank])

# 正常训练
for batch in dataloader:
    loss = model(batch)
    loss.backward()  # 梯度自动同步
    optimizer.step()

张量并行（Tensor Parallelism）

原理

把模型的某一层切分到多张卡上，每张卡算一部分。

最常见的是切分线性层（矩阵乘法）：

原始：Y = X × W

切分 W 成两半：
W = [W1, W2]

卡0 算：Y1 = X × W1
卡1 算：Y2 = X × W2

合并：Y = [Y1, Y2]

Transformer 的张量并行

Transformer 里的 Attention 和 FFN 都可以切：

Attention 切分

Q、K、V 矩阵按 head 切分
卡0：head 0-7
卡1：head 8-15
...

FFN 切分

FFN：Y = ReLU(X × W1) × W2

按列切 W1，按行切 W2
每张卡算一部分，最后 AllReduce 汇总

通信开销

张量并行每一层都要通信（AllReduce 或 AllGather），开销比较大。

所以张量并行通常只在同一个节点内用（走 NVLink），不跨节点。

优点

  • 可以训练超大的单层（比如很宽的 FFN）
  • 显存占用分摊

缺点

  • 通信频繁，需要高带宽互联
  • 实现复杂
  • 不能无限扩展（通信开销会爆炸）

适用场景

  • 单层很大，单卡放不下
  • 节点内多卡（NVLink 连接）
  • 通常和其他并行方式组合使用

流水线并行（Pipeline Parallelism）

原理

把模型按层切分到不同的卡上，数据像流水线一样流过各张卡。

模型：层0 → 层1 → 层2 → ... → 层31

卡0：层 0-7
卡1：层 8-15
卡2：层 16-23
卡3：层 24-31

数据流动：
batch → 卡0 → 卡1 → 卡2 → 卡3 → output

Bubble 问题

流水线并行有个大问题：bubble（气泡）。

时间 →

卡0: [batch0] [batch1] [batch2] [batch3] [等待]  [等待]  [等待]  [等待]
卡1: [等待]   [batch0] [batch1] [batch2] [batch3] [等待]  [等待]  [等待]
卡2: [等待]   [等待]   [batch0] [batch1] [batch2] [batch3] [等待]  [等待]
卡3: [等待]   [等待]   [等待]   [batch0] [batch1] [batch2] [batch3] [等待]

开始和结束时，很多卡在等，效率损失。

解决 Bubble

1. 增加 micro-batch 数量

把一个大 batch 切成很多小 micro-batch：

batch → [micro0, micro1, micro2, micro3, ...]

更多 micro-batch = 更少 bubble 占比

2. 1F1B 调度

交错安排前向和反向传播：

卡0: F0 F1 F2 F3 B0 B1 B2 B3
卡1:    F0 F1 F2 B0 B1 B2 F3 B3
卡2:       F0 F1 B0 B1 F2 F3 B2 B3
卡3:          F0 B0 F1 B1 F2 B2 F3 B3

F = 前向，B = 反向。减少 bubble。

3. Interleaved Schedule

更复杂的调度策略，进一步减少 bubble。

优点

  • 通信开销小（只在层边界通信）
  • 可以跨节点（不需要 NVLink）
  • 显存分摊

缺点

  • Bubble 导致效率损失
  • 实现复杂（调度、梯度累积）
  • 需要精心设计切分点

适用场景

  • 模型层数多，需要切分
  • 跨节点训练
  • 和其他并行组合使用

ZeRO：数据并行的优化

数据并行的问题是每张卡要存完整模型。ZeRO（Zero Redundancy Optimizer）解决这个问题。

ZeRO 的三个阶段

ZeRO Stage 1：切分优化器状态

原来：每张卡存完整的 optimizer states
ZeRO-1：每张卡只存 1/N 的 optimizer states

ZeRO Stage 2：切分优化器状态 + 梯度

ZeRO-2：每张卡只存 1/N 的 optimizer states 和梯度

ZeRO Stage 3：切分优化器状态 + 梯度 + 参数

ZeRO-3：每张卡只存 1/N 的一切
需要时从其他卡取

显存节省

方式7B 模型每卡显存（估算）
普通数据并行80-120GB
ZeRO Stage 150-60GB
ZeRO Stage 230-40GB
ZeRO Stage 310-20GB

ZeRO-3 可以用数据并行训练原本放不下的模型。

代价

ZeRO 需要更多通信。Stage 3 每次前向/反向都要收集参数。

适合通信带宽高的环境。

3D 并行

实际训练超大模型，通常同时用三种并行：

- 数据并行（DP）：跨多个流水线
- 张量并行（TP）：节点内多卡
- 流水线并行（PP）：跨节点

例如 1024 张卡：
- TP = 8（节点内 8 卡张量并行）
- PP = 16（16 个节点流水线并行）
- DP = 8（8 路数据并行）

8 × 16 × 8 = 1024

为什么这么配

  • TP 放节点内：通信密集，需要 NVLink
  • PP 跨节点：通信相对少，可以走网络
  • DP 最外层：扩展性好

选择策略

瓶颈解决方案
单层太大张量并行
模型层太多流水线并行
想提速数据并行
显存不够ZeRO

实际配置示例

7B 模型，8 卡 A100

- 数据并行：8 路
- 或 ZeRO Stage 2/3
- 不需要模型并行

70B 模型，8 卡 A100

- 张量并行：8 路（模型切到 8 张卡）
- 或 ZeRO Stage 3

70B 模型，64 卡（8 节点）

- 张量并行：8（节点内）
- 数据并行：8（跨节点）

175B 模型，1024 卡

- 张量并行：8（节点内）
- 流水线并行：16
- 数据并行：8

小结

三种并行方式：

并行方式切分对象通信开销适用场景
数据并行数据中（AllReduce）模型能放单卡
张量并行层内参数高单层很大，节点内
流水线并行层低层数多，跨节点

组合使用：

  • 小模型：纯数据并行
  • 中等模型：数据并行 + ZeRO
  • 大模型：3D 并行（DP + TP + PP）

关键认知：没有万能的并行方式，要根据模型大小、硬件配置、通信带宽来选择组合。

下一篇讲 NCCL：GPU 之间具体怎么同步数据的。