GPU D2D 原理

当下的大模型训练/推理，多 GPU 之间，除了使用高带宽的 nvlink 通信之外，在不支持 nvlink 的 GPU 上，也可以使用 PCIE 进行互联通信。

一种常见的网络拓扑：8 卡 GPU 的机器通过单边 4 卡通过 PCIE SW 连接到单个 CPU，使用这种架构的常见机器包括 A100(A800) PCIE、4090、L4 等等。两个 CPU 之间再通过 QPI 连接。

常见的跨 GPU 通信算子如AllReduce、AlltoAll 都涉及到最基本的操作：两个设备间的D2D拷贝。一般来说，PCIE 4.0 的协议单 lane 提供 16 GT/s 的比特传输带宽，那么 16 lanes 则是 32GB/s 的单向带宽。上面的拓扑结构中，每块 GPU 到 PCIE Switch 的上下行带宽分别为 32GB/s，PCIE switch 到 CPU0 的单向带宽为 32GB/s。于是，做 infra 加速的同学应该关注以下两个关键的问题：

• 实际情况中，任 2 个 GPU 之间的通信能达到理论的 32GB/s 带宽吗？同 socket 和跨 socket 的通信带宽会有什么差异？
• 在同 PCIE 拓扑结构下，不同 GPU 卡型是否会影响 P2P 的通信带宽？

1 基本概念

1.1 GPU Direct P2P for PCIE

在 D2D 通信中， 如果配置了高速互联，比如 A100 SXM （NVLINK）, H100 SXM（NVLINK），或者 AMD 的 fullmesh 卡间互联，D2D 通信是不需要经过 CPU 的，卡本身会直接允许 GPU Direct peer-to-peer access，走高速互联如 NVLINK 通道通信。

GPU Direct P2P 是 NVIDIA GPU 中的一项功能，允许 CUDA 程序访问并将数据从一个 GPU 的内存传输到另一个 GPU 的内存，而无需通过连接到 CPU 的共享系统内存池。 这允许程序员从 CUDA 内核中直接访问另一个 GPU 的内存地址，而不需要显示的内存拷贝，也不需要 CPU 来参与调度。

这个能力一般用来加速 D2D 的拷贝，常常在 A100 SXM 和 H100 SXM 这些机器中使用。但是，通过 PCIE 连接的两块 GPU 是否也可以利用这项能力呢？答案是：yes！

2 实验

那么，尝试在卡本身不支持 peer access 的多卡机器上，通过 CUDA 程序来实现机内任意两卡的 D2D 通信。传输量足够大（4G 以上）能打满带宽。

环境：使用双卡的 RTX2080Ti，没有安装 nvlink 桥；CPU 是 Intel 10940X

RTX2080Ti 和 CPU 之间通过 PCIE gen3x16 连接。

2.1 Naive CUDA API call

可以直接通过 CUDA API cudaMemcpy(Async) 的 cudaMemcpyDeviceToDevice 选项来实现同机内跨卡 D2D 的拷贝。实现方式非常简单：在 srcDevice 上调用 cudaMemcpyAsync。cudaMemcpyAsync 调用:

for (size_t i = 0; i < devicePairs.size();, , i) {
  int fromDevice = devicePairs[i].first;
  // Ensure device is set to fromDevice
  cudaSetDevice(
      fromDevice);  // Set device to fromDevice where stream[i] resides

  checkCudaErrors(cudaEventCreate(&startEvents[i]));
  checkCudaErrors(cudaEventCreate(&stopEvents[i]));

  checkCudaErrors(cudaEventRecord(startEvents[i], streams[i]));
  checkCudaErrors(cudaMemcpyAsync(d_dsts[i], d_srcs[i], size,
                                  cudaMemcpyDeviceToDevice, streams[i]));
  checkCudaErrors(cudaEventRecord(stopEvents[i], streams[i]));
}

// Synchronize all streams to ensure all operations for the current size are
// completed
for (size_t i = 0; i < devicePairs.size();, , i) {
  int fromDevice = devicePairs[i].first;
  cudaSetDevice(fromDevice);  // Ensure device is set before synchronizing

  checkCudaErrors(cudaStreamSynchronize(streams[i]));

  float milliseconds = 0;
  checkCudaErrors(
      cudaEventElapsedTime(&milliseconds, startEvents[i], stopEvents[i]));
  totalMilliseconds[i] += milliseconds;

  checkCudaErrors(cudaEventDestroy(startEvents[i]));
  checkCudaErrors(cudaEventDestroy(stopEvents[i]));
}

