NCCL源码解析1——生成ncclUniqueId

70阅读　0评论2024-11-23　lvyilong316
分类：LINUX

NCCL源码解析1——生成ncclUniqueId

——lvyilong316

NCCL 通信库通过 ncclGetUniqueId 函数生成并获取 「ncclUniqueId」。本质上，ncclUniqueId 是一个大小为 128 字节数组，如下所示。

#define NCCL_UNIQUE_ID_BYTES 128

typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } ncclUniqueId;

ncclGetUniqueId逻辑如下，ncclGetUniqueId 首先调用 ncclInit 函数初始化网络组件，接着调用 bootstrapGetUniqueId 函数生成 ncclUniqueId。

点击(此处)折叠或打开

  1. ncclResult_t ncclGetUniqueId(ncclUniqueId* out) {
  2.   NCCLCHECK(ncclInit());
  3.   NCCLCHECK(PtrCheck(out, "GetUniqueId", "out"));
  4.   ncclResult_t res = bootstrapGetUniqueId((struct ncclBootstrapHandle*)out);
  5.   TRACE_CALL("ncclGetUniqueId(0x%llx)", (unsigned long long)hashUniqueId(*out));
  6.   return res;
  7. }

l  ncclInit

我们首先来看 ncclInit 初始化函数。NCCL 使用两套网络，一个是 bootstrap 网络，主要用于初始化时交互一些简单的信息，如机器 IP 地址和端口等，且大部分只在初始化阶段执行一次，因此通常采用 TCP 实现。另一个是通信网络，用于实际通信数据的交互，因此优先使用 GdrCopy 和 RDMA。ncclInit 函数主要用于初始化 bootstrap 网络的 IP 地址信息。该函数调用 initOnceFunc 函数，并通过pthread_once保证 initOnceFunc 在多线程环境中只执行一次。

点击(此处)折叠或打开

  1. static ncclResult_t ncclInit() {
  2.   pthread_once(&initOnceControl, initOnceFunc);
  3.   return initResult;
  4. }

在 initOnceFunc 函数中，首先通过 initEnv 函数读取 ~/.nccl.conf 或者 /etc/nccl.conf 文件中的 NCCL 环境变量，这和在用户脚本中通过 export 设置环境变量是等价的。

点击(此处)折叠或打开

  1. static void initOnceFunc() {
  2.   initEnv(); //读取 ~/.nccl.conf 或者 /etc/nccl.conf 文件中的 NCCL 环境变量
  3.   initGdrCopy(); //加载libgdrapi.so，初始化GDR相关操作函数
  4.   // Always initialize bootstrap network
  5.   NCCLCHECKGOTO(bootstrapNetInit(), initResult, exit);

  7.   initNvtxRegisteredEnums();
  8. exit:;
  9. }

下面，给出 nccl.conf 文件的一个简单示例：

NCCL_SOCKET_IFNAME=eth0

NCCL_DEBUG=WARN

NCCL_RINGS=0 1 2 3|3 2 1 0

接着，调用 initGdrCopy 函数，而initGdrCopy又调用ncclGdrInit初始化 Gdr相关函数。具体来说，ncclGdrInit函数加载 libgdrapi.so动态库，并初始化相关的操作函数，如 gdr_open、gdr_close等。接着，判断 GdrCopy 库和驱动版本的兼容性，当前 NCCL 要求 GDRAPI 2.1 及其之后版本。

点击(此处)折叠或打开

  1. ncclResult_t initGdrCopy() {
  2.   if (ncclParamGdrCopyEnable() == 1) {
  3.     ncclGdrCopy = ncclGdrInit();
  4.   }
  5.   return ncclSuccess;
  6. }

之后，调用 bootstrapNetInit 初始化 bootstrap 网络。bootstrapNetInit 函数用于初始化 bootstrap 网络。如下所示：

