2. 分布式训练基础：数据并行、模型并行、流水线并行的原理与适用场景

在讨论 GPU 集群上的任务中断之前，必须先建立对分布式训练三种基本并行范式——数据并行（Data Parallelism）、模型并行（Model Parallelism）和流水线并行（Pipeline Parallelism）的清晰认知。因为绝大多数训练中断（如 NCCL 超时、OOM、通信死锁）都与并行策略的选择和实现方式直接相关。理解它们的原理与边界，是后续排查故障的“地图”。

2.1 核心原理：三种并行方式如何切分训练要素

一次训练迭代包含三个核心要素：训练数据（Batch）、模型参数（Weights）和梯度计算（Backward）。三种并行方式本质上是对这三个要素的不同切分策略。

  • 数据并行（DP）：每个 GPU 持有完整的模型副本，但只处理整个 Batch 的一个分片。前向计算完成后，各 GPU 通过 AllReduce 通信操作同步梯度，保证所有副本的权重更新一致。
  • 模型并行（MP）：将单个模型层内的参数矩阵按行或按列切分到多个 GPU 上。每个 GPU 只负责一部分矩阵乘法运算，计算过程中需要频繁的 AllGather 或 ReduceScatter 通信来交换中间结果。
  • 流水线并行（PP）：将模型按层切分为多个 Stage，每个 GPU（或一组 GPU）负责连续的若干层。数据以 micro-batch 的形式依次流过各个 Stage，类似工厂流水线。

下图直观展示了三种并行方式在 GPU 上的资源分配与数据流差异。图中左侧为数据并行（每卡持有完整模型，数据分片），中间为模型并行（单层参数被切分到两卡），右侧为流水线并行（模型按层切分为 Stage 1 和 Stage 2，数据依次流过）。

三种并行策略的 GPU 资源与数据流对比 数据并行 (DP) GPU 0 完整模型副本 数据分片 A GPU 1 完整模型副本 数据分片 B AllReduce 梯度同步 ✓ 每卡持有完整模型 ✓ 数据切分，计算并行 ✗ 通信量与卡数成正比 模型并行 (MP) 输入 X GPU 0 权重 W 上半 GPU 1 权重 W 下半 输出 Y ✓ 解决单卡显存不足 ✓ 每层计算需通信 ✗ 通信频率极高 ✗ 扩展性差（>8卡） 流水线并行 (PP) Stage 1 GPU 0: 层 1-4 micro-batch 1 Stage 2 GPU 1: 层 5-8 等待中 激活 时间 → ✓ 通信量小（仅层间激活） ✓ 适合超大模型 ✗ 存在流水线气泡

2.2 方法流程：从单卡到多卡的实施路径

在实际训练框架（如 PyTorch）中，三种并行方式的实现路径差异显著，理解其底层通信模式是判断故障的关键。

数据并行（DP/DD）是最常用的起点。PyTorch 中 DistributedDataParallel 的实现流程为：

  1. 每个进程加载模型并初始化 torch.distributed 环境。
  2. 使用 DistributedSampler 对数据集进行分片，确保每个 GPU 看到不同的数据。
  3. 前向传播独立计算，得到各自的 loss。
  4. 反向传播时，梯度通过 Ring-AllReduce 算法在所有 GPU 间同步。该算法将梯度切块，在 GPU 环上依次累加和分发，通信量为 2 × (N-1) / N × 模型大小。

模型并行（MP）通常需要手动切分权重矩阵。以 Transformer 的 MLP 层为例，假设输入为 X，权重为 W，输出为 Y = X·W。若将 W 按列切分为 [W1, W2] 并分别放在 GPU0 和 GPU1 上，则前向计算变为：

# GPU 0 持有 W1 (输入维度, 输出维度/2)
Y1 = X @ W1   # 形状: (batch, 输出维度/2)

# GPU 1 持有 W2 (输入维度, 输出维度/2)
Y2 = X @ W2   # 形状: (batch, 输出维度/2)

# 通过 AllGather 操作合并 Y1 和 Y2，得到完整的 Y
Y = AllGather([Y1, Y2], dim=-1)

流水线并行（PP）的实现则更关注调度策略。以经典的 GPipe 为例，它将一个 Batch 切分为 m 个 micro-batch，流程如下：

