跳到主要内容
AIInfra前置基础

第6章:集合通信基础

掌握分布式训练的通信原语(AllReduce、AllGather 等)、Ring/Tree 通信算法和 NCCL 的使用与调优

集合通信 AllReduce NCCL Ring算法

集合通信是分布式训练的”血管系统”——并行策略再聪明,通信跑不动一切都是空谈。本文从点对点通信讲起,系统介绍 7 种集合通信原语、Ring/Tree AllReduce 算法、通信计算 Overlap、NCCL 用法与调优,并把通信量公式串成一个能直接用来分析并行策略开销的工具。

📑 目录


1. 通信的两种形式:点对点 vs 集合

类型参与方例子用途
点对点两个进程Send / Recv流水线并行的层间传递
集合通信一组进程AllReduce数据并行的梯度聚合

集合通信本质上是点对点通信的”调度模式”——一个 AllReduce 会被拆成若干次 Send/Recv,只是 NCCL 帮你封装好了最优的调度算法。


2. 七种集合通信原语

设有 NN 个进程,数据总量为 VV 字节。

2.1 Broadcast(广播)

一个进程把数据发给所有其他进程:

Before:        After:
P0: [A]        P0: [A]
P1: [ ]        P1: [A]
P2: [ ]        P2: [A]
P3: [ ]        P3: [A]

通信量:每个 receiver 收 VV 字节,sender 总发 (N1)V(N-1) V 字节。

2.2 Reduce(归约)

所有进程的数据汇总到一个进程,执行 sum/max/min 等操作:

Before:                  After:
P0: [a]                  P0: [a+b+c+d]
P1: [b]                  P1: [b]
P2: [c]                  P2: [c]
P3: [d]                  P3: [d]

2.3 AllReduce(全归约)

最常用的原语——每个进程都得到归约结果:

Before:                  After:
P0: [a]                  P0: [a+b+c+d]
P1: [b]                  P1: [a+b+c+d]
P2: [c]                  P2: [a+b+c+d]
P3: [d]                  P3: [a+b+c+d]

数据并行训练的梯度聚合就是 AllReduce(对每个参数的梯度求和后除以 N)。

2.4 Scatter(分散)

一个进程把数据拆成 N 份,每份发给一个进程:

Before:                  After:
P0: [A,B,C,D]           P0: [A]
P1: [ ]                 P1: [B]
P2: [ ]                 P2: [C]
P3: [ ]                 P3: [D]

2.5 Gather(收集)

Scatter 的反向——所有进程的数据收集到一个进程:

Before:                  After:
P0: [a]                  P0: [a,b,c,d]
P1: [b]                  P1: [b]
P2: [c]                  P2: [c]
P3: [d]                  P3: [d]

2.6 AllGather(全收集)

每个进程都得到所有进程拼接后的数据:

Before:                  After:
P0: [a]                  P0: [a,b,c,d]
P1: [b]                  P1: [a,b,c,d]
P2: [c]                  P2: [a,b,c,d]
P3: [d]                  P3: [a,b,c,d]

ZeRO-3 的参数 AllGather、张量并行的输出聚合都用它。

2.7 ReduceScatter(归约后分发)

先 Reduce,然后把结果按块分发(每个进程拿到归约结果的一段):

Before:                  After:
P0: [a0,a1,a2,a3]       P0: [a0+b0+c0+d0]
P1: [b0,b1,b2,b3]       P1: [a1+b1+c1+d1]
P2: [c0,c1,c2,c3]       P2: [a2+b2+c2+d2]
P3: [d0,d1,d2,d3]       P3: [a3+b3+c3+d3]

ZeRO 的梯度切分用它。重要等式:AllReduce = ReduceScatter + AllGather,Ring AllReduce 就是按这个等式实现。


3. Ring AllReduce:带宽最优

3.1 算法直觉

把 N 个 GPU 排成一个环,每张卡只和左右两个邻居通信。算法分两个阶段,各 N1N - 1 步:

Phase 1:Reduce-Scatter

