PyTorch 训练框架内核深度解析

第16章进程组与 NCCL 通信

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

"ProcessGroup is a base class that captures collective and point to point communication in a fixed set of processes. ... Every function that performs I/O is executed asynchronously by a thread pool owned by the ProcessGroup (by default). They return an object that can be used to wait for completion or error."

—— torch/csrc/distributed/c10d/ProcessGroup.hpp 中 ProcessGroup 类上方的注释

本章要点

init_process_group(backend='nccl') 是分布式入口：每个进程调一次，通过 rendezvous 机制（如 TCPStore）找到所有 peer，NCCL 通信器默认在第一次 collective 时才建立（§16.23）
ProcessGroup 是对外的进程组对象：真正干活的 ProcessGroupNCCL（GPU）、ProcessGroupGloo（CPU）、ProcessGroupMPI 都继承自 Backend，按设备类型挂在 ProcessGroup 下面。GPU 训练基本都用 NCCL
6 种核心 collective：all_reduce、all_gather、reduce_scatter、broadcast、reduce、all_to_all
NCCL 算子也走 dispatcher：注册成 c10d::allreduce_ 等 ATen 算子，让 profiler / torch.compile 等基础设施能看到通信
Work 对象是异步通信句柄：work = all_reduce(tensor, async_op=True) 立即返回，work.wait() 等完成
barrier() 让所有进程同步：常用于训练开始 / 结束的对齐点

16.1 分布式训练的核心问题

多卡训练有几种基本模式：

Data Parallel (DDP)：每张卡有完整模型，每个 batch 切到不同卡，反向后 AllReduce 平均梯度
Tensor Parallel：把单个张量切到多卡（如 attention 的 QKV 投影按列切）
Pipeline Parallel：把模型按层切到多卡，流水线执行
Sharded Data Parallel (FSDP)：参数 / 梯度 / optimizer state 都切到多卡（ZeRO-3）

无论哪种模式，核心都是"多 GPU 之间高效通信"。这一章拆通信底座；第 17 章 DDP、第 18 章 FSDP 在它之上展开。

16.2 init_process_group：握手与 rendezvous

每个分布式进程的第一行：

import os
import torch.distributed as dist

dist.init_process_group(
    backend='nccl',
    init_method='env://',     # 从环境变量读 master_addr / master_port
    world_size=8,
    rank=int(os.environ['RANK']),
)

torch/distributed/distributed_c10d.py:1605 的 init_process_group 做的事：

解析 init_method（tcp://、env://、file:// 三种）
创建 Store 对象（TCPStore / FileStore），所有进程通过它做"first contact"
创建 ProcessGroup 实例，并按设备类型挂上后端（ProcessGroupNCCL 或 ProcessGroupGloo）；NCCL 之后建通信器时，由 rank 0 生成 ncclUniqueId 写进 Store、其他 rank 从 Store 读取（ProcessGroupNCCL::broadcastUniqueNCCLID），peer 之间的地址交换由 NCCL 自己的 bootstrap 完成
默认 lazy：不传 device_id 时这一步推迟到第一次 collective（§16.23）
设置全局默认 group，后续 dist.all_reduce(...) 用它

Store 是分布式协调的"小型 KV 数据库"。它不传训练数据，只用于握手 + barrier 等小信息。生产里用 TCPStore（基于 TCP 的简易实现）。

16.3 ProcessGroup 与 Backend：两层结构

v2.11 里用户拿到的 ProcessGroup（torch/csrc/distributed/c10d/ProcessGroup.hpp）是一层外壳，具体后端都继承自 Backend（torch/csrc/distributed/c10d/Backend.hpp），按设备类型登记在 ProcessGroup 里（如 cpu → Gloo、cuda → NCCL）：

后端	实现类（`: public Backend`）	特点
nccl	`ProcessGroupNCCL`	NVIDIA NCCL，GPU 上的最优选择
gloo	`ProcessGroupGloo`	Facebook 自家库，CPU + GPU 都行，CPU 集合通信
mpi	`ProcessGroupMPI`	OpenMPI / MPICH 接入，HPC 场景
ucc	`ProcessGroupUCC`	NVIDIA UCC，下一代统一通信

两层的接口大致（简化自 ProcessGroup.hpp / Backend.hpp）：

class ProcessGroup : public torch::CustomClassHolder {
  // 不是纯虚函数：默认实现去 dispatcher 里找 c10d::allreduce_ 并调用
  virtual c10::intrusive_ptr<Work> allreduce(
      std::vector<at::Tensor>& tensors,
      const AllreduceOptions& opts = AllreduceOptions());
  // ... allgather / broadcast / reduce / scatter / gather / send / recv 同理
};

class Backend : public torch::CustomClassHolder {
  // 默认实现直接报 "Backend xxx does not support allreduce"，由子类覆写
  virtual c10::intrusive_ptr<Work> allreduce(
      std::vector<at::Tensor>& tensors,
      const AllreduceOptions& opts = AllreduceOptions());
};

c10d 算子在 CPU / CUDA 等 dispatch key 上的 kernel（torch/csrc/distributed/c10d/Ops.cpp）再调 process_group->getBackend(DEV)->allreduce(...)，落到具体后端。每个方法返回 Work 对象（异步句柄）。同步语义靠 work->wait() 触发；async 模式下立即返回让 CPU 继续。在 CPU 上用 gloo 实跑，type(dist.group.WORLD) 是 ProcessGroup，dist.group.WORLD._get_backend(torch.device("cpu")) 才是 ProcessGroupGloo。

16.4 6 种核心 collective 操作

graph TB
    subgraph AllReduce["all_reduce: 所有 rank 都拿到 sum/avg/max..."]
        A1[rank 0: 1 2 3] --> AR[result: 10 14 18]
        A2[rank 1: 4 5 6] --> AR
        A3[rank 2: 5 7 9] --> AR
        AR --> B1[rank 0: 10 14 18]
        AR --> B2[rank 1: 10 14 18]
        AR --> B3[rank 2: 10 14 18]
    end

Collective	输入	输出	用途
`all_reduce`	每 rank 一个 tensor	每 rank 都拿到 reduce 结果	DDP 梯度同步
`all_gather`	每 rank 一个 tensor	每 rank 拿到所有 rank tensor 的拼接	FSDP unshard 参数
`reduce_scatter`	每 rank 一个 tensor	每 rank 拿到 reduce 结果的 1/N 切片	FSDP 分片梯度
`broadcast`	一个 src rank 有 tensor	所有 rank 拿到副本	模型初始化广播 weight
`reduce`	每 rank 一个 tensor	一个 dst rank 拿到 reduce 结果	收集训练 metrics
`all_to_all`	每 rank N 个分片	每 rank 拿到来自所有 rank 的对应分片	MoE 路由、Ulysses 式序列并行

all_reduce 是最常用的（DDP 的核心）。all_gather 与 reduce_scatter 在 FSDP / ZeRO 里大量用。第 17 / 18 章会展开各自场景。

16.5 NCCL 算子也是 ATen 算子

PyTorch 为了让 collective 与 autograd / torch.compile 协作，把 NCCL 操作注册成 ATen 算子：

// torch/csrc/distributed/c10d/Ops.cpp:25 (TORCH_LIBRARY(c10d, m) 里)
m.def(
    "allreduce_(Tensor[] tensors, "
    "__torch__.torch.classes.c10d.ProcessGroup process_group, "
    "__torch__.torch.classes.c10d.ReduceOp reduce_op, "
    "Tensor? sparse_indices, bool async_op=True, int timeout=-1) "
    "-> (Tensor[], __torch__.torch.classes.c10d.Work)");

这意味着 dist.all_reduce(tensor) 在内部走的是 dispatcher（第 5 章）—— 经过和普通 ATen 算子一样的 dispatch 路径，能被 torch.compile 看到。

这套设计带来几个好处：

torch.compile 能追踪 collective：collective 作为算子进图，编译器可以重排它与计算的相对位置来争取 overlap（§14.30 的 reorder pass）
Profiler 能记录 collective：自动出现在 chrome trace 里
dispatcher 可拦截：用户写 mode 监控所有 collective

注意 c10d::allreduce_ 这类原地算子没有注册 autograd 公式：在 CPU 上用 gloo 实跑，y = x * 2; dist.all_reduce(y) 之后 y.grad_fn 仍是 MulBackward0，通信不进反向图。可微的 collective 是 §16.12 的 functional collectives（*_autograd 版本）。

第 17 章 DDP 章会看到反向 AllReduce 是怎么自动触发的 —— 是 DDP 挂在梯度累加上的 autograd hook 在反向时调 ProcessGroup::allreduce，再经 dispatcher 走到 c10d::allreduce_。

