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

“DDP 的天才在于：让通信发生在 backward 进行中，而不是 backward 之后。当最后一层梯度算完时，前面层的梯度同步已经在飞了。”

—— PyTorch dev podcast

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（distributed.py:330），但真正的工作在 C++ 端的 Reducer（torch/csrc/distributed/c10d/reducer.cpp，2509 行）。Reducer 持有：

• bucket 列表：每个 bucket 包含若干 parameter，bucket 的总 size 接近 bucket_cap_mb（默认 25MB）
• 每个 parameter 的 grad_accumulator 钩子：当那个 param 的 grad 算完时触发
• Work 列表：进行中的 NCCL collective handles
• expect_grad_count + 完成计数：跟踪哪些 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 (param : model.parameters()) {
    auto grad_accumulator = param.grad_accumulator();   // 第 7 章 §7.6 的 AccumulateGrad
    grad_accumulator->add_post_hook([this, param](){
        this->mark_variable_ready(param);
    });
}

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

• bucket 里所有 param 都 ready → 立即发 AllReduce
• 否则继续等

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 掉，只需 wait 一下最早的 bucket。

17.4 环形 AllReduce 算法

NCCL 内部对 AllReduce 默认用 Ring AllReduce 算法。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 自动按拓扑选最优。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（PyTorch 1.8+）让 param.grad 直接是 bucket 的 view（共享内存），省掉一次 copy：

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

实测节省 5-10% 训练时间，几乎没有副作用。生产代码里强烈建议开启。

17.6 find_unused_parameters：不规则模型支持

某些模型在不同 batch 走不同分支，导致某些 param 在某些 batch 不参与 forward / backward。DDP 默认认为”所有 param 每 batch 都有 grad”，遇到没 grad 的 param 会卡住等。

find_unused_parameters=True 让 DDP 在 forward 后扫描 autograd 图，标记哪些 param 这次不会有 grad，bucket 等到的时候直接当 “ready”。

代价：

• forward 后做一次 graph 遍历（小开销）
• 失去一些 bucket overlap 优化空间

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

17.7 静态图模式：static_graph=True

static_graph=True 是新引入的优化：DDP 假设”forward 图每 batch 都一样”（绝大多数训练满足）。这样：

• 第 1 个 iter：正常跑，记录 graph 结构
• 第 2 个 iter 起：直接复用记录的结构，跳过 unused param 检测、bucket 重排等

实测能让大模型训练再快 5-15%。代价是不能动态改模型（动态改了就重新 fall back 到非静态模式）。

static_graph 是 v1.11+ 才稳定，对 70B 训练几乎是必开优化。

17.8 与 NCCL Async Error Handling 的协作

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

TORCH_NCCL_ASYNC_ERROR_HANDLING=1 让 NCCL 在某 rank 异常时让其他 rank 也快速 abort：

• 检测到通信器异常 → 标记本 rank 异常
• 主程序定期 query 异常状态 → 抛异常退出
• 配合 _set_static_graph / _set_construction_logging 等机制让 stack trace 完整

生产分布式训练必开，否则一次 OOM 让整个集群挂半小时（直到 default 30 分钟 timeout）。

17.8.5 register_comm_hook：通信压缩与算法替换

DDP 默认 hook 是”AllReduce 求和后除以 world_size”。register_comm_hook 让用户替换这个默认行为：

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

# bf16 通信压缩 (省一半带宽, 用 fp32 unscale 保精度)
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/）：

通信压缩在跨节点带宽受限时收益巨大。Llama 训练在 100 Gbps ethernet 上用 bf16 hook 比 fp32 通信快 1.8x（半带宽 ≈ 一半时间）。代价是数值精度可能小损（bf16 通常无可见影响）。

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

17.8.6 _DDPSink：autograd 反向的边界

distributed.py:242 的 _DDPSink 是一个 autograd.Function 子类：

class _DDPSink(Function):
    @staticmethod
    def forward(ctx, ddp_module, *inputs):
        ctx.ddp_module = ddp_module
        return inputs

    @staticmethod
    def backward(ctx, *grad_outputs):
        # 反向开始时通知 DDP "我要开始 backward 了"
        ddp_module._post_backward_setup()
        return (None,) + grad_outputs

DDP 在 forward 末尾把输出包一层 _DDPSink.apply(...)，让反向第一步先进入 _DDPSink.backward。这样 DDP 能在反向真正开始前做一些设置（如重置 bucket 状态、检查未使用 param）。

