NCCL 初始化通信的关键

ncclGetUniqueId这个函数生成一个全局通信标识符。这个 ID 只需要在一个进程中生成（通常 rank=0），再通过 MPI 或 socket 发送 给其他进程。

ncclCommInitRank需要输入参数id，即所有进程共享的ID

那么流程就很清楚了，一个进程初始化ID，通过MPI的接口广播给所有进程，进程通过这个ID建立NCCL通信组然后再进行GPU通信，比纯OpenMPI更快。

#include <mpi.h>
#include <nccl.h>
#include <cuda_runtime.h>
#include <stdio.h>

int main(int argc, char* argv[]) {
    MPI_Init(&argc, &argv);

    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    cudaSetDevice(rank);

    ncclUniqueId id;
    ncclComm_t comm;

    // rank 0 生成 unique id，并广播给其他进程
    if (rank == 0) ncclGetUniqueId(&id);
    MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);

    ncclCommInitRank(&comm, size, id, rank);

    float* sendbuf;
    float* recvbuf;
    cudaMalloc(&sendbuf, sizeof(float));
    cudaMalloc(&recvbuf, sizeof(float));

    float val = (float)(rank + 1); // 1.0, 2.0
    cudaMemcpy(sendbuf, &val, sizeof(float), cudaMemcpyHostToDevice);

    ncclAllReduce(sendbuf, recvbuf, 1, ncclFloat, ncclSum, comm, 0);

    float result;
    cudaMemcpy(&result, recvbuf, sizeof(float), cudaMemcpyDeviceToHost);
    printf("[Rank %d] Result = %f\n", rank, result);

    ncclCommDestroy(comm);
    cudaFree(sendbuf);
    cudaFree(recvbuf);
    MPI_Finalize();
    return 0;
}