2、通信原理解析：AllReduce、AllGather、ReduceScatter 原理与对比

多卡训练的核心，说白了就是「怎么让一堆显卡高效地交换数据」。

我刚开始接触分布式训练时，总觉得通信不就是发发消息嘛，有啥难的？结果第一次跑 64 卡训练，通信开销直接占了 60% 的训练时间。嗯，从那以后我再也不敢小看通信原语了。

今天咱们就掰开揉碎，聊聊三个最核心的集合通信操作：AllReduce、AllGather 和 ReduceScatter。

2.1 为什么需要集合通信？

先想一个问题：你用 8 张卡训练一个模型，每张卡算完自己的梯度后，下一步该干嘛？

对，得把 8 份梯度加起来，然后平均，再分发给每张卡去更新参数。这个过程，就是集合通信要干的事。

我个人习惯把集合通信分成两类：

• 规约类：把多份数据合并成一份（比如求和、求平均）
• 分发类：把一份数据复制给所有人

而 AllReduce、AllGather、ReduceScatter 就是这两类操作的「集大成者」。

2.2 AllReduce：最常用的「求和+广播」

AllReduce 是分布式训练中使用频率最高的通信原语。它的作用很简单：每张卡都拿到所有卡的规约结果。

举个例子：你有 4 张卡，每张卡有一个数值 [1, 2, 3, 4]。AllReduce 求和后，每张卡都得到 [10]。

2.2.1 经典实现：Ring AllReduce

最主流的实现方式是 Ring AllReduce。它分两步走：

1. ReduceScatter 阶段：数据分块，在环上做规约，每张卡只保留一部分结果
2. AllGather 阶段：把规约后的结果广播给所有卡

我在项目中遇到过一个问题：刚开始用 Ring AllReduce 时，总觉得它慢。后来发现是卡间拓扑没对齐——物理环和逻辑环不一致，导致跨节点通信带宽被浪费了。

2.2.2 代码示例：PyTorch 中的 AllReduce

import torch
import torch.distributed as dist

# 初始化进程组
dist.init_process_group(backend='nccl')
rank = dist.get_rank()

# 每张卡生成自己的梯度
tensor = torch.tensor([rank + 1.0]).cuda()

# AllReduce 求和
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(f"Rank {rank}: {tensor.item()}")  # 所有卡都输出 10.0

2.3 AllGather：把所有人的数据「拼」在一起

AllGather 做的事情很直观：每张卡把自己的数据发给所有人，最后每张卡都拥有完整的数据集。

你想想看，这就像 4 个人每人手里有一块拼图，AllGather 之后，每个人手里都有一套完整的拼图。

2.3.1 什么时候用 AllGather？

• 模型并行中的参数同步：每张卡持有部分参数，需要拿到完整参数做前向计算
• 数据并行中的 BatchNorm 统计量同步：每张卡算自己的均值和方差，AllGather 后计算全局统计量
• MoE（混合专家）模型中的专家路由信息收集

2.3.2 代码示例

# 每张卡持有自己的数据
local_data = torch.tensor([rank * 2, rank * 2 + 1]).cuda()

# AllGather：收集所有卡的数据
gathered = [torch.zeros(2).cuda() for _ in range(dist.get_world_size())]
dist.all_gather(gathered, local_data)

# 现在 gathered 包含所有卡的数据
print(f"Rank {rank}: {gathered}")

2.4 ReduceScatter：规约+分发的「二合一」

ReduceScatter 是 AllReduce 的前半段。它把数据分块，对每块做规约，然后把结果分散到各张卡上。

说白了就是：先求和，再分家。

2.4.1 为什么需要 ReduceScatter？

在 Ring AllReduce 中，ReduceScatter 是第一步。但单独使用 ReduceScatter 的场景也很多：

• 梯度分片：每张卡只维护一部分梯度的全局规约结果
• ZeRO 优化器：把优化器状态分片到各卡，ReduceScatter 用于收集梯度分片
• 流水线并行：不同 stage 之间传递部分规约结果

2.4.2 代码示例

# 每张卡有 4 个元素
local_data = torch.tensor([rank + 1.0] * 4).cuda()

# ReduceScatter：规约后分散
output = torch.zeros(2).cuda()  # 每张卡只拿 2 个结果
dist.reduce_scatter(output, [local_data[:2], local_data[2:]], op=dist.ReduceOp.SUM)

print(f"Rank {rank}: {output}")  # 每张卡拿到部分规约结果

2.5 三大原语对比

我整理了一张对比表，方便你快速查阅：

原语	输入	输出	通信量	典型场景
AllReduce	每卡一份数据	每卡一份规约结果	O(N * K)	梯度同步（数据并行）
AllGather	每卡一份数据	每卡拥有全部数据	O(N * K)	参数同步、BN 统计量
ReduceScatter	每卡一份数据	每卡持有部分规约结果	O(K)	梯度分片、ZeRO 优化器

注意看通信量这一列。AllReduce 和 AllGather 都是 O(N*K)，而 ReduceScatter 只有 O(K)。这就是为什么 ZeRO 系列优化器能省带宽——它用 ReduceScatter 替代了 AllReduce。

2.6 核心逻辑框架图

下面这张 SVG 图展示了三种原语的数据流向和关系：

2.7 选型建议

最后，我根据实际项目经验，给你几个选型建议：

• 数据并行梯度同步：直接用 AllReduce，简单高效。NCCL 后端已经优化到极致了。
• ZeRO 系列优化器：用 ReduceScatter 替代 AllReduce，能省 30%-50% 的通信量。
• 模型并行参数同步：用 AllGather，但要注意卡数不要太多，否则通信量会爆炸。
• 混合并行场景：建议把通信原语拆开用，不要一股脑全用 AllReduce。我见过有人把 AllReduce 当万能药，结果通信拓扑乱成一锅粥。

好了，通信原理解析就到这里。记住一句话：选对原语，比调优更重要。