这是 PyTorch 用 autograd.Function 给 DDP 注入”反向 hook 的根”的精妙手法。第 7 章 §7.8 我们看到 autograd.Function 让用户写自定义反向；DDP 用同一接口在框架级别注入逻辑。

17.8.7 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 distributed.py:276）。当某 rank dataloader 耗尽提前 break，Join 让它继续模拟参与 AllReduce（发送 0 梯度），让其他 rank 不 hang。所有 rank 都耗尽时退出 with 块。

这个 API 让”不均衡 sharding”训练能正常工作。实战里 90% 训练不需要它（dataloader 设 drop_last=True），但跨 rank 数据天然不均的场景（如 federated learning）必备。

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

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

1. 找到这个 variable 所属的 bucket（每个 param 在初始化时被分配了 bucket_index）
2. 把 grad 数据拷贝到 bucket 的连续 buffer（如果 gradient_as_bucket_view=False）；否则因为 grad 已是 bucket view，跳过
3. bucket 的 pending 计数 -1
4. 如果 bucket pending == 0（这是 bucket 里最后一个 variable）→ 调 mark_bucket_ready
5. mark_bucket_ready 立即发起这个 bucket 的 AllReduce（launch_bucket）

整套流程在 C++ 端做，主要原因是 hook 触发频次极高（每秒几百次）—— Python 解释器开销 + GIL 抢占会拖死整个 backward。

prepare_for_backward（distributed.py:_DDPSink.backward 间接触发）在反向开始前把所有 bucket 的 pending 计数重置 + 清掉上次的状态，让本次反向能干净地开始计数。

17.8.9 bucket 的”反向重排”：rebuild_buckets

DDP 初始化时按 model.parameters() 顺序构造 bucket。但反向计算的顺序与构造顺序相反 —— 最后一层 forward 是最先 backward 的。如果 bucket 还按 forward 顺序排，第一个 bucket 就是 layer 0 的 param，要等所有反向都跑完才 ready，完全没 overlap 机会。

DDP 用 should_rebuild_buckets() (reducer.cpp:526/719) + rebuild_buckets 解决：

1. 第一轮 backward 时按”实际收到 grad 的顺序”记录 rebuilt_param_indices_
2. 第一轮跑完后重新构造 bucket：按反向顺序重排，最后一层的 param 进 bucket 0、倒数第二层进 bucket 1 …
3. 第二轮起 bucket 顺序就是反向序，layer N 的 grad 一算完 bucket 0 立即满 → 立即 AllReduce → 与 layer N-1 的 backward 计算 overlap

这个 rebuild 让”第一轮慢、第二轮起飞”的现象有了源码解释：第一轮 buckets 顺序不对、几乎没 overlap，吞吐低；第二轮起 buckets 已重排、overlap 充分，吞吐稳定。生产代码看 DDP 训练前几个 step 比后续慢很多就是这个机制。

static_graph=True（§17.7）的关键好处之一就是：第一轮就用反向序构造 bucket（因为图固定可预测），跳过第二轮 rebuild 的开销。

17.8.10 DDP __init__ 详细流程

distributed.py:330 的 DistributedDataParallel.__init__ 做的事远不止 wrap：

1. broadcast 模型权重：rank 0 的初始化权重通过 _broadcast_coalesced 广播到所有 rank，保证训练起点一致。如果不做这步，每个 rank 的随机初始化不同，DDP 训练就退化成”几个独立模型平均梯度”，完全错
2. 创建 Reducer C++ 实例：传入 model.parameters()、process_group、bucket_size 等
3. 注册 backward hook：给每个 leaf parameter 的 grad_accumulator 挂上 mark_variable_ready
4. 注册 forward pre/post hook：用于 _sync_buffers（buffer 同步）和 _DDPSink 包装
5. 如果开 find_unused_parameters=True：注册额外的 unused parameter 检测逻辑
6. 如果开 mixed_precision：配 _MixedPrecision policy，让 forward 自动 cast 到 lower precision

整个 __init__ 涉及 reducer.cpp + distributed.py 一千多行代码协调。用户写一行 ddp_model = DDP(model) 背后是 20+ 个内部函数协同。理解这条流程让你看到 stack trace 时能立刻定位错误来源。

17.8.11 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.8.12 prepare_for_backward：反向开始前的准备

reducer.cpp:1527 的 prepare_for_backward 在每次 forward 完成后被调（通过 _DDPSink 间接触发）。它做的事：

