分布式GPU上计算长向量模的方法

当向量分布在多个GPU卡上时，计算向量模(2-范数)需要以下步骤：

在每个GPU上计算本地数据的平方和
跨GPU通信汇总所有平方和
在根GPU上计算总和的平方根

实现方法

下面是一个完整的CUDA示例代码，使用NCCL进行多GPU通信：

cpp 复制代码

#include <iostream>
#include <cmath>
#include <cuda_runtime.h>
#include <nccl.h>

#define CHECK_CUDA(call) { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \
                  << cudaGetErrorString(err) << std::endl; \
        exit(EXIT_FAILURE); \
    } \
}

#define CHECK_NCCL(call) { \
    ncclResult_t res = call; \
    if (res != ncclSuccess) { \
        std::cerr << "NCCL error at " << __FILE__ << ":" << __LINE__ << ": " \
                  << ncclGetErrorString(res) << std::endl; \
        exit(EXIT_FAILURE); \
    } \
}

// CUDA核函数：计算局部平方和
__global__ void compute_local_square_sum(const float* vec, float* partial_sum, size_t n) {
    extern __shared__ float shared_mem[];
    
    unsigned int tid = threadIdx.x;
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    float sum = 0.0f;
    if (i < n) {
        float val = vec[i];
        sum = val * val;
    }
    
    // 归约到共享内存
    shared_mem[tid] = sum;
    __syncthreads();
    
    // 块内归约
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared_mem[tid] += shared_mem[tid + s];
        }
        __syncthreads();
    }
    
    // 第一个线程写入结果
    if (tid == 0) {
        partial_sum[blockIdx.x] = shared_mem[0];
    }
}

// 计算向量模
float distributed_vector_norm(int ngpus, size_t total_elements, size_t local_elements, const float* local_vec, cudaStream_t stream, ncclComm_t comm) {
    // 1. 每个GPU计算本地平方和
    const int block_size = 256;
    const int grid_size = (local_elements + block_size - 1) / block_size;
    
    float* d_partial_sums;
    CHECK_CUDA(cudaMalloc(&d_partial_sums, grid_size * sizeof(float)));
    
    // 调用核函数计算局部平方和
    compute_local_square_sum<<<grid_size, block_size, block_size * sizeof(float), stream>>>(
        local_vec, d_partial_sums, local_elements);
    
    // 2. 在设备上完成最终归约
    float* d_local_sum;
    CHECK_CUDA(cudaMalloc(&d_local_sum, sizeof(float)));
    
    // 使用CUDA的归约函数完成设备上的最终归约
    void* d_temp_storage = nullptr;
    size_t temp_storage_bytes = 0;
    cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_partial_sums, d_local_sum, grid_size, stream);
    CHECK_CUDA(cudaMalloc(&d_temp_storage, temp_storage_bytes));
    cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_partial_sums, d_local_sum, grid_size, stream);
    
    // 3. 跨GPU通信汇总所有平方和
    float* d_global_sum;
    CHECK_CUDA(cudaMalloc(&d_global_sum, sizeof(float)));
    
    // 使用NCCL进行all reduce操作
    CHECK_NCCL(ncclAllReduce((const void*)d_local_sum, (void*)d_global_sum, 
                            1, ncclFloat, ncclSum, comm, stream));
    
    // 4. 计算平方根（只在root GPU上获取结果）
    float global_sum = 0.0f;
    int root = 0;
    int rank;
    CHECK_NCCL(ncclCommUserRank(comm, &rank));
    
    if (rank == root) {
        CHECK_CUDA(cudaMemcpyAsync(&global_sum, d_global_sum, sizeof(float), 
                                  cudaMemcpyDeviceToHost, stream));
        CHECK_CUDA(cudaStreamSynchronize(stream));
    }
    
    // 清理
    CHECK_CUDA(cudaFree(d_temp_storage));
    CHECK_CUDA(cudaFree(d_partial_sums));
    CHECK_CUDA(cudaFree(d_local_sum));
    CHECK_CUDA(cudaFree(d_global_sum));
    
    return (rank == root) ? sqrtf(global_sum) : 0.0f;
}

int main(int argc, char* argv[]) {
    // 初始化
    int ngpus;
    CHECK_CUDA(cudaGetDeviceCount(&ngpus));
    
    // 初始化NCCL
    ncclComm_t comm;
    ncclUniqueId id;
    if (rank == 0) ncclGetUniqueId(&id);
    MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);
    
    CHECK_NCCL(ncclCommInitRank(&comm, ngpus, id, rank));
    
    // 假设总向量大小为1亿元素
    size_t total_elements = 100000000;
    size_t local_elements = total_elements / ngpus;
    
    // 分配和初始化本地向量
    float* d_local_vec;
    CHECK_CUDA(cudaMalloc(&d_local_vec, local_elements * sizeof(float)));
    
    // 初始化向量数据（这里简单设置为全1，实际应用中应填充真实数据）
    float init_val = 1.0f;
    CHECK_CUDA(cudaMemset(d_local_vec, init_val, local_elements * sizeof(float)));
    
    // 创建CUDA流
    cudaStream_t stream;
    CHECK_CUDA(cudaStreamCreate(&stream));
    
    // 计算向量模
    float norm = distributed_vector_norm(ngpus, total_elements, local_elements, d_local_vec, stream, comm);
    
    int rank;
    CHECK_NCCL(ncclCommUserRank(comm, &rank));
    if (rank == 0) {
        std::cout << "Vector norm: " << norm << std::endl;
        std::cout << "Expected norm: " << sqrtf(total_elements) << std::endl;
    }
    
    // 清理
    CHECK_CUDA(cudaFree(d_local_vec));
    CHECK_CUDA(cudaStreamDestroy(stream));
    CHECK_NCCL(ncclCommDestroy(comm));
    
    return 0;
}

关键点说明

数据分布：向量被均匀分布在多个GPU上，每个GPU处理一部分数据。
本地计算：
- 使用CUDA核函数计算本地数据的平方和
- 使用块内归约优化性能
- 使用CUB库进行设备端最终归约
跨GPU通信：
- 使用NCCL进行all-reduce操作，汇总所有GPU的平方和
- NCCL针对多GPU通信进行了优化
结果计算：
- 只在根GPU上计算最终结果的平方根
- 其他GPU可以忽略结果或用于后续计算

编译说明

编译此代码需要：

CUDA工具包
NCCL库
CUB头文件（通常包含在CUDA工具包中)

编译命令示例：

bash 复制代码

nvcc -o distributed_norm distributed_norm.cu -lnccl

性能优化建议

对于非常大的向量，可以考虑使用更高效的内存访问模式
根据GPU架构调整块大小和网格大小
使用CUDA图来捕获整个计算流程，减少启动开销
考虑使用FP16或TF32计算来提升吞吐量（如果精度允许）

这种方法可以高效地计算分布在多个GPU上的大型向量的模，适用于大规模科学计算和机器学习应用。