深入理解 PyTorch DDP：为什么更少的通信，不一定意味着更快的训练？

起因
#

上一篇讨论完 Checkpointing 和 Offloading 之后，结尾留了一句：单卡优化的代价已经不可接受时，下一步自然是考虑让多张 GPU 分担训练状态。

于是我开始看 DDP（DistributedDataParallel）。一开始觉得 DDP 的核心思想很直白——每张卡各算各的梯度，然后 All-Reduce 同步一下，用平均梯度更新参数。看起来最朴素的做法就是把所有梯度打包成一个大 tensor 做一次 All-Reduce，简单干净。

但我发现 PyTorch DDP 实际上不是这么做的。它把梯度分成了多个 Bucket，分批 All-Reduce。更让我困惑的是：分批意味着更多次通信，通信次数多了，累计的通信时间反而更长——为什么这样做训练反而更快？

这篇想把这个问题从头想通。

每张卡都有完整的模型
#

先简单回顾一下 Data Parallelism 的基本设定，不展开太多。

假设有 4 张 GPU。DDP 的做法是：每张卡各持有一份完整的模型副本，每个 step 各自拿到不同的 mini-batch 数据，独立跑 forward 和 backward。backward 结束后，4 张卡各自算出了一组梯度。

问题是：如果每张卡用自己的梯度独立更新参数，4 份模型很快就会分道扬镳。所以需要一步同步——All-Reduce，让所有卡拿到梯度的平均值，再用这个一致的平均梯度各自更新参数。因为更新规则相同、梯度相同、参数初始值也相同，更新后的参数自然保持一致。

4 GPU DDP 全景：不同数据、完整模型副本、All-Reduce 同步梯度（教学示意）

有一点容易忽略：DDP 不分片 Model States。 每张卡上都保存着完整的参数、梯度、optimizer states。之前算过，7B 模型在 BF16 mixed precision + AdamW 下的 Model States 是 112 GB。DDP 4 卡训练意味着 4 张卡各持有 112 GB——总共 448 GB Model States 分布在集群中，但单卡并没有因此省下任何显存。DDP 扩展的是训练吞吐量（每 step 处理更多数据），不是单卡显存。

这一章不展开更多。后面的重点是：这个 All-Reduce 应该怎么做，才能让训练跑得快。

一个参数一个参数地同步，代价有多大
#

最直接的做法：backward 算完一个参数的梯度，就立刻对这个参数做一次 All-Reduce。

听起来很及时，但有个问题。

每次发起一次集合通信，不管传输的数据量多小，都有一个固定开销——建立通信、同步各 Rank 的进度、启动数据传输。用一个简化的通信成本模型来表示：

$$ T_{\text{comm}} \approx \alpha + \frac{S}{B} $$

其中 $\alpha$ 是每次通信的固定开销（latency），$S$ 是传输的数据量，$B$ 是有效带宽。

需要说明：这是极度简化的模型。真实的 Collective Communication 成本与算法（ring、tree）、网络拓扑、GPU 数量都有关，不是一个简单的线性公式能概括的。但 $\alpha$ 这个固定开销的存在是真实的，用它来理解 Bucketing 的动机已经够了。

假设一个模型有 200 个参数 tensor。如果逐个 All-Reduce，就要发起 200 次通信，每次都付一遍 $\alpha$。而大部分参数 tensor 很小——一个 bias 可能只有几 KB，花在固定开销上的时间远超实际数据传输时间。

这跟写磁盘的道理类似。如果要写 1000 个小文件，每个 4 KB，即使总数据量只有 4 MB，大量时间花在了文件系统的元数据操作上。把它们合并成一个 4 MB 的顺序写，速度可能快一个数量级。网络批量发送也是同理——攒一批数据一起发，摊薄每次通信的固定开销。

DDP 里的 Gradient Bucketing 就是这个思路：不逐个参数 All-Reduce，而是把多个参数的梯度拼接到同一个 Bucket（连续内存缓冲区），Bucket 填满后统一做一次 All-Reduce。

逐参数 All-Reduce vs Gradient Bucketing（教学示意）

但 DDP 的 Bucketing 和磁盘 buffer 有一个关键区别。磁盘 buffer 随时可以写，数据到了就填进去。而 DDP 的梯度不是同时产生的——backward 是按计算图逆序执行的，不同参数的梯度在不同时刻变为 ready。这个时序差异，决定了 Bucket 不是"越大越好"。

梯度不是同时 ready 的
#

用一个四层 MLP 来想这个问题。Forward 的计算顺序是：

$$ L_1 \rightarrow L_2 \rightarrow L_3 \rightarrow L_4 \rightarrow \text{loss} $$