1. 重置所有 bucket 的 pending 计数
2. 清掉上次反向遗留的 grad（如果用户没 zero_grad）
3. 如果 find_unused_parameters=True：从 forward 输出反向 BFS，标记哪些 parameter 不在路径上 → 这些 param 的 grad 直接当 ready
4. 检查 graph 一致性（每次 forward 的 autograd graph 是否一致）—— static_graph=True 时严格检查、否则容忍

这套准备让反向阶段干净开始。如果中途某个 hook 触发异常（如 reducer 检测到 graph 变化但没开 static_graph），reducer.cpp:838-868 的 error message 会精确告诉用户该开什么 flag。

17.8.13 set_static_graph 的 invariants

reducer.cpp:2180 set_static_graph 内部检查几个 invariants：

1. 每次 forward 调用的 module 顺序必须相同
2. 每次 forward 的输出张量数必须一致
3. autograd 图结构必须一致（不能某个 step 走 if 分支某个 step 走 else）

这些检查在第一个 iteration 跑完时完成，第二个 iteration 起跳过 find_unused_parameters 等开销。如果你的 model 实际有动态行为（if x.shape[0] > 0: ...），开 static_graph=True 训练几步后会抛 RuntimeError —— DDP 检测到 graph 变化、无法享受 static graph 优化。

static_graph 还能让 DDP 跳过 bucket rebuild 的第一轮成本（§17.8.9），是大模型训练几乎必开的优化。

17.8.14 DDP × GradScaler：mixed precision 的协作

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

scaler = GradScaler()
ddp_model = DDP(model)

