AllReduce

AllReduce(全归约) 是分布式训练中最核心的集合通信(Collective Communication) 原语。它的作用是：$N$ 个进程各持有一份等长数据，通过归约操作（通常是求和）得到聚合结果，并将完整结果分发到每个进程上。在标准 数据并行(DDP) 训练中，AllReduce 用于跨 GPU 同步梯度。

集合通信原语一览

理解 AllReduce 之前，先了解它的组成部件。以下所有操作中，每个参与者称为一个 rank：

操作	输入	输出	说明
Reduce	每个 rank 各持一份数据	一个指定 rank 持有归约结果	多对一归约
Broadcast	一个 rank 持有数据	所有 rank 持有相同数据	一对多广播
Scatter	一个 rank 持有完整数据	每个 rank 持有 $1/N$ 切片	切分分发
Gather	每个 rank 各持 $1/N$ 切片	一个 rank 持有拼接结果	多对一拼接
All-Gather	每个 rank 各持 $1/N$ 切片	所有 rank 持有拼接结果	多对多拼接
Reduce-Scatter	每个 rank 各持一份完整数据	每个 rank 持有 $1/N$ 的归约结果	归约 + 切分
All-to-All	每个 rank 持有 $N$ 份不同数据	每个 rank 收到来自所有 rank 各一份	转置式交换

AllReduce 在逻辑上等价于 Reduce + Broadcast。但在实际高效实现中，它被分解为 Reduce-Scatter + All-Gather——这正是 Ring AllReduce 的核心思路。

朴素方法与瓶颈

最简单的实现是选一个主节点（类似参数服务器）：

1. 所有 GPU 将数据发送给主节点 → 主节点接收 $(N-1) \times D$ 数据
2. 主节点做归约
3. 主节点将结果广播回所有 GPU → 发送 $(N-1) \times D$ 数据

$D$ 为每个 GPU 的数据量。瓶颈在于主节点的带宽：吞吐量与 GPU 数 $N$ 成正比，扩展性差。对于一个 3 亿参数模型（1.2GB 梯度），10 个 GPU + 1GB/s 带宽，仅通信就需要约 10.8 秒。

Ring AllReduce

Ring AllReduce(环形全归约) 是目前最广泛使用的带宽最优 AllReduce 算法，由百度研究院推广至深度学习领域。

基本设定

$N$ 个 GPU 排成一个逻辑环，每个 GPU 有一个左邻居和一个右邻居。每个 GPU 的数据被均分为 $N$ 个 chunk。

阶段一：Reduce-Scatter（$N-1$ 轮）

每轮中，每个 GPU 向右邻居发送一个 chunk，同时从左邻居接收一个 chunk 并做累加。经过 $N-1$ 轮后，每个 GPU 恰好持有一个已完成归约的 chunk——即所有 GPU 该位置数据之和。

阶段二：All-Gather（$N-1$ 轮）

每轮中，每个 GPU 向右邻居发送一个已完成归约的 chunk，从左邻居接收一个，直接覆盖本地对应位置。经过 $N-1$ 轮后，每个 GPU 拥有全部 $N$ 个完整 chunk，即完整的归约结果。

通信量分析

每个 GPU 在两个阶段各发送 $N-1$ 个 chunk，每个 chunk 大小 $D/N$：

$\text{每 GPU 总发送量} = 2 \times (N-1) \times \frac{D}{N} = \frac{2(N-1)}{N} \cdot D$

当 $N$ 较大时近似 $2D$，与 GPU 数量无关。这就是 Ring AllReduce 被称为带宽最优(bandwidth-optimal) 的原因——每条链路的带宽利用率达到理论上限。瓶颈仅取决于环中最慢的连接，而非 GPU 总数。

详细的步骤数值示例见 Ring AllReduce 数值示例。

其他 AllReduce 算法

Ring AllReduce 是大消息场景下的最优选择，但它的延迟为 $O(N)$ 轮。对于小消息，其他算法可能更合适：

算法	每 GPU 通信量	延迟（轮数）	特点
Ring AllReduce	$\frac{2(N-1)}{N}D$（带宽最优）	$2(N-1)$	大消息最优
Tree AllReduce（双二叉树）	$\sim 2D$	$O(\log N)$	小消息延迟更低
递归倍增/减半(Recursive Halving-Doubling)	$\sim 2D$	$O(\log N)$	延迟最优
Rabenseifner 算法	$\frac{2(N-1)}{N}D$	$O(\log N)$	兼顾带宽和延迟

NCCL 的拓扑感知选择

NVIDIA 的 NCCL(NVIDIA Collective Communications Library) 是 GPU 集群上集合通信的事实标准。它会自动感知硬件拓扑并选择最优算法：

• 大消息（如梯度同步的大张量）→ Ring AllReduce
• 小消息 → Tree AllReduce（双二叉树，降低延迟）
• 自动检测 NVLink / NVSwitch / PCIe / InfiniBand 拓扑，构造最优通信环或树

在分布式训练中的应用

在 DDP 中的角色

标准 数据并行(DDP) 训练中，AllReduce 用于梯度同步：

1. 每个 GPU 独立计算自己 mini-batch 子集上的梯度
2. AllReduce 对所有 GPU 的梯度求和（再除以 $N$ 得到平均）
3. 每个 GPU 用相同的平均梯度更新参数

通信量：一次 AllReduce，约 $2 \times 2\Psi = 4\Psi$ 字节（$\Psi$ 为参数量，fp16 梯度）。

在 FSDP 中的分解

FSDP 将 AllReduce 的两个组成部分拆开使用：

• 前向阶段用 All-Gather 聚合参数（从各 rank 收集参数分片拼成完整参数）
• 反向阶段用 Reduce-Scatter 分发梯度（各 rank 的梯度求和后切分，每个 rank 只保留自己负责的部分）

这种分解使得参数可以在计算后立即释放，是 FSDP 实现显存节省的关键。

计算-通信重叠

反向传播是逐层计算梯度的，因此可以在输出层梯度就绪时立即启动该层的 AllReduce/Reduce-Scatter，与更早层的梯度计算并行。现代框架（DDP、FSDP）都会自动做这种重叠。

相关页面

• 数据并行总览见 data-parallelism
• FSDP 的通信模式详解见 fsdp
• LLM 并行策略总览见 llm-parallelism
• 专家并行 中的 All-to-All 通信