PyTorch 训练框架内核深度解析

第17章 DDP：环形 AllReduce 与梯度桶

作者杨艺韬 · 12,832 字 · 发布于 2026-04-28 · 更新于 2026-09-25

"DistributedDataParallel will bucket parameters into multiple buckets so that gradient reduction of each bucket can potentially overlap with backward computation."

—— bucket_cap_mb 参数的官方文档（torch/nn/parallel/distributed.py:571）

本章要点

DDP 把模型复制到每张卡，每张卡跑独立 forward + backward：反向时通过 AllReduce 同步梯度
关键优化是 communication / computation overlap：用 autograd backward hook 让"梯度算完立刻发 AllReduce"，不等整个 backward 结束
梯度桶 (gradient bucket) 把多个小梯度合并成大消息：减少 NCCL launch 次数，bucket_cap_mb 不传时默认 25 MiB（torch/nn/parallel/distributed.py:833）
环形 AllReduce 是底层算法：N 个 rank 形成环，每步传 1/N 数据，整体 2(N-1)/N × tensor_size 通信量
Reducer C++ 类是核心：管理 buckets、监听 grad 完成、触发 AllReduce、与 autograd Engine 协作
find_unused_parameters=True 用于"某些层在某些 batch 跳过"的场景：DDP 自动检测哪些梯度该等

17.1 数据并行的基本构想

DDP 的工作模式简单：

graph TB
    Data[整个 batch]
    Data --> S1[切片 1] --> R0[rank 0: model copy]
    Data --> S2[切片 2] --> R1[rank 1: model copy]
    Data --> S3[切片 3] --> R2[rank 2: model copy]
    R0 --> G0[grad_0]
    R1 --> G1[grad_1]
    R2 --> G2[grad_2]
    G0 --> AR["AllReduce<br/>每个 rank 拿到 (g0+g1+g2)/3"]
    G1 --> AR
    G2 --> AR
    AR --> Step[每个 rank 用同样平均梯度更新参数]
    style AR fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

每个 rank：本地 forward → 本地 backward → AllReduce 梯度 → step。同步梯度让所有 rank 的参数永远一致，等同于"用 N×batch_size 的大 batch 训练"。

实现关键是第 3 步怎么高效 AllReduce。朴素实现："等 backward 全跑完，再对所有参数一次 AllReduce" —— 慢。DDP 的优化在于让 AllReduce 在 backward 进行中开始，与下一层计算 overlap。

17.2 核心数据结构

DDP 在 Python 端是 torch.nn.parallel.DistributedDataParallel（torch/nn/parallel/distributed.py:330），但真正的工作在 C++ 端的 Reducer（torch/csrc/distributed/c10d/reducer.cpp，2509 行）。Reducer 持有：

bucket 列表：每个 bucket 包含若干 parameter，累计字节数一旦达到上限就封口（所以可能略超），上限是 bucket_cap_mb（默认 25 MiB），重建后的第一个 bucket 另用 1 MiB 的小上限（kDefaultFirstBucketBytes，torch/csrc/distributed/c10d/reducer.hpp:29）
每个 parameter 的 grad_accumulator 钩子：当那个 param 的 grad 算完时触发
每个 bucket 的 future_work：这个 bucket 进行中的 collective（comm hook 返回的 Future）
每个 bucket 的 pending 计数：还差几个 grad 没到，归零时这个 bucket 可以发 AllReduce

graph LR
    subgraph Buckets["梯度桶（按 25MB 分组）"]
        B1[Bucket 1<br/>params 1-15<br/>~25 MB]
        B2[Bucket 2<br/>params 16-30<br/>~25 MB]
        B3[Bucket 3<br/>params 31-45<br/>~25 MB]
    end
    subgraph Hooks["每 param 的 hook"]
        P1[param 1.grad ready]
        P2[param 2.grad ready]
        Pn[param N.grad ready]
    end
    P1 --> B1
    P2 --> B1
    Pn --> B3
    B1 --> R[Reducer]
    R -->|bucket 满了发 AllReduce| NCCL[NCCL all_reduce]
    style B1 fill:#dbeafe
    style B2 fill:#fef3c7
    style B3 fill:#dcfce7

17.3 backward hook：让 AllReduce 在 backward 中触发

Reducer 在 __init__ 时给每个参数挂一个 hook：

// 简化版伪代码
for (index, param : params_) {
    auto grad_accumulator = param.grad_accumulator();   // 第 7 章 §7.6 的 AccumulateGrad
    grad_accumulator->add_post_hook([this, index](){
        this->autograd_hook(index);   // 内部再调 mark_variable_ready
    });
}

AccumulateGrad 是 leaf 张量的反向终点（第 7 章 §7.6）。挂在它上面的 post_hook 在"梯度累积完成"时触发。每次 mark_variable_ready 检查这个 param 所属的 bucket：

bucket 里所有 param 都 ready，且排在它前面的 bucket 都已发出 → 发 AllReduce（bucket 严格按序发，见 §17.12）
否则继续等

17.3.1 Bucket 顺序的智慧

bucket 顺序按 反向计算的顺序 编排（与 forward 顺序相反）。这样最后一层 param 的 grad 先算完 → 第一个 bucket 先发 AllReduce → 与前面层的 backward 计算 overlap：

gantt
    title backward × AllReduce overlap 时间线
    dateFormat X
    axisFormat %s
    section compute stream
    backward layer N      :a1, 0, 2
    backward layer N-1    :a2, 2, 4
    backward layer N-2    :a3, 4, 6
    backward layer N-3    :a4, 6, 8
    section comm stream
    AllReduce bucket1     :crit, b1, 2, 5
    AllReduce bucket2     :crit, b2, 4, 7
    AllReduce bucket3     :crit, b3, 6, 9

通信 stream 与计算 stream 并行执行 —— bucket 1 的 AllReduce 在 backward layer N-1 / N-2 进行时已经在跑，时间被完全 hide 在计算后面。

最理想情况下，所有 AllReduce 都在 backward 进行中完成，到 backward 结束时通信大部分已经 overlap 掉，只需等最后发出的那个 bucket。

17.4 环形 AllReduce 算法

AllReduce 最经典的实现是 Ring AllReduce（NCCL 会按拓扑和消息大小在 Ring、Tree 等算法间挑选）。N 个 rank 形成环，分两阶段：

Reduce-Scatter 阶段（N-1 步）：每步每个 rank 把 1/N 数据传给下一个 rank、累加上一个 rank 传来的部分。N-1 步后每个 rank 拥有 1/N 的最终 reduce 结果。

All-Gather 阶段（N-1 步）：每个 rank 把自己那 1/N 传一圈，最后所有 rank 拥有完整结果。

总通信量：每个 rank 发出 2 × (N-1)/N × tensor_size 字节，与 N 几乎无关（N 大时趋近 2× tensor）。这让 AllReduce 在大集群下仍然高效 —— 1024 卡和 8 卡的单卡通信量几乎一样。

graph LR
    R0[rank 0] -->|1/N 块| R1[rank 1]
    R1 -->|1/N 块| R2[rank 2]
    R2 -->|1/N 块| R3[rank 3]
    R3 -->|1/N 块| R0
    style R0 fill:#dbeafe
    style R1 fill:#fef3c7
    style R2 fill:#dcfce7
    style R3 fill:#fce7f3

NCCL 还有 Tree AllReduce（实现是 double binary tree，小消息、大规模时延迟更低）等算法，默认自动按拓扑和消息大小选择。NCCL_ALGO=ring 强制 Ring，NCCL_ALGO=tree 强制 Tree。

17.5 `gradient_as_bucket_view`：减少一次拷贝

朴素实现里，bucket 是个独立张量、每次梯度算完拷到 bucket 再 AllReduce：

param.grad → copy → bucket → AllReduce

gradient_as_bucket_view=True（v1.7 起）让 param.grad 直接是 bucket 的 view（共享内存），省掉一次 copy：

param.grad (view of bucket) → AllReduce in place

省下的是"每个 bucket 一次全量 copy"的时间与那份额外的显存（docstring 说省下的峰值显存等于全部梯度的大小，且要从第二轮起才看得到）。副作用是 grad 成了 view，不能再对它调 detach_()。生产代码里强烈建议开启。