for batch in loader:
    optimizer.zero_grad()
    with autocast(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 再平均”。

但是有个微妙点：inf 检测必须在 unscale 后做，且所有 rank 必须看到相同的 inf 检测结果。否则一个 rank 跳过 step（有 inf）、其他 rank 正常 step，模型会失同步。scaler.step(optimizer) 内部用 AllReduce 同步 found_inf 状态确保一致 —— 这是 DDP × GradScaler 协作的关键工程细节。

17.8.15 DDP × torch.compile

DDP 与 torch.compile 的协作有顺序讲究：

# 推荐: 先 DDP wrap, 再 compile
ddp_model = DDP(model)
compiled_ddp = torch.compile(ddp_model)

为什么这个顺序？ 因为 DDP 的反向 hook 必须挂在 leaf parameter 的 AccumulateGrad 上（第 7 章 §7.6）。如果先 compile 再 DDP wrap，compile 后的反向是编译产物（看不到 leaf hook 注册点），DDP 拿不到反向触发信号、AllReduce 不发。

反过来先 DDP 再 compile：compile 看到带 hook 的 model，编译时把 hook 调用编入产物。AllReduce 还能正常触发。这是 v2.1+ “DDP × compile” 兼容性的工程基础。

实测 70B 模型 DDP × compile 比纯 DDP 快 1.3-1.5x，主要收益是 compile 把 forward / backward 的 dispatcher 开销消除。

17.8.16 register_comm_hook 的内部实现

§17.8.5 介绍了 comm hook 的用法。底层机制（distributed.py:1987 register_comm_hook）：

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

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

def bf16_compress_hook(state, bucket):
    grad = bucket.buffer()
    grad_bf16 = grad.to(torch.bfloat16)              # 压缩
    fut = dist.all_reduce(grad_bf16, op=dist.ReduceOp.SUM, async_op=True).get_future()
    def decompress(fut):
        return fut.value()[0].to(torch.float32)      # 解压
    return fut.then(decompress)

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

PowerSGD（低秩压缩）、TopK gradient sparsification 等学术 idea 都是通过 register_comm_hook 在 PyTorch 上实现的。第三方实现复用 DDP 的整个 backward hook + bucket 体系，只关注”我自己的 reduce 算法”。

17.8.17 DDP 多 device 模式（已 deprecated）

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

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

v1.x 已经标记 deprecated，生产代码强制 1 rank = 1 GPU（device_ids=[local_rank]）。这种”每 rank 一个 GPU”模式让 DDP 内部逻辑简单很多 —— 单设备无 GIL、单 NCCL communicator、bucket 顺序确定。

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

17.8.18 dist._verify_params_across_processes

DDP 启动时调 dist._verify_params_across_processes 检查所有 rank 的 model 是否一致。它做的事：

1. 计算每个 parameter 的 hash（shape + dtype + 头几个字节）
2. AllGather 所有 rank 的 hash
3. 检查每个位置是否一致

不一致的常见原因：

• 不同 rank 加载了不同 ckpt
• 用户在 DDP 之前手动改了某个 rank 的 model（如不同 seed 初始化）
• DDP 之前忘了 broadcast 初始权重

这套检查避免”不一致的 model 跑 DDP 训练”这种隐蔽 bug —— 用户能立刻收到 RuntimeError 而不是默默训练几小时后发现 loss 不下降。

17.8.19 性能调优数字

实测数字（H100，70B Llama 训练）：

优化叠加能让训练吞吐提升近 2x。生产代码每开一个 flag 都对应可量化的收益，理解每个 flag 的工程原理（前面各小节）让调优有据可依。

跨节点训练（如 32 卡 4 节点）通信占比通常更高（节点间带宽 < NVLink），HSDP（第 18 章）+ DDP 复合策略能让通信占比降回 20% 以内。

17.8.20 DDP forward 流程的精确分解

distributed.py:_pre_forward + forward + _post_forward 的完整链路：

1. pre forward：检查 grad 状态、可能触发 _sync_buffers 同步 BN running stats
2. 真实 forward：self.module(*inputs, **kwargs) 跑用户 model
3. _DDPSink.apply 包装输出：让反向能进入 DDP 控制流
4. 触发 prepare_for_backward：Reducer 重置 bucket pending 计数
5. post forward：处理 unused parameters detection（如有）

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

17.8.21 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 章 §17.8.7 的 Join API 复杂度 —— 数据天然均匀，不需要 Join。

17.8.22 _register_fused_optim：optimizer step 与 backward 融合

DDP 默认流程是”backward AllReduce 完成 → optimizer.step()“。但 step 之前 grad 已经在每个 rank 上 ready —— 如果 step 立即跑、不等其他 rank，可以与下一次 forward overlap。

distributed.py:_register_fused_optim（实验功能）让 DDP 在每个 bucket AllReduce 完成的瞬间立即更新对应参数：

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

每个 bucket 完成 AllReduce 后立即触发 optimizer.step 处理这个 bucket 的参数 —— 整个 step 与剩余 backward 重叠。理论上能再省 5-10% 训练时间。

实际生产应用少，因为：

• 与 LR scheduler 的协作复杂（schedule.step() 要在所有 step 之后）
• 与 GradScaler 的 inf 检测冲突（要等所有 grad 都 ready 才能确认 inf）
• 用户 optimizer 不一定支持”per-bucket update”

实验代码可见这套机制的潜力。生产代码用第 18 章 FSDP-2 的 optim_in_backward 是替代路径 —— 思想一致但兼容性更好。

17.8.23 PowerSGD：低秩压缩通信

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook 实现 PowerSGD（论文 ICML 2019）。算法核心：把 grad 分解成低秩矩阵 G ≈ P @ Q.T，只通信 P 和 Q（远小于 G）。

具体步骤（hook 内部）：

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，通信量减少 1024 倍！

代价：精度损失（低秩近似有误差）+ 计算开销（矩阵乘法）。实际收益取决于网络带宽与 GPU compute 的相对速度。在带宽受限场景（跨节点 100 Gbps ethernet）能让 DDP 训练加速 1.5-3x；在 NVLink 内带宽充足时反而拖慢。

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

17.8.24 NCCL_BLOCKING_WAIT vs ASYNC_ERROR_HANDLING

第 16 章 §16.7.6 提过 ASYNC_ERROR_HANDLING。还有个相关的 NCCL_BLOCKING_WAIT：

实战推荐：生产环境必开 TORCH_NCCL_ASYNC_ERROR_HANDLING=1。NCCL_BLOCKING_WAIT 是更老的机制，会让所有 NCCL 调用同步、性能下降。ASYNC_ERROR_HANDLING 是 v1.8+ 的现代方案 —— 同步成本接近零、出错快速失败。

如果你看到训练任务”hang 在某个 collective 几十分钟”，几乎肯定是没开 ASYNC_ERROR_HANDLING + 默认 NCCL timeout 30 分钟。

17.8.25 elastic 训练与 DDP

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

• 任一 rank 崩溃 → 整个 job 重启（从最近 ckpt 恢复）
• 节点 join / leave 时整个 group 重新 rendezvous
• 与 Kubernetes 等编排系统集成做自动恢复

elastic 不改变 DDP 的语义，只是让”训练 job 整体寿命”更长。生产 70B 训练动辄几天到几周，期间 GPU 故障 / 网络抖动是必然事件 —— 没有 elastic 重启机制就要人工介入恢复。

17.8.26 DDP 与 MoE 的特殊性

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

DDP 默认行为会卡在那些 expert 的 AllReduce 等不到（pending 计数永远不归零）。

解法：

• 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.8.27 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)