Backward 反过来，沿计算图逆序传播梯度：

$$ \text{loss} \rightarrow g_4 \rightarrow g_3 \rightarrow g_2 \rightarrow g_1 $$

也就是说，$g_4$（第四层的梯度）最先计算完成，$g_1$（第一层的梯度）最后完成。

四层 MLP 的 forward 与 backward 梯度产生顺序（教学示意）

现在把梯度分配到 Bucket 里。

假设 Bucket 越大，能装的参数越多。极端情况：只用一个超大 Bucket 装下所有参数的梯度。这样通信次数降到了 1 次，$\alpha$ 只付一遍。但问题是：这个 Bucket 必须等到 $g_1$（最后一个梯度）ready 之后才能发起 All-Reduce——也就是说，通信只能在整个 backward 结束之后才开始。

反过来，如果把 $g_4$ 单独放一个小 Bucket，那 $g_4$ 一算完，这个 Bucket 就可以立刻开始 All-Reduce——此时 backward 还在继续计算 $g_3$、$g_2$、$g_1$。

这就引出了一个关键的可能性：All-Reduce 能不能和 backward 计算同时进行？

答案是可以的。Bucket 里的梯度一旦同步完成，它们的用途是给 optimizer 做参数更新。而 backward 继续计算后续层的梯度时，并不需要等前面层的 All-Reduce 结果。每个参数的梯度同步服务的是后续的 optimizer step，不是当前的 backward 计算。

通信计算重叠
#

这是全文最核心的部分。

先用一个具体的数字例子。

没有 overlap 的情况
#

假设：

Backward 总耗时 80 ms
模型的梯度被分成 4 个 Bucket，每 20 ms 有一个 Bucket 变为 ready
每个 Bucket 的 All-Reduce 耗时 15 ms
通信通道同时只能跑一路 All-Reduce
计算和通信使用不同的硬件资源，可以并发

如果不做任何 overlap——等 backward 全部完成后再依次 All-Reduce：

$$ T = 80 + 4 \times 15 = 140\ \text{ms} $$

有 overlap 的情况
#

同样的 4 个 Bucket，但一旦某个 Bucket ready，就立刻发起 All-Reduce（如果通信通道空闲的话）：

Bucket 1：20 ms ready，All-Reduce 20–35 ms
Bucket 2：40 ms ready，All-Reduce 40–55 ms
Bucket 3：60 ms ready，All-Reduce 60–75 ms
Bucket 4：80 ms ready，All-Reduce 80–95 ms

Backward 在 80 ms 结束，最后一个 All-Reduce 在 95 ms 结束。

$$ T = 95\ \text{ms} $$

注意：4 次 All-Reduce 的累计通信时间仍然是 $4 \times 15 = 60$ ms，一点没少。但其中 45 ms 发生在 backward 还在计算的时候——被计算"隐藏"了。只有最后一个 Bucket 的 15 ms 通信暴露在 backward 结束之后。

无 overlap 与有 overlap 的时间线对比（教学示意，数字为假设值）

我觉得这里有两个概念值得记住：

Hidden Communication：被计算时间完全覆盖的通信。在上面的例子里，Bucket 1、2、3 的 All-Reduce 都发生在 backward 还在进行的时间段内，属于 hidden communication。

Exposed Communication：超出计算时间、暴露在关键路径上的通信。Bucket 4 的 15 ms 就是 exposed communication——backward 已经结束了，只能干等通信完成。

通信更多，训练却更快
#

现在对比两个方案。

方案 A：所有梯度合并成一个超大 Bucket，只做一次 All-Reduce。

Backward = 80 ms
All-Reduce = 50 ms（数据量大，但只付一次 $\alpha$）
没有 overlap 的空间——通信只能在 backward 结束后开始

$$ T_A = 80 + 50 = 130\ \text{ms} $$

方案 B：分成 4 个 Bucket，累计 All-Reduce = 60 ms（数据量一样大，但多付了 3 次 $\alpha$）。

利用 overlap，总耗时 95 ms（如上面算的）

$$ T_B = 95\ \text{ms} $$

通信总耗时：$T_{\text{comm},B} = 60\ \text{ms} \gt T_{\text{comm},A} = 50\ \text{ms}$

训练 step 耗时：$T_{\text{step},B} = 95\ \text{ms} \lt T_{\text{step},A} = 130\ \text{ms}$

方案 B 的通信量更大、通信次数更多、累计通信时间更长，但端到端训练速度更快。

方案 A（大 Bucket）vs 方案 B（小 Bucket）：通信更多但训练更快（教学示意）