17.6 `find_unused_parameters`：不规则模型支持

某些模型在不同 batch 走不同分支，导致某些 param 在某些 batch 不参与 forward / backward。DDP 默认认为"所有 param 每 batch 都有 grad"，没 grad 的 param 所在的 bucket 永远凑不齐，这个 bucket（及排在它后面的）的 AllReduce 根本不会发出。各 rank 情况相同时，下一轮 forward 会直接报错（gloo 双进程实跑，节选）：

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`, and by
making sure all `forward` function outputs participate in calculating loss.
...
Parameter indices which did not receive grad for rank 0: 2 3

各 rank 用到的参数不一样时，有的 rank 发了 AllReduce、有的没发，就会互相等到超时。

find_unused_parameters=True 让 DDP 在 forward 后扫描 autograd 图，标记哪些 param 这次不会有 grad，反向里第一个 hook 触发时就把它们直接标成 "ready"。

代价：

每轮 forward 后做一次 graph 遍历（参数确实全用到时，reducer 会打一条 find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters 的警告提醒你关掉）
不做 bucket rebuild（§17.13），bucket 顺序一直是按定义顺序反过来的近似

默认 False。只有 transformer-with-MoE / 多 task 共享 backbone 等场景才需要。

17.7 静态图模式：`static_graph=True`

static_graph=True（v1.11 起成为构造参数，之前是 v1.9 引入的 _set_static_graph()）让 DDP 假设"用到哪些参数、反向怎么走"在整个训练里不变。源码里的实际行为：

第 1 个 iter：反向中不发任何 AllReduce，只数每个参数的 hook 触发了几次；等反向结束，由 _DDPSink 排进 Engine 回调队列的 _delay_all_reduce（Reducer::delay_all_reduce）一次性发出全部 bucket，同时记下哪些参数没用到
第 2 个 iter 起：按第一轮的计数判断 ready，不再遍历 autograd 图找 unused 参数

按 docstring，它的价值主要是支持老模式不支持的场景：reentrant backward、同一段被多次 activation checkpoint、带 unused 参数的 checkpoint、在 forward 之外用到参数；以及在有 unused 参数时省掉每轮的图遍历。开了它 find_unused_parameters 设不设都无所谓。它不省 bucket rebuild：gloo 实跑，static_graph 下 bucket 在第 2 轮反向记下顺序、第 3 轮 forward 开头重建，比默认模式还晚一轮（§17.13）。代价是图结构一变就报错（§17.17）。

17.8 与 NCCL Async Error Handling 的协作

DDP + NCCL 的死锁场景：某个 rank 的 backward 在某层报错，它不发 AllReduce，其他 rank 在那层 AllReduce 上等到 timeout。

TORCH_NCCL_ASYNC_ERROR_HANDLING 控制 ProcessGroupNCCL 的 watchdog 线程发现异常或超时后怎么收场：

watchdog 轮询未完成的 collective，发现通信器出错或超过进程组 timeout
按模式处理：v2.11 读取这个变量时的默认值是 3（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:911），模式枚举 ErrorHandlingMode 定义在 torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:165：3 是 SkipCleanUp（不清理 communicator、直接拉掉进程），1 是 TearDown（先 abort communicator 再拉掉进程），0 是 NoHandling（不处理）
配合 flight recorder（第 16 章 §16.28）的 dump 让"卡在哪个 collective"可追

所以 v2.2 起不设这个变量也会快速失败（v2.1 默认是 1，更早默认是 0）；要等的是 NCCL 进程组的默认 timeout —— 10 分钟（kProcessGroupNCCLDefaultTimeout，torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:155），可用 init_process_group(timeout=...) 调。

17.9 `register_comm_hook`：通信压缩与算法替换

DDP 默认行为是"先除以 world_size、再 AllReduce 求和"（Reducer 的 mark_variable_ready_dense（torch/csrc/distributed/c10d/reducer.cpp:353）在把 grad 拷进 bucket 时顺带乘 1/world_size；内置的 allreduce_hook 也是先 div_ 再 all_reduce，注释说是为了防 fp16 溢出）。register_comm_hook 让用户替换这个默认行为：

import torch.distributed.algorithms.ddp_comm_hooks.default_hooks as default

# bf16 通信压缩 (fp32 梯度通信量减半, 结果再拷回 fp32 的 bucket)
ddp_model.register_comm_hook(state=None, hook=default.bf16_compress_hook)

# fp16 压缩
ddp_model.register_comm_hook(state=None, hook=default.fp16_compress_hook)

PyTorch 内置的 hook 集合（torch/distributed/algorithms/ddp_comm_hooks/）：

Hook	作用
`allreduce_hook`	`torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py:35`，默认行为（无压缩）
`bf16_compress_hook`	`torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py:116`，把梯度先转 bf16 再 AllReduce，通信量减半
`fp16_compress_hook`	`torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py:96`，同上但 fp16，要小心数值范围
`powerSGD_hook`	低秩压缩。PowerSGD 论文在 16 卡上实测的压缩倍数是 72×~310×（随 rank 而变），通信量本身降两个数量级；但端到端收益远小于此——ResNet18 通信时间减 54%、整体训练时间减 24%，LSTM 分别是 90% 与 55%。压缩比不等于加速比，瓶颈会转移到编解码和计算上
`quantization_pertensor_hook`（`torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py:46`）/ `quantization_perchannel_hook`（`torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py:122`）	量化后通信

通信压缩在跨节点带宽受限时收益最明显：通信量减半，在"通信是瓶颈"的场景下通信时间大致也减半；但如果通信本来就被计算完全 overlap 掉了，压缩只是白付一次 cast 的代价。代价是数值精度有损失，影响大小要在自己的任务上验证。另外 bf16_compress_hook 要求 CUDA 11+ 且 NCCL 2.10+（或 XPU + XCCL）：CPU 上用 gloo 注册会直接抛 TypeError: BF16 all reduce communication hook required CUDA 11+ and NCCL 2.10+ or XPU and XCCL（实跑）。

实现机制：comm hook 在 Reducer 拿到 bucket 准备 AllReduce 时被调用，user hook 返回一个 Future 对象，告诉 Reducer "我会异步算完通信"。这套异步 future 接口让自定义 hook 仍能 overlap。

17.10 `_DDPSink`：autograd 反向的边界

torch/nn/parallel/distributed.py:242 的 _DDPSink 是一个 autograd.Function 子类，真实实现只做两件很具体的事：

class _DDPSink(Function):
    @staticmethod
    def forward(ctx, ddp_weakref, *inputs):
        # set_materialize_grads(False) 让 None 梯度保持 None、不被填成零
        ctx.set_materialize_grads(False)
        ctx.ddp_weakref = ddp_weakref
        ret = inputs
        if ddp_weakref()._ddp_sink_clone:
            ret = tuple(inp.clone() if isinstance(inp, torch.Tensor) else inp
                        for inp in inputs)
        return ret

    @staticmethod
    def backward(ctx, *grad_outputs):
        # static graph 训练的第一轮：把延迟 all-reduce 排进 autograd 引擎的回调队列
        ddp_weakref = ctx.ddp_weakref()
        reducer = ddp_weakref.reducer
        static_graph = ddp_weakref.static_graph
        delay_ar_enqueued = (
            static_graph and ddp_weakref._static_graph_delay_allreduce_enqueued
        )
        if static_graph and not delay_ar_enqueued:
            Variable._execution_engine.queue_callback(reducer._delay_all_reduce)
            ddp_weakref._static_graph_delay_allreduce_enqueued = True
        return (None, *grad_outputs)

注意三点，都容易被想当然地记错：

_DDPSink 不是每次 forward 都套上去的。torch/nn/parallel/distributed.py:1654 那个 if 写得很清楚：只有 find_unused_parameters and not static_graph，或者 static_graph 且延迟 all-reduce 还没排队时，才走 _DDPSink.apply（torch/nn/parallel/distributed.py:1676）。常规的 static_graph=False, find_unused_parameters=False 训练根本不经过它
它的 forward 侧价值是 set_materialize_grads(False) —— 让"没参与 loss 的那些输出"在反向里保持 None 梯度，reducer 才能据此判断哪些 param 该跳过，而不是收到一堆零
它的 backward 侧只负责 static graph 第一轮那次 queue_callback(reducer._delay_all_reduce)，不做"重置 bucket 状态"之类的事

这仍然是 PyTorch 用 autograd.Function 给 DDP 在反向链上插一个可控节点的手法。第 7 章 §7.8 我们看到 autograd.Function 让用户写自定义反向；DDP 用同一接口在框架级别注入逻辑。

17.11 `Join` API：处理 rank 间数据不均

实际数据集每个 rank 的样本数可能不同（如某 rank dataset 比别人少）。少 batch 的 rank 提前结束训练循环，但其他 rank 还在跑 —— 它们的 AllReduce 永远等不到这个 rank 的参与，hang 住整个集群。

torch.distributed.algorithms.join（PyTorch 1.10+）解决这个问题：

from torch.distributed.algorithms import Join

with Join([ddp_model]):
    for batch in loader:
        ...   # 某 rank 提前 break 也没事

机制：进入 Join 上下文时每个 rank 注册自己的 JoinHook（DDP 的是 _DDPJoinHook at torch/nn/parallel/distributed.py:276）。当某 rank dataloader 耗尽提前 break，Join 让它继续模拟参与 AllReduce（发送 0 梯度），让其他 rank 不 hang。所有 rank 都耗尽时退出 with 块。

这个 API 让"不均衡 sharding"训练能正常工作。实战里大多数训练不需要它（DistributedSampler 会把各 rank 补齐到同样多的样本，见 §17.25），但跨 rank 数据天然不均的场景（如 federated learning）必备。

17.12 `mark_variable_ready`：单变量到 bucket 的 C++ 实现

torch/csrc/distributed/c10d/reducer.cpp:353 的 mark_variable_ready_dense（与 :441 mark_variable_ready_sparse 双胞胎）是 hook 触发后真正的工作函数。流程：

找到这个 variable 所属的 bucket（每个 param 在初始化时被分配了 bucket_index）
把 grad 数据拷贝到 bucket 的连续 buffer（如果 gradient_as_bucket_view=False）；否则因为 grad 已是 bucket view，跳过
bucket 的 pending 计数 -1
如果 bucket pending == 0（这是 bucket 里最后一个 variable）→ 调 mark_bucket_ready
mark_bucket_ready（torch/csrc/distributed/c10d/reducer.cpp:1025）按 bucket 序号严格依次发起 AllReduce：如果这个 bucket 前面还有没发出的 bucket，就先不发，等前面的就绪后一并从 next_bucket_ 往后发（all_reduce_bucket，声明在 torch/csrc/distributed/c10d/reducer.hpp:523，调用点在 torch/csrc/distributed/c10d/reducer.cpp:1045）

整套流程在 C++ 端做：hook 每个参数每轮都要触发一次，放在 Python 里会叠加解释器开销与 GIL 争用。

Reducer::prepare_for_backward（torch/csrc/distributed/c10d/reducer.cpp:1527）负责把所有 bucket 的 pending 计数重置、清掉上次的状态，让本次反向能干净地开始计数。它的调用点在 DDP 的 forward 末尾而不是反向里 —— torch/nn/parallel/distributed.py:1646 在 find_unused_parameters and not static_graph 时传入 _find_tensors(output) 的结果（用来标记这次哪些输出参与了 loss），:1648 在其余情况下传空列表。

17.13 bucket 的"反向重排"：rebuild_buckets

DDP 初始化时先按 model.parameters() 顺序算 bucket 划分，再把 bucket 列表整个反过来交给 Reducer（torch/nn/parallel/distributed.py:1257 的 list(reversed(bucket_indices))，注释说是为了近似反向产出梯度的顺序）—— 最后定义的层落在 bucket 0。但定义顺序不一定是真实执行顺序，所以默认配置（find_unused_parameters=False，或 static_graph=True）下第一轮干脆不分桶：上限传 sys.maxsize，全部参数进一个 bucket（torch/nn/parallel/distributed.py:1224-1225），第一轮反向结束才发一次 AllReduce。

DDP 用 should_rebuild_buckets()（torch/csrc/distributed/c10d/reducer.hpp:148：(static_graph_ || !find_unused_parameters_) && !has_rebuilt_bucket_）+ rebuild_buckets 按真实顺序重来：

第一轮 backward 时按"实际收到 grad 的顺序"记录 rebuilt_param_indices_
下一轮 forward 开头（_pre_forward 里的 self.reducer._rebuild_buckets()）按这个顺序重新分桶：最先产出梯度的参数进 bucket 0，上限 1 MiB；之后每个 bucket 上限 bucket_cap_mb
此后 bucket 顺序就是反向序，layer N 的 grad 一算完 bucket 0 就满 → 立即 AllReduce → 与 layer N-1 的 backward 计算 overlap

gloo 双进程实跑（12 个 nn.Linear(1024, 1024)，每个 weight 4 MiB）：第一轮只有 1 个约 48 MiB 的 bucket；第二轮 forward 开头重建成 3 个，字节数依次是 4198400、29388800、16793600，bucket 0 是最后一层的 weight 与 bias（一个 4 MiB 的 weight 就超过了 1 MiB 上限，所以 bucket 0 只装这一层）。find_unused_parameters=True 时不重建，一开始就按"首个 bucket 1 MiB、其余 25 MiB"分好再反过来：同一模型得到 16797696、29388800、4194304 三个 bucket，1 MiB 上限的那个落在最后（装第一层的 weight）。

这个 rebuild 解释了"第一轮通信几乎不 overlap"：第一轮只有一个大 bucket，要等整个反向跑完才发。

static_graph=True（§17.7）并不跳过 rebuild：第一轮所有 AllReduce 推迟到反向结束，第二轮才记录梯度顺序，重建发生在第三轮 forward 开头（同一模型实跑印证）。

17.14 DDP `init` 详细流程

torch/nn/parallel/distributed.py:655 的 DistributedDataParallel.__init__ 做的事远不止 wrap：

校验并广播模型状态（init_sync=True 时，默认开）：先 _verify_param_shape_across_processes 核对各 rank 参数个数与形状（§17.22），再 _sync_module_states 把 rank 0 的参数（以及 broadcast_buffers=True 时的 buffer）经 _broadcast_coalesced 广播到所有 rank，保证训练起点一致。如果不做这步，每个 rank 的随机初始化不同，DDP 训练就退化成"几个独立模型平均梯度"，完全错
创建 Reducer C++ 实例（_ddp_init_helper）：传入参数列表、反转后的 bucket 划分（§17.13）、process_group、bucket 上限等
注册 backward hook：在 Reducer 的 C++ 构造函数里给每个参数的 grad_accumulator 挂 post hook，回调 autograd_hook（torch/csrc/distributed/c10d/reducer.cpp:189）
forward 期间的逻辑不靠 module hook：buffer 同步、bucket rebuild、prepare_for_backward、_DDPSink 包装都写在 DDP.forward 调的 _pre_forward / _post_forward 里（§17.24）
find_unused_parameters=True 不在 __init__ 里多注册什么：它影响的是初始分桶（§17.13）和每轮 prepare_for_backward 里的图遍历
如果开 mixed_precision：配 _MixedPrecision policy，注册 forward pre hook 把参数拷成低精度副本（输入在 _pre_forward 里 cast），并挂一个 AllReduce 后再 upcast 的 comm hook

整个 __init__ 涉及 reducer.cpp + distributed.py 一千多行代码协调。理解这条流程让你看到 stack trace 时能立刻定位错误来源。

17.15 sync_buffers：BatchNorm running stats 的同步

DDP 默认 broadcast_buffers=True：每次 forward 前把 rank 0 的 buffer（如 BN 的 running_mean / running_var）广播到其他 rank。为什么需要？

BN 在 train 模式下，每个 rank 各自更新自己的 running stats（基于本地 batch）。多个 rank 的 running stats 会慢慢漂移（不同 rank 看到的数据分布略不同），最终模型在 eval 时使用的统计量取决于哪个 rank 的 buffer，结果非确定性。

broadcast_buffers=True 强制每 forward 同步一次，让所有 rank 的 buffer 严格一致。代价是每 forward 多一次 broadcast collective（buffer 通常很小，开销可忽略）。

如果你的训练有 SyncBatchNorm（专门的同步 BN 实现，跨 rank 计算 batch 统计量），可以关闭 broadcast_buffers —— SyncBN 已经保证 stats 同步。

17.16 `prepare_for_backward`：反向开始前的准备

torch/csrc/distributed/c10d/reducer.cpp:1527 的 prepare_for_backward 在每次 forward 完成后被调 —— 由 _post_forward 直接调用（§17.12），不经过 _DDPSink。它做的事：

重置所有 bucket 的 pending 计数（reset_bucket_counting）
清空本轮的记账：grad_ready_order_indices_、"已标记 unused"标志、本轮已 ready 的参数集合
仅当 find_unused_parameters=True 且非 static graph：从 forward 输出沿 autograd 图遍历（search_unused_parameters，用栈做的深度优先），没被遍历到的 AccumulateGrad 对应的参数记为 unused，反向中第一个 hook 触发时把它们直接标 ready

它本身不做图一致性检查。一致性问题在反向里暴露：同一参数在一次反向里被标记 ready 两次时，torch/csrc/distributed/c10d/reducer.cpp:838-868 的报错（Expected to mark a variable ready only once）会列出可能原因，并建议试 _set_static_graph() 或关掉 find_unused_parameters。

17.17 `set_static_graph` 做了什么、在哪里报错

torch/csrc/distributed/c10d/reducer.cpp:2180 的 set_static_graph 本身很短：检查 num_iterations_ == 0（必须在训练开始前调），置 static_graph_ = true，初始化 local_used_map_。它不检查任何 invariant。

"图不变"的约束是在运行时撞出来的。gloo 双进程实跑，一个 a → b 两层的模型：

前两轮都用 b、第三轮不用 b：第四轮 forward 开头报 Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your training graph has changed in this iteration, e.g., one parameter is used in first iteration, but then got unused in the second iteration. this is not compatible with static_graph set to True.
前两轮不用 b、第三轮用上 b：第三轮反向里就报 Your training graph has changed in this iteration, e.g., one parameter is unused in first iteration, but then got used in the second iteration. this is not compatible with static_graph set to True.（autograd_hook 里的 REDUCER_CHECK）

所以模型有动态分支（if x.shape[0] > 0: ...）、用到的参数集合会变时，不能开 static_graph=True。

17.18 DDP × GradScaler：mixed precision 的协作

GradScaler（第 20 章 §20.2.2）解决 fp16 下溢，DDP 解决多卡同步。两者协作：

scaler = torch.amp.GradScaler("cuda")
ddp_model = DDP(model)

for batch in loader:
    optimizer.zero_grad()
    with torch.autocast("cuda", dtype=torch.float16):
        loss = ddp_model(batch).sum()
    scaler.scale(loss).backward()       # backward + AllReduce 一起跑
    scaler.unscale_(optimizer)           # 把 grad 除回去
    scaler.step(optimizer)               # 检查 inf, OK 就 step
    scaler.update()

关键：scaler.scale(loss).backward() 让反向算梯度时是放大版（避免下溢），DDP 的 AllReduce 同步的是放大梯度。所有 rank 都按同样 scale 因子放大、AllReduce 求平均、再 unscale —— 数学上等价于"先 unscale 再平均"。

但是有个微妙点：所有 rank 必须看到相同的 inf 检测结果。否则一个 rank 跳过 step（有 inf）、其他 rank 正常 step，模型会失同步。DDP 下这一点是自然满足的：AllReduce 之后各 rank 的梯度完全相同，任一 rank 的 inf/NaN 也会经求和传到所有 rank，各自 unscale_ 查出的 found_inf 必然一致。torch/amp/grad_scaler.py 里并没有任何集合通信；需要额外 AllReduce found_inf 的是梯度分片的 FSDP，所以才有 torch/distributed/fsdp/sharded_grad_scaler.py 的 ShardedGradScaler。

17.19 DDP × torch.compile

常见写法是先 DDP wrap、再 compile：

ddp_model = DDP(model)
compiled_ddp = torch.compile(ddp_model)

真正的问题不在顺序，而在 overlap。 Dynamo + AOTAutograd 会把整个反向编成一个融合的函数，从 autograd Engine 看它是"原子"执行的：所有参数的梯度同时就绪，DDP 的 hook 要等整段反向跑完才一起触发，通信与计算的 overlap 就没了（torch/_dynamo/backends/distributed.py 里 DDPOptimizer 的 Note 原话）。

v2.0 起 Dynamo 内置的 DDPOptimizer 解决这个问题：Dynamo 编译的代码若正运行在 DDP 的 forward 里（DDP._inside_ddp_forward 会登记 _active_ddp_module，torch/_dynamo/convert_frame.py:2291 读它），就按 DDP 的 bucket_bytes_cap 从后往前累加参数大小，在 bucket 边界处主动切图，各子图分别编译。反向时每个子图的梯度一出来 hook 就能触发，overlap 回来了。开关是 torch._dynamo.config.optimize_ddp（默认 True，即 "ddp_optimizer"）。

用 backend="aot_eager"、gloo 双进程实跑（6 个 nn.Linear(1024, 1024)，bucket_cap_mb=5）：torch.compile(DDP(m)) 与 DDP(torch.compile(m)) 两种顺序都被 DDPOptimizer 切成 4 个子图（TORCH_LOGS=ddp_graphs 可见），各 rank 梯度也都正确同步。注意被编译的必须是用户自己定义的 module 类：直接拿 nn.Sequential 套 DDP 再 compile，Dynamo 在 DDP forward 处断图后不会接着编译 torch 内部的 Sequential.forward，实跑里一个图都没编出来。

DDP × compile 的收益来源是 compile 消掉 forward / backward 的 Python 与 dispatcher 开销，与 DDP 的通信 overlap 是两件独立的事，可以叠加。

17.20 `register_comm_hook` 的内部实现

§17.9 介绍了 comm hook 的用法。底层机制（torch/nn/parallel/distributed.py:1987 register_comm_hook）：

user hook 接收 (state, bucket) 参数 + 返回 Future 对象
DDP 在 bucket ready 时不直接调 NCCL，而是调 user hook
user hook 内部决定怎么 reduce（可以压缩、可以做特殊算法、可以同时跑多个 collective）
user hook 返回的 Future 解析为 reduced grad，DDP 用它继续后续步骤

具体例子（bf16 压缩 hook 的简化实现）：

def bf16_compress_hook(state, bucket):
    grad = bucket.buffer()
    grad_bf16 = grad.to(torch.bfloat16).div_(dist.get_world_size())  # 压缩, 并先除 world_size
    fut = dist.all_reduce(grad_bf16, op=dist.ReduceOp.SUM, async_op=True).get_future()
    def decompress(fut):
        return grad.copy_(fut.value()[0])            # 解压: 就地拷回 fp32 的 bucket
    return fut.then(decompress)

返回的 fut 让 DDP 等 AllReduce 完成、然后调 decompress 拿到 fp32 grad 继续训练。整套接口让任意自定义压缩 / 优化算法能无缝接入 DDP。

PyTorch 自带的 PowerSGD（低秩压缩）就是以 comm hook 的形式实现的（torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py），TopK 稀疏化之类的算法也可以走同一接口。第三方实现复用 DDP 的整个 backward hook + bucket 体系，只关注"我自己的 reduce 算法"。

17.21 DDP 单进程多 device 模式（已移除）

PyTorch 历史上 DDP 还支持"一个 rank 控制多 GPU"模式（device_ids=[0, 1]），思想类似 DataParallel。但这种模式：

单线程 dispatch，多 GPU 受 GIL 限制
反向 AllReduce 顺序复杂
与 NCCL 多 device communicator 协作有 corner case

v1.6–v1.8 在这种模式下会打 "Single-Process Multi-GPU is not the recommended mode for DDP" 的警告，v1.9 起直接抛 ValueError: device_ids can only be None or contain a single element.（v2.11 的 torch/nn/parallel/distributed.py __init__ 里仍是这句，CPU 实跑可见），只能 1 rank = 1 GPU（device_ids=[local_rank]）。这种"每 rank 一个 GPU"模式让 DDP 内部逻辑简单很多 —— 进程内没有多份 replica 争 GIL、单 NCCL communicator、bucket 顺序确定。

如果遇到老代码用 device_ids=[0, 1] 的 DDP，迁移到 torchrun --nproc_per_node=2 + 每个 rank 自己控一个 GPU 是标准做法。

17.22 `dist._verify_params_across_processes`

DDP 启动时调 dist._verify_params_across_processes（Python 侧经 torch/distributed/utils.py:281 的 _verify_param_shape_across_processes 转调，C++ 实现是 torch/csrc/distributed/c10d/reducer.cpp:2347 的 verify_params_across_processes）检查所有 rank 的 model 结构是否一致。它做的事：

先 allgather 各 rank 的 parameter 个数，任一 rank 对不上就报 "DDP expects same model across all ranks"
再把各 parameter 的 sizes 与 strides（每个 param 贡献 2 * ndim 个 int64）从 rank 0 broadcast 出去，各 rank 与本地对比，不符就报 params[i] in this process with sizes [...] appears not to match sizes of the same param in process 0.

注意它比对的是结构（个数、形状、stride），不是权重的具体数值 —— 权重一致靠 __init__ 里紧接着的那次 broadcast（§17.14 第 1 步）保证。所以不同 rank 加载了不同 ckpt、用了不同 seed 初始化，这里都查不出来，也不需要查（会被 rank 0 的权重覆盖）。

它能抓住的是模型构造本身随 rank 变了，例如按 rank 决定要不要某一层、某层的宽度。gloo 双进程实跑，rank 0 的 Linear 带 bias、rank 1 不带时，报 RuntimeError: DDP expects same model across all ranks, but Rank 0 has 2 params, while rank 1 has inconsistent 1 params.

这套检查避免"不一致的 model 跑 DDP 训练"这种隐蔽 bug —— 用户在构造 DDP 时就收到 RuntimeError，而不是在后续 collective 里出现形状错乱或 hang。

17.23 各调优开关省的是哪段时间

本机无 GPU，这里不给倍数，只给每个 flag 省的是哪一段时间——照着这张表用 profiler 一项项验证：

配置	省掉什么
`gradient_as_bucket_view=True`	每个 bucket 一次 grad → buffer 的全量 copy，以及那份额外显存
`static_graph=True`	有 unused 参数时每轮的 autograd 图遍历（不省 bucket rebuild，§17.7 / §17.13）
bf16 comm hook	AllReduce 的通信量减半（只有在通信没被完全 overlap 掉时才转化成时间）
DDP × compile	forward / backward 的 Python 与 dispatcher 开销（与通信无关，可叠加）

跨节点训练（如 32 卡 4 节点）通信占比通常更高（节点间带宽远低于机内 NVLink），HSDP（第 18 章）这类"机内分片、机间复制"的复合策略就是冲这个来的。

17.24 DDP forward 流程的精确分解

torch/nn/parallel/distributed.py 里 forward = _pre_forward → _run_ddp_forward → _post_forward，完整链路：

pre forward：需要同步梯度时调 reducer.prepare_for_forward()；只在训练期第一次成功时做一次 _rebuild_buckets()（§17.13）；broadcast_buffers=True 时 _sync_buffers 同步 BN running stats 等 buffer
真实 forward：self.module(*inputs, **kwargs) 跑用户 model
post forward · prepare_for_backward：Reducer 重置 bucket pending 计数；find_unused_parameters=True 且非 static graph 时传入 _find_tensors(output)，顺带做 unused 参数检测
post forward · _DDPSink.apply：仅在 find_unused_parameters=True 且非 static graph，或 static graph 第一轮时包装输出（§17.10）

如果用户 model 输出包含多个 tensor（如 (logits, hidden_states)），_find_tensors(output) 会扁平化收集所有 tensor 给 prepare_for_backward。这套机制让多输出模型也能正确识别"哪些张量参与反向"。

17.25 DistributedSampler：DDP 数据切分的"黄金搭档"

DDP 训练每 rank 必须看到不同 batch，否则等于 N 个 rank 训同一份数据 —— DDP 完全失效。torch.utils.data.distributed.DistributedSampler 是标准切分器：

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(dataset, sampler=sampler, batch_size=B)

for epoch in range(epochs):
    sampler.set_epoch(epoch)    # 关键: 每 epoch 重置 seed, 让 shuffle 跨 epoch 不同
    for batch in loader:
        ...

set_epoch(epoch) 是新手常忘的关键调用 —— 不调的话每 epoch 都用同一个 random shuffle 顺序。

DistributedSampler 算法：把 dataset 索引 [0, 1, 2, ..., N-1] 切成 world_size 份，rank i 拿到第 i, i+world_size, i+2*world_size, ... 份。这种 strided 切法让数据在 rank 间均匀分布。

如果 len(dataset) % world_size != 0，DistributedSampler 会用前几个样本 padding（保证每 rank 数据量一致）。这避免了 §17.11 的 Join API 复杂度 —— 数据天然均匀，不需要 Join。

17.26 `_register_fused_optim`：optimizer step 与 backward 融合

DDP 默认流程是"backward 里所有 AllReduce 完成 → optimizer.step()"。但每个参数的梯度一 reduce 完就已是最终值 —— 它的更新不必等其余 bucket，可以与仍在进行的其他 bucket 通信重叠。

_register_fused_optim（v1.11 起，私有接口）让 DDP 在参数梯度 reduce 完成后立即更新它：

ddp_model._register_fused_optim(torch.optim.SGD, lr=0.01)

第一个参数 optim 是位置参数（写成 optimizer_class=... 会报 missing 1 required positional argument: 'optim'，gloo 实跑）。注册后一次 backward() 就把参数更新了，不用再调 optimizer.step()（实跑印证）。docstring 称这可能加速训练，并可能降低峰值显存（一次只需一个参数的 optimizer state）（入口是 torch/nn/parallel/distributed.py:2097 的 _register_fused_optim）。

docstring 里列出的限制：

每个 DDP 实例只能注册一次
不能与自定义 comm hook（register_comm_hook）组合
不支持梯度累积和 no_sync

另一条相近的路径是 torch/distributed/optim/apply_optimizer_in_backward.py 的 _apply_optimizer_in_backward（同样是私有接口）：DDP 在 _lazy_init → _setup_in_backward_optimizers 里识别它，把 optimizer step 与 DDP 的 AllReduce 串起来。

17.27 PowerSGD：低秩压缩通信

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook 实现 PowerSGD（Vogels、Karimireddy、Jaggi，PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization，NeurIPS 2019，arXiv:1905.13727）。算法核心：把 grad 分解成低秩矩阵 G ≈ P @ Q.T，只通信 P 和 Q（远小于 G）。

具体步骤（hook 内部，简化伪代码：真实实现对 bucket 里每个 ≥2 维的梯度张量各自做分解，1 维张量走普通 AllReduce，另有 error feedback、Q 的 warm start、前 start_powerSGD_iter 轮先用普通 AllReduce、结果除以 world_size 等细节）：

def powerSGD_hook(state, bucket):
    grad = bucket.buffer().reshape(M, N)   # 示意: 真实实现逐个张量处理

    # 第一次跑: 初始化随机 Q (rank=R << min(M,N))
    if state.iter == 0:
        Q = torch.randn(N, R)

    # 1. P = grad @ Q
    P = grad @ Q
    dist.all_reduce(P)               # 通信 P (M*R 元素)

    # 2. orthogonalize P
    P, _ = torch.linalg.qr(P)

    # 3. Q = grad.T @ P
    Q = grad.T @ P
    dist.all_reduce(Q)               # 通信 Q (N*R 元素)

    # 4. reconstruct: grad ≈ P @ Q.T
    return P @ Q.T

通信量从 M*N 降到 (M+N)*R。M=N=4096、R=4 时压缩比是 4096×4096 /(2×4096×4) = 512 倍。

但压缩比不等于加速比：§17.9 的表里已经引过论文的实测——ResNet18 上通信量压了约两个数量级（72×~243×），通信时间（含编解码）只减 54%、端到端训练时间只减 24%，因为计算本身没变、编解码又有开销。带宽越紧、compute 越富余，这套越划算；机内带宽本就充足时未必划算。

PowerSGD 是 DDP comm hook 设计带来的工程红利 —— 算法研究者不用改 PyTorch 主仓就能把新算法接入生产训练。

17.28 TORCH_NCCL_BLOCKING_WAIT vs ASYNC_ERROR_HANDLING

第 16 章 §16.9 提过 ASYNC_ERROR_HANDLING。还有个相关的 TORCH_NCCL_BLOCKING_WAIT（旧名 NCCL_BLOCKING_WAIT 仍认）：

环境变量	行为
`TORCH_NCCL_BLOCKING_WAIT=1`	`wait()` 在 CPU 侧阻塞直到完成或 timeout 后抛错；此时不创建 watchdog 线程
`TORCH_NCCL_ASYNC_ERROR_HANDLING`	`wait()` 不阻塞 CPU，watchdog 监控；v2.11 默认 3（`SkipCleanUp`，直接拉掉进程），1 是 `TearDown`（先 abort communicator），0 是不处理

两者 v2.11 下的默认组合是 BLOCKING_WAIT=0 + ASYNC_ERROR_HANDLING=3（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:909-911），也就是说默认已经会在出错或超时后快速失败。BLOCKING_WAIT 是更老的机制（v1.3 起），会让每次 wait() 都阻塞 CPU、性能下降；ASYNC_ERROR_HANDLING 是 v1.7 起的方案，不阻塞 CPU、出错快速失败。

如果你看到训练任务"hang 在某个 collective"，先查是否把 TORCH_NCCL_ASYNC_ERROR_HANDLING 显式设成了 0，以及 timeout 设了多大：NCCL 进程组默认 10 分钟（kProcessGroupNCCLDefaultTimeout），gloo 等其他后端才是 30 分钟（kProcessGroupDefaultTimeout，torch/csrc/distributed/c10d/ProcessGroup.hpp:21）；这两个值都是 PyTorch 设的，不是 NCCL 的（第 16 章 §16.9）。

17.29 elastic 训练与 DDP

torchrun --nproc_per_node=8 --rdzv_backend=c10d ... 启动的就是 elastic 训练。它在 DDP 之上加一层 fault tolerance：

任一 worker 失败 → 整个 worker group 重启（次数上限 --max_restarts，默认 0 即不重启；从最近 ckpt 恢复要训练脚本自己做）
节点 join / leave 时整个 group 重新 rendezvous
与 Kubernetes 等编排系统集成做自动恢复

elastic 不改变 DDP 的语义，只是让"训练 job 整体寿命"更长。大模型训练动辄几天到几周，期间 GPU 故障 / 网络抖动难以避免 —— 没有 elastic 重启机制就要人工介入恢复。

17.30 DDP 与 MoE 的特殊性

Mixture-of-Experts (MoE) 模型每 token 只激活部分 expert。DDP 默认假设"所有 rank 看到的所有 expert 都参与 backward"，但 MoE 实际上不同 rank 不同 token 激活不同 expert，导致某些 expert 在某些 rank 上没被使用 → 没收到 grad。

DDP 默认行为下，这些 expert 所在 bucket 的 pending 计数永远不归零，AllReduce 发不出去：各 rank 情况一致时下一轮报 Expected to have finished reduction in the prior iteration（§17.6），各 rank 不一致时则互相等到超时。

解法：

find_unused_parameters=True：DDP 自动检测哪些 expert 这次没用、当作 ready
自定义 sampler 让所有 rank 都激活所有 expert（牺牲 expert specialization）
用 expert parallel（每 expert 在固定 rank）+ 跨 expert 用 all-to-all 通信

MoE 训练在生产里通常用第三条 —— 与 DDP 解耦，用专门的 expert parallel framework（如 Tutel / DeepSpeed-MoE）。理解 DDP 的"全 rank 同步"模型让你看清为什么 MoE 不能简单套 DDP。

17.31 SyncBatchNorm：跨 rank 同步 BN

普通 nn.BatchNorm2d 在每 rank 各自计算 batch 统计量。多 rank 的 batch 统计量不同 → 训练时模型在每 rank 看到的"输入分布"不一致 → 收敛差异。

nn.SyncBatchNorm（或 convert_sync_batchnorm 函数）把 BN 的 mean / var 计算改成跨 rank：

# 一行把所有 BN 替换成 SyncBN
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
ddp_model = DDP(model)

底层机制（torch/nn/modules/_functions.py）：SyncBN forward 先用 torch.batch_norm_stats 算本地 mean / invstd，连同样本数拼成一个小张量做一次 all_gather，再由 batch_norm_gather_stats_with_counts 算出全局 mean / invstd 做 normalization；backward 再做一次 all_reduce 同步 sum_dy / sum_dy_xmu。代价是每个 BN 层前向、反向各多一次小 collective。

CV 类大模型训练（每 rank batch 小、需要全局统计量）几乎必用 SyncBN。LLM 类训练（每 rank batch 大 + 用 LayerNorm 不是 BN）不需要。

convert_sync_batchnorm 内部是模型 surgery（第 9 章 §9.8.3）：递归遍历找 BN 替换成 SyncBN。操作要在 DDP wrap 之前做 —— wrap 后再改 model 会破坏 DDP 注册的 hook。

17.32 DDP × activation_checkpoint 协作

第 7 章 §7.5.4 讲过 activation_checkpoint。它与 DDP 协作的微妙点：

use_reentrant=True 的 checkpoint 用 autograd.Function，反向时重新跑一次被 checkpoint 的那段 forward 取激活，再在里面嵌套调一次 backward。被重算的只是那段子函数，不是 DDP.forward；麻烦出在嵌套的 reentrant backward 里参数的 AccumulateGrad hook 照样触发，DDP 的计数就可能对不上。gloo 双进程实跑：同一个 nn.Linear 被连续 checkpoint 两次，find_unused_parameters=False 时能跑；find_unused_parameters=True 时报 Expected to mark a variable ready only once（§17.16 那段报错，原因 2 就写的是"Reused parameters in multiple reentrant backward passes"）。DDP 给出的规避办法是 static_graph=True（§17.7 docstring 列的前三条场景都与此相关）。

use_reentrant=False 用 saved_tensors_hooks 实现，不走嵌套 backward → 同样的模型加 find_unused_parameters=True 实跑也能正常训练。注意它并不是默认值：v2.11 的 torch/utils/checkpoint.py:352 里 use_reentrant 默认还是 None，:480 走一次 "should be passed explicitly" 的告警之后回落到 True（第 18 章 §18.11 也点了这条）。

实战建议：新代码强制 use_reentrant=False。老代码遇到 "DDP 反向异常"时，checkpoint 的 use_reentrant 参数值得先查。

第 18 章 §18.11 详细讲了 FSDP × checkpoint 的协作（更复杂，因为 FSDP 还要管 unshard）。DDP 这边相对简单 —— 只要 reentrant backward 不让同一参数的 hook 重复触发就 OK。

17.33 `find_unused_parameters` 的实现机制

§17.6 讲过用法，这里展开实现。find_unused_parameters=True 时：

forward 完后，DDP 从 forward 输出沿 autograd 图遍历（search_unused_parameters，栈式深度优先）
标记所有"通往输出路径上的"parameter 为 used
没标记的 parameter 在 backward 不会有 grad
反向里第一个 hook 触发时，reducer 把这些 parameter 直接标成 ready，扣减所在 bucket 的 pending
它们在 bucket 里的那一段填 0 参与 AllReduce（保证集合通信参与）；各 rank 的"本地是否用到"表另做一次 AllReduce，全局都没用到的参数，其 .grad 保持不动

这套机制让"动态分支"模型能在 DDP 下训练，代价是：

forward 后多一次图遍历
每次反向都要重新分析（不能缓存，因为分支可能变）
不做 bucket rebuild（§17.13），bucket 顺序只是"定义顺序反过来"的近似，与真实反向顺序可能对不齐
多一次 local-used 表的 AllReduce

开 find_unused_parameters 要在每次 forward 之后多做一次 autograd 图遍历，bucket 也不按真实反向顺序重排，只在真的需要时开。static_graph=True 时 docstring 说设不设 find_unused_parameters 都无所谓：unused 集合在第一轮定下后就不再遍历图。

17.34 DDP 多进程的工程理由

DDP 用多进程而不是多线程。为什么？

GIL 阻塞：Python 多线程跑同一进程，CPU bound 操作（autograd 调度、Python forward 逻辑）被 GIL 串行化
NCCL 单进程多 device 模型问题：NCCL 早期版本一个进程管多 device 容易出问题
fork + CUDA 不兼容：CUDA context 在 fork 后不可用，必须用 spawn
每进程独立内存空间：Out of memory / segfault 不影响其他进程

多进程模型让每 rank 独立、CUDA-friendly、容错好。代价是进程间通信只能走 IPC（NCCL / shared memory），不能直接共享 Python 对象。

torchrun 启动多进程的标准方式：

torchrun --nproc_per_node=8 --master_addr=... train.py

每个进程 RANK / WORLD_SIZE / LOCAL_RANK 通过环境变量传入，init_process_group(backend='nccl', init_method='env://') 自动读取这些环境变量握手。

历史上 PyTorch 还有 torch.multiprocessing.spawn(...) 启动方式，与 torchrun 等价但 API 更原始。新代码用 torchrun。

17.35 distributed gradient clipping

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) 是单卡 grad clipping。DDP 下要不要改？

取决于调用时机：

DDP backward 内部 AllReduce 完成后 grad 已经是全局平均
此时调 clip_grad_norm_ 算的是全局 grad norm，正确
不需要额外的"分布式 grad norm"实现

但要注意：所有 rank 必须用相同 clip 因子。如果某 rank 算出 norm=10、另一个算出 norm=11，clip 后两者梯度不一致，下次 step 后 model 就失同步。

clip_grad_norm_ 在 DDP 下的标准用法：

loss.backward()                                           # 含 AllReduce
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # 每 rank 算同样 norm
optimizer.step()

因为 grad 已经是全局平均（所有 rank 的 grad 数值一致），每 rank 算的 norm 也一致 —— 自然同步。

FSDP 下情况不同：FSDP 的 grad 是分片的，每 rank 算 local norm 后还要 AllReduce 才得全局 norm。第 18 章 §18.18 讲过，FSDP 提供 clip_grad_norm_ 自家版本处理这个。

17.36 DDP 在 PyTorch 历史上的演进

以下按各 tag 的源码核对（git show <tag>:<路径>）：

v0.2 DistributedDataParallel（DDP，多进程）首次出现，Python 实现里已经按 bucket 合并归约；DataParallel（DP，单进程多线程）在 v0.1.x 就有

v1.1 Reducer 改成 C++ 实现（reducer.cpp），同时加入 find_unused_parameters

v1.7 gradient_as_bucket_view；PyTorch 侧的 NCCL 异步错误处理（NCCL_ASYNC_ERROR_HANDLING）

v1.8 register_comm_hook 公开 API（v1.7 是私有的 _register_comm_hook）+ PowerSGD hook

v1.9 _set_static_graph()；单进程多卡模式改为直接报错（§17.21）

v1.10 Join 上下文管理器、bf16_compress_hook

v1.11 static_graph 成为构造参数、_register_fused_optim

v2.0 Dynamo 进入主仓，自带 DDPOptimizer 与 torch.compile 协作（§17.19）

v2.3 FSDP-2（fully_shard）以 prototype 形态出现（第 18 章 §18.5），大模型场景转向分片方案，DDP 仍是模型能放进单卡时的主选

理解这条演进让你看 DDP 源码时能识别"这段是 v1.x 老代码、那段是 v2.0+ 新功能"。

17.37 几条工程经验

1. gradient_as_bucket_view=True + static_graph=True 是大模型训练的标配组合

2. bucket_cap_mb 调优：默认 25 MiB 是折中。bucket 调大能减少 collective 次数，调小能增加 overlap 机会，具体取值按 §17.39 的方法用 profiler 量出来

3. 监控通信开销占比：用 profiler 看 NCCL kernel 中没被计算 overlap 掉的那部分占 step 时间多少。占比明显时说明通信是瓶颈，考虑通信压缩、调 bucket 或改善网络

4. find_unused_parameters 慢就调成 False：很多用户随手开但不需要

5. broadcast_buffers=True 默认开启：每 forward 把 rank 0 的 buffer（如 BN running stats）广播到其他 rank。可以关掉省一点时间，但要确保 buffer 同步逻辑用户自己处理

6. compile(model) 之后再 wrap DDP vs wrap DDP 再 compile：v2.11 两种顺序都会走 DDPOptimizer 按 bucket 切图（§17.19 实跑），关键是别关掉 torch._dynamo.config.optimize_ddp

7. multi-node 训练前先 dist.barrier() + small AllReduce 测试连通性：避免 forward 跑完才发现 NCCL 配置错

8. torch.distributed.algorithms.ddp_comm_hooks 提供 communication hook：用它做 fp16 通信压缩、PowerSGD 等高级优化，第 18 章 FSDP 也用类似机制

17.38 DDP overlap 为什么达不到理想值

理论上 DDP 让 communication 与 computation 完全 overlap、训练吞吐接近"纯单卡 × N"。但实测往往达不到 100% 线性扩展。原因：

最后一个 bucket 无 overlap 空间：最后一个 bucket（最靠近输入的那几层）的 AllReduce 必须等所有 backward 跑完才发。它发完才能跑 optimizer step。这部分 AllReduce 天然不能与 backward overlap。

带宽限制让某些 bucket 排队：如果 bucket_cap_mb 太大（如 100MB），单次 AllReduce 时间长、GPU compute 已经空闲等通信。如果太小（如 5MB），bucket 数量多、launch overhead 累积。

heterogeneous step time：每 rank 的 backward 速度可能微差（CPU 占用、温度 throttling），慢 rank 的 AllReduce 让快 rank 等。

这三条决定了实际扩展效率与理论线性扩展之间的差距有多大：机内 NVLink 通常离线性较近，跨节点因为带宽低得多、差距会明显拉开。具体到自己的集群，用"1 卡吞吐 × N vs N 卡实测吞吐"直接量。

调优思路：用 profiler（第 21 章）抓 trace、看 NCCL kernel 与 backward kernel 的重叠程度。如果发现"某段 NCCL kernel 完全单独跑"、说明 overlap 没拿到、要查为什么。

17.39 bucket_cap_mb 调优指南

方向（不是标定值，具体值必须实测）：模型越小、单次 AllReduce 越容易落进"延迟主导"区（第 16 章 §16.18），bucket 反而要调小才能换来更多 overlap 机会；模型越大、grad 张量本身越大，bucket 可以调大摊平每次 collective 的固定开销。默认 25 MiB 是对大多数模型的折中。

逻辑：bucket 越大、AllReduce 越大、单次开销摊平到更多 grad，但要等更多 grad 才 ready。bucket 越小、launch 越频繁、overhead 累积。

经验法则：单个 AllReduce 的耗时应该与几个 layer 的 backward 耗时同量级，overlap 才自然发生。真正的调法是拿 profiler 量出"一个 bucket 的 AllReduce 用多久"和"一个 layer 的 backward 用多久"，再按这个比例反推 bucket_cap_mb，而不是照抄某个数字。

eager DDP 的实际分桶用 ddp_model._get_ddp_logging_data() 查：bucket_sizes 是初始划分，重建后多出 rebuilt_bucket_sizes 与 rebuilt_per_bucket_param_indices（§17.13 的数字就是这样取的）。TORCH_LOGS=ddp_graphs 只在 torch.compile 走 DDPOptimizer 时有输出，打印的是按 bucket 切出来的各个子图（§17.19）。

17.40 TorchElastic 实操

elastic 启动的命令模板：

torchrun \
  --nnodes=4 \
  --nproc_per_node=8 \
  --rdzv_id=my_job \
  --rdzv_backend=c10d \
  --rdzv_endpoint=master_node:29500 \
  --max_restarts=3 \
  train.py

关键参数：

--nnodes=4 --nproc_per_node=8 → 32 卡训练
--rdzv_* → rendezvous 配置（节点协商 master 地址）
--max_restarts=3 → 单次 job 失败最多重启 3 次

elastic 还支持 --nnodes=MIN:MAX 形式的弹性节点数（默认 1:1）—— 节点增减时重新 rendezvous。这套机制让大规模训练在云上能自动恢复，而不必每次人工重启。

17.41 Reducer 怎么等所有 AllReduce 完成

DDP 反向触发 AllReduce 后，Reducer 在每个 bucket 的 future_work 里持有这次通信返回的 Future（底层是第 16 章 §16.6 的 Work）。但 optimizer.step() 之前必须等所有 AllReduce 完成，否则 step 用的是部分 reduced 部分未 reduced 的 grad。

实现机制：最后一个 bucket 发出后，mark_variable_ready 通过 Engine::queue_callback 往本次反向的回调队列里塞一个 finalize_backward（torch/csrc/distributed/c10d/reducer.cpp:1703），它在反向收尾时依次对每个 bucket 调 future_work->wait()，再把结果写回各参数的 .grad。这套同步对用户透明，也解释了为什么"loss.backward() 返回后 grad 一定是 reduced 的" —— DDP 已经在 backward 内部 wait 完所有 collective。

对 CUDA 张量，这里的 wait() 按 CUDA Future 的语义是让当前 stream 等待通信完成的事件，而不是阻塞 CPU 线程（按源码判断、未实跑）。

17.42 跨专栏关联

第 7-8 章 Autograd：DDP 的 backward hook 注册在 AccumulateGrad post_hook 上，理解第 7 章 §7.6 才能理解 DDP 怎么"在 backward 中触发通信"
第 16 章 ProcessGroup：DDP 的 AllReduce 走 ProcessGroupNCCL.allreduce
《vLLM 推理内核深度解析》专栏第 14 章分布式：vLLM 的 TP 用 AllReduce 同步注意力输出，思想与 DDP AllReduce 一致但用法不同（TP 是同一前向中的多 GPU 协作，DDP 是不同 batch 的协同）

17.43 DDP 与 LR Scheduler 协作的细节

DDP 不直接管 lr_scheduler，但有几个隐藏陷阱：

1. step-based scheduler vs epoch-based：

# step-based: 每 step 调一次
for batch in loader:
    train_step(batch)
    scheduler.step()         # ← 每 rank 调一样次, 同步

# epoch-based: 每 epoch 调一次
for epoch in range(epochs):
    for batch in loader: ...
    scheduler.step()         # 每 rank 调一次, 自然同步

只要每 rank 调 scheduler.step() 的次数相同，lr 就一致。如果用 ReduceLROnPlateau 这种基于 metric 触发的 scheduler，必须用同步过的 metric（如 dist.all_reduce(val_loss)），否则不同 rank 触发时机不同，lr 失同步。

2. warmup 期间 lr 异常小：DDP 每 rank 看到的是相同 lr，warmup 期间 lr 极小（~1e-7）→ optimizer step 几乎无更新 → 看起来"训练没动"。这是正常现象，warmup 完成后 loss 才开始下降。

3. linear scaling rule 与 DDP：经典做法是 lr = base_lr × world_size（让大 batch 等价 N × 小 batch 训练）。LLM 训练通常不用 linear scaling，而是用 cosine schedule + 经验调出的 base_lr。

17.44 DDP 与 NCCL 版本兼容性

DDP 通过 ProcessGroupNCCL（第 16 章）调 NCCL 库。NCCL 版本演进给 DDP 带来一些工程影响：

NCCL 2.7：加入 ncclSend / ncclRecv 点对点原语，pipeline parallel 的激活值传递从此有了原生支持（第 16 章 §16.15）
NCCL 2.9 起：collective 可以被 CUDA stream capture 录进 CUDA Graph（第 15 章 §15.8.2）
ncclCommSplit：split_group（§16.11）的底层，让子组创建不必从头重新握手

具体版本要求以你安装的 wheel 里带的 NCCL 为准 —— torch.cuda.nccl.version() 直接打出来。生产环境手动升级 NCCL 时要小心：NCCL 与 GPU driver 有兼容性约束，显式锁版本（如 nvidia-nccl-cu12==<版本>）是稳妥做法。

17.45 DDP 错误诊断速查表

实战 DDP 训练常见错误 + 诊断思路：

症状	可能原因	诊断
`RuntimeError: Expected to have finished reduction in the prior iteration`	上次反向有参数没收到 grad（没参与 loss）	开 `find_unused_parameters=True`；未用参数集合固定时也可 `static_graph=True`；设 `TORCH_DISTRIBUTED_DEBUG=INFO` 看参数名
训练 hang 在某 collective	某 rank 提前死、其他 rank 等	确认 `TORCH_NCCL_ASYNC_ERROR_HANDLING` 没被设成 0，配合 flight recorder 看哪个 rank 卡住
loss 不下降 / NaN	rank 间 model 不一致	检查是否用了 `init_sync=False`（此时 DDP 不广播初始权重）
`RuntimeError: DDP expects same model across all ranks, but Rank ...`	各 rank 构造出的模型结构不同	检查按 rank 分支的模型构造代码（§17.22）
单 step 时间忽快忽慢	数据不均 / 慢 rank	用 distributed view 看每 rank 时间
NCCL watchdog 报 `Watchdog caught collective operation timeout`	某 collective 超过进程组 timeout（NCCL 默认 10 分钟）	找出没到的 rank；需要时用 `init_process_group(timeout=...)` 调整

17.46 一个完整的 DDP 训练 step 时间分解

把整章内容串起来，一个健康的 DDP step 在 profiler 里应该长这样（本机无 GPU，只给结构，具体占比自己量）：

阶段	该看到什么
forward（compute）	GPU 时间线密实，没有等 CPU 的空档
backward（compute）	同上，并且 comm stream 上有 NCCL kernel 与它并排跑
AllReduce 中被 overlap 掉的部分	在 backward 的阴影里，不额外占墙钟时间
AllReduce 中没能 overlap 的部分	主要是最后一个 bucket（§17.38），backward 结束后才能发
optimizer.step	用 `fused=True`（第 10 章 §10.6.2）可以进一步压缩
dataloader	正常情况下应被 worker 池 + pin_memory 完全隐藏

优化的着力点是"没能 overlap 的那段 AllReduce"：bf16 comm hook（§17.9）把它的通信量减半，PowerSGD（§17.27）能压得更狠但要付精度与计算的代价。先量出这段占多少，再决定值不值得。

17.47 DDP 与第 8 章 Engine 的协作

第 8 章讲过 autograd Engine 怎么调度反向 DAG。DDP 完全建立在 Engine 上：

DDP 不替换 Engine，只在 AccumulateGrad 后 hook（§17.14）
Engine 调度反向 DAG 完全不知道 DDP 存在 —— hook 触发 AllReduce 是从 Engine 视角的"普通副作用"
AllReduce 本身是异步，不阻塞 Engine 继续调度后面的 Node

这种"DDP 寄生在 Engine 之上"的设计让 Engine 与分布式训练解耦。但 Engine 侧的大改动也未必能直接兼容：C++ Reducer 挂的 hook 不支持 Compiled Autograd（第 13 章 §13.27），torch/csrc/distributed/c10d/reducer.cpp:189 起注册 hook 的地方会报 Compiled autograd is not compatible with C++ DDP Reducer, please use torch._dynamo.config.optimize_ddp="python_reducer".；切到 python_reducer 后，DDP 改用 Python 侧的 register_post_accumulate_grad_hook + functional collective 做 AllReduce（_register_accum_grad_hook）。

第 8 章 §8.18 的 Engine 设计启示对应到 DDP：DDP 用 hook 而不是改 Engine 的方法实现并行训练 —— 可扩展性的胜利。

17.48 设计启示

DDP 的几条核心思想：

第一：通信与计算重叠是分布式性能的最大杠杆：任何"backward 全跑完才同步"的设计都会让多卡训练吞吐崩盘

第二：hook 让通信非侵入：用户写的 model 不需要知道 DDP 存在，DDP 通过 backward hook 在 autograd 体系外接入通信

第三：bucket 减少 launch 开销：N 个小 collective 不如 1 个大 collective，这条规则在所有 collective 框架（NCCL / MPI / Gloo）都成立

第四：Reducer 在 C++ 实现而不是 Python：每个参数每轮都要触发一次 hook，C++ 实现避免 Python 解释器开销 + GIL 争用

下一章拆 FSDP —— 当模型大到单卡放不下，DDP 不够用了。FSDP 把参数 / 梯度 / optimizer state 都切到多卡，用通信换显存。