底层机制：SyncBN forward 时调 dist.all_reduce 收集所有 rank 的 sum / sum_squared，算出全局 mean / var，再用全局值做 normalization。代价是每个 BN forward 多两次小 collective（sum 和 sum_squared）。

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

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

17.8.28 DDP × activation_checkpoint 协作

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

use_reentrant=True（旧默认）的 checkpoint 用 autograd.Function，反向时重新跑一次 forward 取激活。重 forward 期间触发的算子调用经过 dispatcher → 经过 DDP 的 forward hook → DDP 误以为是新的 forward → 错乱。

use_reentrant=False（v2.0+ 新默认）用 saved_tensors_hooks 实现，不触发 forward hook → 与 DDP 兼容。

实战建议：新代码强制 use_reentrant=False。老代码遇到 “DDP 反向异常”看 checkpoint 的 use_reentrant 参数往往是根因。

第 18 章 §18.6.9 详细讲了 FSDP × checkpoint 的协作（更复杂，因为 FSDP 还要管 unshard）。DDP 这边相对简单 —— 只要 hook 不被重 forward 误触发就 OK。

17.8.29 find_unused_parameters 的实现机制

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

1. forward 完后，DDP 从 forward 输出反向 BFS（用 autograd 图）
2. 标记所有”通往输出路径上的”parameter 为 used
3. 没标记的 parameter 在 backward 不会有 grad
4. DDP 在 reducer 里把这些 parameter 的 bucket pending 直接当 ready
5. AllReduce 用 0 张量代替 grad（保证集合通信参与，但效果是 no-op）

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

• forward 后多一次 BFS（小开销）
• 每次反向都要重新分析（不能缓存，因为分支可能变）
• 失去 bucket overlap 的部分优化（pending 提前归零让 AllReduce 提前发，但顺序可能与 backward 计算不对齐）

实测开 find_unused_parameters 比关时慢 5-15%，只在真的需要时开。static_graph=True 与它互斥（static_graph 假设 graph 不变）。

17.8.30 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.8.31 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.6.9 简单提过，FSDP 提供 clip_grad_norm_ 自家版本处理这个。

17.8.32 DDP 在 PyTorch 历史上的演进

v0.4 引入 torch.nn.parallel.DataParallel（DP，单进程多线程）+ DistributedDataParallel（DDP，多进程）共存

v1.0-1.3 DDP 早期：bucket 概念引入、find_unused_parameters

v1.5-1.7 重大重构：reducer 移到 C++（性能 10x）、static_graph 引入

v1.10+ comm hook API + powerSGD 默认实现

v2.0+ 与 torch.compile 集成、_register_fused_optim 实验

v2.4+ FSDP-2 推出，DDP 在大模型场景被取代但仍是中小模型主选

理解这条演进让你看 DDP 源码时能识别”这段是 v1.x 老代码、那段是 v2.0+ 新功能”。reducer.cpp 几千行里有相当一部分是为兼容老 API 留的，不是核心算法。

17.9 几条工程经验

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

2. bucket_cap_mb 调优：默认 25MB 适合大多数场景。如果模型层非常多（如 70B 有几千层）但每层小，调到 50-100MB 减少 launch；如果每层就很大（如 attention QKV），降到 10MB 提高 overlap

3. 监控通信开销占比：用 profiler 看 NCCL kernel 占总时间比例。> 30% 说明通信瓶颈，考虑 FSDP / 加 IB 网卡

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.x 推荐先 DDP 再 compile，让 compile 看到 DDP 的反向 hook

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

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

17.9.5 DDP overlap 的实测分析

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

最后一个 bucket 无 overlap 空间：bucket 0（最靠近 loss 的 layer）的 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 等。

实测 typical 数字（H100, NVLink）：

• 完美 overlap 上限：100% 线性扩展（理论值）
• 实际 8 卡 DDP：90-95% 扩展效率
• 跨节点 32 卡：75-85%（跨节点带宽限制）

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