这个反直觉的结果背后是一个更一般的概念：关键路径（Critical Path）。

端到端 step time 取决于最长的依赖链——从 step 开始到所有必要工作（包括 backward 和梯度同步）完成的最长路径。优化一段已经被隐藏的通信（比如方案 B 里 Bucket 1 的 All-Reduce 从 15 ms 压缩到 10 ms），不会改善 step time，因为那段通信不在关键路径上。

这个道理在之前讨论 Offloading 的时候也出现过——CPU-GPU 数据搬运能否被计算隐藏，取决于搬运是否落在关键路径上。

需要补充一个重要的限定。上面的例子假设计算和通信可以理想地并发，互不干扰。真实系统中，GPU 的计算单元和网络 DMA 引擎确实可以并行工作，但两者共享 GPU 内存带宽、PCIe/NVLink 带宽等资源。在负载很高的情况下，通信和计算可能互相拖慢。所以 overlap 不等于通信完全免费——它减少了暴露在关键路径上的通信时间，但并非零成本。

PyTorch DDP 内部怎么做到的
#

上一节讲了 Bucketing 和 overlap 的原理。这一节看 PyTorch DDP 怎么把这套机制实现出来。

以一次 loss.backward() 调用为线索，把关键步骤串起来。

初始化阶段（DistributedDataParallel.__init__）：

DDP 构造时，在本地创建一个 Reducer 对象
根据模型参数和 bucket_cap_mb 配置，把参数分配到若干个 Bucket
为每个参数的梯度累积节点（AccumulateGrad）注册一个 Autograd Hook

Backward 阶段（loss.backward() 触发）：

PyTorch 的 Autograd Engine 开始沿计算图反向传播，逐个计算参数的梯度
每当一个参数的梯度计算完成，对应的 Hook 被触发，通知 Reducer：“这个参数的梯度 ready 了”
Reducer 检查这个参数所在的 Bucket——如果 Bucket 内所有参数的梯度都已 ready，就把这个 Bucket 标记为可发送
按照预定的 Bucket 顺序，发起异步 All-Reduce

Backward 结束后：

loss.backward() 返回前，DDP 会等待所有 Bucket 的 All-Reduce 完成
之后 optimizer 可以直接使用同步后的梯度更新参数

Autograd → Hook → Reducer → All-Reduce 执行流程（教学示意）

核心分工是：Autograd Engine 负责计算梯度，Hook 负责通知，Reducer 负责协调通信。

关于 Autograd Engine、backward 和梯度的关系，之前的文章已经详细讨论过，这里不重复。

Bucket 的通信顺序
#

这里有一个容易混淆的地方。需要区分三个不同的"顺序"：

Bucket Ready 顺序：某个 Bucket 内所有梯度 ready 的时间顺序。对于简单的顺序模型，最后一层参数所在的 Bucket 通常最先 ready
All-Reduce 发起顺序：Reducer 实际启动 All-Reduce 的顺序。这个顺序是确定的、所有 Rank 必须一致的
All-Reduce 完成顺序：各 Bucket 的 All-Reduce 实际完成的先后。因为不同 Bucket 的数据量可能不同，完成顺序不一定和发起顺序一致

这三个顺序不是同一回事。

为什么 All-Reduce 的发起顺序必须所有 Rank 一致？因为 All-Reduce 是集合通信——所有参与者必须就"现在同步的是哪个 Bucket"达成一致。如果 Rank 0 先发起 Bucket 3 的 All-Reduce，Rank 1 先发起 Bucket 1 的 All-Reduce，两边对不上，就会死锁或产生错误结果。

PyTorch DDP 的做法是：Bucket 的编号在初始化时按参数注册的逆序确定（大致对应 backward 的梯度产生顺序），All-Reduce 严格按 Bucket 编号从小到大（即 backward 中最先 ready 的 Bucket 优先发起）。所有 Rank 使用相同的模型结构和参数顺序，因此 Bucket 划分和发起顺序天然一致。

双 Rank 的 Bucket Ready 与 All-Reduce 顺序约束（教学示意）

那 Bucket Ready 的顺序呢？对于简单的顺序模型（如 MLP、标准 Transformer），backward 的执行顺序几乎是确定的——即使不同 Rank 的执行速度有快有慢，梯度产生的相对顺序不会变。但如果计算图存在独立的并行分支（例如两个不共享中间结果的子网络），不同 Rank 上的分支执行先后可能不同，导致 Bucket Ready 的顺序不一致。这没关系——Reducer 不要求 Bucket Ready 的顺序一致，它只要求 All-Reduce 的发起顺序一致。即使 Bucket 3 比 Bucket 1 更早 ready，Reducer 也会等到 Bucket 1 发起并完成后再轮到 Bucket 3。