16.6 Work 对象与异步通信

work = dist.all_reduce(tensor, async_op=True) 立即返回。tensor 此时还没真正被 reduce，但 CUDA stream 上已经排了一个 NCCL kernel。work.wait() 让当前 stream 等 NCCL 完成。

Async 模式让 communication 与 computation overlap：

work = dist.all_reduce(grad, async_op=True)
# 在 work 完成前 CPU 已经继续做其他事
do_other_work()
work.wait()    # 必要时再 sync

这种 overlap 是 DDP / FSDP 性能的核心。如果每个 collective 都同步等，多卡训练的通信开销会让吞吐崩盘。第 17 章会展开 DDP 怎么用 backward hook 触发 async AllReduce + 与下一层 backward 计算 overlap。

16.7 NCCL 之外的几个细节

ProcessGroup 不是单例：可以创建多个 group。最常用的子 group：

# 数据并行 + 张量并行混合时, 创建两个子 group
dp_group = dist.new_group(ranks=[0, 4, 8, 12])    # 数据并行组
tp_group = dist.new_group(ranks=[0, 1, 2, 3])     # 张量并行组

不同 group 在不同 NCCL 通信器上跑，互不干扰。这是 hybrid parallel（同时 DP + TP + PP）的基础。

barrier() 用于强同步：

dist.barrier()      # 所有 rank 等到这一行才继续

底层就是这样（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:5196 的 ProcessGroupNCCL::barrier）：在选定的 device 上建一个 at::zeros({1}, kFloat) 张量、对它做一次 allreduce，同步模式下再 currentStream.synchronize() 做一次 CPU 侧等待。常用于：

训练开始前等所有 rank 完成 dataloader 初始化
训练结束前对齐 final metrics
ckpt 保存前确保所有 rank 数据一致

但 barrier() 在生产代码里慎用 —— 它强制 GPU 同步，破坏 communication / computation overlap。每个 barrier 等于损失一次 overlap 机会。

16.8 ProcessGroupNCCL 内部：communicator 缓存

ProcessGroupNCCL.cpp 6036 行的核心数据结构是 NCCL communicator 缓存。NCCL 创建 communicator (ncclCommInitRank) 是重量级操作 —— 要跨进程握手、分配 GPU buffer。如果每次 collective 都新建一个 comm，开销无法接受。

PyTorch 按 设备（点对点按 rank 对）缓存：

// ProcessGroupNCCL.hpp:1348
std::unordered_map<std::string, std::shared_ptr<NCCLComm>> devNCCLCommMap_;
// collective 的 key 是设备号字符串（getKeyFromDevice，如 "0"）；
// 点对点的 key 是两端 rank 组成的 "低:高"（如 "1:2"）

第一次某个 ProcessGroup 在某个设备上做 collective 时，初始化一个 NCCLComm 缓存到 map。后续命中直接复用。

这套缓存让"在同一个 group 上反复做 collective"只有第一次付握手代价。但代价是 NCCLComm 对象长期占 GPU 内存，子组越多（如 3D parallel）累计占用越可观，具体大小取决于 NCCL 版本与配置。

abortCommsFromMap()（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1397）在 group 被 abort 时遍历 map，对每个缓存的 communicator 调 abort()（内部走 ncclCommAbort）。

16.9 Heartbeat / Watchdog：异常感知

ProcessGroupNCCL 有两条后台线程，构造时创建（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:970-972），类定义在 torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:609（HeartbeatMonitor）与 :689（Watchdog）：

Watchdog 线程轮询每个未完成的 WorkNCCL：collective 出错或超过进程组 timeout（NCCL 默认 10 分钟，torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:155 的 kProcessGroupNCCLDefaultTimeout）时，通过 Store 通知其他 rank 做 Flight Recorder dump（§16.28），再按 TORCH_NCCL_ASYNC_ERROR_HANDLING 的模式处理——v2.11 默认值是 3（SkipCleanUp：不清理 communicator、直接拉掉进程），1 是 TearDown（先 abort communicator 再拉掉进程）
HeartbeatMonitor 线程盯的是 Watchdog 自己的心跳：Watchdog 每轮循环递增心跳计数，监控线程按 heartbeatTimeoutInSec_ 判超时，默认 8 分钟（torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1709：getCvarInt(TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC, 60 * 8)）
心跳停了 → 认为 Watchdog 卡死（报错信息里写的典型原因是 CudaEventDestroy 之类 NCCL/CUDA API 调用挂住）
触发 dumpDebuggingInfo 把 Flight Recorder 记录、栈信息 dump 出来
最后 terminateProcess 用 LOG(FATAL) 调 std::abort() 直接结束进程（TORCH_NCCL_ENABLE_MONITORING=0 可关掉这一步，默认开）

这套机制让分布式训练能"快速失败" —— 一个 rank 卡住后，进程在 timeout 后被主动拉掉，交给外层 elastic 重启（§16.17），而不是无限期 hang 下去。DDP 视角的讨论见第 17 章 §17.8。

死锁的 dump 文件含每个 rank 在哪个 collective 上等，是诊断"hang 在 backward 哪一层"的金钥匙。

16.10 abort 与 recovery

NCCL 失败有几种：

collective timeout（某 rank 跑得太慢）
网络断（IB / ethernet 拔了）
GPU 故障（ECC error）

PyTorch 处理：

检测到失败 → WorkNCCL::abort() (torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:860) → 调 ncclComm_->abort() 让所有该 comm 的 op 立即返回错误
用户代码看到异常 + NCCL 错误信息（默认的 error handling 模式下进程随后被拉掉，见 §16.9）
ncclCommAbort 之后这个 comm 不能再用 —— 必须创建新 comm 或重新 init_process_group

c10d 本身不做"自动 recovery"：abortCommsFromMap 旁的注释写明，aborted 的 communicator 留在缓存里不删，由应用自己销毁并重建进程组来恢复。生产代码靠 elastic 训练（torchelastic）：失败后整个 job 重启，从 ckpt 恢复继续训练。这条限制让大规模训练对存储 ckpt 频次有要求 —— 失败重启的代价 = 上次 ckpt 到现在的所有计算。

16.11 split_group / new_subgroups：子组管理

torch/distributed/distributed_c10d.py:5198 的 split_group 是 v2.5 加入的 API，让 ProcessGroup 在子组上分裂（签名是 split_group(parent_pg=None, split_ranks=None, timeout=None, pg_options=None, group_desc=None)）：

# world_size=8 拆成 2 个 4 卡子组
# 前提：默认进程组用 init_process_group(..., device_id=...) 绑定了设备
sub_group = torch.distributed.split_group(parent_pg=None, split_ranks=[[0,1,2,3], [4,5,6,7]])

new_group 与 split_group 都要求父组里所有 rank 以相同参数调用。区别在底层：split_group 从父组已建好的 NCCL communicator 上做 ncclCommSplit，省掉一次全新的 bootstrap；new_group 则为子组另建通信器。split_group 在 v2.11 仍标为 experimental，只支持 NCCL 和自定义插件后端，且默认进程组必须已绑定设备 —— 在 CPU 上用 gloo 实跑，它直接抛 RuntimeError: No device associated with the default pg, not safe to split any process groups。

new_subgroups_by_enumeration（torch/distributed/distributed_c10d.py:5689）批量创建子组，返回 (当前 rank 所在子组, 全部子组列表)：

# 一次创建 4 个子组: [0,1] [2,3] [4,5] [6,7]
cur_group, sub_groups = torch.distributed.new_subgroups_by_enumeration(
    ranks_per_subgroup_list=[[0,1],[2,3],[4,5],[6,7]]
)

混合并行（DP + TP）经常需要十几个子组（同一 rank 同时属于 dp_group / tp_group / pp_group）。这些 API 让构造更直接。DeviceMesh（第 18 章 §18.7）是更高级的封装，把多个 group 表示成一个 mesh 对象。

16.12 Functional Collectives：torch.compile 友好的新 API

torch.distributed._functional_collectives（1714 行，v2.1 起随发行版提供）是新一代集合通信 API，与传统 dist.all_reduce 等的根本区别是 返回新 tensor 而非 inplace 修改：

# 老 API: inplace 修改 tensor, 返回 Work 句柄
work = dist.all_reduce(tensor, async_op=True)
work.wait()
# tensor 被修改

# 新 API: 返回新 tensor (是个 AsyncCollectiveTensor)
import torch.distributed._functional_collectives as funcol
new_tensor = funcol.all_reduce(tensor, "sum", group)
# new_tensor 是 AsyncCollectiveTensor wrapper, 用到时自动 wait