# 假设 4 个 Stage，m = 4 个 micro-batch
for micro_batch in range(m):
    # Stage 0 处理第 micro_batch 个数据
    output_0 = stage0(micro_batch)
    # 将 output_0 发送给 Stage 1
    send(output_0, to=stage1)

# Stage 1 在收到第一个 micro-batch 后立即开始计算
# 而 Stage 0 继续处理下一个 micro-batch
# 这种调度方式称为 "1F1B" (One Forward One Backward)

2.3 适用场景与选择策略

三种并行方式并非互斥，现代大模型训练通常采用 混合并行（如 Megatron-Turing 的 3D 并行）。下表总结了各自的适用场景与典型配置：

并行方式 核心解决痛点 典型适用场景 通信瓶颈 推荐规模
数据并行 (DP) 训练速度慢（吞吐量不足） 模型可放入单卡显存，但需要更多数据加速收敛 AllReduce 梯度同步，通信量随模型增大而增大 8~256 卡
模型并行 (MP) 单卡显存无法容纳模型参数 单层参数量极大（如 GPT-3 175B 的 attention 层） 每层前向/反向都需 AllGather/ReduceScatter，通信频率极高 2~8 卡（单节点内）
流水线并行 (PP) 模型层数过多导致显存不足 深层 Transformer（>40 层），且单层可放入单卡 仅 Stage 间传输激活值，通信量小 4~64 卡（跨节点）
选择建议：当模型参数量小于单卡显存的 50% 时，优先使用数据并行；当单层参数量超过单卡显存时，必须使用模型并行；当模型层数超过 30 层且单层可放入单卡时，考虑流水线并行。实际训练中，数据并行是基础，模型并行和流水线并行通常叠加在数据并行之上。

2.4 常见误区与故障复盘

在分布式训练任务中断的案例中，以下三个误区最为常见，它们直接导致 NCCL 超时或显存溢出（OOM）：

误区一：盲目增大 Batch Size 导致通信瓶颈。
在数据并行中，将 Batch Size 从 32 增大到 64 时，若 GPU 数量不变，每个 GPU 的本地 Batch 翻倍，但 AllReduce 的通信量不变。然而，若为了“充分利用”而将 GPU 数量从 8 卡扩展到 32 卡，通信量将增加 4 倍。当集群使用千兆以太网而非 InfiniBand 时，网络带宽成为瓶颈，导致 NCCL 超时中断。
误区二：模型并行中忽略通信算子开销。
模型并行需要频繁的 AllGather 操作。例如，在 8 卡上进行 MLP 层切分时，每个 Transformer 层需要 2 次 AllGather 和 2 次 ReduceScatter。若模型有 96 层，则单次迭代需要 384 次集合通信。如果节点间使用 PCIe 而非 NVLink，通信延迟将急剧上升，表现为 GPU 利用率波动和任务卡死。
误区三：流水线并行中 micro-batch 数量设置不当。
流水线并行存在“气泡”（bubble）——即某些 Stage 在等待前序 Stage 数据时处于空闲状态。若 micro-batch 数量 m 远小于 Stage 数量 p，气泡占比会很高。例如，当 p=8、m=4 时，气泡率约为 (p-1)/(m+p-1) ≈ 47%，GPU 利用率极低。更严重的是，若 micro-batch 切分导致单个 micro-batch 的激活值超过显存，会直接触发 OOM 中断。

复盘案例：某团队训练 13B 参数模型，使用 32 卡 A100（每卡 80GB）。最初采用纯数据并行，但模型权重 + 优化器状态 + 梯度需要约 120GB 显存，远超单卡容量，导致 OOM。随后改为“数据并行 + 流水线并行”（每卡放置 4 层，共 8 个 Stage，4 路数据并行）。但在训练 2000 步后频繁出现 NCCL 超时。排查发现：流水线并行中 Stage 间传输激活值使用的是 torch.distributed.send/recv，而数据并行的梯度同步使用 AllReduce。两种通信操作共享同一 NCCL 通信域，当 AllReduce 流量过大时，阻塞了 send/recv 操作，造成死锁。最终通过为 PP 和 DP 创建独立的通信组（new_group）解决了问题。

← 上一章 📖 返回目录 下一章 →