4. 通信模式
4.1 从点对点通信到集合通信
在单机多卡训练中,GPU 之间通过 PCIe 或 NVLink 交换数据,通信模式相对简单。但在多节点集群中,动辄数十甚至数百个 GPU 需要协同完成梯度同步、参数分发、数据加载等任务。如果每个 GPU 都与其他 GPU 建立独立的点对点连接,通信复杂度将呈平方级增长,且极易造成网络拥塞。为此,业界定义了集合通信(Collective Communication)原语:一组进程(通常对应一个 GPU 或一个 rank)共同参与、按特定拓扑完成数据的聚合、分发或交换。
集合通信的核心价值在于:将多对多的通信模式抽象为少数几个原语,由底层运行时(如 NCCL、RCCL、Gloo)针对不同网络拓扑做优化。上层框架(PyTorch、TensorFlow)只需调用这些原语,即可在千卡规模下保持较高的通信效率。
4.2 核心集合通信原语
多节点训练中最常用的集合通信原语包括以下四种,它们构成了梯度同步和参数分发的基石。
| 原语 | 语义 | 典型用途 | 数据量特征 |
|---|---|---|---|
| AllReduce | 所有 rank 提供一份数据,归约(求和/求均值等)后结果分发给所有 rank | 梯度同步(数据并行训练的核心) | 通信量 = 2×(N-1)/N × 数据量,近似 2×数据量 |
| AllGather | 每个 rank 提供一份数据,收集后拼接,结果分发给所有 rank | 收集各 rank 的统计量、权重碎片、嵌入表 | 通信量 = (N-1)/N × 数据量,近似 1×数据量 |
| ReduceScatter | 先归约,再将结果按 rank 分片,每个 rank 只保留自己那份 | ZeRO 优化器状态分片、分布式训练前奏 | 通信量 = (N-1)/N × 数据量 |
| Broadcast | 单一 rank 的数据发送给所有其他 rank | 参数初始化、学习率调度广播 | 通信量 = 1×数据量 |
4.3 AllReduce 的实现算法与演进
AllReduce 的实现并非只有一种固定方式。不同算法在通信量、延迟和带宽利用率上各有取舍,理解它们有助于在实际训练中定位性能瓶颈。
4.3.1 Ring AllReduce(环形归约)
Ring AllReduce 是 NCCL 最经典的算法,适用于节点内 NVLink 高速互联或节点间环形拓扑。其核心思想是:将数据切分为 N 个 chunk(N 为 rank 数),每个 rank 只与相邻的两个 rank 通信。整个过程分为两步:
- ReduceScatter 阶段:每个 rank 将自己的数据切分为 N 块,沿环向下一 rank 发送第 i 块,同时从上一 rank 接收第 i-1 块并累加。经过 N-1 轮后,每个 rank 持有某个 chunk 的全局归约结果。
- AllGather 阶段:每个 rank 将自己持有的完整 chunk 沿环继续转发,经过 N-1 轮后,所有 rank 获得全部 chunk 的归约结果。
// 伪代码:Ring AllReduce(N 个 rank,数据切为 N 块)
// 阶段1:ReduceScatter
for step in 0..N-2:
send_chunk = (rank - step) mod N
recv_chunk = (rank - step - 1) mod N
send(rank+1, data[send_chunk])
recv(rank-1, tmp)
data[recv_chunk] += tmp // 累加
// 阶段2:AllGather
for step in 0..N-2:
send_chunk = (rank - step + 1) mod N
recv_chunk = (rank - step) mod N
send(rank+1, data[send_chunk])
recv(rank-1, data[recv_chunk])
Ring AllReduce 的通信量约为 2×(N-1)/N × 数据总量,当 N 较大时接近 2 倍数据量。其优势在于带宽利用率高,每个 rank 在每一轮都在同时收发,没有空闲等待;劣势是延迟随 N 线性增长,在小数据量场景下(如模型较小时)延迟开销占比高。
4.3.2 Tree AllReduce(树形归约)
Tree AllReduce 采用二叉树或双二叉树结构。数据从叶子节点向上归约到根节点,再由根节点向下广播。其通信量约为 2×log(N) × 数据量(每层传输一份完整数据),延迟为 O(log N)。
NCCL 在实际部署中常采用双二叉树(Double Binary Tree)结构:将 rank 组织成两棵独立的树,一棵负责数据的前半部分,另一棵负责后半部分,从而在归约和广播阶段都保持所有链路忙碌,避免单棵树中父节点成为瓶颈。
- 数据量 < 512KB:使用 Tree 算法(延迟优先)
- 数据量 ≥ 512KB:使用 Ring 算法(带宽优先)
在 NCCL 2.12+ 中,还引入了 NVLS(NVLink Sharp) 和 CollNet 等针对特定硬件的优化。
4.4 AllGather 与 ReduceScatter 的配合:ZeRO 场景
在大模型训练中,AllGather 和 ReduceScatter 往往成对出现。以微软 ZeRO(零冗余优化器)为例,它将优化器状态、梯度和参数分片到各 rank 上,训练过程中需要频繁执行集合通信:
- 前向传播:每个 rank 需要完整的参数权重 → 执行 AllGather 收集被分片的参数。
- 反向传播:每个 rank 计算完梯度后,需要对梯度做分片归约 → 执行 ReduceScatter,只保留自己负责的梯度分片。
- 优化器更新:各 rank 独立更新自己分片上的参数,无需通信。
这种模式下,通信量从传统数据并行的 2×模型大小(AllReduce)降低为 2×模型大小(AllGather + ReduceScatter 各一次),但通信次数翻倍。因此,在 ZeRO 场景中,通信延迟的优化比带宽更重要,这也是为什么 ZeRO 通常配合高带宽、低延迟的节点内互联(NVLink)使用。
# PyTorch 中 ZeRO 风格的集合通信调用示例
import torch.distributed as dist
# 前向:收集完整参数(假设 param 被分片在 rank 0..N-1)
dist.all_gather(param_full_list, param_shard)
# 反向:梯度分片归约(grad 为完整梯度,grad_shard 只保留本 rank 负责的部分)
dist.reduce_scatter(grad_shard, grad_list, op=dist.ReduceOp.SUM)
4.5 案例:8 节点 64 卡 AllReduce 实测分析
某训练集群使用 8 台 8×A100 服务器,节点间通过 8×200Gbps InfiniBand 互联。在一次千亿参数模型训练中,梯度张量大小为 1.2GB,NCCL 选择了 Ring 算法。实测通信耗时约 18ms,而理论计算如下:
理论通信时间 = 2 × 数据量 / 有效带宽
= 2 × 1.2GB / (8 × 200Gbps × 0.9)
= 2.4GB / 180GB/s
≈ 13.3ms
实测 18ms 与理论值存在约 35% 的差距,排查后发现:
- 节点内 NVLink 带宽虽高,但 Ring 算法在跨节点时需经过节点间链路,跨节点链路成为瓶颈。
- NCCL 默认使用 2 个通信通道(channel),未充分利用 8 条 IB 链路。
- CPU 频率波动导致 NCCL 内核启动延迟增加。
通过设置 NCCL_IB_HCA 指定全部 8 个 IB 设备、增加 channel 数至 4,并将梯度切分为更细粒度的 chunk 以提升流水线并行度,最终通信耗时降至 14.2ms,接近理论值。
NCCL_ALGO 和 NCCL_PROTO 环境变量做 A/B 测试。
4.6 集合通信实现中的关键参数与调优
NCCL 作为最主流的集合通信库,其行为可通过环境变量精细控制。以下参数在多节点调优中最为关键:
| 环境变量 | 作用 | 推荐设置 |
|---|---|---|
NCCL_ALGO |
强制指定算法(Ring/Tree/CollNet) | 留空让 NCCL 自动选择 |
NCCL_PROTO |
选择协议(LL/LL128/Simple) | 默认 Simple;小数据量可试 LL128 |
NCCL_IB_HCA |
指定使用的 IB 设备列表 | 如 mlx5_0,mlx5_1,... |
NCCL_SOCKET_IFNAME |
指定 TCP/IP 通信网卡 | 如 eth0 或 ib0 |
NCCL_NTHREADS |
每个 channel 的线程数 | 默认 64;可调至 128 或 256 |
NCCL_MAX_NCHANNELS |
限制最大 channel 数 | 默认自动;可增大以利用多链路 |
4.7 误区与复盘
loss.backward() 后立即调用 optimizer.step(),通信尚未完成就会阻塞计算。正确做法是使用 PyTorch 的 torch.distributed.all_reduce 的异步版本,或依赖 DDP 内部的梯度钩子(hook)实现通信与反向计算的重叠。
if dist.get_rank() == 0 时也要确保所有分支都执行相同次数的集合通信。
1. 先用
nccl-tests 的 all_reduce_perf 测量裸通信带宽,排除框架干扰;2. 对比 Ring 与 Tree 在不同数据量下的表现,确定算法选择策略;
3. 检查
NCCL_DEBUG=INFO 日志中的拓扑检测结果,确认 IB 链路和 NVLink 均被正确识别;4. 逐步增加 channel 数和线程数,观察带宽变化直至饱和;
5. 最后在真实训练脚本中验证通信与计算的重叠效果。