64.5 分布式训练
一句话
分布式训练不是「多卡一定更快」，而是把计算、显存、通信三者重新分配：数据并行省时间但复制状态，模型并行省显存但增加通信，流水并行省显存但产生气泡——每一种并行都在用某一种资源换另一种资源。
四种并行
| 并行 | 切什么 | 解决什么 | 新增成本 | 典型 |
|---|---|---|---|---|
| 数据并行（DP） | 切 batch | 吞吐、训练时间 | 每卡一份模型 + 梯度 AllReduce | DDP、FSDP（分片） |
| 模型并行（MP） | 切层内矩阵 | 单卡放不下模型 | 层内通信（AllReduce/AllGather） | Megatron-LM 张量并行 |
| 流水并行（PP） | 切层 | 单卡放不下深度 | pipeline bubble、微批调度 | GPipe、Megatron PP |
| ZeRO / 分片 | 切优化器状态/梯度/参数 | 显存 | 参数 AllGather、通信量增加 | DeepSpeed ZeRO-1/2/3、FSDP |
🔑 选并行先看瓶颈：单卡放得下模型 → 数据并行；放不下 → 模型/流水/ZeRO；通信带宽不够 → 减少并行度、梯度累积、通信重叠。不要一上来就上 3D 并行。
数据并行的通信账
每个 step：
- 各卡用自己的 micro-batch 前向/反向，得到梯度；
- 对所有卡的梯度做 AllReduce（求平均/求和）；
- 各卡用相同梯度更新，保持参数一致。
Ring AllReduce 的通信量：
每个 step 搬运字节数 ≈ 2 × (N-1)/N × P
N 很大时趋近 2P
P 是梯度总量。也就是说，通信量与卡数无关，存在下界 2P；卡越多，计算被切得越小，通信占比越高。
代码：1GB 梯度、10GB/s，加卡能快多少
package main
import "fmt"
const (
paramsGB = 1.0 // 梯度/参数总量（GB）
bandwidth = 10.0 // 单卡有效带宽（GB/s）
latencyMS = 0.01 // 单跳延迟（ms）
computeMS = 100.0 // 单卡计算时间（ms）
)
func ringMS(n int) float64 {
if n <= 1 {
return 0
}
// ring all-reduce：2(N-1) 步，每步搬 P/N
bytesPerStep := paramsGB / float64(n)
return 2*float64(n-1)*(bytesPerStep/bandwidth)*1000 + 2*float64(n-1)*latencyMS
}
func main() {
base := computeMS
fmt.Println("=== 数据并行 ring all-reduce：1GB 梯度、10GB/s、单卡计算 100ms ===")
fmt.Printf("%-6s %-14s %-12s %-12s %-10s %s\n",
"GPUs", "ring(ms)", "step(ms)", "speedup", "eff(%)", "bottleneck")
for _, n := range []int{1, 8, 64, 512} {
ring := ringMS(n)
step := base + ring
speedup := float64(n) * base / step
eff := speedup / float64(n) * 100
bottleneck := "compute"
if ring > base {
bottleneck = "comm"
}
fmt.Printf("%-6d %-14.1f %-12.1f %-12.2f %-10.1f %s\n",
n, ring, step, speedup, eff, bottleneck)
}
fmt.Println("\n=== 如果带宽只有 2GB/s ===")
for _, n := range []int{8, 64, 512} {
ring := 2*float64(n-1)*(paramsGB/float64(n)/2.0)*1000 + 2*float64(n-1)*latencyMS
step := base + ring
fmt.Printf(" N=%-4d ring=%.1fms step=%.1fms speedup=%.2fx efficiency=%.1f%%\n",
n, ring, step, float64(n)*base/step, (float64(n)*base/step)/float64(n)*100)
}
fmt.Println("\n结论：ring 的通信量有下界 2P/B（与 N 无关），本例约 200ms；")
fmt.Println(" compute 只有 100ms 时，加卡也突破不了约 3x 的加速上限。")
}
本机实测（macOS，Apple M1，Go 1.27.1）：
=== 数据并行 ring all-reduce：1GB 梯度、10GB/s、单卡计算 100ms ===
GPUs ring(ms) step(ms) speedup eff(%) bottleneck
1 0.0 100.0 1.00 100.0 compute
8 175.1 275.1 2.91 36.3 comm
64 198.1 298.1 21.47 33.5 comm
512 209.8 309.8 165.25 32.3 comm
=== 如果带宽只有 2GB/s ===
N=8 ring=875.1ms step=975.1ms speedup=0.82x efficiency=10.3%
N=64 ring=985.6ms step=1085.6ms speedup=5.90x efficiency=9.2%
N=512 ring=1008.3ms step=1108.3ms speedup=46.20x efficiency=9.0%
结论：ring 的通信量有下界 2P/B（与 N 无关），本例约 200ms；
compute 只有 100ms 时，加卡也突破不了约 3x 的加速上限。
读数：
| 卡数 | ring | step | 加速比 | 效率 | 说明 |
|---|---|---|---|---|---|
| 1 | 0 | 100ms | 1.00x | 100% | 基线 |
| 8 | 175ms | 275ms | 2.91x | 36% | 通信已超过计算 |
| 64 | 198ms | 298ms | 21.5x | 34% | 加卡几乎全被通信吃掉 |
| 512 | 210ms | 310ms | 165x | 32% | ring 趋近 2P/B 下界 |
🔑 这组数字解释了一个常见困惑：为什么从 8 卡加到 64 卡，训练时间没怎么降？ 因为单卡计算只有 100ms，而 AllReduce 下界约 200ms；加卡把计算切小，但通信不降，效率卡在 30% 左右。要提速，先提带宽/减少通信量（梯度压缩、ZeRO 分片、通信重叠），而不是继续加卡。
带宽更差时，加卡可能变慢
2GB/s 时，8 卡 ring 875ms，step 975ms，加速比 0.82x——比单卡还慢。原因：计算 100ms，通信 875ms，加卡带来的计算收益远小于通信成本。
⚠️ 分布式训练的第一个反直觉结论：在网络差的环境里，多卡可能比单卡慢。上多卡前先测「单卡 step 时间 + AllReduce 时间」，算清楚 scaling efficiency，再决定并行度。
显存不够时：ZeRO / 模型并行 / 流水并行
| 方案 | 切什么 | 显存节省 | 通信 | 适用 |
|---|---|---|---|---|
| ZeRO-1 | 优化器状态 | 约 4x | 低 | 大多数数据并行 |
| ZeRO-2 | + 梯度 | 约 8x | 中 | 大模型、单卡放不下优化器 |
| ZeRO-3 / FSDP | + 参数 | 约 Nx | 高（AllGather） | 单卡放不下参数 |
| 张量并行 | 层内矩阵 | 单层切分 | 高（每层 AllReduce） | 单机多卡、NVLink |
| 流水并行 | 层 | 按 stage 切 | 中（激活传递） | 多机、深模型 |
🔑 ZeRO 的本质是「用通信换显存」：把优化器状态/梯度/参数分片到各卡，需要时再 AllGather。分片越多，显存越省，通信越贵；ZeRO-3 不是默认最优，要按模型大小与带宽选。
流水并行的气泡
流水并行把模型按层切成 stage，micro-batch 依次流过；但首尾 stage 会空闲，形成 pipeline bubble。气泡比例近似：
bubble ≈ (stages - 1) / (micro_batches + stages - 1)
例如 8 个 stage、8 个 micro-batch，气泡约 7/15 ≈ 47%；要压到 10% 以下，micro-batch 数要远大于 stage 数。这就是为什么流水并行常配「1F1B / interleaved schedule」和大量 micro-batch。
工程细节：梯度累积、通信重叠、straggler
| 问题 | 做法 | 代价 |
|---|---|---|
| 小 batch 通信太频繁 | 梯度累积 N 步再 AllReduce | 有效 batch 变大，收敛动态变 |
| 通信等待 | 计算/通信重叠（bucket、异步） | 实现复杂，显存增加 |
| 慢卡拖累 | straggler 检测、弹性训练、重调度 | 容错复杂 |
| 断点恢复 | checkpoint（参数+优化器+调度器+RNG） | 存储/恢复时间 |
| 数据不均 | 按 token 数而非样本数分片 | 需要长度统计 |
⚠️ checkpoint 只存模型参数不够：优化器状态（Adam 的 m/v）、学习率调度器、梯度累积步数、RNG 状态都要存，否则恢复后训练曲线会跳。大模型 checkpoint 可能几百 GB，恢复时间要算进 RTO。
生产落地：指名 → 可抄 → 可验 → 可退
指名：PyTorch DDP（数据并行基线）、FSDP（参数分片）、DeepSpeed ZeRO（显存优化）、Megatron-LM（张量/流水并行）；调度用 Slurm/K8s + Ray。
可抄（PyTorch DDP，未验证：本机未安装 GPU/PyTorch）：
model = DistributedDataParallel(model, device_ids=[local_rank])
for step, batch in enumerate(loader):
loss = model(batch)
loss.backward() # DDP 在 backward 中 AllReduce 梯度
optimizer.step()
optimizer.zero_grad(set_to_none=True)
# 8 卡单机
torchrun --nproc_per_node=8 train.py
# 多机：--nnodes=2 --node_rank=0 --master_addr=... --master_port=...
可验：先跑 100 step，记录单卡 step 时间、AllReduce 时间、scaling efficiency；效率 < 50% 先查网络/带宽/梯度大小。可退：并行度与 ZeRO stage 作为配置，可回退到单卡/小并行；checkpoint 保留最近 N 个，恢复时校验参数/优化器/RNG。
反模式
| 反模式 | 后果 |
|---|---|
| 不看通信量直接上多机 | 加卡不加速，甚至变慢（本节 0.82x） |
| 小模型上 ZeRO-3 | 通信开销大于显存收益 |
| 流水并行 micro-batch 太少 | 气泡 50%+，GPU 大量空闲 |
| checkpoint 只存参数 | 恢复后优化器状态丢失，训练曲线跳 |
| 不固定数据分片顺序 | 各卡重复/漏样本 |
| 用平均 loss 掩盖慢卡 | straggler 拖垮整体，误判扩容有效 |
| 梯度累积后不调学习率 | 有效 batch 变了，收敛异常 |
| 不监控 GPU 利用率/NCCL 带宽 | 不知道瓶颈在算还是通信 |
自检
Q1 数据并行、模型并行、流水并行各切什么？分别解決什么问题？
▸ 答案：数据并行切 batch，每卡一份模型，解决吞吐/训练时间；模型并行切层内矩阵（张量并行），解决单卡放不下单层；流水并行切层，解决单卡放不下整个模型；ZeRO 切优化器状态/梯度/参数，解决显存。它们都在用通信换计算或显存。
Q2 Ring AllReduce 的通信量为什么与卡数无关？
▸ 答案：ring 算法分 scatter-reduce 与 all-gather 两阶段，每阶段每卡搬运约 (N-1)/N × P，合计约 2P；卡数增加时每步搬的块变小、步数变多，总量趋近 2P。所以加卡不能降低通信总量，只能降低单卡计算量。
Q3 为什么 8 卡只加速 2.91x，而不是 8x？
▸ 答案：单卡计算 100ms，8 卡 ring 通信约 175ms，step = 100 + 175 = 275ms；加速比 = 8×100/275 ≈ 2.91。通信占比 64%，加卡收益被通信吃掉。要提升，先提带宽、减少梯度体积、通信重叠或换并行策略。
Q4 带宽 2GB/s 时，8 卡为什么比单卡还慢？
▸ 答案：ring 通信约 875ms，step 975ms；单卡 step 100ms。8 卡的计算收益远小于通信成本，加速比 0.82x。分布式训练的前提是「计算时间 > 通信时间」，否则加卡是负收益。
Q5 ZeRO-1/2/3 分别切什么？为什么 ZeRO-3 不一定最优？
▸ 答案：ZeRO-1 切优化器状态，ZeRO-2 再切梯度，ZeRO-3 再切参数。切得越多显存越省，但通信越多（参数 AllGather）；ZeRO-3 在带宽不足时会被通信拖慢，小模型上尤其不划算。要按模型大小、带宽、显存瓶颈选 stage。
Q6 流水并行的气泡怎么算？怎么减小？
▸ 答案：气泡比例约 (stages-1)/(micro_batches+stages-1)；stage 越多、micro-batch 越少，气泡越大。减小方法：增加 micro-batch 数、1F1B/interleaved schedule、减少 stage 数、把 stage 切得更均衡。
Q7 大模型 checkpoint 要存什么？只存参数会怎样？
▸ 答案：参数、优化器状态（Adam m/v）、学习率调度器、梯度累积步数、RNG 状态、数据加载器位置。只存参数会导致恢复后优化器动量丢失、学习率跳变、数据重复/漏样本，训练曲线不连续，甚至无法复现。