为什么这是 torch.compile 友好？因为 inplace 操作让 functionalize（第 13 章 §13.4）很难处理 —— 编译器不知道 collective 把 tensor 改了。新 API 全程纯函数式，collective 在图里就是普通算子，Inductor 可以调整它与计算的先后顺序来争取 overlap（§14.30）。

AsyncCollectiveTensor 是懒同步张量：它包着 collective 的输出张量，但不立即 wait；下次张量被算子使用时（如 + 1）触发 unwrap 并 wait。这套 lazy wait 让 communication / computation overlap 自动发生，不需要用户手写 async_op + wait。

API 列表（_functional_collectives.py）：

函数	行号	对应老 API
`all_reduce`	155	`dist.all_reduce`
`all_gather_tensor`	177	`dist.all_gather_into_tensor`
`all_gather_tensor_autograd`	218	+ autograd 支持的版本（反向自动加 reduce_scatter）
`reduce_scatter_tensor`	251	`dist.reduce_scatter_tensor`
`all_to_all_single`	457	`dist.all_to_all_single`
`all_reduce_coalesced`	334	多 tensor 一次 reduce

*_autograd 后缀的版本支持 collective 可微 —— 反向时自动产生对应的反向 collective（all_gather 的反向是 reduce_scatter、reduce_scatter 的反向是 all_gather）。这让 Tensor Parallel 训练里 collective 也能参与 autograd，反向链不断。

DTensor（第 18 章 §18.8）内部就是用 functional collectives 实现 placement 转换。v2.11 里 dist.all_reduce 等传统 API 并没有被标记为弃用，两套 API 并存：要进 torch.compile 图、要可微的场景用 funcol，普通 eager 代码用传统 API 即可。

16.13 TCPStore 协议：rendezvous 怎么真正做

init_method='env://' 的背后是 TCPStore。它是个迷你 KV server，所有 rank 通过它做"first contact"。源码 torch/csrc/distributed/c10d/TCPStore.cpp。

启动顺序：

sequenceDiagram
    participant R0 as rank 0 (master)
    participant R1 as rank 1
    participant R2 as rank 2
    participant R3 as rank 3
    R0->>R0: 启动 TCPStore server (listen on master_port)
    R1->>R0: connect to master_addr:master_port
    R2->>R0: connect
    R3->>R0: connect
    Note over R0,R3: 所有 rank 连上 master
    Note over R0,R3: 第一次 collective 时（lazy init）
    R0->>R0: ncclGetUniqueId() + store.set("0", id)
    R1->>R0: store.get("0")（阻塞到 key 出现）
    R2->>R0: store.get("0")
    R3->>R0: store.get("0")
    Note over R0,R3: 所有 rank 拿到同一个 ncclUniqueId, 调 ncclCommInitRank<br/>peer 地址交换由 NCCL 自己的 bootstrap 完成

图里的 key 是 ProcessGroupNCCL::broadcastUniqueNCCLID 用的通信器计数（std::to_string(ncclCommCounter_++)），实际还带着各进程组的 PrefixStore 前缀；点对点通信则用 "低 rank:高 rank" 作 key。

支持的 KV 操作：

set(key, value)：写入
get(key)：读取（阻塞等到 key 存在）
add(key, delta)：原子加法（用于 barrier 计数）
wait(keys)：等多个 key 同时存在
compare_set(key, expected, desired)：CAS 操作

实现细节：

协议二进制：每条请求以 1 字节 QueryType 开头，后面跟长度前缀的 key / value（torch/csrc/distributed/c10d/TCPStore.cpp 的 SendBuffer）；客户端建连后先发一次带 magic number 的 VALIDATE 再 PING，没有用 HTTP 等高级协议
libuv 后端：TCPStoreLibUvBackend.cpp（1594 行）是 server 端的实现，v2.4 起成为默认（TCPStoreOptions::useLibUV = true，env:// 也可用 USE_LIBUV=0 退回旧实现），比早期纯 socket 的实现更能扛住大集群同时握手
超时机制：每个 wait 操作支持独立超时，避免无限阻塞
建连重试：客户端建连 / 校验失败时按指数退避重试，直到 store 的 timeout 用完

专门做一个 store 而不用 Redis / etcd 的好处是：rendezvous 阶段不需要用户先搭外部服务。TCPStore.cpp 本体只有 760 行，够用即可。

FileStore（基于所有进程都能访问的文件）是另一种实现，常见于单机或有共享文件系统的场景。HashStore 是进程内的内存版本，多用于测试。生产几乎都用 TCPStore。

16.14 NCCL 内部算法：Ring vs Tree vs CollNet

NCCL 的 AllReduce 不是简单的"全发全收"。NCCL 内部有几种算法，不同拓扑下选择不同：

Ring AllReduce（最常用）：

graph LR
    R0((rank 0)) -->|chunk 0| R1((rank 1))
    R1 -->|chunk 1| R2((rank 2))
    R2 -->|chunk 2| R3((rank 3))
    R3 -->|chunk 3| R0

每个 rank 把自己的 tensor 切成 N 块，每步发送一块给下一个 rank、从前一个 rank 收到一块。先走 N-1 步 reduce-scatter（每个 rank 攒齐一块的规约结果），再走 N-1 步 all-gather，共 2(N-1) 步后所有 rank 都拿到完整 reduce 结果。

复杂度：每 rank 收发 2(N-1) 块、即 2(N-1)/N 倍张量大小的数据量，N 变大时趋近 2 倍、与 N 无关，带宽利用接近最优
缺点：延迟与 N 线性相关，大集群不友好

Tree AllReduce：

构造 binary tree，叶子向上 reduce 到 root，root 再向下 broadcast。延迟 = O(log N)，但 bandwidth 利用率较低。

NCCL 默认根据 message size + topology 自动选：小 message 用 tree（延迟敏感）、大 message 用 ring（带宽敏感）。NCCL_ALGO=Ring 强制用 ring，调试时可用。

CollNet：

借助交换机上的 SHARP（NVIDIA 的 in-network compute）做 in-network reduction —— reduce 的一部分算力从 GPU 挪到网络里，减少 GPU 之间要搬的数据量。需要 IB 硬件 + SHARP 组件支持。

实战看法：日常训练让 NCCL 自动选；性能调优时 NCCL_DEBUG=INFO 看选了哪个 algo + 实际 bandwidth、再决定是否手动 force。

16.15 send / recv：点对点通信

除了 collective，PyTorch 也暴露点对点 API：

# 阻塞版本
dist.send(tensor, dst=1)        # rank 0
dist.recv(tensor, src=0)        # rank 1

# 异步版本
work = dist.isend(tensor, dst=1)
work = dist.irecv(tensor, src=0)
work.wait()

底层走 ncclSend / ncclRecv（NCCL 2.7+ 支持），两个 rank 在同一 stream 里点对点传输。Pipeline Parallel 训练核心 —— stage i 把激活值发给 stage i+1：

# Pipeline Parallel 简化版
def stage_forward(x, stage_id):
    if stage_id == 0:
        x = layer1(x)
        dist.send(x, dst=1)            # 发给 stage 1
    elif stage_id == 1:
        x = torch.empty_like(...)
        dist.recv(x, src=0)            # 从 stage 0 收
        x = layer2(x)
        dist.send(x, dst=2)
    ...

batch_isend_irecv 一次批量提交多个点对点操作，让 NCCL 内部并发：

ops = [
    dist.P2POp(dist.isend, x_to_next, next_rank),
    dist.P2POp(dist.irecv, x_from_prev, prev_rank),
]
works = dist.batch_isend_irecv(ops)
for w in works:
    w.wait()

这套 API 是 Megatron-LM、DeepSpeed Pipeline 训练的基础。理解它让你看清 Pipeline Parallel 的"激活值传递"不是魔法 —— 就是 isend/irecv 配合 stream 同步。

16.16 GPU Direct RDMA 与跨节点拓扑

跨节点 collective 的最大开销是"GPU 显存 → CPU 内存 → 网卡"的拷贝。NCCL 用 GPU Direct RDMA (GDR) 跳过 CPU 直接 DMA：

graph LR
    G1[GPU 1<br/>node A] --gdr.--> N1[NIC<br/>node A]
    N1 --IB/RoCE--> N2[NIC<br/>node B]
    N2 --gdr.--> G2[GPU 2<br/>node B]
    style G1 fill:#dcfce7
    style G2 fill:#dcfce7
    style N1 fill:#fef3c7
    style N2 fill:#fef3c7

要求：

