3. NVLink多对通信模式测试

任务目标

在包含4张GPU卡（2对NVLink连接）的节点上，实现多GPU全交换通信模式测试：每张GPU同时向其余3张GPU发送数据并接收数据，测量聚合带宽与端到端延迟，并分析通信规模扩展时的性能变化趋势。

环境依赖

• NVIDIA GPU（Volta及以上架构，支持NVLink）
• CUDA Toolkit ≥ 11.0（包含nvcc与cuda_runtime.h）
• GPU间通过NVLink互连（可用nvidia-smi topo -m确认拓扑）
• Linux操作系统

实现方案

使用CUDA P2P（Peer-to-Peer）API实现GPU间直接通信，避免经过主机内存拷贝。核心流程：

1. 查询并启用所有GPU对的P2P访问能力
2. 每张GPU分配发送缓冲区和接收缓冲区（接收缓冲区需容纳来自其他3张GPU的数据）
3. 使用cudaMemcpyPeerAsync在默认流中发起全交换通信
4. 同步后统计耗时，计算聚合带宽与平均延迟

关键完整代码

#include <cstdio>
#include <cuda_runtime.h>
#include <vector>
#include <chrono>

#define CHECK_CUDA(call) do { \
    cudaError_t err = (call); \
    if (err != cudaSuccess) { \
        printf("CUDA error at %s:%d - %s\n", __FILE__, __LINE__, \
               cudaGetErrorString(err)); \
        return -1; \
    } \
} while(0)

// 全交换模式：每张GPU向其他所有GPU发送data_size字节
int run_all_to_all(int num_gpus, size_t data_size) {
    // 1. 启用所有GPU对的P2P
    for (int i = 0; i < num_gpus; i++) {
        CHECK_CUDA(cudaSetDevice(i));
        for (int j = 0; j < num_gpus; j++) {
            if (i == j) continue;
            int can_access = 0;
            CHECK_CUDA(cudaDeviceCanAccessPeer(&can_access, i, j));
            if (can_access) {
                CHECK_CUDA(cudaDeviceEnablePeerAccess(j, 0));
            } else {
                printf("Warning: GPU %d cannot directly access GPU %d\n", i, j);
            }
        }
    }

    // 2. 分配缓冲区
    std::vector<char*> send_bufs(num_gpus);
    std::vector<char*> recv_bufs(num_gpus);
    for (int i = 0; i < num_gpus; i++) {
        CHECK_CUDA(cudaSetDevice(i));
        CHECK_CUDA(cudaMalloc(&send_bufs[i], data_size));
        // 接收缓冲区需容纳来自其他num_gpus-1张卡的数据
        CHECK_CUDA(cudaMalloc(&recv_bufs[i], data_size * (num_gpus - 1)));
        // 初始化发送数据（非零值便于验证）
        CHECK_CUDA(cudaMemset(send_bufs[i], 0xAB, data_size));
        CHECK_CUDA(cudaMemset(recv_bufs[i], 0, data_size * (num_gpus - 1)));
    }

    // 3. 执行全交换通信
    auto start = std::chrono::high_resolution_clock::now();
    for (int src = 0; src < num_gpus; src++) {
        CHECK_CUDA(cudaSetDevice(src));
        for (int dst = 0; dst < num_gpus; dst++) {
            if (src == dst) continue;
            // 目标GPU上接收偏移：按发送方编号排列
            size_t offset = (src < dst) ? src : src - 1;
            CHECK_CUDA(cudaMemcpyPeerAsync(
                recv_bufs[dst] + offset * data_size, dst,
                send_bufs[src], src,
                data_size));
        }
    }
    // 等待所有GPU完成
    for (int i = 0; i < num_gpus; i++) {
        CHECK_CUDA(cudaSetDevice(i));
        CHECK_CUDA(cudaDeviceSynchronize());
    }
    auto end = std::chrono::high_resolution_clock::now();
    double elapsed_ms = std::chrono::duration<double, std::milli>(end - start).count();

    // 4. 计算性能指标
    double total_bytes = (double)num_gpus * (num_gpus - 1) * data_size;
    double agg_bw_GBps = total_bytes / (elapsed_ms * 1e-3) / 1e9;
    double avg_latency_us = elapsed_ms * 1000.0 / (num_gpus * (num_gpus - 1));

    printf("=== All-to-All NVLink Test (num_gpus=%d, data_size=%.2f MB) ===\n",
           num_gpus, data_size / 1e6);
    printf("Total data transferred: %.2f MB\n", total_bytes / 1e6);
    printf("Elapsed time: %.3f ms\n", elapsed_ms);
    printf("Aggregate bandwidth: %.2f GB/s\n", agg_bw_GBps);
    printf("Avg per-transfer latency: %.2f us\n", avg_latency_us);

    // 5. 验证数据完整性（仅检查第一对通信）
    CHECK_CUDA(cudaSetDevice(1));
    char* host_buf = new char[data_size];
    CHECK_CUDA(cudaMemcpy(host_buf, recv_bufs[1], data_size, cudaMemcpyDeviceToHost));
    bool ok = true;
    for (size_t i = 0; i < data_size; i++) {
        if (host_buf[i] != 0xAB) { ok = false; break; }
    }
    printf("Data verification: %s\n", ok ? "PASS" : "FAIL");
    delete[] host_buf;

    // 6. 清理
    for (int i = 0; i < num_gpus; i++) {
        CHECK_CUDA(cudaSetDevice(i));
        for (int j = 0; j < num_gpus; j++) {
            if (i != j) cudaDeviceDisablePeerAccess(j);
        }
        CHECK_CUDA(cudaFree(send_bufs[i]));
        CHECK_CUDA(cudaFree(recv_bufs[i]));
    }
    return ok ? 0 : -1;
}

