分布式GPU上計算長向量模的方法

當向量分布在多個GPU卡上時，計算向量模(2-范數)需要以下步驟：

1. 在每個GPU上計算本地數據的平方和
2. 跨GPU通信匯總所有平方和
3. 在根GPU上計算總和的平方根

實現方法

下面是一個完整的CUDA示例代碼，使用NCCL進行多GPU通信：

#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));// 初始化NCCLncclComm_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;
}

關鍵點說明

1. 數據分布：向量被均勻分布在多個GPU上，每個GPU處理一部分數據。

2. 本地計算：

   • 使用CUDA核函數計算本地數據的平方和
   • 使用塊內歸約優化性能
   • 使用CUB庫進行設備端最終歸約

3. 跨GPU通信：

   • 使用NCCL進行all-reduce操作，匯總所有GPU的平方和
   • NCCL針對多GPU通信進行了優化

4. 結果計算：

   • 只在根GPU上計算最終結果的平方根
   • 其他GPU可以忽略結果或用于后續計算

編譯說明

編譯此代碼需要：

• CUDA工具包
• NCCL庫
• CUB頭文件（通常包含在CUDA工具包中)

編譯命令示例：

nvcc -o distributed_norm distributed_norm.cu -lnccl

性能優化建議

1. 對于非常大的向量，可以考慮使用更高效的內存訪問模式
2. 根據GPU架構調整塊大小和網格大小
3. 使用CUDA圖來捕獲整個計算流程，減少啟動開銷
4. 考慮使用FP16或TF32計算來提升吞吐量（如果精度允許）

這種方法可以高效地計算分布在多個GPU上的大型向量的模，適用于大規模科學計算和機器學習應用。