从 Nsight System 的 timeline 观察到，cudaMemcpyAsync 实际会以 CPU 为中转，转为一个 D2H 和 H2D 的 memory copy。这个 CPU memory 是 pagable 的，这意味着:

• 拷贝操作无法和其他 CPU 操作异步。Timeline 上无论是多个 devicePair 的 kernel 下发，或者是多轮迭代执行，都无法在 timeline 上同步执行。
• 拷贝过程 CPU 的参与：GPU 无法直接从可分页主机内存访问数据，需要 CPU 参与分配一块临时 pinned memory 才能完成到 GPU 的拷贝，这无疑是低效的。

2.2 优化

1. 基于 CPU 中转的 Double Buffer 流式传输

由于上述实现中，pagable memory 拷贝时非常低效的，并且阻碍了我们同时 launch 多个链路的通信。 我们有没有办法自己实现 pinned memory 的传输，同时能使多路通信同时发生（比如0-1-2-3的环状通信）呢？

为此实现了一个基于 Double buffer 的流式传输代码实现。为了实现流式传输，我们将原数据分成多个相等大小的 chunk，每个 chunk 选择 chunk_size 左右的大小（根据实际的总负载指定，比如传输 4GB，chunk_size 设为 64MB）。在 host memory 上，申请两块 chunk_size 大小的 pinned memory 作为 buffer。双 buffer 是为了实现前一份 chunk 的 H2D 和 下一份 chunk 的 D2H 并发执行。

因为前一份 chunk 的 D2H 结束后，pinned buffer 才能被覆写以接受下一份 chunk，我们需要利用 event 来进行时序的控制。代码实现核心逻辑如下：

template<typename T>  
void nonPeerD2DCopyWithDoublePinned(const T* d_src, int srcDevice,  
                                   T* d_dst, int dstDevice,  
                                   size_t numElements,  
                                   cudaStream_t srcStream,  
                                   cudaStream_t dstStream,  
                                   PinnedMemoryPool<T>& memPool) {  
    const size_t CHUNK_SIZE = memPool.getSize();  
    const size_t numChunks = (numElements + CHUNK_SIZE - 1) / CHUNK_SIZE;  
    
    int currentBuffer = 0;  
    size_t offset = 0;  
    
    // 启动第一次传输  
    if (numChunks > 0) {  
        size_t currentChunkSize = std::min(CHUNK_SIZE, numElements);  
        checkCudaErrors(cudaSetDevice(srcDevice));  
        checkCudaErrors(cudaMemcpyAsync(memPool.getBuffer(currentBuffer),  
                                      d_src,  
                                      currentChunkSize * sizeof(T),  
                                      cudaMemcpyDeviceToHost,  
                                      srcStream));  
        checkCudaErrors(cudaEventRecord(memPool.getSrcEvent(currentBuffer), srcStream));  
    }  
    
    // 处理所有完整的块  
    for (size_t chunk = 1; chunk < numChunks; ++chunk) {  
        int nextBuffer = 1 - currentBuffer;  
        size_t nextOffset = chunk * CHUNK_SIZE;  
        size_t currentChunkSize = std::min(CHUNK_SIZE, numElements - offset);  
        size_t nextChunkSize = std::min(CHUNK_SIZE, numElements - nextOffset);  
        
        checkCudaErrors(cudaSetDevice(dstDevice));  
        checkCudaErrors(cudaStreamWaitEvent(dstStream, memPool.getSrcEvent(currentBuffer)));  
        checkCudaErrors(cudaMemcpyAsync(d_dst + offset,  
                                      memPool.getBuffer(currentBuffer),  
                                      currentChunkSize * sizeof(T),  
                                      cudaMemcpyHostToDevice,  
                                      dstStream));  
        checkCudaErrors(cudaEventRecord(memPool.getDstEvent(currentBuffer), dstStream));  
        
        checkCudaErrors(cudaSetDevice(srcDevice));  
        checkCudaErrors(cudaStreamWaitEvent(srcStream, memPool.getDstEvent(nextBuffer)));  
        checkCudaErrors(cudaMemcpyAsync(memPool.getBuffer(nextBuffer),  
                                      d_src + nextOffset,  
                                      nextChunkSize * sizeof(T),  
                                      cudaMemcpyDeviceToHost,  
                                      srcStream));  
        checkCudaErrors(cudaEventRecord(memPool.getSrcEvent(nextBuffer), srcStream));  
        
        offset = nextOffset;  
        currentBuffer = nextBuffer;  
    }  
    
    // 处理最后一块数据  
    if (numChunks > 0) {  
        size_t currentChunkSize = std::min(CHUNK_SIZE, numElements - offset);  
        checkCudaErrors(cudaSetDevice(dstDevice));  
        checkCudaErrors(cudaStreamWaitEvent(dstStream, memPool.getSrcEvent(currentBuffer)));  
        checkCudaErrors(cudaMemcpyAsync(d_dst + offset,  
                                      memPool.getBuffer(currentBuffer),  
                                      currentChunkSize * sizeof(T),  
                                      cudaMemcpyHostToDevice,  
                                      dstStream));  
        checkCudaErrors(cudaEventRecord(memPool.getDstEvent(currentBuffer), dstStream));  
    }  
}  