IB / RoCE 网络（不是普通 ethernet）
nvidia-peermem 或 GPU Direct module 加载
GPU 与 NIC 在同一 PCIe root（否则跨 root 拷贝）

GDR 省掉的是"显存 → 主机内存 → 网卡"这一趟额外拷贝，跨节点带宽因此能贴近网卡的线速而不是被 PCIe/主机内存拖住。这是大模型训练能扩到几千卡的硬件基础（具体数字取决于网卡代次与拓扑，用 nccl-tests 的 all_reduce_perf 在自己集群上量）。

诊断：NCCL_DEBUG=INFO 的建连日志形如 Channel 00/0 : 0[0] -> 8[0] [send] via NET/IB/0/GDRDMA，带 /GDRDMA 后缀表示 GDR 生效；只有 via NET/IB/0 而没有这个后缀，则数据经主机内存中转。这是大规模训练性能调优的常见诊断点。

PyTorch 这层不做 GDR 配置，全靠 NCCL + driver。但工程师必须懂这个 stack —— 看 NCCL log 不懂等于黑盒。

16.17 TorchElastic：失败恢复机制

PyTorch 不支持 ProcessGroup recovery（§16.10），那大模型训练几千卡几小时一次失败怎么办？答案是 torchrun + TorchElastic（v1.9+ 内置）。

机制：

torchrun --rdzv_backend=c10d --rdzv_endpoint=master:29400 启动训练
内置 elastic agent 监控所有 worker 进程
某 worker 失败 / 节点死 → agent 检测到 → 杀掉所有 worker、所有节点重启 worker
worker 重启后从最近 ckpt 加载 → 继续训练

关键设计："快速失败 + 快速重启"。不试图救活单个 worker（NCCL 限制），而是把失败成本降到最低（重启时间 + ckpt 回滚的步数）。

rdzv_backend=c10d 用 PyTorch 自家的 rendezvous 协议（基于 TCPStore）；rdzv_backend=etcd 用 etcd（适合超大集群）。前者足够大多数场景。

实战配置：

torchrun \
  --nnodes=8 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  --rdzv_id=$JOB_ID \
  --max_restarts=5 \
  train.py

max_restarts=5 让 elastic agent 最多重启整个 job 5 次。生产代码里这个数字根据训练时长 + 历史失败率调（如训练 7 天历史每天平均失败 2 次 → max_restarts=20）。

理解这套"failure recovery 在外层做、内层不管"的工程思路：c10d 这一层只负责快速失败，恢复交给外层 elastic 一刀切地处理。

16.18 collective 吞吐随 message size 的变化

NCCL collective 的延迟 / 吞吐与 message size 强相关，曲线形状对任何集群都成立（绝对值必须在自己的集群上用 NVIDIA 官方的 nccl-tests（all_reduce_perf -b 8 -e 4G -f 2 -g 8）实测，本机无 GPU，本节不给数）：

message size 区间	主导因素	有效带宽
很小（KB 量级）	延迟主导 —— 每次 collective 的固定 startup cost 摊不开	远低于链路上限
中等（MB 量级）	进入 ring algorithm 的高效区	快速爬升
很大（几百 MB 以上）	带宽主导	饱和到链路上限（NVLink / IB 的理论值打个折）

工程意义：

DDP gradient bucketing：把多个小 grad tensor 合并成一个 bucket 再 AllReduce（第 17 章，默认 25 MB）—— 让每次 AllReduce 都落在高效区
FSDP all_gather chunk：太小延迟主导、太大显存压力，中间取值是经验最优
避免大量极小的 collective：如果你训练里有这种调用，bundle 成大 collective 收益明显

跨节点走 IB 时曲线形状相同，只是饱和值更低（IB 带宽通常低于机内 NVLink）。理解这条曲线让你看 DDP / FSDP 的设计选择不是凭空的，是基于硬件性能曲线的精算。

16.19 NCCL 调优环境变量大全

实战调 NCCL 性能必备的环境变量：

变量	作用	常用值
`NCCL_DEBUG`	日志详细度	INFO（看每次 collective）/ TRACE（更细）
`NCCL_DEBUG_SUBSYS`	限制日志子系统	INIT,COLL,GRAPH（不要打 ALL，太吵）
`NCCL_SOCKET_IFNAME`	选择网卡	eth0 / ens6 / ib0
`NCCL_IB_HCA`	InfiniBand HCA 选择	mlx5_0:1（指定第一个 IB 端口）
`NCCL_IB_DISABLE`	关闭 IB 走 socket（调试用）	0（默认走 IB）/ 1（强制 socket）
`NCCL_P2P_DISABLE`	关闭 GPU 间 P2P	0 / 1（PCIe 跨 root 时调试）
`NCCL_NTHREADS`	NCCL 内部线程数	256（多线程加速 collective）
`NCCL_BUFFSIZE`	内部缓冲区大小	4194304（4 MB，大 collective 调大）
`NCCL_ALGO`	强制算法	Ring / Tree / CollNet（autotuning fail 时用）
`NCCL_PROTO`	强制协议	Simple / LL / LL128（小 message 用 LL 更快）
`NCCL_MAX_NCHANNELS`	最大 channel 数	16-32（更多 channel 让多 GPU 间并发）
`TORCH_NCCL_ASYNC_ERROR_HANDLING`	PyTorch 侧：watchdog 发现 collective 出错 / 超时后怎么处理	v2.11 默认 3（直接拉掉进程）；1 = 先 abort communicator 再拉掉进程

（注意最后一项是 PyTorch 的环境变量，不是 NCCL 的；collective 超时也是在 PyTorch 这层设的 —— init_process_group(timeout=...)。不传时按后端取默认值：NCCL 是 10 分钟（torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:155 的 kProcessGroupNCCLDefaultTimeout），Gloo 等其他后端是 30 分钟（torch/csrc/distributed/c10d/ProcessGroup.hpp:21 的 kProcessGroupDefaultTimeout），分流逻辑在 torch/distributed/distributed_c10d.py 的 _get_default_timeout。）

实战调优顺序：

第一步开 NCCL_DEBUG=INFO 看 NCCL 选了什么算法、哪个网卡、bandwidth 多少
如果 bandwidth 远低于硬件理论值：检查 NCCL_IB_HCA / NCCL_SOCKET_IFNAME 是否选对网卡
小 collective 多场景：开 NCCL_PROTO=LL 让小 message 走低延迟协议
大 collective 主导场景：调 NCCL_BUFFSIZE=8388608 + NCCL_MAX_NCHANNELS=32 提升吞吐
跨节点拓扑诡异：开 NCCL_TOPO_DUMP_FILE=/tmp/nccl_topo.xml dump 拓扑文件，肉眼检查

理解这些 flag 让你有杠杆调 NCCL 性能，而非只能 "不行就重启"。

16.20 monitored_barrier：可调试的同步点

普通 dist.barrier() 不告诉你"哪个 rank 卡了"。dist.monitored_barrier（v1.9 起提供）解决这个：

from datetime import timedelta

try:
    dist.monitored_barrier(timeout=timedelta(minutes=10))
except RuntimeError as e:
    # 报错信息含哪些 rank 没到达 barrier
    print(e)

实现：只有 Gloo 后端支持（ProcessGroupGloo 里实现，其他后端直接抛 ValueError: monitored_barrier is only implemented for GLOO backend.），内部用 send / recv 让每个 rank 通知 rank 0 "我到了"。rank 0 等待所有 rank、超时后报出没到的 rank。在 CPU 上 4 个进程实跑、让 rank 3 不进 barrier、timeout 设 3 秒，rank 0 抛出的 RuntimeError 开头是 [Rank 0]: Rank 3 failed to pass monitoredBarrier in 3000 ms。它比普通 barrier 多了逐 rank 收发，但调试 hang 时极其有用。

实战：NCCL 训练要用它，得另建一个 gloo 子组（dist.new_group(backend="gloo")）传给 group=。长跑训练偶尔 hang，先在 epoch 边界加 monitored_barrier，下次 hang 时立刻能看到"rank 47 没到达"——直接登 rank 47 节点 py-spy 看 stack。比干等 NCCL collective 超时再抛错高效得多。

但 monitored_barrier 不能放训练 hot path（每 step 都加 = 性能崩盘）。生产建议放在 epoch / checkpoint / 异常恢复等低频节点。

16.21 ProcessGroupGloo：CPU 通信与调试用途

NCCL 是 GPU 王者，但CPU 训练 / 调试场景用 Gloo：

dist.init_process_group(backend='gloo')   # 不需要 GPU

何时用 Gloo：

