【NCCL】5 GPU 间链路 Preconnect 机制

NCCL GPU 间链路 Preconnect 机制详解

详细分析 NCCL 中 GPU 间链路的 preconnect（预连接）机制。这是一个关键的优化，用于在实际通信前建立好所有必要的连接。

Preconnect 机制概览

RING
TREE
PAT
NVLS
ncclPrepareTasks
标记需要连接的算法
algoNeedConnect数组
ncclCollPreconnect
根据算法类型
ncclTransportRingConnect
ncclTransportTreeConnect
ncclTransportPatConnect
ncclNvlsBufferSetup
ncclTransportP2pConnect
ncclTransportP2pSetup
建立物理连接

详细流程

第一阶段：标记需要连接的算法

1. 在 ncclPrepareTasks() 中标记

ncclResult_t ncclPrepareTasks(struct ncclComm* comm, bool* algoNeedConnect, bool* needConnect, ...) {
  // 遍历所有任务
  task = ncclIntruQueueHead(&planner->collTaskQueue);
  while (task != nullptr) {
    // 如果是运行时连接模式，且该算法尚未初始化
    if (comm->runtimeConn && comm->initAlgoChannels[task->algorithm] == false) {
      // 标记该算法需要连接
      comm->initAlgoChannels[task->algorithm] = true;
      algoNeedConnect[task->algorithm] = true;  // ← 标记需要连接
      *needConnect = true;
    }
    task = task->next;
  }
}

第二阶段：执行集合通信预连接

2. 在 groupLaunch() 中调用

static ncclResult_t groupLaunch(...) {
  // 准备任务并预连接
  do {
    comm = cliqueHead;
    do {
      NCCLCHECKGOTO(ncclPrepareTasksAndCollPreconnect(comm, simInfo, &asyncCollJobs), ret, fail);
      comm = comm->groupNext[ncclGroupTaskTypeCollective];
    } while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
    
    // 异步执行连接任务
    NCCLCHECKGOTO(asyncJobLaunch(&asyncCollJobs, groupAbortFlag), ret, fail);
  } while (cliqueHead != nullptr);
}

3. ncclPrepareTasksAndCollPreconnect()

static ncclResult_t ncclPrepareTasksAndCollPreconnect(...) {
  bool needConnect = false;
  bool algoNeedConnect[NCCL_NUM_ALGORITHMS];
  memset(algoNeedConnect, 0, sizeof(bool) * NCCL_NUM_ALGORITHMS);
  
  // 调用 ncclPrepareTasks 标记需要连接的算法
  NCCLCHECK(ncclPrepareTasks(comm, algoNeedConnect, &needConnect, simInfo));
  
  // 如果需要连接，创建异步任务
  if (comm->cuMemSupport && needConnect) {
    struct ncclPreconnectJob* job;
    job->base.func = ncclCollPreconnectFunc;  // ← 异步执行的函数
    job->algoNeedConnect = algoNeedConnect;   // ← 传递需要连接的算法
    ncclIntruQueueEnqueue(asyncCollJobs, &job->base);
  }
}

4. ncclCollPreconnect() - 核心分发函数

static ncclResult_t ncclCollPreconnect(struct ncclComm* comm, bool* algoNeedConnect) {
  // 遍历所有算法
  for (int i = 0; i < NCCL_NUM_ALGORITHMS; ++i) {
    if (algoNeedConnect[i]) {
      switch (i) {
        case NCCL_ALGO_RING:
          NCCLCHECK(ncclTransportRingConnect(comm));  // ← Ring 算法连接
          break;
        case NCCL_ALGO_TREE:
          NCCLCHECK(ncclTransportTreeConnect(comm));  // ← Tree 算法连接
          break;
        case NCCL_ALGO_PAT:
          NCCLCHECK(ncclTransportPatConnect(comm));   // ← PAT 算法连接
          break;
        case NCCL_ALGO_NVLS:
          NCCLCHECK(ncclNvlsBufferSetup(comm));       // ← NVLS 缓冲区设置
          break;
        // ... 其他算法
      }
    }
  }
}

第三阶段：PAT 算法的连接过程

5. ncclTransportPatConnect() - PAT 特定连接