点击(此处)折叠或打开

  1. ncclResult_t bootstrapNetInit() {
  2.   if (bootstrapNetInitDone == 0) {
  3.     pthread_mutex_lock(&bootstrapNetLock);
  4.     if (bootstrapNetInitDone == 0) {
  5.       /* 如果设置了 NCCL_COMM_ID 环境变量（NCCL_COMM_ID的格式为“nccl_socket:ip:port”，其中ip指定要使用的IP地址，port指定要使用的端口号），则查找一个和该环境变量中指定的 IP 地址处于同一子网的网卡作为 booststrap 网络通信所使用的网卡 bootstrapNetIfAddr。 */
  6.       const char* env = ncclGetEnv("NCCL_COMM_ID");
  7.       if (env) {
  8.         union ncclSocketAddress remoteAddr;
  9.         if (ncclSocketGetAddrFromString(&remoteAddr, env) != ncclSuccess) {
  10.           WARN("Invalid NCCL_COMM_ID, please use format: : or []: or :");
  11.           pthread_mutex_unlock(&bootstrapNetLock);
  12.           return ncclInvalidArgument;
  13.         }
  14.         if (ncclFindInterfaceMatchSubnet(bootstrapNetIfName, &bootstrapNetIfAddr, &remoteAddr, MAX_IF_NAME_SIZE, 1) <= 0) {
  15.           WARN("NET/Socket : No usable listening interface found");
  16.           pthread_mutex_unlock(&bootstrapNetLock);
  17.           return ncclSystemError;
  18.         }
  19.       } else {
  20.         int nIfs = ncclFindInterfaces(bootstrapNetIfName, &bootstrapNetIfAddr, MAX_IF_NAME_SIZE, 1);
  21.         if (nIfs <= 0) {
  22.           WARN("Bootstrap : no socket interface found");
  23.           pthread_mutex_unlock(&bootstrapNetLock);
  24.           return ncclInternalError;
  25.         }
  26.       }
  27.       char line[SOCKET_NAME_MAXLEN+MAX_IF_NAME_SIZE+2];
  28.       sprintf(line, " %s:", bootstrapNetIfName);
  29.       ncclSocketToString(&bootstrapNetIfAddr, line+strlen(line));
  30.       INFO(NCCL_INIT, "Bootstrap : Using%s", line);
  31.       bootstrapNetInitDone = 1;
  32.     }
  33.     pthread_mutex_unlock(&bootstrapNetLock);
  34.   }
  35.   return ncclSuccess;
  36. }

    首先，如果设置了 NCCL_COMM_ID 环境变量，则查找一个和该环境变量中指定的 IP 地址处于同一子网的网卡作为 booststrap 网络通信所使用的网卡 bootstrapNetIfAddr。

否则，使用 ncclFindInterfaces 函数选择一个合适的网卡。这里，如果用户通过 NCCL_SOCKET_IFNAME 环境变量指定了特定的网卡名称，则选择该网卡。

否则网卡选择的具体顺序如下：

l  首先，选择名称以 "ib" 为起始字符串的网卡；如果没找到，

l  接着，选择与 NCCL_COMM_ID 中指定 IP 地址处于同一子网的网卡；否则，

l  其次，选择名称不以 "docker" 和 "lo" 为起始字符串的网卡；否则，

l  再次，选择名称以 "docker" 为起始字符串的网卡；否则，

l  {BANNED}最佳后，选择名称以 "lo" 为起始字符串的网卡。

当不指定环境变量时，就是上述红色字体部分选择出的网卡。这里就确定了bootstrap网络的接口名称bootstrapNetIfName和接口地址bootstrapNetIfAddr。

{BANNED}最佳后，通过 initNvtxRegisteredEnums 函数为基于 nvtx 的工具注册额外的信息。具体来说，注册 reduce 类通信的算子类型。

    回到ncclGetUniqueId函数，调用完ncclInit后，则调用bootstrapGetUniqueId。

l  bootstrapGetUniqueId

我们看bootstrapGetUniqueId的调用方式如下，其返回参数out即ncclUniqueId结构，同时可以转换为ncclBootstrapHandle结构，所以其实ncclUniqueId的本质就是ncclBootstrapHandle。

点击(此处)折叠或打开

  1. ncclResult_t res = bootstrapGetUniqueId((struct ncclBootstrapHandle*)out);

所以，我们先来看一下ncclBootstrapHandle这个结构体，如下面的代码片段所示。ncclBootstrapHandle 包含一个魔术值和记录 socket 地址的数据结构。也就是说，ncclUniqueId 是一个包含魔术值和root进程的 socket 地址的数据结构，包括 root 进程的监听地址和端口。

点击(此处)折叠或打开

  1. struct ncclBootstrapHandle {
  2.   uint64_t magic;
  3.   union ncclSocketAddress addr;
  4. };
  5. /* Common socket address storage structure for IPv4/IPv6 */
  6. union ncclSocketAddress {
  7.   struct sockaddr sa;
  8.   struct sockaddr_in sin;
  9.   struct sockaddr_in6 sin6;
  10. };

    下面看bootstrapGetUniqueId的代码逻辑