单机多进程 CPU 训练：没 GPU 时唯一选择
调试分布式 bug：在没有 GPU 集群的笔记本上跑 4 进程模拟分布式
CPU tensor 的 collective：metrics 收集等场景，避免来回 cuda → cpu 拷贝

Gloo 走的是 CPU + 主机网络这条路，没有 NVLink / GDR 可用，GPU 上的吞吐与 NCCL 完全不是一个量级。绝不要在生产 GPU 训练用 Gloo。

特殊用途：hybrid backend。某些 PyTorch 代码同时用 NCCL（GPU collective）+ Gloo（CPU collective）。init_process_group(backend='cpu:gloo,cuda:nccl') 创建混合后端。但实战很少这么用，绝大多数用统一 NCCL。

16.22 NCCL 性能在 chrome trace 里的可视化

torch.profiler 自动捕获 NCCL collective 到 chrome trace：

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
) as prof:
    for batch in loader:
        ...
        prof.step()

打开 tensorboard --logdir log 或 chrome://tracing 加载 trace.json，能看到：

NCCL kernel：在 GPU stream 那一行显示为 kernel 条
NCCL launch：在 CPU thread 上对应 launch 时间
collective metadata：tensor size、rank 列表、reduce_op 都在 tooltip 里
同步点：work.wait() 显示为 stream sync 标记

调优时看的关键：

collective 与 computation overlap：理想情况是 NCCL kernel 与 forward / backward kernel 并发执行（不同 stream）
collective 自身吞吐：message_size / time = 实际 bandwidth，与 §16.18 基准对比
rank 间不平衡：看不同 rank 的 trace 拼起来，发现某 rank 总是慢几 ms 是诊断 stragglers 关键

第 21 章 profiler 章会展开。理解 NCCL 在 trace 里长什么样，能很快判断"通信是不是瓶颈"。

16.23 NCCL communicator lazy init

§16.8 提到 communicator 缓存，但第一次 comm 什么时候创建？默认是 lazy init：init_process_group 只创建 ProcessGroup 对象 + Store 握手，不立即创建 NCCL communicator。传了 device_id 时则会绑定设备、立即建好通信器（eager init），§16.11 的 split_group 就依赖这一点。

第一次某个 collective 在某个 device set 上调用时才触发：

// 简化版 ProcessGroupNCCL::allreduce
auto comm = getNCCLComm(devices);    // 从缓存找, 没找到才 ncclCommInitRank
ncclAllReduce(input, output, count, ..., comm, stream);

为什么 lazy？因为：

多子组场景：用户创建 8 个 sub_group 但实际只用 4 个 → eager init 浪费 4 个 comm 的握手时间 + GPU 显存
错误恢复：comm 创建失败时（如网络配置错），训练第一个 collective 才报错而不是 init 时报错 —— 让用户能在更具体上下文里看错
CUDA context 时机：comm init 要 CUDA context，lazy 让用户先设好 device 再 init

代价：第一次 collective 明显更慢（要握手）。生产里这只是一次性开销，摊到整个训练上影响可忽略。

理解 lazy init 让你看到偶发"训练第一个 step 慢"的原因不是模型，是 NCCL 第一次握手。dist.barrier() 在训练开始前调一次能强制 NCCL init 完成 —— 让正式训练第一步快。

16.24 nccl init 失败诊断流程

NCCL init 失败是分布式训练最常见的入门坑。报错典型：ncclInternalError、ncclInvalidArgument、ncclSystemError、ncclRemoteError。诊断流程：

flowchart TD
    Err[ncclXXXError]
    Err --> A{INFO log 看选了什么<br/>NIC / GPU / algo}
    A --> B1{NIC 是预期的吗?}
    B1 -->|否| F1[NCCL_SOCKET_IFNAME=eth0<br/>强制选对 NIC]
    B1 -->|是| C1{GPU 通讯走 NVLink 吗?}
    C1 -->|否| F2[nvidia-smi topo -m<br/>看物理拓扑]
    C1 -->|是| C2{所有 rank 都连上了?}
    C2 -->|否| F3[检查防火墙 / port<br/>master_addr 路由通否]
    C2 -->|是| C3{IB 配置对吗?}
    C3 -->|否| F4[NCCL_IB_HCA=mlx5_0:1<br/>NCCL_IB_GID_INDEX=3]
    C3 -->|是| C4[NCCL_DEBUG=TRACE<br/>看更细日志]
    style Err fill:#fee2e2
    style F1 fill:#dcfce7
    style F2 fill:#dcfce7
    style F3 fill:#dcfce7
    style F4 fill:#dcfce7

诊断标准动作清单：

nvidia-smi：所有节点都看到 GPU 吗？驱动版本一致吗？
nvidia-smi topo -m：GPU 之间的拓扑是 NVLink 还是 PCIe？跨 NUMA 吗？
ibstat（IB 集群）：每个 HCA 都 Active 吗？
ifconfig / ip addr：能看到预期网卡吗？master_addr 在哪个网络？
跨节点 ping：master_addr 各节点都能 ping 通吗？
netstat -an | grep <master_port>：master 节点 port 是否真的 listen 了？

多数"NCCL init 失败"都能在这份清单里找到原因。理解这个 flow 让你不会盲目改代码、瞎试 NCCL flag。

16.25 DeviceMesh：把多个 PG 抽象成网格

第 18 章 FSDP 章会展开 DeviceMesh，这里先建立桥梁。torch/distributed/device_mesh.py 的 DeviceMesh 类是"多维 ProcessGroup 容器"（v2.0 时放在私有的 torch.distributed._tensor 包里，v2.2 起移到公开路径 torch/distributed/device_mesh.py）。

为什么需要它？hybrid parallel（DP × TP × PP）创建几十个子组、每个 rank 同时属于多个组、用户要管理一堆 group 对象 —— 代码丑且容易错。

DeviceMesh 把这些组织成多维网格：

import torch
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh

# 8 卡, 2 维网格: dp=2, tp=4
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

# 拿子 group
dp_group = mesh.get_group("dp")    # rank 0 上拿到的是 [0,4] 这个组
tp_group = mesh.get_group("tp")    # rank 0 上拿到的是 [0,1,2,3] 这个组

底层实现：DeviceMesh 内部用一个 torch.Tensor 记录 rank 排布（mesh.mesh 返回的就是 Tensor），每个 cell 是 rank id。子组在构造 mesh 时（_init_process_groups）就按每个维度切出的 rank list 建好：默认进程组已绑定设备时走 split_group，否则逐个调 new_group；get_group(dim_name) 只是按名字取回已建好的组。

DeviceMesh 还提供：

mesh.get_local_rank("dp")：当前进程在 dp 维上的索引
mesh.size(1) / mesh["tp"].size()：tp 维的大小（size() 只收整数维度下标，传 "tp" 会抛 TypeError）
mesh.get_coordinate()：当前 rank 在 mesh 中的多维坐标

DTensor（第 18 章）的所有 placement 操作都基于 mesh，FSDP-2 / TP 的入口也接收 mesh 参数。

理解 mesh 是这一章 → 18 章的桥梁。本章讲的"ProcessGroup 是通信原语"，下一层 mesh 把多个 PG 组织成"逻辑视角"，让分布式代码可读性大幅提升。

16.26 ProcessGroup 在 PT2 时代的演进

PT2（v2.0+）时代 ProcessGroup 的几次大改造：

方向	代表产物（v2.11 树上的位置）	意义
functional collectives（§16.12）	`torch/distributed/_functional_collectives.py`（1714 行）	编译友好，collective 能进图
TCPStore 的 libuv 后端	`torch/csrc/distributed/c10d/TCPStoreLibUvBackend.cpp`（1594 行）	大集群启动更稳
DeviceMesh	`torch/distributed/device_mesh.py`	hybrid parallel 标准 API
Heartbeat / Watchdog	`torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp:609` 的 `HeartbeatMonitor`	生产可观测
`split_group`	`torch/distributed/distributed_c10d.py:5198`	从父 communicator 分裂子组，省一次 bootstrap
Flight Recorder	`torch/csrc/distributed/c10d/FlightRecorder.{hpp,cpp}` / `torch/csrc/distributed/c10d/FlightRecorderCuda.cpp`	出错 / 超时时 dump 最近的 collective 记录
symmetric memory	`torch/csrc/distributed/c10d/symm_mem/`	NVLink switch 集群优化

整体趋势：

可编译性：从 inplace API 转向 functional API，让 collective 进 graph
可观测性：从"NCCL 黑盒 timeout"转向"详细日志 + 状态 dump"
可恢复性：从"挂了重启"转向"快速失败 + elastic 恢复"
可扩展性：从"手写一堆 group"转向"DeviceMesh 多维抽象"

