NCCL 集合通信--Collective Operations

深度解析：NVIDIA Collective Operations：AllReduce、Broadcast、Reduce与AllGather详解

原创

已于 2022-08-01 10:44:32 修改 · 2.6k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#网络 #c++ #开发语言

于 2022-07-25 15:22:19 首次发布

本文详细介绍了NVIDIA Collective Operations中的AllReduce、Broadcast、Reduce和AllGather四种通信操作，包括它们的原理、示例和RingAllReduce、RingReduceScatter、RingAllGather等实际实现策略。这些技术在深度学习和分布式计算中至关重要。

集合通信 Collective Operations

AllReduce

在这里插入图片描述

ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream)

N 个设备参加的AllReduce，没有root，初始状态下，所有的N个设备上sendbuffer 长度为 count（函数中定义 size_t count），N个sendbuffer中的内容记为 in_0,in_1,…,in_N-1，运行结束后，所有sendbuffer中内容经过op操作，使得N个设备上的recv buffer中有相同的结果副本，长度均为count，即out。
以op操作为ncclSum为例， out[i] 来自于所有sendbuffer的 in_k[i] 的和，即 out[i] = sum(in_k[i])

Broadcast

在这里插入图片描述

ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream)

N个设备参与的broadcast，root 是函数中指定的 int root（以上图为例则是2)，初始状态下，rank 为 root（上图中2）的设备上的send buffer 长度为count。运行结束后，所有参与broadcast 的N个设备上的recv buffer 长度为 count，且所有N个rank 的 recv buffer 部分的数据均来自于 root设备的send buffer 的拷贝。

Reduce

在这里插入图片描述

ncclResult_t ncclReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream)

N个设备参与的reduce ，root 是函数中指定的 int root（上图中为2)，初始状态下，N个设备上的send buffer 长度为 count，记为数组 in_0,in_1,…in_N-1。运行结束后，rank 为 root （上图中为2）的设备的recv buffer长度为count，记为数组out 。其中 recv buffer的数据来自于N个的send buffer经op操作后的结果。
以op操作为 ncclSum为例，out[i] 为各个send buffer 中 in_X[i] 的和，即out[i] = sum(in_X[i])

AllGather

在这里插入图片描述

ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount, ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream)

N个设备参与AllGather，没有root，初始状态下，N个设备上的send buffer 长度为 sendcount，记为数组 in_0,int_1,…,in_N-1，运行结束后，N个设备上的recv buffer 内容均相同，长度均为 $N * se n d co u n t$ ，记为数组out 。其中每个设备上的recv buffer中的 $o u t [Y * se n d co u t + i]$ 来自于rank为 Y 的设备的send buffer中的 in_Y[i] , 即 out[Y*count + i] = in_Y[i]

ReduceScatter

在这里插入图片描述

ncclResult_t ncclReduceScatter(const void* sendbuff, void* recvbuff, size_t recvcount, ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream)

N 个设备参与的ReduceScatter，没有root设备。初始状态下，N个设备的sendbuffer 长度均为 N*recvcount。运行结束后，N个设备的recv buffer长度为 recvcount，记为数组out_0,out_1,…,out_N-1。其中，N个 send buffer中数据经过op 操作得到 reduced result，将其散布在 N 个设备中，每个设备仅包含 reduced result中一部分。
以op操作为 ncclSum为例， rank为Y的设备中数据 out_Y[i] 为 reduced result 中 Yrecvcount+i 位置的数据，即 out_Y[i] = sum(in_X[Ycount + i])

集合通信的实现

Ring AllReduce

各个设备首尾相连，形成单向的环。每个环上处理一部分数据 block，NCCL 在 luanch kernel 时，会确定 block的数量，一个block对应一个环。

一个循环中AllReduce的实现：

一共有 k 各设备，将每个设备上的数据划分成等大的 k 个 chunk
step 0 当前设备的 rank 为 ringIx 则将（ringIx + k -1 ) mod N 号 chunk 传给下一个设备

在这里插入图片描述

 // step 0: push data to next GPU
      chunk = modRanks(ringIx + nranks-1);
      offset = calcOffset(chunk);
      nelem = min(realChunkSize, size-offset);
      prims.send(offset, nelem);

进行 K - 2 次循环，j 从2 遍历道 N-1；每次循环将接收该设备的（ringIx + k - j ) mod N 作为 chunk_id ，进行reduce之后，传给下一个设备

在这里插入图片描述

 // k-2 steps: reduce and copy to next GPU
      for (int j=2; j<nranks; ++j) {
   
   
        chunk = modRanks(ringIx + nranks-j)

最低0.47元/天解锁文章