点击(此处)折叠或打开

  1. ncclResult_t bootstrapGetUniqueId(struct ncclBootstrapHandle* handle) {
  2.   memset(handle, 0, sizeof(ncclBootstrapHandle));

  4.   const char* env = ncclGetEnv("NCCL_COMM_ID");
  5.   if (env) {
  6.     INFO(NCCL_ENV, "NCCL_COMM_ID set by environment to %s", env);
  7.     /* 首先判断是不是设置了 「NCCL_COMM_ID」 环境变量。NCCL_COMM_ID 中包含 IP 地址和端口号，比如"ip:port"。*/
  8.     if (ncclSocketGetAddrFromString(&handle->addr, env) != ncclSuccess) {
  9.       WARN("Invalid NCCL_COMM_ID, please use format: : or []: or :");
  10.       return ncclInvalidArgument;
  11.     }
  12.     handle->magic = NCCL_MAGIC;
  13.   } else {
  14.     NCCLCHECK(getRandomData(&handle->magic, sizeof(handle->magic)));
  15.     memcpy(&handle->addr, &bootstrapNetIfAddr, sizeof(union ncclSocketAddress));
  16.     NCCLCHECK(bootstrapCreateRoot(handle, false));
  17.   }

  19.   return ncclSuccess;
  20. }

bootstrapGetUniqueId函数首先判断是不是设置了NCCL_COMM_ID环境变量。NCCL_COMM_ID 中包含 IP 地址和端口号，比如"ip:port"。上面我们提到，ncclUniqueId 是一个包含魔术值和 root 的 socket 地址的数据结构。那么，如果设置了 NCCL_COMM_ID 环境变量，那么我直接使用 NCCL_COMM_ID中的 IP 地址和端口初始化 ncclUniqueId 数据结构，并使用默认值 NCCL_MAGIC初始化魔术值。

反之，函数通过读取 /dev/urandom获取随机魔术值。接着，使用前面bootstrapNetInit

函数中获取的bootstrapNetIfAddr继续初始化 ncclUniqueId 结构。这里bootstrapNetIfAddr 是在 ncclInit ->bootstrapNetIfAddr函数中构建的用于 bootstrap 网络的 IP 地址。另外，需要注意的是，bootstrapCreateRoot 函数可能会进一步初始化 ncclUniqueId，以包含监听的端口信息。

接下来，我们来看一看 bootstrapCreateRoot 函数, 该函数在所有进程间建立环形可达的通信链路。为了便于理解后面的代码，我们先这里简单介绍下环形可达通信链路的含义。具体来说，每个进程都建立一个监听端，并且知道其后序进程的监听地址。这样，就构成了一个如下图所示的环形通信链。例如，rank=1 的进程知道 rank=2 的进程的监听地址，这样，rank=1 的进程可以和 rank=2 的后序进程通信。


点击(此处)折叠或打开

  1. ncclResult_t bootstrapCreateRoot(struct ncclBootstrapHandle* handle, bool idFromEnv) {
  2.   struct ncclSocket* listenSock;
  3.   struct bootstrapRootArgs* args;
  4.   pthread_t thread;

  6.   NCCLCHECK(ncclCalloc(&listenSock, 1));
  7.   NCCLCHECK(ncclSocketInit(listenSock, &handle->addr, handle->magic, ncclSocketTypeBootstrap, NULL, 0));
  8.   NCCLCHECK(ncclSocketListen(listenSock));
  9.   NCCLCHECK(ncclSocketGetAddr(listenSock, &handle->addr));

  11.   NCCLCHECK(ncclCalloc(&args, 1));
  12.   args->listenSock = listenSock;
  13.   args->magic = handle->magic;
  14.   NEQCHECK(pthread_create(&thread, NULL, bootstrapRoot, (void*)args), 0);
  15.   ncclSetThreadName(thread, "NCCL BootstrapR");
  16.   NEQCHECK(pthread_detach(thread), 0); // will not be pthread_join()'d
  17.   return ncclSuccess;
  18. }

 bootstrapCreateRoot 函数首先调用 ncclSocketInit 函数初始化 listenSock。接着，调用 ncclSocketListen 函数绑定并在 listenSock 监听。这里使用的是handle->addr，即前文所获取的bootstrapNetIfAddr。此外，ncclSocketListen 还通过 bind 函数确定了监听端口。ncclSocketGetAddr 函数将监听的网络地址和端口拷贝到 ncclUniqueId 数据结构中。也就是说，此时 ncclUniqueId 包含了 root 进程监听的 IP 地址和端口。（这里其实主要是端口信息，地址在调用bootstrapCreateRoot之前已经对handle->addr进行了赋值）

在实际使用时，用户需要将 ncclUniqueId 广播到所有进程。因此，其它进程都可以通过 root 进程的监听地址连接 root 进程。这样，其它进程可以将它们自己的监听地址发送到 root 进程，从而构建上面提到的环形通信链。

{BANNED}最佳后，bootstrapCreateRoot 函数还启动 bootstrapRoot 线程，在所有进程间建立环形通信链。

l  bootstrapRoot 线程

bootstrapRoot 在所有进程间建立环形通信链，注意这里是bootstrap网络的环形通信。首先，bootstrapRoot 通过下面的代码片段接收其它所有进程连接，获取与 root 进程连接的 sockaddr 和与其它进程连接的 sockaddr，分别保存在 rankAddressesRoot和 rankAddresses中。

