NCCL：GPU 之间怎么同步数据

前面讲了分布式训练要同步梯度、传输激活值。具体怎么传？

在 NVIDIA GPU 上，用的是 NCCL。

NCCL 是什么

NCCL（NVIDIA Collective Communications Library）是 NVIDIA 开发的 GPU 集合通信库。

「集合通信」是指多个节点之间的数据交换操作，比如：

把所有节点的数据汇总（AllReduce）
把一个节点的数据广播给所有节点（Broadcast）
把数据分发给各节点（Scatter）

NCCL 针对 NVIDIA GPU 和高速网络优化，是目前最快的 GPU 通信库。

基本通信操作

AllReduce

最常用的操作。把所有 GPU 的数据汇总，结果返回给所有 GPU。

GPU 0: [1, 2, 3]
GPU 1: [4, 5, 6]
GPU 2: [7, 8, 9]

AllReduce (sum)
↓

GPU 0: [12, 15, 18]
GPU 1: [12, 15, 18]
GPU 2: [12, 15, 18]

用于梯度同步：所有 GPU 的梯度求和取平均。

Broadcast

一个 GPU 的数据广播给所有 GPU。

GPU 0: [1, 2, 3]  ← 源
GPU 1: []
GPU 2: []

Broadcast
↓

GPU 0: [1, 2, 3]
GPU 1: [1, 2, 3]
GPU 2: [1, 2, 3]

用于初始化：主节点的模型参数广播给所有节点。

AllGather

每个 GPU 贡献一部分数据，最后所有 GPU 都得到完整数据。

GPU 0: [A]
GPU 1: [B]
GPU 2: [C]

AllGather
↓

GPU 0: [A, B, C]
GPU 1: [A, B, C]
GPU 2: [A, B, C]

用于 ZeRO：收集分散在各 GPU 的参数。

ReduceScatter

先 Reduce，再 Scatter。每个 GPU 得到结果的一部分。

GPU 0: [1, 2, 3]
GPU 1: [4, 5, 6]
GPU 2: [7, 8, 9]

ReduceScatter (sum)
↓

GPU 0: [12]      # 第一列的和
GPU 1: [15]      # 第二列的和
GPU 2: [18]      # 第三列的和

用于 ZeRO：梯度求和后分散存储。

AllToAll

每个 GPU 都向其他所有 GPU 发送数据。

GPU 0: [A0, A1, A2]  发送 A1 给 GPU1，A2 给 GPU2
GPU 1: [B0, B1, B2]  发送 B0 给 GPU0，B2 给 GPU2
GPU 2: [C0, C1, C2]  发送 C0 给 GPU0，C1 给 GPU1

AllToAll
↓

GPU 0: [A0, B0, C0]
GPU 1: [A1, B1, C1]
GPU 2: [A2, B2, C2]

用于张量并行里的数据重排。

通信算法

Ring AllReduce

经典算法。GPU 组成环形，数据沿环传递。

GPU0 → GPU1 → GPU2 → GPU3 → GPU0

分两个阶段：

Reduce-Scatter：数据沿环传递并累加，每个 GPU 最终持有 1/N 的完整结果
AllGather：完整结果沿环传递，每个 GPU 得到所有数据

优点：

通信量固定，和 GPU 数量无关
带宽利用率高

缺点：

延迟随 GPU 数量增加
不适合超大规模

Tree AllReduce

用树形结构通信。

        GPU0
       /    \
    GPU1    GPU2
   /    \
GPU3    GPU4

先从叶子向根 Reduce，再从根向叶子 Broadcast。

优点：

延迟低（log N 级别）
适合延迟敏感场景

缺点：

带宽利用不均匀（根节点压力大）

Double Binary Tree

两棵树交织，平衡负载。大规模训练常用。

NCCL 的选择

NCCL 会自动根据硬件配置选择最优算法。

# 查看 NCCL 选择的算法
export NCCL_DEBUG=INFO

通信和计算重叠

好的分布式训练会让通信和计算同时进行，减少等待。

梯度分桶

把模型参数分成多个桶，算完一个桶的梯度就开始同步，不用等所有梯度算完。

计算: [桶0 梯度] [桶1 梯度] [桶2 梯度] [桶3 梯度]
通信:           [桶0 同步] [桶1 同步] [桶2 同步] [桶3 同步]

PyTorch DDP 自动做这个优化。

异步通信

发起通信后不等待完成，继续做其他计算。

# 同步通信
dist.all_reduce(tensor)  # 等待完成

# 异步通信
handle = dist.all_reduce(tensor, async_op=True)
# 做其他事情
handle.wait()  # 需要结果时再等待

NCCL 使用

PyTorch 里使用 NCCL

import torch.distributed as dist

# 初始化，指定 nccl 后端
dist.init_process_group(backend='nccl')

# AllReduce
tensor = torch.tensor([1.0, 2.0]).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

# Broadcast
if rank == 0:
    tensor = torch.tensor([1.0, 2.0]).cuda()
else:
    tensor = torch.zeros(2).cuda()
dist.broadcast(tensor, src=0)

# AllGather
tensor_list = [torch.zeros(2).cuda() for _ in range(world_size)]
dist.all_gather(tensor_list, tensor)

直接使用 NCCL

#include <nccl.h>

ncclComm_t comm;
ncclCommInitRank(&comm, nRanks, id, myRank);

ncclAllReduce(sendbuff, recvbuff, count, datatype,
              ncclSum, comm, stream);

一般用 PyTorch 就够了，不用直接调 NCCL。

常见问题和调优

1. NCCL 超时

NCCL WARN Timeout waiting for GPU

原因：

某个 GPU 卡住了
网络问题
GPU 之间速度差异大

解决：

# 增加超时时间
export NCCL_TIMEOUT=1800  # 秒

2. 性能不达预期

# 开启 NCCL 调试
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

看日志分析：

用的什么算法
通信带宽多少
是否有重试

3. 跨节点通信慢

# 指定网卡
export NCCL_SOCKET_IFNAME=eth0

# 使用 InfiniBand
export NCCL_IB_DISABLE=0
export NCCL_IB_GID_INDEX=3

4. 多卡拓扑问题

NCCL 会自动检测 GPU 拓扑，但有时候检测不准。

# 手动指定拓扑
export CUDA_VISIBLE_DEVICES=0,1,2,3
export NCCL_P2P_LEVEL=NVL  # 强制走 NVLink

通信带宽测试

nccl-tests

NVIDIA 官方提供的测试工具：

git clone https://github.com/NVIDIA/nccl-tests
cd nccl-tests
make

# 测试 AllReduce
./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8

输出示例：

#       size         count    type   redop     time   algbw   busbw
         8             2   float     sum    0.01    0.00    0.00
        16             4   float     sum    0.01    0.00    0.00
       ...
  134217728      33554432   float     sum    4.21   31.87   55.77