NCCL (音 “Nickel”) 是一个独立的库，包含用于 GPU 的标准通信例程。它已针对使用 PCIe、NVLink、NVswitch 以及使用 InfiniBand Verbs 或 TCP/IP 套接字的平台进行了优化，以实现高带宽。NCCL 支持单个节点或跨多个节点安装的任意数量的 GPU，并且可以在单进程或多进程（例如 MPI）应用程序中使用。

2017 年百度将 HPC 领域的 Ring AllReduce 算法引入机器学习领域，支持 GPU 的告诉通信。本文主要分析 NCCL 的 Reduce-Scatter/All-Gather 的算法逻辑。

Reduce-Scatter

Algorithm

Ring Reduce-Scatter 的核心思想是将 $N$ 个 GPU 连成环，将数据切分成 $0 \sim N-1$ 个 block，对于第 $i$ 个 GPU，在第 $k$ 次时发送第 $k + i - 1$ 个 block 到下一个 GPU 并在其上做规约操作，总共只需要进行 $N - 1$ 次发送就能完成规约¹。

这样对于每个 GPU，都会发送和接受 $N-1$ 次 block，如果带宽为 $B$ ，则其通信复杂度为 $$O(\frac{\frac{x}{N} * (N-1)}{B}) = O(\frac{x}{B})$$ 注意到与 GPU 个数 $N$ 无关，只取决于整个环中最低的带宽 $B$

Implementation

Reduce-Scatter 主要实现在 src/device/reduce_scatter.h 中，在机器学习中常常发生在反向传播需要进行梯度同步时。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
template<typename T, typename RedOp, typename Proto>
__device__ __forceinline__ void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
	ncclRing *ring = &ncclShmem.channel.ring;
	int const *ringRanks = ring->userRanks;
	const int nranks = ncclShmem.comm.nRanks;
	size_t count;
	size_t gridOffset; // start position of first data
	size_t channelCount; // number elements of all data
	size_t chunkCount; // number of elements in a chunk
	ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), &count, &gridOffset, &channelCount, &chunkCount);
	size_t offset;
	size_t dataOffset;
	uint32_t nelem;
	int rankDest;

	Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
	prims(tid, nthreads, &ring->prev, &ring->next, work->sendbuff, work->recvbuff, work->redOpArg);

	for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
	    // 计算当前处理的元素数量，确保不超过通道总数
	    nelem = min(chunkCount, channelCount - elemOffset);
	    // 计算数据在全局数组中的偏移量
	    dataOffset = gridOffset + elemOffset;
	    // /////////////// begin ReduceScatter steps ///////////////
	    // 步骤 0: 将数据推送到下一个 GPU
	    rankDest = ringRanks[nranks-1]; // 获取目标 GPU 的排名
	    offset = dataOffset + rankDest * count; // 计算目标 GPU 的数据偏移量
	    prims.send(offset, nelem); // 发送数据到目标 GPU
	
	    // 步骤 k-2: 对数据进行归约并复制到下一个 GPU
	    for (int j=2; j<nranks; ++j) { // 对后面的 GPU
	        rankDest = ringRanks[nranks-j]; // 获取当前目标 GPU 的排名
	        offset = dataOffset + rankDest * count; // 计算目标 GPU 的数据偏移量
	        prims.recvReduceSend(offset, nelem); // 执行接收、归约并发送数据
	    }
	    
	    // 步骤 k-1: 归约这个缓冲区和数据，产生最终结果
	    rankDest = ringRanks[0]; // 获取第一个 GPU 的排名
	    offset = dataOffset + rankDest * count; // 计算目标 GPU 的数据偏移量
	    prims.recvReduceCopy(offset, dataOffset, nelem, /*postOp=*/true); // 执行接收、归约并复制数据
	}
}

注意

Ring Reduce 实现的过程中，19 行遍历了每个 GPU 负责下的 channel（不确定是否是被切成了 channelCount 个小的 chunk，还是数组）。然后对于每个 GPU ，都是以自己为头建的环，因此 nranks-j 都是和自己距离为 j 的 GPU ，而 rank 0 就是自己。因此可以看到，虽然只需要 $N-1$ 步即可完成聚合，但是多发送了一步，以保证每块 GPU 负责的数据部分在原位置更新为聚合后的值。

简而言之，每个 GPU 的逻辑是两块：

负责的数据块，逻辑为先发送（步骤 0）、等待最后收到然后本地规约（步骤 k-1）
是其余的数据快，逻辑为收到后规约然后立刻发送到下个 GPU（步骤 k-2）

All-Gather

Algorithm

Allgather 中，只需要把 Reduce-Scatter 中使用的算子中的规约部分去除即可。

Implementation

Reduce-Scatter 主要实现在 src/device/all_gather.h 中，在机器学习中常常发生在反向传播需要进行参数同步时。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
    // 输入缓冲区指针
    T *inputBuf = (T*)work->sendbuff;
    // 输出缓冲区指针
    T *outputBuf = (T*)work->recvbuff;

    for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
        // /////////////// begin AllGather steps ///////////////
        // 计算当前处理的元素数量
        nelem = min(chunkCount, partCount - elemOffset);
        // 计算当前数据块的偏移量
        dataOffset = partOffset + elemOffset;

        // 步骤 0: 将数据推送到下一个 GPU
        rankDest = ringRanks[0]; // 获取目标 GPU 的排名
        offset = dataOffset + rankDest * count; // 计算目标偏移量

        // 检查是否为就地操作
        if (inputBuf + dataOffset == outputBuf + offset) { // 如果输入和输出缓冲区重叠
            prims.directSend(dataOffset, offset, nelem); // 直接发送数据
        } else {
            prims.directCopySend(dataOffset, offset, nelem); // 否则复制后发送数据
        }

        // 步骤 k-2: 将数据复制到下一个 GPU
        for (int j=1; j<nranks-1; ++j) { // 从 rank 1 到 rank nranks-2
            rankDest = ringRanks[nranks-j]; // 获取当前目标 GPU 的排名
            offset = dataOffset + rankDest * count; // 计算目标偏移量

            prims.directRecvCopyDirectSend(offset, nelem); // 接收并直接发送数据
        }

        // 将缓冲区的最终副本复制到目标
        rankDest = ringRanks[1]; // 获取第二个 GPU 的排名
        offset = dataOffset + rankDest * count; // 计算目标偏移量

        // 最终的等待和复制
        prims.directRecv(offset, offset, nelem); // 从目标接收数据
    }
}

[AllReduce Blog 图解]https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/ ↩︎

Last modified on 2024-12-03