All-Reduce、All-Gather 和 Reduce-Scatter 都是分布式训练中的集合通信操作collective communication。它们不是两个 GPU 之间的点对点通信而是由一个通信组中的所有 GPU 共同参与。假设每张 GPU 上都有一个形状相同的张量GPU0: x0 GPU1: x1 GPU2: x2 GPU3: x3All-Reduce 聚合每张卡都得到完整结果。执行求和形式的 All-Reduce Sumdist.all_reduce(x, opdist.ReduceOp.SUM)后会得到GPU0: x0 x1 x2 x3 GPU1: x0 x1 x2 x3 GPU2: x0 x1 x2 x3 GPU3: x0 x1 x2 x3All-Gather收集所有分片每张卡都得到完整张量。它可以理解为每张 GPU 把自己的分片广播给其他所有 GPU。执行 All-Gather 后GPU0:[X0, X1, X2, X3]GPU1:[X0, X1, X2, X3]GPU2:[X0, X1, X2, X3]GPU3:[X0, X1, X2, X3]在FSDP中平时每张 GPU 上只保存一部分参数GPU0: parameter shard P0 GPU1: parameter shard P1 GPU2: parameter shard P2 GPU3: parameter shard P3某一层即将执行 forward 时通过 All-Gather 临时恢复完整参数GPU0:[P0, P1, P2, P3]GPU1:[P0, P1, P2, P3]GPU2:[P0, P1, P2, P3]GPU3:[P0, P1, P2, P3]Reduce-Scatter先聚合再把结果切分到不同 GPU可以理解为Reduce Scatter两个操作的组合。假设每张 GPU 上都有一个完整张量但把它逻辑上分成 4 段GPU0:[a0, b0, c0, d0]GPU1:[a1, b1, c1, d1]GPU2:[a2, b2, c2, d2]GPU3:[a3, b3, c3, d3]首先对相同位置进行 ReduceAa0 a1 a2 a3 Bb0 b1 b2 b3 Cc0 c1 c2 c3 Dd0 d1 d2 d3然后把结果 Scatter 到各 GPU。形式化地第jjj张GPU得到yj∑ixi(j)y_j\sum_ix_i^{(j)}yj∑ixi(j)其中xi(j)x_i^{(j)}xi(j)表示第iii张GPU输入的第jjj的分片。GPU0: A GPU1: B GPU2: C GPU3: D在FSDP backward 时每张 GPU 都可能计算出当前层的完整梯度GPU0: full gradient G0 GPU1: full gradient G1 GPU2: full gradient G2 GPU3: full gradient G3但 FSDP 不希望每张 GPU 都长期保存完整聚合梯度因此执行 Reduce-Scatter。这样既完成了数据并行中的梯度同步又让梯度继续保持分片状态。GPU0: shard0of(G0 G1 G2 G3)GPU1: shard1of(G0 G1 G2 G3)GPU2: shard2of(G0 G1 G2 G3)GPU3: shard3of(G0 G1 G2 G3)All-Reduce Reduce-Scatter All-Gather这是三者之间最重要的关系。假设目标是让所有 GPU 都得到x0x1x2x3x_0x_1x_2x_3x0x1x2x3而目前有GPU0: x0 GPU1: x1 GPU2: x2 GPU3: x3第一步Reduce-Scatter把最终求和结果切成 4 个分片。GPU0: sum(x)[shard0]GPU1: sum(x)[shard1]GPU2: sum(x)[shard2]GPU3: sum(x)[shard3]第二步All-Gather收集所有聚合后的分片。GPU0: complete sum(x)GPU1: complete sum(x)GPU2: complete sum(x)GPU3: complete sum(x)NCCL 中常见的 Ring All-Reduce本质上就可以看成Ring Reduce-Scatter Ring All-Gather