NCCL源码详解3：通信器初始化ncclCommInitRank() 含视频教程

爱串门的小马驹

已于 2025-02-19 10:44:42 修改

阅读量2.9k

点赞数 23

分类专栏：集合通信文章标签：人工智能分布式信息与通信

于 2024-05-24 20:06:16 首次发布

本文链接：https://blog.csdn.net/lianghuaju/article/details/138965028

版权

集合通信专栏收录该内容

15 篇文章

订阅专栏

上上节课我们Nvidia NCCL的官方案例中

NCCL源码解读1：NCCL使用/调用案例 Example : One Device per Process or Thread-CSDN博客

已经讲了，第四步：4、通信初始化如何获取唯一ID UniqueId

NCCL源码详解2：通信初始化如何获取唯一ID UniqueId，ncclGetUniqueId()中ncclInit()、bootstrapGetUniqueId()包含视频教程-CSDN博客

今天讲，第六步是：6、初始化NCCL通信器。

初始化NCCL通信器的源码ncclCommInitRank()

一、ncclCommInitRank()源码详解

哈哈哈，视频教程在这

哔哩哔哩_bilibili

什么是通信器？

通信器：NCCL通信器是一个用于管理多GPU环境中进程间通信的集合，通过定义通信域来支持并行计算任务的执行。一旦通信器被创建，进程就可以使用这个通信器来进行集合通信操作/原语(Broadcast、AllReduce等)，例如数据同步和数据传输等。

简单点说：每个GPU知道我要和哪些GPU通信，并且能够进行集合通信。

哈哈哈，通信器的定义是我瞎编的，官网没找到定义。

ncclCommInitRank()核心逻辑

1、加载CUDA驱动

2、获取当前CUDA设备ID

3、根据rank号和CUDA设备ID等，完成NCCL通信器初始化

源码速递

源码位置：nccl-master\src\init.cc

ncclResult_t ncclCommInitRank(ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank) {  

  // 1、加载CUDA驱动/  
  (void)ncclCudaLibraryInit();  
  
  // 声明一个变量来存储当前CUDA设备的ID  
  int cudaDev;  

  // 初始化NCCL配置结构体，使用默认的配置  
  ncclConfig_t config = NCCL_CONFIG_INITIALIZER;  

  // 2、获取当前CUDA设备ID   
  CUDACHECK(cudaGetDevice(&cudaDev));  
  
  // 存储参数方便可视化和分析 
  NvtxParamsCommInitRank payload{myrank, nranks, cudaDev};  
  NVTX3_FUNC_WITH_PARAMS(CommInitRank, CommInitRankSchema, payload)  

  //3、根据rank号和CUDA设备ID等，完成NCCL通信器初始化/
  // 调用ncclCommInitRankDev函数来初始化NCCL通信器  
  // ncclCommInitRankDev是一个更底层的函数，它允许指定CUDA设备ID和NCCL配置    
  NCCLCHECK(ncclCommInitRankDev(newcomm, nranks, commId, myrank, cudaDev, &config));  
  
  // 如果所有操作都成功完成，则返回ncclSuccess表示成功  
  return ncclSuccess;  

}

ncclCommInitRank()中最最核心的显然是第3步ncclCommInitRankDev(),下面我们向下挖一挖。

二、ncclCommInitRankDev()源码解读

ncclCommInitRankDev()核心逻辑

1、确保NCCL驱动已经初始化

2、确保CUDA驱动已经初始化

3、检查参数/指针是否为空，是否在有效范围内

4、配置NCCL通信器的一些属性，是否阻塞，通信通道数量等

5、分配一个作业对象 job,并设置作业对象的各个成员变量

6、使用 ncclAsyncLaunch 异步启动 ncclCommInitRankFunc 函数来初始化通信。异步启动通信器初始化

7、退出和错误处理

源码速递

源码位置：nccl-master\src\init.cc