理解这条演进，能看清 PyTorch 原生分布式栈（FSDP-2、DTensor、DeviceMesh）为什么能承担越来越多原本由外部框架自己实现的工作 —— 底座是这套 PG 抽象。

16.27 一段总结：把 PG 看成"分布式 dispatcher"

把全章串起来：ProcessGroup 在分布式训练里的地位 = dispatcher 在单机训练里的地位。

单机	分布式
dispatcher 把 ATen op 路由到 backend	ProcessGroup 把 collective 路由到通信库
多个 backend（CPU、CUDA、XPU、MPS）	多个后端（NCCL、Gloo、MPI、UCC）
dispatcher 是抽象基类	ProcessGroup 是抽象基类
op 注册让算子统一接口	collective 注册成 ATen op 让接口统一
autograd 集成	functional collectives 让 autograd 能跟踪
compile 友好	functional collectives 让 collective 能进 compile 的图

这种"dispatcher 思路在分布式延伸"的设计是 PyTorch 工程哲学的体现 —— 抽象层一致，让上层代码无差异写。理解了这个类比，你看 c10d 源码不会再觉得"通信库怎么这么复杂" —— 它就是另一套 dispatcher。

16.28 NCCL Flight Recorder：生产诊断的瑞士军刀

NCCL Flight Recorder 是分布式训练 hang / crash 的核心诊断工具。它的环形 trace buffer 在 v2.2 已随 TORCH_NCCL_TRACE_BUFFER_SIZE 出现（当时在 TraceUtils.h），v2.6 起实现拆到独立的 torch/csrc/distributed/c10d/FlightRecorder.hpp 与 .cpp，ProcessGroupNCCL.cpp 通过 FlightRecorderCUDA::get() 在每次 collective 调用前后写入 entry（v2.11 里在 ProcessGroupNCCL.cpp 的 :398 reset_all / :411 dump / :711 getEntry / :2424 retire_id 等处调用）。机制：每个 rank 维护一个循环 buffer记录最近 N 次 collective 的状态：

record entry:
  - collective_id (递增序号)
  - collective_type (allreduce / allgather / ...)
  - input_size, output_size
  - device, stream
  - state (scheduled / started / completed)
  - start_time, end_time
  - traceback (Python stack 截图)

v2.11 里它默认就是开的，下面几个变量写出来的都是默认值（TORCH_NCCL_TRACE_BUFFER_SIZE 设成 0 才会关掉）：

export TORCH_NCCL_TRACE_BUFFER_SIZE=2000   # 保留最近 2000 条（默认值）
export TORCH_NCCL_DUMP_ON_TIMEOUT=1        # 超时 / 出错时 dump（默认开）

某个 rank 的 Watchdog 发现超时或出错时，会经 Store 通知其他 rank，各 rank 各自把自己的 buffer dump 到磁盘。默认路径是 $XDG_CACHE_HOME/torch/（未设时为 ~/.cache/torch/）下的 comm_lib_trace_rank_<rank>，可用 TORCH_NCCL_DEBUG_INFO_TEMP_FILE 改前缀（torch/csrc/distributed/c10d/FlightRecorder.cpp）：

comm_lib_trace_rank_0
comm_lib_trace_rank_1
...

文件内容是 pickle 序列化的记录（FlightRecorder::dump 最后 pickle_str），含每个 collective 的 state。把所有 rank 的 trace 摆在一起对比，能看出：

哪个 collective rank A 等到了 rank B 没到
rank B 的 trace 显示卡在哪个 collective
collective ID 不一致暴露 "rank 间 collective 顺序不同" bug

典型用法：训练 hang 住 → flight recorder dump 出所有 rank 的 buffer → 对齐各 rank 的 collective 序号，找出"多数 rank 停在序号 N，某个 rank 已经走到 N+1"的那个 rank → 顺着它的 traceback 找出该 rank 走了不同代码分支的原因。

这是 v2.x 分布式调试的一项重要工程进步。没有 flight recorder 时同样的 bug 要靠 py-spy 逐 rank 看 stack，排查费时得多。理解它让你在生产 hang 时第一时间想到查 flight recorder log。

16.29 NCCL kernel 的 stream 模型

NCCL 不是"同步 API"，而是把 kernel launch 到 stream 上：

ncclAllReduce(input, output, count, dtype, op, comm, stream)

最后一个参数 stream 是 cudaStream_t。NCCL 内部 launch 一个 GPU kernel 到这个 stream。返回时不等 kernel 完成，仅完成 launch（几 us）。

这种设计让 NCCL 与 PyTorch 的 stream 模型无缝集成：

async_op=False（默认）时，ProcessGroupNCCL 直接把 NCCL kernel 发到调用方的当前 stream 上
async_op=True 时改用 ProcessGroupNCCL 自己的内部 NCCL stream，先让它等当前 stream 上的输入就绪，从而能与后续计算并发（ProcessGroupNCCL.cpp 的 collective() 里 "in asyncOp=false [default] mode, we use currentStream as ncclStream" 那段）
WorkNCCL::wait() 默认不阻塞 CPU：它让当前 stream 等 NCCL 的结束 event（§16.36）；只有开了 TORCH_NCCL_BLOCKING_WAIT 或传了 timeout 时才在 CPU 上轮询等待

并发例子：

# stream A: 计算
with torch.cuda.stream(stream_A):
    out1 = model.forward(input1)

# stream B: 通信（与 stream A 并发）
with torch.cuda.stream(stream_B):
    work = dist.all_reduce(grad_buffer, async_op=True)

# 同步点
torch.cuda.synchronize()    # 等所有 stream 完成

DDP 内部就用这套：backward hook 触发的 AllReduce 以 async 方式发出、跑在 ProcessGroupNCCL 的内部 NCCL stream 上，与后续 backward 计算（在计算 stream 上）并发。第 17 章会展开。

用户侧控制 stream 的正规方式是 with torch.cuda.stream(s): 加上 Stream.wait_stream / Tensor.record_stream；生产代码很少手动控制，默认行为就够用。

16.30 collective 与 sparse tensor 的协作

PyTorch 支持 sparse tensor（COO / CSR），但 NCCL 的集合通信只处理稠密缓冲区。sparse 张量调 dist.all_reduce 时，dispatcher 按 SparseCPU / SparseCUDA key 另走一条注册（torch/csrc/distributed/c10d/Ops.cpp：SparseCPU 进后端的 allreduce，由后端按 layout 分流；SparseCUDA 进后端的 allreduce_sparse），各后端做法不同：

Gloo：AsyncSparseAllreduceWork（torch/csrc/distributed/c10d/ProcessGroupGlooDetail.hpp）先 allgather 各 rank 的元数据，再分别 allgather indices 和 values，最后在本地求和，结果仍是 sparse。CPU 上 2 进程实跑，各 rank 贡献一个非零元，all_reduce 后得到含两个非零元的 sparse 张量
NCCL：标准构建里 ProcessGroupNCCL::allreduce_sparse 直接报错 NCCL does not support all_reduce with sparse tensors. Please use dense tensors instead.；只有 NCCLX 构建（IS_NCCLX）才走一条"按完整 shape 分配稠密输出、做完再 to_sparse()"的路径

所以在 GPU 上要对 sparse 梯度做规约，常见做法是自己先 to_dense() 再 all_reduce。代价：dense 版本要按完整 shape 分配显存，与稀疏度无关。一个 95% 稀疏的张量，稀疏表示只存 5% 的 (index, value)，转 dense 后却要占满整个 shape —— 越稀疏，这条路径的浪费越大。只在张量本身不大时才合理。生产场景里 sparse collective 不常用，但这条机制说明了"后端不支持的类型"是怎么在 PyTorch 这层分流的。

16.31 monitor heartbeat 配置实战

§16.9 提了 Watchdog 与 HeartbeatMonitor，这里给具体配置 cheat sheet（v2.11 里多数项的默认值已经是"开"，写出来是为了知道能调什么）：

# 错误处理模式（默认 3：直接拉掉进程；1：先 abort communicator 再拉掉进程）
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

# Watchdog / dump
export TORCH_NCCL_DUMP_ON_TIMEOUT=1          # 超时 / 出错时 dump flight recorder trace（默认开）
# 注意：TORCH_NCCL_BLOCKING_WAIT=1 会让 wait() 在 CPU 上阻塞并自己检查 timeout，
#       同时不再启动 watchdog 线程 —— 与上面这套机制二选一，一般不开

# Heartbeat
export TORCH_NCCL_ENABLE_MONITORING=1        # 默认开
export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=600  # watchdog 无心跳多久后拉掉进程 (不设时默认 8 分钟)
export TORCH_NCCL_DESYNC_DEBUG=1             # 超时时额外分析是哪些 rank 没跟上（默认关）

