3、分布式训练基础：数据并行、模型并行、流水线并行、混合并行、通信原语、集合通信
当模型参数量从百万级增长到千亿级，单张 GPU 的显存和算力都无法承载完整训练任务。分布式训练的本质，是把「模型计算」和「训练数据」切分到多张 GPU 上，再通过通信把各卡的结果重新拼合。理解数据并行、模型并行、流水线并行、混合并行，以及它们背后的通信原语与集合通信，是掌握 DeepSpeed 优化技术的前置知识。
一、四种并行范式：切什么，怎么切
分布式训练的切分维度只有两个：数据维度和模型维度。由此衍生出四种基本范式。
- 数据并行（Data Parallel, DP）：每张卡持有完整模型副本，输入不同的数据分片，各自前向、反向计算梯度，再通过 AllReduce 求平均梯度后同步更新。优点是实现简单、扩展性好；缺点是每张卡都要存完整模型、优化器状态和激活值，显存冗余严重。
- 模型并行（Model Parallel, MP）：把模型本身切开，不同层或同一层的不同张量分到不同卡上。张量并行（TP）切分单层内的权重矩阵，流水线并行（PP）按层切分。优点是单卡显存压力小；缺点是层间依赖强，通信频繁，容易出现 GPU 空转。
- 流水线并行（Pipeline Parallel, PP）：把模型按层分成若干 stage，每个 stage 放在不同卡上，微批次（micro-batch）像流水线一样依次流过各 stage。它用「气泡（bubble）」换取显存节省，需要精心设计 micro-batch 数量与调度策略。
- 混合并行（Hybrid Parallel）：实际大模型训练几乎都是 DP + TP + PP 的组合。例如 Megatron-LM 的 3D 并行：节点内做 TP，节点间做 PP，最外层做 DP。DeepSpeed 的 ZeRO 则是在数据并行基础上进一步切分优化器状态、梯度和参数。
二、通信原语与集合通信：分布式训练的「神经系统」
并行策略决定了「怎么切」，通信原语决定了「怎么合」。点对点通信（Send/Recv）用于流水线并行的 stage 间传递激活值；集合通信（Collective Communication）用于多卡协同，常见原语包括：
| 原语 | 语义 | 典型用途 | 通信量 |
|---|---|---|---|
| Broadcast | 一卡数据复制到所有卡 | 参数初始化、超参同步 | O(N) |
| AllReduce | 所有卡数据归约后结果广播给所有卡 | 数据并行梯度平均 | O(2N) |
| AllGather | 收集所有卡的分片并拼成完整张量 | ZeRO 参数收集、TP 输出拼接 | O(N) |
| ReduceScatter | 归约后把结果按卡分散 | ZeRO 梯度分片、TP 反向 | O(N) |
| AllToAll | 每张卡向其他卡发送不同数据 | MoE 专家路由、序列并行 | O(N) |
关键结论：AllReduce 可分解为 ReduceScatter + AllGather，这正是 Ring AllReduce 的实现基础。DeepSpeed 的 ZeRO 正是利用这一分解，把梯度归约和参数收集拆开，从而在通信量不变的前提下大幅降低单卡显存。
三、方法流程：一次混合并行训练的前向与反向
以 3D 并行（TP=2，PP=4，DP=8）为例，一次迭代的典型流程如下：
- 数据加载：DP 组内每张卡读取不同 mini-batch，PP 组内按 micro-batch 切分。
- 前向传播：micro-batch 依次流过 PP 的 4 个 stage；每个 stage 内部，TP 组通过 AllReduce 或 AllGather 合并张量并行结果。
- 反向传播：梯度沿 PP 反向回传，TP 组内做 ReduceScatter 或 AllReduce 归约梯度。
- 梯度同步：DP 组内对完整梯度执行 AllReduce，得到全局平均梯度。
- 参数更新：优化器更新参数，进入下一轮迭代。
# 以 DeepSpeed 配置为例，描述 3D 并行关键参数
{
"train_batch_size": 1024,
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8,
"zero_optimization": {
"stage": 2,
"allgather_partitions": true,
"reduce_scatter": true
},
"tensor_parallel": { "tp_size": 2 },
"pipeline_parallel": { "pp_size": 4 }
}
四、案例与场景：不同规模模型的选择
- 小模型（<1B）：单卡放得下，直接用数据并行 + DDP，AllReduce 梯度即可。
- 中等模型（1B～10B）：单卡放不下但层数不多，优先 ZeRO-2/3 数据并行，避免 TP 的频繁通信。
- 大模型（10B～100B）：节点内 TP=8，节点间 PP=4～16，最外层 DP + ZeRO-1，兼顾显存与通信效率。
- 超大模型（>100B）：TP + PP + DP + ZeRO-3 + MoE，配合 AllToAll 做专家路由，通信拓扑与 NVLink/IB 带宽成为瓶颈。
五、误区与复盘
误区一：并行度越高越好。TP 超过单节点 GPU 数会跨机通信，带宽骤降；PP 过大则气泡占比升高。并行度必须与网络拓扑匹配。
误区二：只关注计算不关注通信。大模型训练中通信可占 30%～50% 时间。AllReduce 的 Ring 与 Tree 算法选择、梯度压缩、通信与计算重叠，都是 DeepSpeed 优化的重点。
误区三：忽略 micro-batch 与全局 batch 的关系。global_batch = micro_batch × grad_accum × DP_size，配置不当会导致梯度噪声大或显存溢出。
复盘时建议按「显存占用 → 通信量 → 计算利用率」三步排查：先用 ZeRO 或 PP 解决显存，再用 TP/DP 调整通信，最后用 micro-batch 和重叠策略提升 GPU 利用率。理解这些基础，才能读懂 DeepSpeed 后续的 ZeRO、Offload、通信压缩等优化手段。