CUDA + MPI 异构并行：多 GPU 分布式编程实战

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

__global__ void kernel(float *d_data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) d_data[idx] *= 2.0f;
}

int main(int argc, char **argv) {
    int rank, size;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // 绑定 MPI rank 到本地 GPU
    int dev_count;
    cudaGetDeviceCount(&dev_count);
    int dev_id = rank % dev_count;
    cudaSetDevice(dev_id);

    int n_local = 1024 * 1024;
    size_t bytes = n_local * sizeof(float);

    // GPU 显存分配
    float *d_data;
    cudaMalloc(&d_data, bytes);

    // 主机内存分配（建议锁页内存）
    float *h_data;
    cudaMallocHost(&h_data, bytes);  // page-locked

    // 初始化数据
    for (int i = 0; i < n_local; i++) {
        h_data[i] = (float)(i + rank * n_local);
    }

    // H2D
    cudaMemcpy(d_data, h_data, bytes, cudaMemcpyHostToDevice);

    // 启动核函数
    int threads = 256;
    int blocks = (n_local + threads - 1) / threads;
    kernel<<<blocks, threads>>>(d_data, n_local);
    cudaDeviceSynchronize();

    // D2H
    cudaMemcpy(h_data, d_data, bytes, cudaMemcpyDeviceToHost);

    // ---- MPI 进程间通信 ----
    float local_sum = 0.0f;
    for (int i = 0; i < n_local; i++) local_sum += h_data[i];

    float global_sum;
    MPI_Allreduce(&local_sum, &global_sum, 1, MPI_FLOAT,
                  MPI_SUM, MPI_COMM_WORLD);

    if (rank == 0) {
        printf("Global sum = %.2f\n", global_sum);
    }

    // 清理
    cudaFree(d_data);
    cudaFreeHost(h_data);
    MPI_Finalize();
    return 0;
}

mpicc -cuda hybrid.cu -o hybrid -I/usr/local/cuda/include \
      -L/usr/local/cuda/lib64 -lcudart

# 1. 检查 nvidia_p2p 模块
lsmod | grep nvidia_p2p

# 2. 启用 CUDA+IB 直接访问
export MPIR_CVAR_ENABLE_GPU=1
export MPICH_GPU_SUPPORT_ENABLED=1  # MPICH 系
export OMPI_MCA_btl_openib_allow_cuda_memory=1  # OpenMPI

# 3. 使用 CUDA-aware MPI：直接在显存指针上调用 MPI
float *d_sendbuf, *d_recvbuf;  // 设备指针
cudaMalloc(&d_sendbuf, bytes);
cudaMalloc(&d_recvbuf, bytes);
// ... 填充 d_sendbuf ...
MPI_Send(d_sendbuf, count, MPI_FLOAT, dest, tag, MPI_COMM_WORLD);

#include <nccl.h>

ncclComm_t comm;
int localRank = rank % dev_count;
// 使用 ncclGetUniqueId + ncclCommInitRank 在多进程间建立 NCCL 通信子
ncclUniqueId id;
if (rank == 0) ncclGetUniqueId(&id);
MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);
ncclCommInitRank(&comm, size, id, rank);

// NCCL 集合通信：直接在显存上操作
ncclAllReduce(d_sendbuf, d_recvbuf, count,
              ncclFloat, ncclSum, comm, cudaStreamDefault);

// 同步
ncclCommDestroy(comm);

cudaStream_t stream_compute, stream_comm;
cudaStreamCreate(&stream_compute);
cudaStreamCreate(&stream_comm);

for (int iter = 0; iter < max_iters; iter++) {
    // 启动计算核（stream_compute 上）
    compute_kernel<<<grid, block, 0, stream_compute>>>(d_data, ...);

    // 同时将上一迭代需要通信的数据 D2H（stream_comm 上）
    cudaMemcpyAsync(h_halo, d_halo, halo_bytes,
                    cudaMemcpyDeviceToHost, stream_comm);

    // 更新边界缓冲（stream_compute 完成后）
    update_boundary<<<grid, block, 0, stream_compute>>>(d_data, ...);

    // 等待通信流完成数据准备
    cudaStreamSynchronize(stream_comm);

    // MPI 交换边界
    MPI_Sendrecv(h_halo_send, count, MPI_FLOAT, ...,
                 h_halo_recv, count, MPI_FLOAT, ...);

    // 将收到的边界数据传回 GPU
    cudaMemcpyAsync(d_halo_recv, h_halo_recv, halo_bytes,
                    cudaMemcpyHostToDevice, stream_compute);
}

# rank-to-node mapping：确保相邻 rank 在物理上靠近
mpirun -np 8 --map-by ppr:4:node:PE=8 --bind-to core \
       ./cuda_mpi_app

# 配合 NCCL 拓扑自动检测：
# NCCL 自动选择 NVLink > PCIe P2P > IB 的最快路径

// 严格的 rank-to-GPU 绑定
int ngpus;
cudaGetDeviceCount(&ngpus);
if (ngpus == 0) {
    fprintf(stderr, "No GPU found\n");
    MPI_Abort(MPI_COMM_WORLD, 1);
}
int gpu_id = local_rank % ngpus;
cudaSetDevice(gpu_id);

# PyTorch DDP 内部使用 NCCL 通信后端
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
model = MyModel().cuda(local_rank)
ddp_model = DDP(model, device_ids=[local_rank])

for batch in dataloader:
    loss = ddp_model(batch)
    loss.backward()
    # DDP 自动执行 AllReduce 梯度同步（NCCL）
    optimizer.step()

# PyTorch Elastic Launch（容忍节点增减）
torchrun --nnodes=1:4 --nproc_per_node=8 --max_restarts=3 train.py