memPool 的设计是为了最终做显存的回收，避免显式的 cudaFree 带来的强制同步。从nsys timeline 上我们可以看到成功实现了多个通信链路的并发传输，并且是一个个 chunk 串行流式传输的实现。

这个实现显著提升了 cudaMemcpyAsync 实现的吞吐！并且成功实现了并发。可以看到 0->1, 2->3 的并发传输案例，因为都需要经过PCIE SW-CPU 的通信链路，导致了带宽的竞争，导致单条链路的单向带宽只能达到无并发时的一半。

第一个 D2H 和最后一个 H2D 的拷贝因为没有办法并发，往往会带来一定的延迟和带宽损耗。当总传输量较小，比如当 chunk size 设为64MB，而总传输量只有 256MB 时， double buffer 首尾段的开销就比较大了。如果把 chunk size 调小，那么launch overhead 也会变高。

2：Enable Peer Access

DoubleBuffer 的方案虽然取得了一定优化，还是避免不了需要经过 host memory 中转这件事。因此对于支持 Peer Access 的设备来说，我们应该避免经过 host memory 的拷贝，直接进行设备到设备的传输。这样各个设备之间的通信不需要经过 PCIE-CPU 的通信链路，避免了多卡并行时的竞争。

代码实现上，只需要在前面的基础上增加 cudaDeviceEnablePeerAccess 的使能，再调用cudaMemcpy(Async) 的 cudaMemcpyDeviceToDevice 选项即可。

// Enable peer access between devices if possible and requested
  if (enablePeerAccess) {
    for (const auto& pair : devicePairs) {
      int fromDevice = pair.first;
      int toDevice = pair.second;

      int canAccessPeer = 0;
      checkCudaErrors(
          cudaDeviceCanAccessPeer(&canAccessPeer, fromDevice, toDevice));
      if (canAccessPeer) {
        cudaSetDevice(fromDevice);
        cudaError_t err = cudaDeviceEnablePeerAccess(toDevice, 0);
        if (err == cudaSuccess) {
          std::cout << "Peer access enabled from device " << fromDevice
                    << " to device " << toDevice << std::endl;
        } else {
          std::cout << "Failed to enable peer access from device " << fromDevice
                    << " to device " << toDevice << ": "
                    << cudaGetErrorString(err) << std::endl;
        }
      } else {
        std::cout << "Peer access not supported from device " << fromDevice
                  << " to device " << toDevice << std::endl;
      }

      checkCudaErrors(
          cudaDeviceCanAccessPeer(&canAccessPeer, toDevice, fromDevice));
      if (canAccessPeer) {
        cudaSetDevice(toDevice);
        cudaError_t err = cudaDeviceEnablePeerAccess(fromDevice, 0);
        if (err == cudaSuccess) {
          std::cout << "Peer access enabled from device " << toDevice
                    << " to device " << fromDevice << std::endl;
        } else {
          std::cout << "Failed to enable peer access from device " << toDevice
                    << " to device " << fromDevice << ": "
                    << cudaGetErrorString(err) << std::endl;
        }
      } else {
        std::cout << "Peer access not supported from device " << toDevice
                  << " to device " << fromDevice << std::endl;
      }
    }
  }

// ... 调用cudaMemcpyAsync