broadcast allreduce nccl

norch/csrc/distributed.cpp

@@ -1,27 +1,69 @@
-//#include <nccl.h>
+#include <nccl.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include "distributed.h"
+#include "tensor.h"
+#include "cuda.h"
 #include <mpi.h>
+#include <nccl.h>
+
+int rank;
+int world_size;
+ncclComm_t nccl_comm;
 
 extern "C" {
-  void init_process_group(int rank, int world_size) {
 
-      MPI_Init(NULL, NULL);
-      
-      int id;
+void init_process_group(int env_rank, int env_world_size) {
+
+    rank = env_rank;
+    world_size = env_world_size;
+
+    // init MPI
+    MPI_CHECK(MPI_Init(NULL, NULL));
+
+
+    ncclUniqueId nccl_id;
+
+    if (rank == 0) {
+        ncclGetUniqueId(&nccl_id);
+    }
+
+    // send nccl unique id to all processes
+    MPI_CHECK(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
+
+    // init NCCL communication group
+    NCCL_CHECK(ncclCommInitRank(&nccl_comm, world_size, nccl_id, rank));
 
-      if (rank == 0) {
-          id = 123;
-      }
-      printf("O id é: %d no rank %d/%d", id, rank, world_size);
-      MPICHECK(MPI_Bcast((void *)&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD));
-      printf("\n\n");
-      printf("O id é: %d no rank %d/%d", id, rank, world_size);
-      MPI_Finalize();
-  }
 }
 
-int main(){
-  init_process_group(5, 10);
-}
+void broadcast_tensor(Tensor* tensor) {
+
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+    
+    NCCL_CHECK(ncclBroadcast(tensor->data, tensor->data, tensor->size * sizeof(float), ncclFloat, 0, nccl_comm, stream));
+    cudaStreamSynchronize(stream);
+    cudaStreamDestroy(stream);
+}
+
+void allreduce_mean_tensor(Tensor* tensor) {
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+
+    // Perform NCCL AllReduce operation to calculate the sum of all tensors across all processes
+    NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
+
+    // Divide the sum by the world size to get the mean
+    float scale = 1.0f / world_size;
+    scalar_mul_tensor_cuda(tensor, scale, tensor->data);
+
+    cudaStreamSynchronize(stream);
+    cudaStreamDestroy(stream);
+}
+
+void end_process_group() {
+    MPI_CHECK(MPI_Finalize());
+    NCCL_CHECK(ncclCommDestroy(nccl_comm));
+}
+
+}