# Trace buffer
export TORCH_NCCL_TRACE_BUFFER_SIZE=2000     # 默认值
export TORCH_NCCL_DEBUG_INFO_TEMP_FILE=/tmp/nccl_dump_rank_
export TORCH_NCCL_DEBUG_INFO_PIPE_FILE=/tmp/nccl_pipe_   # 往这个管道写东西可手动触发 dump

（v2.11 里避免 record_stream 已是默认行为，再设 TORCH_NCCL_AVOID_RECORD_STREAMS=1 只会得到一条 deprecated 警告，见 torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:931-933。）

这一组环境变量让分布式训练生产可观测。某个 rank 卡住时：

其他 rank 的 watchdog 在 collective 超过进程组 timeout（NCCL 默认 10 分钟）时发现
经 Store 通知各 rank，各自 dump 到 /tmp/nccl_dump_rank_X
按错误处理模式决定是否先 abort NCCL comm
进程被拉掉退出
TorchElastic 检测到 worker 死亡 → 重启整个 job
工程师看 dump 文件诊断根因

如果连 watchdog 自己都卡在 NCCL/CUDA API 调用里，就轮到 heartbeat 监控线程在 TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC 之后兜底 dump 并结束进程。整套机制让大规模训练能在出错后快速失败并重启，而不是无限期 hang 住。生产工程师必须懂这套配置。

16.32 千卡级训练的工程问题

8-128 卡集群与 1024+ 卡集群有本质差别。后者的特殊问题：

1. rendezvous 不能简单 N 个进程都连一个 master

1024 个 rank 连同一个 TCPStore master → master 要同时应付上千个连接和大量 set/get 请求，启动阶段容易成为瓶颈。v2.4 起默认的 libuv server（§16.13）就是为扛住这种并发握手而来；更大规模时还要靠减少启动期的 store 往返次数。

2. 跨节点 collective 拓扑感知

千卡集群通常是"node 内 NVLink + 跨 node IB"两级网络。Naive ring algorithm 让一个 chunk 跨多次 node 边界，浪费 IB 带宽。hierarchical AllReduce：

第 1 步：node 内 reduce_scatter（NVLink 高带宽）
第 2 步：跨 node AllReduce（每 node 只发 1/8 数据，节省带宽）
第 3 步：node 内 all_gather（NVLink）

NCCL 会根据探测到的拓扑自动选择合适的算法（NCCL_ALGO 可取 Ring / Tree / NVLS 等）。前提是 NCCL 看到的拓扑正确，虚拟化等环境下可能需要用 NCCL_TOPO_FILE=/path/topo.xml 手动提供；拓扑错配会明显拉低吞吐。

3. failure rate 累积

单卡的故障概率再低，乘上几千张卡、再乘上几周的训练时长，"整个 job 全程无故障"的概率就低得没法指望。规模越大，平均无故障时间越短是数学必然。生产必须靠 elastic + 足够频繁的 ckpt —— ckpt 间隔的下界由"重跑这段的代价"和"存 ckpt 本身的开销"两头夹出来。

4. straggler effect

同步训练里每一步都由最慢的那个 rank 决定。某个 rank 每步慢一点（disk I/O 抖动、温度降频），这点延迟会乘上总步数变成可观的墙钟时间，而且卡数越多、出现至少一个慢 rank 的概率越高。解法：straggler detection，找出持续偏慢的 rank、隔离它（重启该 worker）让 elastic 接管。

5. NCCL communicator 累积

1024 卡 + 几个子组（DP、TP、PP）→ 每 rank 持有多个 NCCL comm，每个都占一份显存缓冲，累加起来不可忽视。要监控 + 必要时主动销毁不用的 group。

理解千卡训练的特殊问题让你看 Megatron / DeepSpeed 这类框架为什么要做大量"看似冗余"的优化 —— 在 8 卡环境下确实多余，但 1024 卡环境下每个都是必须。

16.33 NCCL 算法选择的内部决策

NCCL 在每次 collective 调用前自动选 algo + protocol。它内部是按拓扑、带宽、延迟参数给每种 (algo, protocol) 组合估算耗时、取最小者；下图只是帮助建立直觉的示意，分界点并非 NCCL 源码里的固定常量，实际选择以 NCCL_DEBUG=INFO 的输出为准：

flowchart TD
    Call[collective 调用]
    Call --> Size{message size}
    Size -->|< 32 KB| Small[small msg path]
    Size -->|32 KB - 4 MB| Med[medium]
    Size -->|> 4 MB| Large[large]
    Small --> SP{topology}
    SP -->|单 node NVLink| L1[Tree + LL128]
    SP -->|跨 node IB| L2[Tree + LL]
    Med --> MP{topology}
    MP -->|单 node NVLink| L3[Ring + LL128]
    MP -->|跨 node IB| L4[Ring + Simple]
    Large --> LP{topology}
    LP -->|集群有 SHARP| L5[CollNet + Simple]
    LP -->|普通 IB| L6[Ring + Simple]
    LP -->|单 node| L7[Ring + LL128]

各 protocol 的取舍：

LL (Low Latency)：每 8 字节里 4 字节数据 + 4 字节 flag，接收方看 flag 就知道数据是否到齐，不用额外同步；有效载荷只有一半。适合小 message
LL128：进化版，把 128 字节切成 120 valid + 8 marker，平衡延迟与吞吐
Simple：标准 send/recv，适合大 message（吞吐主导）

如果发现 NCCL 选错 algo，可以强制：

NCCL_ALGO=Ring NCCL_PROTO=LL128 python train.py

但慎用：NCCL 自动选的通常是最优的，强制可能反而慢。生产里仅在"NCCL_DEBUG=INFO 显示选了反直觉 algo + 实测性能差"时强制。

理解这套决策让你看 NCCL 内部不是黑盒，而是一个按 size × topology × protocol 估算代价的模型。

16.34 fp16 / bf16 collective 的注意事项

混合精度训练里 gradient 是 fp16 / bf16，AllReduce 也走低精度。但有两个坑：

1. fp16 求和溢出

1024 个 rank 各 reduce 一个 fp16 grad。fp16 max 约 65504，1024 个 ~100 的 grad 求和约 102400，超出 → 出 inf。

解法：pre_div_factor——每 rank 先除 N（grad / 1024）再 AllReduce + sum。等价数学但中间值不溢出。

# 示意伪代码（真实的 dist.all_reduce 是原地修改、不返回张量）
# 错误（fp16 易溢）
grad_global = all_reduce(grad, ReduceOp.SUM) / 1024

# 正确（pre_div）
grad_local = grad / 1024
grad_global = all_reduce(grad_local, ReduceOp.SUM)

ReduceOp.AVG（torch/csrc/distributed/c10d/Types.hpp:38）让后端直接做平均，省掉手写的除法。

2. bf16 精度损失累积

bf16 的尾数只有 7 位，reduce 过程中每一次加法都在这个精度上舍入，参与的 rank 越多、累积误差越大。

解法：在 fp32 上做 reduce——把梯度先 cast 到 fp32、AllReduce 完再 cast 回来。PyTorch 这层没有一个叫 TORCH_NCCL_FP32_REDUCE 的开关；实际做法是在 DDP 的 comm hook（第 17 章 §17.9）或 FSDP 的 MixedPrecision(reduce_dtype=torch.float32) 里指定 reduce 用的 dtype。

实战影响：低精度梯度在大集群上 reduce，误差会随参与 rank 数累积。用 reduce_dtype=float32 是最直接的规避手段，代价是通信量翻倍。

理解这两个坑让你看 DDP / FSDP 源码里的"莫名 cast"不奇怪 —— 都是为了避开数值问题。

16.35 跨数据中心训练：地理分布式新场景

大模型训练的规模不断逼近单一数据中心（DC）的供电与机房容量，跨 DC 训练开始成为讨论的方向。这层网络与 DC 内网络的差异是定性的：

维度	单 DC（IB / RoCE）	跨 DC（WAN）
带宽	高，每卡配专用高速网卡	低得多，且多个作业共享
延迟	微秒级	毫秒级（受物理距离限制）
抖动	小	大
丢包率	几乎 0	偶发

NCCL 的设计前提是 IB / NVLink 这类低延迟网络，直接跨 DC 跑同步 collective 性能会大幅下降。常被讨论的思路：

1. Hierarchical 训练

DC 内做 DP / TP（NCCL 高带宽），DC 间做 PP（低频通信，激活值传递可承受高延迟）。让大延迟通信只发生在低频路径上。

