📚 课程目录第2章分布式训练基础

2. 分布式训练基础

在进入 DeepSpeed 的模型状态分片与内存复用之前，必须先建立一张清晰的“并行地图”。大模型训练之所以需要分布式，是因为单个 GPU 的显存无法同时容纳模型参数、梯度、优化器状态以及激活值。不同的并行策略，本质上是在“切分什么”和“复制什么”之间做权衡。理解数据并行、模型并行、流水线并行和混合并行的差异，是后续理解 ZeRO 分片逻辑的前提。

一、核心原理：四种并行策略的分工

数据并行（Data Parallelism, DP）是最常见的策略。每张 GPU 持有一份完整的模型副本，但只处理不同的数据批次。前向和反向计算在本地完成，随后通过 AllReduce 同步梯度，使所有副本保持一致。它的优点是实现简单、扩展性好；缺点是每张卡都要存完整模型状态，显存开销随模型增大而线性增长。

模型并行（Model Parallelism, MP）又称张量并行，把单个层内的矩阵运算切分到多张 GPU 上。例如把权重矩阵按列或按行拆分，每张卡只计算一部分，再通过 AllGather 或 ReduceScatter 拼接结果。它解决了单层放不下的问题，但通信频繁，对互联带宽要求高。

流水线并行（Pipeline Parallelism, PP）按层切分模型，不同 GPU 负责不同层段。数据被切成多个 micro-batch，在流水线中依次流动，从而让各段 GPU 交替工作。它降低了单卡显存压力，但会引入流水线气泡，需要合理设置 micro-batch 数量来掩盖空闲。

混合并行（Hybrid Parallelism）则是把上述策略组合使用，例如“数据并行 + 张量并行 + 流水线并行”，典型如 Megatron-LM 的 3D 并行。DeepSpeed 常与这些策略配合，在数据并行维度上进一步做 ZeRO 分片，从而在超大规模模型上取得显存与吞吐的平衡。

二、方法流程：从单卡到多卡的决策路径

选择并行策略时，可以按以下顺序判断：

模型能否放入单卡？能，则优先数据并行，简单高效。
单层是否放不下？若单层参数过大，引入张量并行，把矩阵切到多卡。
层数很多但单层可放下？使用流水线并行，按层切分，降低单卡驻留。
规模继续增大？组合 DP、TP、PP 形成混合并行，并在 DP 维度上叠加 ZeRO 分片。

在 DeepSpeed 中，数据并行是默认底座，ZeRO 则是在 DP 组内对优化器状态、梯度和参数进行分片。因此，理解 DP 的“复制”语义，才能理解 ZeRO 的“分片”语义。

三、案例场景：8 卡训练一个 13B 模型

假设有 8 张 80GB GPU，训练 13B 参数模型。若采用纯数据并行，每张卡需保存完整参数、梯度和优化器状态，显存很快耗尽。此时可配置为：TP=2、PP=2、DP=2，即 2×2×2=8 卡。每张卡只持有部分层和部分矩阵，显存压力显著下降。再在 DP=2 的维度上启用 ZeRO-2，进一步分片优化器状态和梯度。

{
  "train_batch_size": 64,
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" }
  },
  "tensor_parallel": { "tp_size": 2 },
  "pipeline_parallel": { "pp_size": 2 }
}

该配置下，通信组被划分为 TP 组、PP 组和 DP 组，各自使用不同的集合通信原语。若 TP 组跨节点，带宽不足会成为瓶颈，因此通常把 TP 放在节点内，PP 和 DP 放在节点间。

四、误区与复盘

常见误区一：并行度越高越好。实际上，TP 过大导致通信量激增，PP 过大导致气泡占比升高，DP 过大则 AllReduce 开销增加。并行度需要与模型结构、网络拓扑匹配。

复盘要点：先判断瓶颈在显存还是吞吐。显存不足优先 TP/PP/ZeRO；吞吐不足则检查通信与气泡。DeepSpeed 的 ZeRO 主要作用于 DP 维度，不能替代 TP 和 PP 对单层和层间的切分。

策略	切分对象	复制对象	主要通信	典型适用
数据并行 DP	数据批次	完整模型	AllReduce	模型可放入单卡
模型并行 MP	层内矩阵	无	AllGather / ReduceScatter	单层过大
流水线并行 PP	模型层段	无	点对点	层数多、单层可放下
混合并行 3D	数据 + 层内 + 层间	按维度组合	组合通信	超大规模模型

掌握这四种并行策略的切分与复制关系后，就能理解 DeepSpeed 为什么选择在数据并行维度上做 ZeRO 分片：它不改变模型结构，却能把冗余的模型状态从每张卡上“拆走”，从而在相同硬件上训练更大的模型。