static ncclResult_t ncclCommInitRankDev(ncclComm_t* newcomm, int nranks, ncclUniqueId commId, int myrank, int cudaDev, ncclConfig_t *config) {
  ncclResult_t res = ncclSuccess; // 初始化结果，成功则为ncclSuccess
  ncclComm_t comm = NULL; // 初始化通信对象指针为NULL
  struct ncclCommInitRankAsyncJob *job = NULL; // 初始化异步任务结构体指针为NULL
  const char* env = ncclGetEnv("NCCL_COMM_ID"); // 获取环境变量NCCL_COMM_ID的值

  if (env && myrank == 0) { // 如果环境变量存在且当前节点是第一个节点
    INFO(NCCL_ENV, "NCCL_COMM_ID set by environment to %s", env); // 输出日志，显示环境变量设置的值
    NCCLCHECKGOTO(bootstrapCreateRoot((struct ncclBootstrapHandle*)&commId, true), res, fail); // 创建根节点
  }

  1、确保NCCL驱动已经初始化
  NCCLCHECKGOTO(ncclInit(), res, fail); // 初始化NCCL库
  if (myrank == 0) showVersion(); // 如果当前节点是第一个节点，显示版本信息

  2、确保CUDA运行时已经初始化
  CUDACHECKGOTO(cudaFree(NULL), res, fail); // 尝试释放NULL指针，如果失败则返回错误

  3、检查参数/指针是否为空，是否在有效范围内
  NCCLCHECKGOTO(PtrCheck(newcomm, "CommInitRank", "newcomm"), res, fail); // 检查通信对象指针是否有效
  NCCLCHECKGOTO(PtrCheck(config, "CommInitRank", "config"), res, fail); // 检查配置结构体指针是否有效
  if (nranks < 1 || myrank < 0 || myrank >= nranks) { // 验证输入参数
    WARN("Invalid rank requested : %d/%d", myrank, nranks); // 如果参数无效，输出警告
    res = ncclInvalidArgument; // 设置结果为无效参数错误
    goto fail; // 跳转到失败处理部分
  }

  3、配置NCCL通信器的一些属性，是否阻塞，通信通道数量等
  NCCLCHECKGOTO(ncclCalloc(&comm, 1), res, fail); // 为通信对象分配内存
  comm->startMagic = comm->endMagic = NCCL_MAGIC; // 设置通信对象的开始和结束魔术数，用于检测通信对象损坏
  NCCLCHECKGOTO(ncclCudaHostCalloc((uint32_t**)&comm->abortFlag, 1), res, fail); // 为通信对象的中断标志分配内存
  NCCLCHECKGOTO(ncclCalloc((uint32_t**)&comm->abortFlagRefCount, 1), res, fail); // 为通信对象的中断标志引用计数分配内存
  *comm->abortFlagRefCount = 1; // 初始化中断标志引用计数为1
  NCCLCHECKGOTO(parseCommConfig(comm, config), res, fail); // 解析通信配置
  /* start with ncclInternalError and will be changed to ncclSuccess if init succeeds. */
  comm->initState = ncclInternalError; // 初始化通信对象的状态为内部错误，若初始化成功则修改为成功
  *newcomm = comm; // 将通信对象赋值给传入的指针

  5、分配一个作业对象 job,并设置作业对象的各个成员变量
  NCCLCHECKGOTO(ncclCalloc(&job, 1), res, fail); // 为异步任务结构体分配内存
  job->comm = comm; // 设置异步任务的通信对象
  job->nranks = nranks; // 设置异步任务的节点数量
  job->commId = commId; // 设置异步任务的唯一ID
  job->myrank = myrank; // 设置异步任务的当前节点排名
  job->cudaDev = cudaDev; // 设置异步任务的CUDA设备号

  6、使用 ncclAsyncLaunch 异步启动 ncclCommInitRankFunc 函数来初始化通信。当这个函数完成时，它将自动调用 free 函数来释放 comm 对象     异步启动通信器初始化
  NCCLCHECKGOTO(ncclAsyncLaunch(&job->base, ncclCommInitRankFunc, NULL, free, comm), res, fail); // 异步启动初始化通信对象的任务

  7、退出和错误处理
exit:
  return ncclGroupErrCheck(res); // 返回最终的结果，检查是否有错误发生
fail:
  if (comm) { // 如果通信对象不为空
    if (comm->abortFlag) ncclCudaHostFree((void *)comm->abortFlag); // 如果存在中断标志，释放其内存
    if (comm->abortFlagRefCount) free(comm->abortFlagRefCount); // 释放中断标志引用计数的内存
    free(comm); // 释放通信对象的内存
  }
  if (newcomm) *newcomm = NULL; // 将通信对象指针设置为NULL
  goto exit; // 跳转到退出部分
}

主要步骤说明

前三步没啥好说的，主要就是检测状态，没有初始化的初始化，甚至不能算核心逻辑。

第4步说明：