Bucket 大小怎么选
#

PyTorch DDP 通过 bucket_cap_mb 参数控制每个 Bucket 的容量上限：

1ddp_model = DistributedDataParallel(
2    model,
3    device_ids=[local_rank],
4    bucket_cap_mb=25,  # 默认 25 MiB
5)

注意几点：bucket_cap_mb 的单位是 MiB，它是每个 Bucket 的容量上限，不代表每个 Bucket 都恰好一样大。参数按模型注册顺序的逆序（大致对应 backward 顺序）依次填入 Bucket，填满一个就开始下一个。最后一个 Bucket 可能远小于上限。

上限值越大，Bucket 越少、每个越大；上限越小，Bucket 越多、每个越小。前面的讨论已经说清楚了两个方向的 trade-off：

Bucket 太小：通信次数多，每次固定开销 $\alpha$ 重复支付，通信可能积压
Bucket 太大：通信启动时间晚，overlap 机会少，通信暴露在关键路径上

用一个简化的性能模型把这个 trade-off 量化一下。

五种配置的对比
#

假设条件（教学用途，不是真实 GPU benchmark）：

梯度总量：80 MiB
Backward 时间：80 ms，匀速产出梯度（每 ms 产出 1 MiB）
每次 All-Reduce 固定开销 $\alpha = 6$ ms
有效传输带宽：5 MiB/ms
通信通道同一时刻只能跑一路 All-Reduce
计算与通信理想并发

每个 Bucket 的 All-Reduce 耗时：

$$ T_{\text{comm}}(b) = 6 + \frac{b}{5}\ \text{ms} $$

其中 $b$ 是 Bucket 大小（MiB）。

Bucket Size	Bucket 数	单次 AR	完成时间
5 MiB	16	7 ms	117 ms
10 MiB	8	8 ms	88 ms
20 MiB	4	10 ms	90 ms
40 MiB	2	14 ms	94 ms
80 MiB	1	22 ms	102 ms

以 Bucket Size = 10 MiB 为例推演一下：8 个 Bucket，每 10 ms 一个 ready。每个 Bucket 的 All-Reduce 耗时 8 ms——通信完成时间（10 + 8 = 18 ms）早于下一个 Bucket 的 ready 时间（20 ms），所以通信不积压，每个 Bucket 都能在下一个到来之前完成。最后一个 Bucket 在 80 ms ready，88 ms 完成。Backward 80 ms 结束，exposed communication 只有 8 ms。

再看 5 MiB 的情况。16 个 Bucket，每 5 ms 一个 ready，但每个 All-Reduce 要 7 ms。Bucket 1 在 5 ms ready，12 ms 完成——但 Bucket 2 在 10 ms 就 ready 了，必须等 Bucket 1 的 All-Reduce 结束才能开始。从 Bucket 2 开始通信就积压了，后面越积越多。最终最后一个 All-Reduce 在 117 ms 才结束。

5 种 Bucket Size 的性能对比（教学模拟，非 GPU benchmark）

性能曲线不是单调的。太小（5 MiB）通信积压严重，太大（80 MiB）没有 overlap 机会。在这组假设下最优点在 10 MiB 附近。

当然，最优 Bucket Size 取决于具体的工作负载。不同模型的梯度分布不一样（有的层参数特别大，有的特别小），backward 的计算时间也不均匀，实际通信带宽和固定开销与硬件环境相关。没有一个普适的最优值——这就是为什么需要 Profiler。

怎么判断 Bucket Size 是否值得调
#

假设一个场景：DDP 训练跑起来了，但感觉效率不太对。应该怎么入手？

核心方法是用 PyTorch Profiler 看实际的时间线。

 1from torch.profiler import profile, schedule, tensorboard_trace_handler
 2
 3with profile(
  activities=[
      torch.profiler.ProfilerActivity.CPU,
      torch.profiler.ProfilerActivity.CUDA,
  ],
  schedule=schedule(wait=1, warmup=1, active=3),
  on_trace_ready=tensorboard_trace_handler("./log"),
  record_shapes=True,
  with_stack=True,
12) as prof:
  for step, batch in enumerate(dataloader):
      loss = model(batch)
      loss.backward()
      optimizer.step()
      optimizer.zero_grad()
      prof.step()

在 TensorBoard 或 Chrome Trace Viewer 里打开 trace，重点看几个东西：