17.9.6 bucket_cap_mb 调优指南

不同模型规模的推荐值：

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

经验法则：单个 AllReduce 时间应该 ≈ 几个 layer 的 backward 时间，让 overlap 自然发生。70B 模型每 layer 几百 ms backward，bucket 应该 ≥ 几百 ms 通信对应的 size（按 NVLink 600GB/s 算约 100MB+）。

TORCH_LOGS=ddp_graphs 输出每个 bucket 的实际 size 与命中算子，调优时拿来对照。

17.9.7 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 还支持 min_nodes / max_nodes 弹性配置 —— 节点抖动时自动扩缩。这套机制让大规模训练在云上能从”运维噩梦”变成”自动恢复”，是生产 70B+ 训练的工程必备。

17.9.9 BarrierEvent vs Reducer 的 Work 句柄

DDP 反向触发 AllReduce 后，Reducer 持有 Work 对象（第 16 章 §16.5）跟踪每个 collective 完成状态。但 optimizer.step() 之前必须等所有 AllReduce 完成，否则 step 用的是部分 reduced 部分未 reduced 的 grad。

实现机制：DDP 在 reducer 里有一个 pending_works 列表，每次新 AllReduce 入列。finalize_backward（每次反向最后调）调每个 work.wait()。这套同步对用户透明，但理解后能解释为什么”loss.backward() 后 grad 一定是 reduced 的” —— DDP 已经在 backward 内部 wait 完所有 collective。

async_op 模式（DDP 内部用）让 wait 不立即阻塞 Python —— wait 排队进 CUDA stream，与下一次 forward kernel 同流水线。这是 PyTorch v1.x 末期引入的优化，让 backward 与下一次 forward 能 overlap。

17.10 跨书关联

• 第 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.9.8 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.10.4 DDP 与 NCCL 版本兼容性

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

• NCCL 2.x → 2.18：原生支持 ncclSend / ncclRecv 让 P2P 通信进入主流（pipeline parallel 受益）
• NCCL 2.20+：支持 ncclCommSplit（§16.7.8 提过的 split_group 底层）+ 高效 group-level abort
• NCCL 2.22+：原生 fp8 reduction 支持，让 fp8 训练梯度同步不需要 cast 回 fp16
• NCCL 2.24+：与 NVIDIA UCC 集成，让 cross-rack 通信走 InfiniBand RDMA 优化

PyTorch 与 NCCL 版本的对应关系（v2.11 默认带 NCCL 2.20+）通常自动维护。但生产环境手动升级 NCCL 时要小心 —— 太新的 NCCL 与某些 GPU driver 不匹配会让训练崩。pip install nvidia-nccl-cu12==2.20.5 这种显式锁版本是稳妥做法。

17.10.5 DDP 错误诊断速查表

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

把这套表打印贴在工位，遇到 DDP issue 时按表查能省 80% 调试时间。

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

把整章内容串起来，看 70B Llama 训练单 step 的时间分布（H100，8 卡 NVLink）：

| 阶段                            | 占比   | 时长   |
| forward (compute)              | 35%    | 700ms  |
| backward (compute)             | 50%    | 1000ms |
| AllReduce (overlap 部分)        | -      | -      |
| AllReduce (无法 overlap 部分)    | 8%     | 160ms  |
| optimizer.step                  | 5%     | 100ms  |
| dataloader / 杂项               | 2%     | 40ms   |
| 单 step 总时间                  | 100%   | 2000ms |

关键观察：

• backward + 部分 AllReduce overlap → AllReduce 剩 8% 不能 overlap 的部分主要是最后一个 bucket
• optimizer step 本身只占 5%，但 fused = True（第 10 章 §10.6.2）能让它从 5% 降到 1-2%
• DataLoader 可忽略（worker 池 + pin_memory 跑得过 GPU）

优化空间：能再压缩的部分主要是 AllReduce 的 8% —— 用 bf16 comm hook（§17.8.5）压一半带宽降到 4%、用 PowerSGD（§17.8.23）能压更狠。但每 100ms 的优化要付精度风险，是否值得看场景。

17.10.7 DDP 与第 8 章 Engine 的协作

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

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

这种”DDP 寄生在 Engine 之上”的设计让 Engine 与分布式训练完全解耦。Engine 改进（如 v2.4+ Compiled Autograd）DDP 自动受益、不需要适配。

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

17.11 设计启示

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 都切到多卡，用通信换显存。