ncclResult_t ncclTransportPatConnect(struct ncclComm* comm) {
  if (comm && comm->nRanks > 1) {
    // PAT 使用二叉树结构，需要连接多个层级
    for (int mask=1; mask<comm->nRanks; mask<<=1) {
      // 计算前驱和后继节点
      int prevPeer = (comm->rank + mask) % comm->nRanks;
      int nextPeer = (comm->rank + comm->nRanks - mask) % comm->nRanks;
      
      // ReduceScatter 方向：prevPeer -> 本节点 -> nextPeer
      for (int c = 0; c < comm->nChannels; c++) {
        NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &prevPeer, 1, &nextPeer, 0), ret, fail);
      }
      NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_TREE], 0), ret, fail);
      
      // AllGather 方向：nextPeer -> 本节点 -> prevPeer（反向）
      for (int c = 0; c < comm->nChannels; c++) {
        NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &nextPeer, 1, &prevPeer, 0), ret, fail);
      }
      NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_TREE], 0), ret, fail);
    }
    INFO(NCCL_INIT, "Connected binomial trees");
  }
}

PAT 连接模式说明：

• PAT 使用**二叉树（binomial tree）**结构
• 对于 N 个节点，需要 log2(N) 层连接
• 每层连接距离为 2^k（k=0,1,2,...）
• 例如 8 个节点：
  • 第1层：距离1（0↔1, 2↔3, 4↔5, 6↔7）
  • 第2层：距离2（0↔2, 1↔3, 4↔6, 5↔7）
  • 第3层：距离4（0↔4, 1↔5, 2↔6, 3↔7）

第四阶段：底层连接建立

6. ncclTransportP2pConnect() - 标记连接

ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, 
                                     int nrecv, int* peerRecv, 
                                     int nsend, int* peerSend, 
                                     int connIndex) {
  struct ncclChannel* channel = &comm->channels[channelId];
  uint64_t mask = 1UL << channel->id;
  
  // 标记需要接收的对等节点
  for (int i=0; i<nrecv; i++) {
    int peer = peerRecv[i];
    if (peer == -1 || peer >= comm->nRanks || peer == comm->rank || 
        channel->peers[peer]->recv[connIndex].connected) continue;
    comm->connectRecv[peer] |= mask;  // ← 标记该通道需要从 peer 接收
  }
  
  // 标记需要发送的对等节点
  for (int i=0; i<nsend; i++) {
    int peer = peerSend[i];
    if (peer == -1 || peer >= comm->nRanks || peer == comm->rank || 
        channel->peers[peer]->send[connIndex].connected) continue;
    comm->connectSend[peer] |= mask;  // ← 标记该通道需要向 peer 发送
  }
  
  return ncclSuccess;
}

7. ncclTransportP2pSetup() - 建立物理连接

这是最核心的函数，负责实际建立连接：

ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex) {
  // 分配临时数据结构
  struct ncclConnect** data;
  struct ncclConnect** recvData;
  struct ncclConnect** sendData;
  
  // 遍历所有对等节点（除了自己）
  for (int i=1; i<comm->nRanks; i++) {
    int recvPeer = (comm->rank - i + comm->nRanks) % comm->nRanks;
    int sendPeer = (comm->rank + i) % comm->nRanks;
    uint64_t recvMask = comm->connectRecv[recvPeer];
    uint64_t sendMask = comm->connectSend[sendPeer];
    
    // 步骤 1: 选择传输方式（P2P/SHM/NET）并设置连接信息
    for (int c=0; c<MAXCHANNELS; c++) {
      if (recvMask & (1UL<<c)) {
        // 选择最优传输方式（P2P CUDA、共享内存、网络等）
        NCCLCHECKGOTO(selectTransport<0>(comm, graph, recvData[p]+recvChannels++, 
                                         c, recvPeer, connIndex, &type), ret, fail);
      }
    }
    for (int c=0; c<MAXCHANNELS; c++) {
      if (sendMask & (1UL<<c)) {
        NCCLCHECKGOTO(selectTransport<1>(comm, graph, sendData[p]+sendChannels++, 
                                         c, sendPeer, connIndex, &type), ret, fail);
      }
    }
    
    // 步骤 2: 通过 Bootstrap 交换连接信息
    if (sendPeer == recvPeer) {
      // 同一个对等节点（发送和接收）
      NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, recvPeer, bootstrapTag, data[p], 
                                  sizeof(struct ncclConnect)*(recvChannels+sendChannels)), ret, fail);
      NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, recvPeer, bootstrapTag, data[p], 
                                  sizeof(struct ncclConnect)*(recvChannels+sendChannels)), ret, fail);
    } else {
      // 不同的对等节点
      if (recvChannels) NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, recvPeer, bootstrapTag, 
                                                     recvData[p], sizeof(struct ncclConnect)*recvChannels), ret, fail);
      if (sendChannels) NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, sendPeer, bootstrapTag, 
                                                     sendData[p], sizeof(struct ncclConnect)*sendChannels), ret, fail);
      if (sendChannels) NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, sendPeer, bootstrapTag, 
                                                     sendData[p], sizeof(struct ncclConnect)*sendChannels), ret, fail);
      if (recvChannels) NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, recvPeer, bootstrapTag, 
                                                     recvData[p], sizeof(struct ncclConnect)*recvChannels), ret, fail);
    }
    
    // 步骤 3: 建立实际连接（可能需要多轮）
    bool allChannelsConnected = false;
    while (!allChannelsConnected) {
      allChannelsConnected = true;
      for (int c=0; c<MAXCHANNELS; c++) {
        if (sendMask & (1UL<<c)) {
          struct ncclConnector* conn = comm->channels[c].peers[sendPeer]->send + connIndex;
          if (conn->connected == 0) {
            // 调用传输层的 connect 函数
            NCCLCHECKGOTO(conn->transportComm->connect(comm, sendData[p] + sendDataOffset, 
                                                       1, comm->rank, conn), ret, fail);
            if (ret == ncclSuccess) {
              conn->connected = 1;
              // 将连接信息复制到设备内存
              CUDACHECKGOTO(cudaMemcpyAsync(&comm->channels[c].devPeersHostPtr[sendPeer]->send[connIndex], 
                                            &conn->conn, sizeof(struct ncclConnInfo), 
                                            cudaMemcpyHostToDevice, hostStream), ret, fail);
            } else if (ret == ncclInProgress) {
              allChannelsConnected = false;  // 需要继续等待
            }
          }
        }
        // 接收通道类似处理...
      }
    }
  }
  
  // 步骤 4: 同步所有节点，确保连接完成
  for (int i = 1; i < comm->nRanks; i++) {
    // 通过 Bootstrap 进行最终同步
    NCCLCHECKGOTO(bootstrapSend/Recv(...), ret, fail);
    // 清除连接标记
    comm->connectRecv[recvPeer] = comm->connectSend[sendPeer] = 0UL;
  }
}