第4步：4、配置NCCL通信器的一些属性，是否阻塞，通信通道数量等

直接看parseCommConfig()函数的源码会有点懵逼，那就直接看parseCommConfig()的最终结果，其实就是赋值了comm->config

源码速递

源码位置：nccl-master\src\init.cc

static ncclResult_t parseCommConfig(ncclComm_t comm, ncclConfig_t *config) {

  /*其它代码略*/


  //通信器配置赋值
  /* assign config to communicator */
  comm->config.blocking = internalConfigPtr->blocking;
  comm->config.cgaClusterSize = internalConfigPtr->cgaClusterSize;
  comm->config.minCTAs = internalConfigPtr->minCTAs;
  comm->config.maxCTAs = internalConfigPtr->maxCTAs;
  comm->config.netName = internalConfigPtr->netName;
  comm->config.splitShare = internalConfigPtr->splitShare;

  NCCLCHECKGOTO(envConfigOverride(comm), ret, fail);

exit:
  return ret;
fail:
  goto exit;
}

第6步说明：

异步启动通信器初始化

ncclCommInitRankFunc()核心逻辑：

1、获取 CUDA 设备和架构信息，初始化 CUDA 内核

2、是否有父通信器

a、有，从父通信器分裂出来子通信器，并初始化

b、无，直接为其分配内存，并初始化

3、设置通信器的CUDA架构版本和哈希值。

4、始化当前通信器的传输层。

5、加载调整器插件。调整器用于动态调整通信算法，以优化性能。

6、更新通信器状态为成功，表示通信器初始化成功

源码速递

源码位置：nccl-master\src\init.cc

// 定义了一个函数ncclCommInitRankFunc，用于初始化NCCL通信器的特定排名  
static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) {  
    
  // 将传入的通用作业指针转换为特定于ncclCommInitRankAsyncJob的指针  
  struct ncclCommInitRankAsyncJob* job = (struct ncclCommInitRankAsyncJob*)job_;  
    
  // 获取作业中指定的通信器  
  ncclComm_t comm = job->comm;  
    
  // 初始化结果码为成功  
  ncclResult_t res = ncclSuccess;  
    
  // 初始化变量用于保存CUDA设备的计算能力的主版本和次版本  
  int archMajor, archMinor;  
    
  // 初始化变量用于保存内核栈的最大大小（字节为单位）  
  size_t maxLocalSizeBytes = 0;  
    
  // 获取作业中指定的CUDA设备ID  
  int cudaDev = job->cudaDev;  
    
  // 初始化一个指向父通信器排名的指针（初始为NULL）  
  int* parentRanks = NULL;  
  /  
  // 初始化一个变量用于保存CUDA设备的架构值（由主版本和次版本组成）  
  int cudaArch;  
  
  //
  1、获取 CUDA 设备和架构信息，初始化 CUDA 内核
  // 设置CUDA设备为指定的设备ID  
  CUDACHECKGOTO(cudaSetDevice(cudaDev), res, fail);  
    
  // 获取CUDA设备的主计算能力版本  
  CUDACHECKGOTO(cudaDeviceGetAttribute(&archMajor, cudaDevAttrComputeCapabilityMajor, cudaDev), res, fail);  
    
  // 获取CUDA设备的次计算能力版本  
  CUDACHECKGOTO(cudaDeviceGetAttribute(&archMinor, cudaDevAttrComputeCapabilityMinor, cudaDev), res, fail);  
    
  // 计算CUDA设备的架构值（例如，7.5对应750）  
  cudaArch = 100*archMajor + 10*archMinor;  
  
  // 为指定设备初始化NCCL内核  
  NCCLCHECK(ncclInitKernelsForDevice(cudaArch, &maxLocalSizeBytes));  
    
  // 如果内核栈的最大大小大于0且允许设置栈大小，则设置CUDA设备的栈大小限制  
  // 这可以避免在加载时重新配置CUDA内存（例如，NVSHMEM问题）  
  if (maxLocalSizeBytes > 0 && ncclParamSetStackSize() == 1) {  
    TRACE(NCCL_INIT, "Setting cudaLimitStackSize to %zi", maxLocalSizeBytes);  
    CUDACHECKIGNORE(cudaDeviceSetLimit(cudaLimitStackSize, maxLocalSizeBytes));  
  }  

  /
  /*2、是否有父通信器
        a、有，从父通信器分裂出来子通信器，并初始化
        b、无，直接为其分配内存，并初始化*/
  // 如果指定了父通信器  
  if (job->parent) {  
    // 为父通信器的排名分配内存  
    NCCLCHECKGOTO(ncclCalloc(&parentRanks, job->parent->nRanks), res, fail);  
      
    // 获取分裂信息，如新的排名数、当前排名、父通信器的排名等  
    NCCLCHECKGOTO(commGetSplitInfo(comm, job->parent, job->color, job->key, &job->nranks, &job->myrank, parentRanks), res, fail);  
      
    // 如果颜色为负，表示只是参与了一个allgather操作，但不需要创建新的通信器  
    // 我们已经完成了需要做的部分，因此跳转到exit标签  
    if (job->color == NCCL_SPLIT_NOCOLOR) goto exit;  
      
    // 生成一个唯一的通信器ID，该ID基于父通信器的哈希值  
    snprintf((char*)&job->commId, sizeof(job->commId), "%016lx-%d", job->parent->commHash, job->color);  
      
    // 分配通信器资源  
    NCCLCHECKGOTO(commAlloc(comm, job->parent, job->nranks, job->myrank), res, fail);  
      
    // 初始化分裂后的通信器  
    NCCLCHECKGOTO(bootstrapSplit((struct ncclBootstrapHandle*)&job->commId, comm, job->parent, job->color, job->key, parentRanks), res, fail);  
  } else {  
    // 如果没有指定父通信器，则为新的通信器分配资源  
    NCCLCHECKGOTO(commAlloc(comm, NULL, job->nranks, job->myrank), res, fail);  

    // 初始化新的通信器  

    NCCLCHECKGOTO(bootstrapInit((struct ncclBootstrapHandle*)&job->commId, comm), res, fail);  

}  
  
///
3、设置通信器的CUDA架构版本和哈希值。/

comm->cudaArch = cudaArch;  

// 根据通信器的唯一ID计算哈希值，并设置为通信器的哈希值  

comm->commHash = getHash(job->commId.internal, NCCL_UNIQUE_ID_BYTES);  

// 如果指定了父通信器，打印通信器的分裂信息  

if (job->parent) {  

  INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx parent %p color %d key %d commId 0x%llx - Init START",  

  comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId));  

} else {  

  // 如果没有指定父通信器，打印通信器的初始化信息  

  INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx commId 0x%llx - Init START",  

  comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, (unsigned long long)hashUniqueId(job->commId));  

}  
  