int main() {
    int num_gpus = 0;
    CHECK_CUDA(cudaGetDeviceCount(&num_gpus));
    if (num_gpus < 2) {
        printf("Need at least 2 GPUs, found %d\n", num_gpus);
        return -1;
    }
    printf("Found %d GPUs\n", num_gpus);

    // 测试不同数据规模下的扩展性
    std::vector<size_t> sizes = {8 << 20, 16 << 20, 32 << 20}; // 8, 16, 32 MB
    for (size_t sz : sizes) {
        int ret = run_all_to_all(num_gpus, sz);
        if (ret != 0) return ret;
        printf("\n");
    }
    return 0;
}

编译与运行

nvcc -o nvlink_alltoall nvlink_alltoall.cu -std=c++11
./nvlink_alltoall

验证方法与判定标准

• 正确性验证：程序输出Data verification: PASS，表明GPU 1正确收到了来自GPU 0的数据（内容为0xAB），证明P2P传输路径有效。
• 性能指标：观察不同数据规模下的聚合带宽。在4张A100（每对NVLink约600GB/s单向）环境下，聚合带宽应随数据规模增大趋近NVLink总带宽上限（3对链路×600GB/s≈1.8TB/s单向，实际受限于PCIe交换机拓扑可能略低）。
• 扩展性分析：比较8MB、16MB、32MB三种规模的聚合带宽。若带宽基本持平，说明通信已饱和；若带宽随规模上升，说明小数据量时延迟开销占主导。

扩展性分析要点

全交换模式的通信量为O(N²)，当GPU数量从2增至4时，总数据量增至6倍（每卡3收3发）。NVLink拓扑中，若GPU间为全连接（如DGX系列），聚合带宽应随GPU数线性扩展；若为环形或混合拓扑，远端GPU通信需经过中间GPU转发（通过NVLink Switch或PCIe），导致带宽下降。可通过nvidia-smi topo -m查看实际拓扑，结合本程序输出的聚合带宽判断是否存在转发瓶颈。