从零实现集合通信原语

问题

在 N 个 rank 上做朴素 allreduce，需要把张量发送 N 次到根节点，再广播 N 次回来。每个 rank 的带宽按 O(N) 扩展，根节点成为瓶颈，墙钟时间下限是最慢链路乘以 N。Ring allreduce 把它摊平成 2(N-1) 个大小为 T/N 的块，因此每个 rank 的字节数降到 2T(N-1)/N，与集群规模无关。Tree allreduce 在小 N 和高延迟链路上占优，因为深度是 \log_2(N) 跳而不是 2(N-1) 跳。为集群形状选错拓扑，最慢的 GPU 就会决定每步时间。

本路线中你会读到的每个分布式训练框架都依赖这四个原语。PyTorch DDP 对每个参数 bucket 做一次 allreduce 来同步梯度。ZeRO 用 reduce_scatter 分片优化器状态，用 allgather 广播更新后的参数。FSDP 把完整的 forward 变成 allgather 加 reduce_scatter。Pipeline 并行需要 broadcast 在阶段组之间传递激活。如果你无法实现这四种集合操作，就无法推理训练为何停滞、为什么第 3 个 rank 出现梯度不匹配，或者为什么换拓扑后 pipeline 气泡翻倍。

概念

flowchart LR
  R0[rank 0] --> Q01[(queue 0 to 1)]
  Q01 --> R1[rank 1]
  R1 --> Q12[(queue 1 to 2)]
  Q12 --> R2[rank 2]
  R2 --> Q23[(queue 2 to 3)]
  Q23 --> R3[rank 3]
  R3 --> Q30[(queue 3 to 0)]
  Q30 --> R0

两趟 ring allreduce

把张量切成 N 个等大的块，索引 0..N-1。每个 rank 拥有索引等于其 rank 的块。第 1 趟 reduce-scatter 运行 N-1 步。在第 s 步，rank r 把块 (r - s) mod N 发给 rank (r + 1) mod N，并从 rank (r - 1) mod N 接收块 (r - s - 1) mod N，把收到的块累加到本地副本。N-1 步之后，rank r 拥有块 r 的完整和。第 2 趟 allgather 再运行 N-1 步，让完成的块沿环旋转，直到每个 rank 持有每个块的完整和。

用 Queue 网格替代 NCCL

NCCL 运行在 PCIe 和 NVLink 上，带有硬件卸载的归约。在 CPU 上你没有这些。每条环边的 multiprocessing.Queue 提供有序的点对点交付，只有一个生产者和一个消费者。归约发生在用户空间，因此要付出 Python 开销，但线路模式与 NCCL ring allreduce 完全相同。只要在 queue 版本上把正确性推理清楚，集群行为自然随之而来。

对照 gloo 验证

每个原语都配一个单元测试，在相同 world size 下，用 gloo 后端初始化的 torch.distributed 对相同张量运行，并比较输出。如果你的 ring allreduce 与 gloo 的偏差超过 float32 epsilon，测试就失败。对照参考实现验证不可妥协；没有它，原语看起来正确，直到真实训练的第 10000 步才暴露问题。

构建它

code/main.py 实现：

• Mesh 类：把 N 个 multiprocessing.Queue 实例连成一个环，并为每个 rank 暴露 send(dst, tensor) 和 recv(src)。
• ring_allreduce(mesh, rank, world_size, tensor)：运行两趟算法。
• broadcast(mesh, rank, world_size, tensor, src)：在对数树上广播。
• allgather(mesh, rank, world_size, tensor)：使用 N-1 次旋转。
• reduce_scatter(mesh, rank, world_size, tensor)：作为 allreduce 的前半部分。
• _gloo_reference(op, world_size, tensor)：把相同输入通过 gloo 后端的 torch.distributed 运行，用于逐字节比较。

运行：

bash复制python3 code/main.py

输出：逐原语的验证表，比较 queue 网格与 gloo 的输出，随后是每个 rank 的字节计数器，证明 2T(N-1)/N 的扩展关系。

生产环境中的模式

有三种模式能让原语达到可交付的强度。

在 allreduce 之前对梯度分桶。 一个 10 亿参数的模型有数万个梯度张量。每个张量一次 allreduce 要付 N 次延迟下限。DDP 把梯度分桶到约 25 MB 的块，每个桶发一次 allreduce；小张量搭大张量的便车。不分桶的话，延迟开销会主导整个 step。

让通信与计算重叠。 反向传播按逆序逐层计算梯度。最后一层的梯度一就绪，就启动它的 allreduce，同时下一层继续计算。PyTorch DDP 用 bucket-ready 钩子实现这种重叠。当网络有余量时，重叠能把可见通信时间减半。

按消息大小选 ring 或 tree，而不是按信仰。 NCCL 附带拓扑检测器：消息超过约 1 MB 用 ring，低于则用 tree。交叉点由带宽与延迟的权衡决定：超过 1 MB 时，带宽项 2T(N-1)/N 占主导，ring 胜出；低于 1 MB 时，log2(N) 的跳数胜出。硬编码一种拓扑会在错误的消息大小上损失吞吐量。

使用它

生产实现：

• PyTorch DDP。 在 backward 之后对分桶的梯度调用 dist.all_reduce。桶大小可调；默认 25 MB 对 100Gbit 以太网是合理值。
• DeepSpeed ZeRO。 发 reduce_scatter 来分片梯度，发 allgather 在 forward 前重建完整参数。本课的原语正是 ZeRO 发出的调用。
• FSDP。 forward 以 allgather 开始，unshard 该层，计算后用 reduce_scatter 归约并丢弃 unshard。同样的原语，不同的调度。

交付它

在第 77 到 81 课中使用这个 queue 网格原语。第 77 课把 allreduce 接进 DDP。第 78 课把 reduce_scatter 接进 ZeRO。第 79 课把 broadcast 接进 pipeline 激活。第 81 课把四者组合成端到端演示。

练习

1. 添加 tree allreduce 变体，并按消息大小在 ring 和 tree 之间切换。测量交叉点。
2. 添加 recv_timeout_ms，让停滞的 rank 抛出截止时间错误，而不是永远挂起。
3. 用 TCP socket 替换 multiprocessing.Queue 实现四个原语。同样的测试，真实的线路。
4. 添加带宽插桩钩子，把每个 rank 的字节计数写入 JSONL。
5. 在 4 个 rank 上比较 ring 与 tree 对 1KB、1MB、16MB 张量的墙钟时间。用实测数据论证交叉点。

关键术语

延伸阅读

• PyTorch Distributed: NCCL collectives
• Horovod ring allreduce paper
• NCCL topology and algorithm selection
• Patarasuk and Yuan, Bandwidth optimal allreduce algorithms
• 第 10 阶段第 05 课 - 分布式训练概览
• 第 19 阶段第 77 课 - 在这些原语之上接入 DDP