Backward Kernel 的时间分布。计算是均匀分布还是集中在某几层？如果某一层特别慢（比如一个很大的矩阵乘），它之后的梯度会集中在短时间内 ready，可能导致通信积压。

NCCL All-Reduce 的启动时间。第一个 All-Reduce 是什么时候开始的？如果 backward 已经过了很久才启动第一次通信，说明 Bucket 可能太大。

通信是否积压。如果连续多个 All-Reduce 之间没有间隔，上一个刚完成下一个立刻开始，说明通信速度跟不上梯度产出速度——可能是 Bucket 太小（固定开销太高），也可能是带宽不够。

Backward 结束后是否有明显的通信尾部。如果 backward 结束后还有大段 All-Reduce 在跑，说明 overlap 不充分——要么是 Bucket 太大导致通信启动晚，要么是通信本身太慢。

不同 Rank 是否存在明显等待。如果某个 Rank 在 All-Reduce 上花的时间远多于其他 Rank，可能存在负载不均衡（不同 Rank 的 backward 速度差异大）。

两种 Profiler 时间线示意：良好 overlap vs 明显通信尾部（教学示意）

把典型的判断整理一下：

通信次数多、每次很短：尝试增大 bucket_cap_mb，减少固定开销的累积
第一段通信启动太晚：尝试减小 bucket_cap_mb，让较小的 Bucket 提前发起
通信积压、连续排队：判断是固定开销太高（减少通信次数）还是带宽不足（不是 Bucket Size 能解决的问题）
通信几乎完全被计算隐藏：Bucket Size 可能已经不是主要瓶颈了，去找其他优化点

还有两个值得知道的 DDP 配置。gradient_as_bucket_view=True 让参数的 .grad 直接指向 Bucket 的内存，避免了梯度从参数 tensor 拷贝到 Bucket 缓冲区的开销。默认 False，但实际使用中通常建议打开。

torch.compile 与 DDP 配合时，编译器可能会改变计算图的结构，影响梯度 ready 的时序。如果发现 overlap 行为在 compile 前后差异很大，值得用 Profiler 重新检查。

这两个点各自展开都可以写很多，这里只提一下，不深入。

写在最后
#

回过头看这篇的主线。

一开始的直觉是"减少通信次数 → 训练更快"。从固定开销的角度看，这个方向没错——通信次数少确实减少了 $\alpha$ 的累积。但通信次数少意味着每个 Bucket 更大，Bucket 更大意味着通信启动更晚，overlap 的空间更小。

最终决定 step time 的不是通信的总耗时，而是暴露在关键路径上的通信时间。

通信次数少，不等于性能好。 100 次完全被隐藏的通信，对 step time 的影响可能是零。

通信时间更长，不等于训练更慢。 方案 B 的累计通信比方案 A 多了 10 ms，但 step time 反而快了 35 ms——因为多出的通信全部被计算隐藏了。

分布式训练的性能优化目标是缩短端到端 Critical Path。 不只是减少通信量。通信粒度、启动时机、和计算的 overlap 程度，都直接影响最终的训练速度。

这些原则其实不只适用于 DDP。之前写过的几篇文章反复出现类似的结构：

Offloading 里，CPU-GPU 数据搬运能否被计算隐藏，取决于传输是否在关键路径上
Arithmetic Intensity 里，计算快不代表性能好——如果瓶颈在数据搬运，再多的算力也发挥不出来
GPU 存储层级里，FlashAttention 的核心也是减少关键路径上的数据搬运

抽象出来就是几条：减少总工作量，摊薄固定开销，通过异步和流水化隐藏等待，然后优化剩下的关键路径。对于 DDP 的 Gradient Bucketing 来说，Bucket Size 的选择本质上就是在"摊薄固定开销"和"最大化 overlap"之间找平衡。

参考资料
#

GPU 显存不够怎么办？Checkpointing 与 Offloading 的工程权衡（本文前篇）
7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？
从 Forward 到 Optimizer：再看一遍 Training Loop
从 CS336 到 ML Infra：理解 Arithmetic Intensity
GPU 的数据存在哪？——从 HBM 到片上存储的存储层级
PyTorch DistributedDataParallel API 文档
PyTorch DDP 设计文档
MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/

起因#

每张卡都有完整的模型#

一个参数一个参数地同步，代价有多大#

梯度不是同时 ready 的#

通信计算重叠#

没有 overlap 的情况#

有 overlap 的情况#

通信更多，训练却更快#

PyTorch DDP 内部怎么做到的#

Bucket 的通信顺序#

Bucket 大小怎么选#

五种配置的对比#

怎么判断 Bucket Size 是否值得调#

写在最后#

参考资料#

相关文章