从零实现集合通信
撑起分布式训练的四个 collective 操作是 allreduce、broadcast、allgather 和 reduce_scatter。训练框架提供的其他所有原语，都是对这四个的封装。用一套
multiprocessing.Queue组成的 mesh 把它们实现一遍，再拿一个参考实现来逐字节核对，这条线后面的内容就只剩接管线了。
类型： Build 语言： Python 前置要求： 阶段19 Track C 第42-49课 预计时间： ~90 分钟
学习目标
- 用两趟（先 reduce-scatter 再 allgather）实现 ring allreduce，并证明每个 rank 的通信量是每元素 2(N-1)/N 字节。
- 在基于
multiprocessing.Queue的点对点 send 之上，构建 broadcast、allgather 和 reduce_scatter。 - 用同样的输入，把每个原语和
torch.distributed的 gloo 参考实现逐字节对比验证。 - 从集群形态、延迟下限、带宽上限三个角度，论证 ring 与 tree 之间该怎么选。
问题背景
N 个 rank 上的朴素 allreduce，会把张量发给一个 root 共 N 次，再广播回来 N 次。每个 rank 的带宽按 O(N) 增长，root 成为瓶颈，墙钟时间下限是最慢那条链路乘以 N。Ring allreduce 把这个过程压平成 2(N-1) 个大小为 T/N 的 chunk，于是每个 rank 的字节数降到 2T(N-1)/N，且与集群规模无关。Tree allreduce 在 N 小、链路高延迟时占优，因为它的深度是 log2(N) 跳，而不是 2(N-1)。给集群形态选错了拓扑，那块最慢的 GPU 就会决定整步的耗时。
这条线你要读的每一个分布式训练框架，都依赖这四个原语。PyTorch DDP 用每个参数 bucket 一次 allreduce 来同步梯度。ZeRO 用 reduce_scatter 来 shard optimizer state，用 allgather 来广播更新后的参数。FSDP 把整个 forward 拆成 allgather 加 reduce_scatter。Pipeline parallel 需要 broadcast 在 stage group 之间传递激活值。如果你实现不了这四个 collective，你就没法推理为什么训练卡住、为什么梯度不匹配偏偏出在 rank 3、为什么换个拓扑 pipeline bubble 就翻倍。
核心概念
flowchart LR
R0[rank 0] --> Q01[(队列 0 到 1)]
Q01 --> R1[rank 1]
R1 --> Q12[(队列 1 到 2)]
Q12 --> R2[rank 2]
R2 --> Q23[(队列 2 到 3)]
Q23 --> R3[rank 3]
R3 --> Q30[(队列 3 到 0)]
Q30 --> R0两趟实现 ring allreduce
把张量切成 N 个等大的 chunk，编号 0..N-1。每个 rank 拥有编号等于自身 rank 的那个 chunk。第一趟是 reduce-scatter，跑 N-1 步。在第 s 步，rank r 把 chunk (r - s) mod N 发给 rank (r + 1) mod N，并从 rank (r - 1) mod N 收到 chunk (r - s - 1) mod N，把收到的 chunk 累加到自己的本地副本上。N-1 步之后，rank r 拥有 chunk r 的完整和。第二趟是 allgather，再跑 N-1 步，把算好的 chunk 沿 ring 转一圈，直到每个 rank 都持有每个 chunk 的完整和。
| 原语 | 每 rank 字节 | 步数 | 何时使用 |
|---|---|---|---|
| Ring allreduce | 2T(N-1)/N | 2(N-1) | T 大、粗管道同构集群 |
| Tree allreduce | T log2(N) | 2 log2(N) | T 小或链路高延迟 |
| Broadcast | T | log2(N) tree | 参数初始化、标量配置 |
| Allgather | T(N-1)/N | N-1 | sharded forward、ZeRO unshard |
| Reduce_scatter | T(N-1)/N | N-1 | ZeRO 梯度 sharding |
用队列 mesh 替代 NCCL
NCCL 跑在 PCIe 和 NVLink 上，归约由硬件 offload。在 CPU 上你没这条件。给 ring 的每条边配一个 multiprocessing.Queue，就有了单生产者、单消费者的有序点对点投递。归约发生在用户态，所以你要付 Python 的开销，但线上的传输模式和 NCCL ring allreduce 完全一致。在队列版本上把正确性推理清楚，集群行为也就顺理成章了。
用 gloo 来验证
每个原语落地时都带一个单元测试，把它的输出和用 gloo backend 初始化、跑在同样 world size 同一张张量上的 torch.distributed 做对比。如果你的 ring allreduce 和 gloo 的偏差超过 float32 epsilon，测试就失败。拿参考实现来验证这件事没得商量；没有它，原语会一直看起来正确，直到真实训练跑到第 10000 步才露馅。
ci-ring-allreduce动手构建
code/main.py 实现了：
Mesh类，把 N 个multiprocessing.Queue实例接成一个 ring，并为每个 rank 暴露send(dst, tensor)和recv(src)。ring_allreduce(mesh, rank, world_size, tensor)，跑两趟算法。broadcast(mesh, rank, world_size, tensor, src)，走对数 tree。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，做逐字节对比。
运行：
python3 code/main.py输出：逐原语的验证表，对比队列 mesh 和 gloo 的输出，后面跟一个逐 rank 的字节计数器，证明 2T(N-1)/N 的扩展规律。
真实世界中的生产模式
有三个模式能把这些原语打磨到可以上线。
allreduce 前先把梯度 bucket 起来。 一个 10 亿参数的模型有几万个梯度张量。每个张量一次 allreduce，就要付 N 次延迟下限的代价。DDP 把梯度 bucket 成约 25 MB 的 chunk，每个 bucket 发一次 allreduce；小张量搭在大张量的便车上。不做 bucketing，延迟开销就会主导整步。
让通信和计算重叠。 Backward 是按反向逐层算梯度的。最后一层的梯度一就绪，就启动它的 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 之后对 bucket 化的梯度调
dist.all_reduce。bucket 大小可调；默认 25 MB 对 100Gbit 以太网是合理的。 - DeepSpeed ZeRO。 发 reduce_scatter 来 shard 梯度，发 allgather 在 forward 前重建完整参数。本课的原语正是 ZeRO 所调用的那些。
- FSDP。 Forward 以 allgather 开头来 unshard 这一层，算完后用 reduce_scatter 归约，再丢弃 unshard。同样的原语，不同的调度。
拿去用
在第 77-81 课里使用这套队列 mesh 原语。第 77 课把 allreduce 接进 DDP。第 78 课把 reduce_scatter 接进 ZeRO。第 79 课把 broadcast 接进 pipeline 激活值。第 81 课把四个 collective 全部组装进端到端 demo。
练习
- 加一个 tree allreduce 变体，按消息大小在 ring 和 tree 之间切换。测出临界点。
- 加一个
recv_timeout_ms，让卡住的 rank 抛一个截止时间错误，而不是永远挂着。 - 把四个原语的
multiprocessing.Queue换成 TCP socket。同样的测试，真实的线路。 - 加一个带宽埋点钩子，让逐 rank 字节计数器把数据记到 JSONL。
- 在 4 个 rank 上，对大小为 1KB、1MB、16MB 的张量对比 ring 与 tree 的墙钟时间。用实测数据论证临界点。
关键术语
| 术语 | 大家怎么说 | 实际含义 |
|---|---|---|
| Allreduce | "跨 rank 求和" | 调用之后每个 rank 都持有同一个归约后的张量 |
| Ring | "那个快拓扑" | N-1 个大小为 T/N 的 chunk 沿环转两圈 |
| Tree | "那个 log 拓扑" | 归约沿二叉树进行；深度是 log2(N) 跳 |
| Allgather | "拼接各分片" | 每个 rank 最终都拿到其他每个 rank 的分片 |
| Reduce_scatter | "把和拆开" | 每个 rank 最终只拿到一个 chunk 的和 |
| Bucket | "融合小张量" | 把 N 个小 allreduce 合并成一个大的 |
延伸阅读
- PyTorch Distributed：NCCL collectives
- Horovod ring allreduce 论文
- NCCL 拓扑与算法选择
- Patarasuk 与 Yuan，带宽最优的 allreduce 算法
- 阶段10 第05课 - 分布式训练概览
- 阶段19 第77课 - 在这些原语之上接起来的 DDP