点击(此处)折叠或打开

  1. static void *bootstrapRoot(void* rargs) {
  2. ......
  3.   /* Receive addresses from all ranks */
  4.   do {
  5.     struct ncclSocket sock;
  6.     NCCLCHECKGOTO(ncclSocketInit(&sock), res, out);
  7.     NCCLCHECKGOTO(ncclSocketAccept(&sock, listenSock), res, out);
  8.     NCCLCHECKGOTO(bootstrapNetRecv(&sock, &info, sizeof(info)), res, out);
  9.     NCCLCHECKGOTO(ncclSocketClose(&sock), res, out);

  11.     if (c == 0) {
  12.       nranks = info.nranks;
  13.       NCCLCHECKGOTO(ncclCalloc(&rankAddresses, nranks), res, out);
  14.       NCCLCHECKGOTO(ncclCalloc(&rankAddressesRoot, nranks), res, out);
  15.     }

  17.     if (nranks != info.nranks) {
  18.       WARN("Bootstrap Root : mismatch in rank count from procs %d : %d", nranks, info.nranks);
  19.       goto out;
  20.     }

  22.     if (memcmp(zero, &rankAddressesRoot[info.rank], sizeof(union ncclSocketAddress)) != 0) {
  23.       WARN("Bootstrap Root : rank %d of %d ranks has already checked in", info.rank, nranks);
  24.       goto out;
  25.     }

  27.     // Save the connection handle for that rank
  28.     memcpy(rankAddressesRoot+info.rank, &info.extAddressListenRoot, sizeof(union ncclSocketAddress));
  29.     memcpy(rankAddresses+info.rank, &info.extAddressListen, sizeof(union ncclSocketAddress));

  31.     ++c;
  32.     TRACE(NCCL_INIT, "Received connect from rank %d total %d/%d", info.rank, c, nranks);
  33.   } while (c < nranks);
  34.   TRACE(NCCL_INIT, "COLLECTED ALL %d HANDLES", nranks);

上述过程如下图所示。

这里可能有人会疑惑，怎么上来就接收呢？发送方是谁呢？不要忘了bootstrapRoot 是一个独立的线程，主线程和其他rank还在并行执行，他们得到ncclUniqueId后就会调用MPI_Bcast广播了。

接着，root 进程分别连接每一个 rank=i 进程，并将其后序 rank=i+1 进程的监听地址发送给它。这样，每个进程就可以连接其后序进程，从而构成一个环形通信链。

点击(此处)折叠或打开

  1. // Send the connect handle for the next rank in the AllGather ring
  2.   for (int r=0; r<nranks; ++r) {
  3.     int next = (r+1) % nranks;
  4.     struct ncclSocket sock;
  5.     NCCLCHECKGOTO(ncclSocketInit(&sock,rankAddressesRoot+r,magic, ncclSocketTypeBootstrap), res, out);
  6.     NCCLCHECKGOTO(ncclSocketConnect(&sock), res, out);
  7.     NCCLCHECKGOTO(bootstrapNetSend(&sock,rankAddresses+next,sizeof(union ncclSocketAddress)), res, out);
  8.     NCCLCHECKGOTO(ncclSocketClose(&sock), res, out);
  9.   }
  10.   TRACE(NCCL_INIT, "SENT OUT ALL %d HANDLES", nranks);

如下图所示，root 进程连接 rank=2 进程，并将 rank=3 进程的监听地址发送给它。

    {BANNED}最佳后回到nccl-test的run函数中，如上所述，在调用完ncclGetUniqueId后，调用MPI_Bcast将生成的ncclUniqueId广播给所有的rank。这里要注意我们是以root rank为视角讲解的，就如同上面所讲的bootstrapRoot 线程首先通过接收其它所有进程连接。所有一定有其他rank向root发送的过程，那么这一点我们将在后面部分讲解。

小结

在{BANNED}中国第一部分，我们主要分析了 ncclGetUniqueId 函数了实现。整体流程如所示：

该函数在 root 进程中执行，根据用户设置或者系统选择决定所使用的 bootstrap 网络的网卡，并初始化 bootstrap 网络。这里，获取的 ncclUniqueId 实际上是包含魔术值和使用上述网卡 IP 地址和绑定端口的信息。通过广播该 ncclUniqueId，所有进程都可以连接 root 进程，并进一步创建环形通信链路。

 

我们将在下一篇分析 ncclCommInitRank 函数，说明非 root 进程如何和 root 进程协作完成 bootstrap 网络的初始化、communicator 的创建和更为复杂的 NCCL 拓扑探测过程。

上一篇：PCIe总线的事物层和链路层
下一篇：NCCL源码解析2——bootstrap网络初始化