把数据切成 N 段,每张卡负责其中一段。第 kk 步,GPU ii 把自己持有的某段累加值发给 GPU i+1i+1:

N=4, 数据切 4 段
Step 0:   GPU 0→1 发段 0
          GPU 1→2 发段 1
          GPU 2→3 发段 2
          GPU 3→0 发段 3

Step 1:   GPU 0→1 发段 3 (含 0+3 累加)
          ...
Step 2:   每张卡持有自己负责段的完整累加结果

Phase 2:All-Gather

每张卡把自己负责的累加段沿环传播,N1N - 1 步后所有卡都拿到完整结果。

3.2 通信量分析

每张卡总共发 2(N1)NV\frac{2(N-1)}{N} \cdot V 字节,接收同等量。

  • 带宽最优:不管 N 多大,每张卡的发送量逼近 2V2V(一个 ReduceScatter 的 VV + 一个 AllGather 的 VV),接近理论下限
  • 延迟随 N 线性增长:总步数 2(N1)2(N - 1),跨机时延迟可能成瓶颈

3.3 工程意义

这就是为什么 NCCL 的默认算法是 Ring,以及为什么 GPU 数越多,Ring 的优势越明显——带宽利用率始终接近 100%。


4. Tree AllReduce:延迟更低

把 N 个 GPU 组成一棵二叉树,Reduce 沿树往上聚合,Broadcast 沿树往下分发。

算法总通信步数每卡数据量适用场景
Ring2(N1)2(N-1)2V\sim 2V大消息 / 大集群
Tree2log2N2 \log_2 N2V\sim 2V小消息 / 小集群
Double Binary Treelog2N\log_2 N + 一些2V\sim 2VNCCL 默认混合

NCCL 2.4+ 内置 Double Binary Tree(两棵二叉树 + 反向),根据消息大小和拓扑自动选择。


5. 通信与计算 Overlap

5.1 为什么要 Overlap

如果通信和计算串行进行:

计算       ████████____________ ████████____________
通信       ____________████████ ____________████████
总时间     ←─────────────→

如果能让通信”藏”在计算背后:

计算       ████████████████████████████████
通信       ____________████████____________████████
总时间     ←──────→  (节省了通信时间)

5.2 PyTorch DDP 的 Bucket 机制

DDP 把梯度按层组合成”桶”(默认 25MB 一桶)。一个桶满了立刻触发 AllReduce,不等所有梯度算完——这样反向传播底层的梯度还在算的时候,顶层的梯度已经在通信。

model = DistributedDataParallel(
    model,
    bucket_cap_mb=25,         # 桶大小
    gradient_as_bucket_view=True,  # 梯度直接作为 bucket 的视图,省内存
)

5.3 ZeRO-3 的 Prefetch

ZeRO-3 的参数是切片的,每层 forward/backward 前都要 AllGather 拿全参数。优化做法:提前一两层 prefetch——上一层在算,下一层的参数已经在通信。


6. NCCL 实战与调优

6.1 最小 NCCL 程序(C++)

#include <nccl.h>

ncclComm_t comm;
ncclCommInitAll(&comm, nGpus, devs);

float *sendbuff, *recvbuff;
cudaMalloc(&sendbuff, count * sizeof(float));
cudaMalloc(&recvbuff, count * sizeof(float));

ncclAllReduce(sendbuff, recvbuff, count,
              ncclFloat, ncclSum, comm, stream);
cudaStreamSynchronize(stream);

ncclCommDestroy(comm);

6.2 PyTorch 中通常不直接用

import torch.distributed as dist

dist.init_process_group(backend='nccl')
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

PyTorch 的 dist.* 接口底层就是 NCCL,直接用即可。

6.3 关键环境变量

变量作用
NCCL_DEBUG=INFO打印通信细节,debug 必备
NCCL_DEBUG_SUBSYS=ALL更详细日志
NCCL_IB_DISABLE=0启用 IB(默认是,显式确认)
NCCL_IB_HCA=mlx5_0,mlx5_1指定使用哪些 IB 网卡
NCCL_SOCKET_IFNAME=eth0控制流走哪个网卡
NCCL_NET_GDR_LEVEL=PHB启用 GPU Direct RDMA(GPU 直接读写网卡内存)
NCCL_TREE_THRESHOLD切换 Tree/Ring 的消息大小阈值
NCCL_NTHREADS通信线程数

