07. 大规模训练时通信瓶颈在哪里？AllReduce、AllGather、ReduceScatter 有什么区别？

简单回答

大规模分布式训练的通信瓶颈主要来自梯度同步（数据并行的 AllReduce）和参数/激活传递（模型并行的 AllGather/ReduceScatter）。AllReduce 是每张 GPU 的数据最终汇聚成所有 GPU 相同的全局结果；AllGather 是把每张 GPU 各自的一部分数据收集成完整数据到所有 GPU；ReduceScatter 是先对所有 GPU 的数据做规约（求和/平均），再把结果分片分发给各 GPU。ZeRO 的通信模式就是把 DDP 的 AllReduce 替换成了 ReduceScatter + AllGather。

详细解答

三种集合通信原语

AllReduce 的语义是：每张 GPU 持有一个向量，执行后所有 GPU 都持有这些向量的聚合结果（通常是逐元素求和）。对于 N 张 GPU，每个元素的值都是所有 GPU 上该元素的和。

$\text{AllReduce: }[v_1,v_2,...,v_N]\to [v_1+v_2+...+v_N,v_1+v_2+...+v_N,...,v_1+v_2+...+v_N]$

DDP 的梯度同步用的就是 AllReduce——所有 GPU 各自算了本地梯度，AllReduce 之后每张 GPU 都得到全局平均梯度。

Ring-AllReduce 是目前最主流的 AllReduce 实现，通信量为 $2(N-1)/N\approx 2$ 倍消息大小（一次 ReduceScatter + 一次 AllGather），和 GPU 数量无关，扩展性好。

ReduceScatter 的语义是：每张 GPU 持有一个向量，执行后每张 GPU 得到全局规约结果的一个分片（而不是完整结果）。第 $i$ 张 GPU 得到的是所有 GPU 在第 $i$ 段的求和。

$\text{ReduceScatter: }[v_1,v_2,...,v_N]\to [{\text{GPU}}_i\text{ 持有 }(v_1+...+v_N)\text{ 的第 }i\text{ 段}]$

ZeRO-2 的梯度同步用的是 ReduceScatter 而不是 AllReduce——每张 GPU 只需要自己负责的那 $1/N$ 份梯度，所以 ReduceScatter 就够了，不需要 AllReduce 把全部结果广播给所有 GPU。

AllGather 的语义是：每张 GPU 持有全局数据的一个分片，执行后所有 GPU 都得到完整的全局数据。

$\text{AllGather: }[{\text{GPU}}_i\text{ 持有第 }i\text{ 段}]\to [\text{所有 GPU 持有完整数据}]$

ZeRO-3 在前向传播时需要某层参数，就用 AllGather 把分散在各 GPU 上的参数片段收集成完整参数，用完再释放。FSDP 的"unshard"操作就是 AllGather。

Ring-AllReduce = ReduceScatter + AllGather

这个等价关系很重要，解释了为什么 ZeRO-2 的通信量和 DDP 相同：

DDP 的 AllReduce 通信量 = ZeRO-2 的 ReduceScatter + AllGather 通信量

ZeRO-2 把后半段（AllGather）省掉了，因为每张 GPU 不需要完整结果，只需要自己那 $1/N$ 份。所以 ZeRO-2 通信量是 DDP 的 一半，显存更低但通信量不增加，是非常合算的升级。

ZeRO-3 还需要额外的 AllGather（前向时取参数），所以通信量比 DDP 更多。

实际瓶颈分析

大规模训练中，通信瓶颈的位置和并行策略密切相关。

数据并行（DDP/ZeRO）的通信：在每个训练步的反向传播结束后做梯度同步（AllReduce 或 ReduceScatter），可以和反向传播的计算做重叠（Gradient Overlap）——当一层的反向算完，立刻启动这一层梯度的通信，同时继续算更深层的反向。PyTorch DDP 默认开启这个重叠，DeepSpeed 和 FSDP 也都支持。

张量并行的通信：每层内部有 AllReduce（Megatron 方案），频率高，每层都有，无法完全和计算重叠。这就是为什么张量并行必须在 NVLink 节点内——NVLink 的带宽（A100 节点约 600 GB/s 双向）才够用，InfiniBand HDR（200 Gbps ≈ 25 GB/s）带宽差太多。

流水线并行的通信：相邻流水线阶段之间传递 Activation 和 Activation 梯度，只发生在边界处，通信量是 $[B,S,H]$ 的 Activation 大小，相对少，跨节点用 InfiniBand 也可以。

通信/计算比（Compute-to-Communication Ratio） 是判断某种并行策略是否高效的核心指标。GPU 算力越强（H100 比 A100 快 3~4 倍），通信量不变但计算更快，通信占比更高，瓶颈从计算转向通信。这也是为什么随着 GPU 算力的提升，网络带宽的重要性越来越高。

面试时可以这样答

先区分三个通信原语。AllReduce 是所有 GPU 汇总，最后每张卡都有相同的完整结果，DDP 梯度同步用这个；ReduceScatter 是汇总后分片，每张卡只拿自己负责的那份，ZeRO-2 梯度同步换成了它，通信量砍了一半；AllGather 是把各卡的片段收集成完整数据，ZeRO-3 前向取参数用的就是它。

实际通信瓶颈在哪里取决于并行策略。数据并行的梯度同步可以和反向计算重叠，在精心设计下通信开销不明显。张量并行每层都有通信，频率最高，必须放在 NVLink 节点内跑，否则带宽不够根本跑不动。流水线并行的通信量少，适合跨节点。

随着 GPU 越来越快（H100 的峰值算力是 A100 的 3 倍多），但网络带宽升级没那么快，通信瓶颈越来越明显。这也是为什么通信和计算的 overlap 设计在大规模训练里越来越重要。

常见追问

1. DDP 的梯度 overlap 具体是怎么实现的？Bucket 机制是什么？
2. 在 InfiniBand 带宽有限的情况下，张量并行跨节点有没有什么优化手段？
3. 如何评估一个训练集群的通信效率是否成为瓶颈？