不同算法的连接模式

1. Ring 算法连接 (ncclTransportRingConnect)

// Ring 拓扑：每个节点连接前驱和后继
for (int c = 0; c < comm->nChannels; c++) {
  struct ncclChannel* channel = comm->channels + c;
  // 连接：prev <- 本节点 -> next
  ncclTransportP2pConnect(comm, c, 
                          1, &channel->ring.prev,  // 接收：从 prev
                          1, &channel->ring.next,  // 发送：到 next
                          0);
}
ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_RING], 0);

连接模式：

Rank 0 <-> Rank 1 <-> Rank 2 <-> ... <-> Rank N-1 <-> Rank 0

2. Tree 算法连接 (ncclTransportTreeConnect)

// Tree 拓扑：每个节点连接父节点和多个子节点
for (int c = 0; c < comm->nChannels; c++) {
  struct ncclChannel* channel = comm->channels + c;
  // 下行连接：本节点 -> 子节点们
  ncclTransportP2pConnect(comm, c, 
                          NCCL_MAX_TREE_ARITY, channel->tree.down,  // 发送到多个子节点
                          1, &channel->tree.up,                      // 接收从父节点
                          0);
  // 上行连接：子节点们 -> 本节点
  ncclTransportP2pConnect(comm, c, 
                          1, &channel->tree.up,                      // 发送到父节点
                          NCCL_MAX_TREE_ARITY, channel->tree.down,  // 接收从多个子节点
                          0);
}
ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_TREE], 0);

连接模式：

        Rank 0 (root)
       /   |   \
   Rank 1 Rank 2 Rank 3
   /  \
Rank 4 Rank 5

3. PAT 算法连接 (ncclTransportPatConnect)

// PAT 使用二叉树（binomial tree）结构
for (int mask=1; mask<comm->nRanks; mask<<=1) {
  // 计算距离为 mask 的对等节点
  int prevPeer = (comm->rank + mask) % comm->nRanks;
  int nextPeer = (comm->rank + comm->nRanks - mask) % comm->nRanks;
  
  // ReduceScatter 方向连接
  for (int c = 0; c < comm->nChannels; c++) {
    ncclTransportP2pConnect(comm, c, 1, &prevPeer, 1, &nextPeer, 0);
  }
  ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_TREE], 0);
  
  // AllGather 方向连接（反向）
  for (int c = 0; c < comm->nChannels; c++) {
    ncclTransportP2pConnect(comm, c, 1, &nextPeer, 1, &prevPeer, 0);
  }
  ncclTransportP2pSetup(comm, &comm->graphs[NCCL_ALGO_TREE], 0);
}

PAT 连接示例（8 个节点）：

层级 0 (mask=1): 每个节点连接距离1的节点
  0↔1, 2↔3, 4↔5, 6↔7

层级 1 (mask=2): 每个节点连接距离2的节点
  0↔2, 1↔3, 4↔6, 5↔7