6.4 排查通信慢的清单

  1. NCCL_DEBUG=INFO 跑一次,看实际选了哪个 channel(NVLink? IB? PCIe?)
  2. nvidia-smi topo -m 确认拓扑,验证物理上能走 NVLink
  3. nccl-testsall_reduce_perf -b 1M -e 1G -f 2 -g 8,看带宽是否接近理论值
  4. 多机时检查 NCCL_IB_HCANCCL_SOCKET_IFNAME 是否走对网卡
  5. 检查 IB 网卡 ibstat 是否 Active,ib_send_bw 测点对点带宽

7. 通信公式速查表

设 N 个 GPU,数据 V 字节(梯度量、激活量等)。

操作每 GPU 发送量总通信量用途
Broadcast (Tree)V\sim VVlog2NV \log_2 N模型参数广播
Reduce (Tree)V\sim VVlog2NV \log_2 N单点聚合
AllReduce (Ring)2V\sim 2V2V(N1)2 V (N-1)DDP 梯度同步
AllGather (Ring)VN1N\sim V \frac{N-1}{N}V(N1)V (N-1)ZeRO-3 参数收集
ReduceScatter (Ring)VN1N\sim V \frac{N-1}{N}V(N1)V (N-1)ZeRO 梯度切分
All-to-AllV\sim VV(N1)V (N-1)MoE expert 路由

速记口诀:AllReduce = AllGather + ReduceScatter,各占一半带宽

7.1 一个具体例子:DDP 训练 LLaMA-7B

  • 模型参数 7B,FP16 梯度 = 14 GB
  • 8 卡 AllReduce(NVLink 900 GB/s):每卡发送约 2×14=282 \times 14 = 28 GB,理论耗时 28/90031\approx 28 / 900 \approx 31 ms
  • 跨机 8×8=64 卡(IB 25 GB/s):每卡发送约 28 GB,理论耗时 28/251.1\approx 28 / 25 \approx 1.1 s

🌟 结论:跨机 AllReduce 比单机慢约 35 倍——这就是为什么大模型训练必须做”梯度聚合 + 计算 overlap”,否则 GPU 大半时间在等通信。


✅ 自我检验清单

  • 七大原语:能默写 Broadcast / Reduce / AllReduce / AllGather / ReduceScatter / Scatter / All-to-All 的数据流示意图
  • AllReduce 公式:能写出 Ring AllReduce 每卡发送量 2(N1)NV\frac{2(N-1)}{N} V 的来源
  • Ring vs Tree:能解释为什么大消息选 Ring、小消息选 Tree
  • AllReduce 分解:能解释 AllReduce = ReduceScatter + AllGather,以及这一等式如何指导 ZeRO 设计
  • 拓扑感知:nvidia-smi topo -m 输出 NV4 / SYS 等标记,能判断哪些卡走 NVLink、哪些走 PCIe / 跨 NUMA
  • DDP Bucket:能解释 PyTorch DDP 的 bucket 机制如何实现通信和计算 overlap
  • Overlap 直觉:给定一个 7B 模型 8 卡 DDP,能口算单步 AllReduce 时间
  • NCCL 调试:能用 NCCL_DEBUG=INFO + nccl-tests 排查通信瓶颈
  • 跨机对比:能解释为什么张量并行不能跨机,而流水线并行可以

📚 参考资料

论文与教程

  • Bringing HPC Techniques to Deep Learning (Baidu, 2017):Ring AllReduce 在深度学习中的应用
  • NCCL: Optimized Primitives for Inter-GPU Communication:NVIDIA 论文

官方文档

中文解读

  • OneFlow:Ring AllReduce 全图解
  • 猛猿:图解集合通信原语
  • 方佳瑞:深入理解 NCCL 通信库