/
4、根据通信器的配置，初始化所有的传输层///  

NCCLCHECKGOTO(initTransportsRank(comm, job->parent), res, fail);  
  
/
5、加载NCCL的调优插件  //

NCCLCHECKGOTO(ncclTunerPluginLoad(&comm->tuner), res, fail);  

if (comm->tuner) {  

  // 如果成功加载了调优插件，调用其初始化函数  

  NCCLCHECK(comm->tuner->init(comm->nRanks, comm->nNodes, ncclDebugLog, &comm->tunerContext));  

}  
  
/
 6、更新通信器状态为成功，表示通信器初始化成功 /
// update communicator state  
comm->initState = ncclSuccess;

//
/到这里核心逻辑就结束了，已经完成通信器初始化，下面是日志跟踪等非核心逻辑///
  // Trace this call for replay tool
  if (job->parent) {
    /* unlink child abort flag. */
    __atomic_store_n(&job->parent->childAbortFlag, NULL, __ATOMIC_RELEASE);
    TRACE_CALL("ncclCommSplit(%p, %d, %d, %p, %d, %d)",
                job->parent, job->color, job->key, comm, comm->rank, comm->nRanks);
  } else {
    TRACE_CALL("ncclCommInitRank(%p, %d, 0x%llx, %d, %d)",
                comm, comm->nRanks, (unsigned long long)hashUniqueId(job->commId), comm->rank, comm->cudaDev);
  }

  if (job->parent) {
    INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx parent %p color %d key %d commId 0x%llx - Init COMPLETE",
    comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId));
  } else {
    INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx commId 0x%llx - Init COMPLETE",
    comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, (unsigned long long)hashUniqueId(job->commId));
  }
exit:
  if (job->newcomm) {
    /* assign it to user pointer. */
    __atomic_store_n(job->newcomm, comm, __ATOMIC_RELEASE);
  }
  free(parentRanks);
  return res;
fail:
  comm->initState = res;
  goto exit;
}