fix buffer init

karpathy · Jun 17, 2024 · 751bec9 · 751bec9
1 parent 6d548ab
commit 751bec9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train_gpt2.cu b/train_gpt2.cu
@@ -165,6 +165,7 @@ void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) {
 void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config, float *unified_buffer) {
 #ifdef MULTI_GPU
     if (multi_gpu_config->num_processes > 1) {
+        if (unified_buffer == NULL) cudaCheck(cudaMallocManaged(&unified_buffer, sizeof(float)));
         ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, 0));
     }
 #endif