2. Pipeline Parallel + WAN

PP 第 i stage 把 activation 发给 stage i+1。如果两个 stage 在不同 DC，每个 micro-batch 前向时跨 DC 发一次 activation、反向时跨 DC 回传一次梯度。PP 的跨 stage 流量是点对点、量相对小，因此常被看作最适合放在跨 DC 链路上的并行维度。

3. 异步训练

不再要求所有 rank 同步。某个 DC 训练慢 → 不等它 → 用 stale 梯度。代价：收敛变慢、超参更难调。这条路线在学术上有大量工作（如各类 local SGD / DiLoCo 变体），工业界的具体采用情况多数没有公开材料，不宜按传闻当事实写。

4. 通信压缩

跨 DC 带宽贵 → 压缩 grad 减少传输量（bf16 / int8 / 稀疏化等，压缩比越高、精度损失通常越大）。PyTorch 这层的接口在 torch.distributed.algorithms.ddp_comm_hooks（第 17 章 §17.9）。

init_process_group 本来就支持 init_method='tcp://...' 并允许把 timeout 调到很大，跨 DC 的 rendezvous 因此不成问题；真正的瓶颈在 NCCL 本身仍假设低延迟、低丢包的网络。跨 DC 的通信层还在演进中。

理解这层让你看到分布式训练可能的下一个 frontier 不在单卡 / 单集群，而在跨 DC（这是趋势判断，不是既成事实）。

16.36 Work、cudaEvent 与 stream 同步的精细协作

WorkNCCL 的 wait() 内部不是 cudaStreamSynchronize，而是配合 cudaEvent 做精确的"何时完成"判定：

// 简化自 torch/csrc/distributed/c10d/ProcessGroupNCCL.{hpp,cpp} 的 WorkNCCL
class WorkNCCL : public Work {
    // NCCL kernel 发出后在 NCCL stream 上 record 的结束 event
    std::shared_ptr<at::cuda::CUDAEvent> ncclEndEvent_;

    void synchronize() {       // wait() 先调它
        synchronizeStream();
    }

    void synchronizeStream() {
        // 不阻塞 CPU：让调用方的当前 stream 在 NCCL kernel 之后
        // 才开始执行后续 op（CUDAEvent::block 即 cudaStreamWaitEvent）
        auto currentStream = at::cuda::getCurrentCUDAStream(device_.index());
        ncclEndEvent_->block(currentStream);
    }
};

只有开了 TORCH_NCCL_BLOCKING_WAIT、给 wait() 传了 timeout，或者是 barrier 的 Work 时，wait() 才会在 CPU 上等 kernel 真正完成。

这种 event-based 设计让一个 stream 的 NCCL collective 可以与另一个 stream 的 computation 精确同步：

# stream_compute 上跑计算
with torch.cuda.stream(stream_compute):
    out = forward(x)

# stream_comm 上跑 collective, 但要等 forward 完成
with torch.cuda.stream(stream_comm):
    # ProcessGroupNCCL 只会让 NCCL stream 等"当前 stream"（这里是 stream_comm），
    # 不知道 stream_compute 的存在 —— 跨 stream 的依赖要自己声明
    stream_comm.wait_stream(stream_compute)
    work = dist.all_reduce(grad, async_op=True)

理解 events / streams 的协作让你看 DDP / FSDP 源码里大量"record_stream / wait_event / current_stream" 调用不困惑 —— 都是为了让 collective 能与 computation 真正并发，又不丢数据依赖关系。

torch.cuda.current_stream(device).wait_stream(other_stream)（v1.0 就已提供）是用户层 API，封装了 event-based 同步。生产代码里几乎都用这层、不直接 cudaEvent。

16.37 自定义 backend 注册：第三方通信库的接入方式

PyTorch 不只支持 NCCL / Gloo / MPI，第三方可以注册自家通信库。torch/distributed/distributed_c10d.py 的 Backend 类支持注册：

import torch.distributed as dist

# 第三方库实现 c10d::Backend 的 C++ 子类后
dist.Backend.register_backend(
    "my_backend",
    create_my_backend,                # 工厂函数, 返回后端（Backend）实例
    extended_api=True,
    devices=["cuda", "cpu"],
)

# 用户用法
dist.init_process_group(backend="my_backend", ...)

实战例子：

Habana HCCL：Intel Gaudi 加速器的通信库，以 hccl 后端名接入
Ascend HCCL：华为昇腾（torch_npu）的通信库，同样以 hccl 接入

要区分的是另一类"换掉 NCCL 下层"的方案：AWS 的 aws-ofi-nccl 是 NCCL 的网络插件（让 NCCL 走 EFA），AMD ROCm 版 PyTorch 在 nccl 这个后端名下用的是 RCCL —— 它们都不经过 register_backend，PyTorch 这层照旧是 ProcessGroupNCCL。

硬件厂商通过这个机制接入 PyTorch，用户代码基本只需改 backend="..." 一个字符串。这种"接口稳、后端可换"的设计是 §16.27 dispatcher 类比的具体体现。

理解这套注册机制让你看到 PyTorch 分布式层没有绑死 NVIDIA，而是在工程上为多硬件留好了接入点。

16.38 几条工程经验

1. NCCL 环境变量：NCCL_DEBUG=INFO 看 NCCL 通信细节；NCCL_SOCKET_IFNAME=eth0 指定网卡；NCCL_IB_DISABLE=1 关闭 InfiniBand 走 socket（调试用）

2. init_method='env://' + torchrun：标准的多进程启动方式。不要手写 mp.spawn

3. timeout 按需调大：NCCL 默认 10 分钟（Gloo 30 分钟），大模型训练里 rank 0 单独加载 ckpt、其他 rank 在 collective 上等的情况可能更久。init_process_group(timeout=timedelta(hours=2))

4. 跨节点通信看 NCCL_IB_HCA / NCCL_IB_GID_INDEX：物理网络配置不对会让 collective 从 IB 退到普通 TCP socket，吞吐掉一个量级 —— NCCL_DEBUG=INFO 里能直接看出走的是哪条路

5. NCCL 死锁的常见原因：

不同 rank 调 collective 顺序不一致（如 rank 0 先 AllReduce 再 Broadcast，rank 1 反过来）
某个 rank 异常退出，其他 rank 卡在 AllReduce 等

6. 别把 TORCH_NCCL_ASYNC_ERROR_HANDLING 设成 0：v2.11 默认值 3 已经会在 NCCL 出错 / 超时后拉掉进程；设成 0（NoHandling）则 watchdog 发现错误后不再主动处理，进程可能卡在出错的 collective 上

7. 避免 CPU 上的 collective：CPU AllReduce 走 Gloo，比 NCCL 慢得多。把 metrics 张量先 .cuda() 再 collective

8. 子 group 创建有顺序约束：所有 rank 必须以同样顺序、同样参数创建子 group。否则各 rank 的 group ID 对不上

16.39 跨专栏关联

《vLLM 推理内核深度解析》专栏第 14 章分布式：vLLM 的 TP 分片用 ProcessGroup AllReduce 同步 attention 输出
第 5 章 dispatcher：理解 collective 也走 dispatcher 让 autograd / compile 集成成为可能
《Tokio 源码深度解析》专栏第 2 章 Future 与 poll 模型：Work 对象与 Tokio Future 思想一致 —— 表示一个尚未完成的异步操作

16.40 设计启示

ProcessGroup 设计的核心思想：

第一：collective 即算子：把通信作为 ATen 算子注册让它享受 PyTorch 生态（compile、profiler；可微的版本由 functional collectives 提供）。这是与早期 MPI 风格"通信和计算分离"完全不同的思路

第二：Work 对象表示异步操作：让 communication 与 computation overlap 成为标准模式，而非 advanced 优化

第三：多后端共享接口：让 NCCL / Gloo / MPI 等通信库可互换。研究新算法时切到 Gloo 调试，生产切回 NCCL

第四：子 group 让 hybrid parallel 简单：不需要新概念，复用 group 抽象就能表达 DP/TP/PP 混合

第五：生产可观测性是分布式系统的基本要求：Heartbeat、Watchdog、Flight Recorder 这些机制看似与"通信原语"无关，但缺了它们千卡训练根本不能稳定跑。这是 PyTorch 团队从研究框架走向工业框架的一个标志性转变 —— 把"出问题怎么诊断"做成一等公民

下一章拆 DDP —— 看 PyTorch 怎么用 backward hook + 梯度桶 + AllReduce 实现"几乎线性扩展"的多卡训练，看本章这套底座是怎么被上层数据并行训练真正用起来的，从抽象通信原语到具体训练加速的最后一公里怎么落地。