层级 2 (mask=4): 每个节点连接距离4的节点
  0↔4, 1↔5, 2↔6, 3↔7

传输层选择机制

selectTransport() - 选择最优传输方式

template <int type>  // type: 0=recv, 1=send
static ncclResult_t selectTransport(...) {
  // 按优先级尝试不同的传输方式
  for (int t=0; t<NTRANSPORTS; t++) {
    struct ncclTransport *transport = ncclTransports[t];
    // 传输方式优先级：
    // 0. p2pTransport   - CUDA P2P（GPU Direct）
    // 1. shmTransport   - 共享内存（同节点）
    // 2. netTransport   - 网络传输（跨节点）
    // 3. collNetTransport - 集合网络
    
    int ret = 0;
    NCCLCHECK(transport->canConnect(&ret, comm, graph, myInfo, peerInfo));
    if (ret) {
      connector->transportComm = transportComm;
      // 调用传输层的 setup 函数
      NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, 
                                     connect, connector, channelId, connIndex));
      return ncclSuccess;
    }
  }
}

传输方式选择逻辑

1. P2P Transport (优先级最高)

   • 条件：同节点 GPU，支持 CUDA P2P
   • 优势：最快，直接 GPU 间内存访问

2. SHM Transport

   • 条件：同节点，不支持 CUDA P2P
   • 优势：通过共享内存通信

3. NET Transport

   • 条件：跨节点
   • 方式：通过网络（InfiniBand、以太网等）

连接信息交换

Bootstrap 协议

使用 bootstrapSend/Recv() 交换连接信息：

// 发送连接信息给对等节点
bootstrapSend(comm->bootstrap, peer, tag, connectData, size);

// 接收对等节点的连接信息
bootstrapRecv(comm->bootstrap, peer, tag, connectData, size);

ncclConnect 结构体包含：

• 传输类型（P2P/SHM/NET）
• 缓冲区地址
• IPC 句柄（CUDA IPC 或共享内存）
• 网络地址（跨节点）

连接状态管理

连接标记位图

// 每个对等节点有两个 64 位位图
comm->connectSend[peer]  // 需要向 peer 发送的通道掩码
comm->connectRecv[peer]  // 需要从 peer 接收的通道掩码

// 例如：connectSend[3] = 0b00000111
// 表示通道 0、1、2 需要向 Rank 3 发送数据

连接状态

struct ncclConnector {
  int connected;              // 0=未连接, 1=已连接
  struct ncclTransportComm* transportComm;  // 传输层接口
  struct ncclConnInfo conn;   // 连接信息（设备端使用）
  // ...
};

运行时连接 vs 初始化连接

初始化时连接（默认）

• 在 ncclCommInitRank() 中：

// 初始化时就连接所有算法
NCCLCHECKGOTO(ncclTransportRingConnect(comm), ret, fail);
if (comm->maxLocalRanks == 1) 
  NCCLCHECKGOTO(ncclTransportPatConnect(comm), ret, fail);

运行时连接（comm->runtimeConn）

• 延迟到第一次使用时才连接
• 在 ncclPrepareTasks() 中检测并标记
• 在 ncclCollPreconnect() 中执行连接

异步连接机制

异步任务执行

// 创建异步连接任务
struct ncclPreconnectJob* job;
job->base.func = ncclCollPreconnectFunc;  // 异步执行的函数
job->algoNeedConnect = algoNeedConnect;

// 加入异步任务队列
ncclIntruQueueEnqueue(asyncCollJobs, &job->base);

// 启动异步任务（可能创建新线程）
asyncJobLaunch(&asyncCollJobs, groupAbortFlag);

异步执行流程

1. 单任务：在主线程直接执行
2. 多任务：创建 pthread 并行执行
3. 等待完成 ：通过原子变量 job->state 轮询状态

总结

Preconnect 的关键作用

1. 性能优化：提前建立连接，避免首次通信时的延迟
2. 资源准备：分配缓冲区、建立 IPC 句柄、设置网络连接
3. 拓扑优化：根据硬件拓扑选择最优传输方式
4. 并行化：支持异步并行连接多个节点

Preconnect 的执行时机

时机	触发点	说明
初始化时	ncclCommInitRank	连接常用算法（Ring、PAT等）
运行时	ncclPrepareTasks	首次使用某算法时连接
P2P操作	p2pTaskAppend	标记需要连接的通道

关键数据结构

• comm->connectSend/Recv[peer]: 连接标记位图
• comm->initAlgoChannels[algo]: 算法是否已初始化
• algoNeedConnect[algo]: 本次需要连接的算法
• ncclConnect: 连接信息交换结构
• ncclConnector: 连接器状态和接口