broadcast allreduce nccl

build/Makefile

@@ -1,7 +1,7 @@
 # Compiler
 CC = g++
 NVCC = nvcc
-MPI_CC = mpiCC
+MPI_CC = mpicc
 
 # Compiler flags
 CFLAGS = -Wall -Wextra -std=c++11
@@ -13,7 +13,7 @@ CUDAFLAGS = -arch=sm_75
 CUDALIBS = -I/usr/local/cuda/include -L/usr/local/cuda/lib64 -lcudart -lcuda
 
 # MPI flags and libraries
-MPILIBS = -lmpi_cxx
+DISTRIBUTEDLIBS = -lmpi_cxx -lnccl
 
 # Directories
 SRCDIR = ../norch/csrc
@@ -31,7 +31,7 @@ MPI_OBJS := $(BUILDDIR)/distributed.o
 
 # Rule to build the target
 $(TARGET): $(OBJS) $(MPI_OBJS) $(CU_OBJS)
-	$(NVCC) --shared -o $(TARGET) $(OBJS) $(MPI_OBJS) $(CU_OBJS) $(CUDALIBS) $(MPILIBS)
+	$(NVCC) --shared -o $(TARGET) $(OBJS) $(MPI_OBJS) $(CU_OBJS) $(CUDALIBS) $(DISTRIBUTEDLIBS)
 
 # Rule to compile C++ source files
 $(BUILDDIR)/%.o: $(SRCDIR)/%.cpp

norch/__init__.py

@@ -1,4 +1,4 @@
-from norch.tensor import Tensor
+from norch.tensor import *
 from .nn import *
 from .optim import *
 from .utils import *

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));
+}
+
+}

norch/csrc/distributed.h

@@ -1,7 +1,7 @@
 #ifndef DISTRIBUTED_H
 #define DISTRIBUTED_H
 
-#define MPICHECK(cmd) do {                          \
+#define MPI_CHECK(cmd) do {                          \
   int e = cmd;                                      \
   if( e != MPI_SUCCESS ) {                          \
     printf("Failed: MPI error %s:%d '%d'\n",        \
@@ -10,14 +10,14 @@
   }                                                 \
 } while(0)
 
-/*#define NCCLCHECK(cmd) do {                         \
+#define NCCL_CHECK(cmd) do {                         \
   ncclResult_t r = cmd;                             \
   if (r!= ncclSuccess) {                            \
     printf("Failed, NCCL error %s:%d '%s'\n",             \
         __FILE__,__LINE__,ncclGetErrorString(r));   \
     exit(EXIT_FAILURE);                             \
   }                                                 \
-} while(0)*/
+} while(0)
 
 extern "C" {
   void init_process_group(int rank, int world_size);

norch/distributed/distributed.py

@@ -1,6 +1,6 @@
 import os
 import ctypes
-from norch import Tensor
+from norch import Tensor, CTensor
 
 def init_process_group(rank, world_size, backend='nccl'):
 
@@ -9,3 +9,16 @@ def init_process_group(rank, world_size, backend='nccl'):
 
     Tensor._C.init_process_group(rank, world_size)
 
+def broadcast_tensor(tensor):
+
+    Tensor._C.broadcast_tensor.argtypes = [ctypes.POINTER(CTensor)]
+    Tensor._C.broadcast_tensor.restype = None
+    
+    Tensor._C.broadcast_tensor(tensor)
+
+def allreduce_mean_tensor(tensor):
+
+    Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
+    Tensor._C.allreduce_mean_tensor.restype = None
+    
+    Tensor._C.allreduce_mean_tensor(tensor)

setup.py

@@ -5,6 +5,8 @@ import subprocess
 with open("README.md", "r", encoding = "utf-8") as fh:
     long_description = fh.read()
 
+
+
 apt_dependencies = [
     'nvidia-cuda-toolkit'
 ]
@@ -14,12 +16,25 @@ apt_get_dependencies = [
     'libopenmpi-dev'
 ]
 
+apt_nccl = [
+    'libnccl2=2.21.5-1+cuda12.2',
+    'libnccl-dev=2.21.5-1+cuda12.2 '
+]
+
+subprocess.check_call(['sudo', 'apt-get', 'update'])
+
 for package in apt_dependencies:
     subprocess.check_call(['sudo', 'apt', 'install', package])
 
 for package in apt_get_dependencies:
     subprocess.check_call(['sudo', 'apt-get', 'install', 'y', package])
 
+subprocess.check_call(['wget', 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb '])
+subprocess.check_call(['sudo', 'dpkg', '-i', 'cuda-keyring_1.0-1_all.deb'])
+
+for package in apt_nccl:
+    subprocess.check_call(['sudo', 'apt', 'install', package])
+
 class CustomInstall(install):
     def run(self):
         subprocess.call(['make', '-C', 'build'])

train.py

@@ -1,4 +1,5 @@
 import os
+import norch
 import norch.distributed as dist
 
 def main():
@@ -9,5 +10,17 @@ def main():
 
     dist.init_process_group(rank, world_size)
 
+    tensor = norch.Tensor([1,2,3])
+    tensor = rank * tensor
+    print(f"BEFORE on rank {rank}: ", tensor, '\n\n')
+
+    dist.broadcast_tensor(tensor)
+
+    print(f"AFTER BROADCAST on rank {rank}: ", tensor, '\n\n')
+
+    dist.allreduce_mean_tensor(tensor)
+
+    print(f"AFTER ALLREDUCE MEAN on rank {rank}: ", tensor, '\n\n')
+
 if __name__ == "__main__":
     main()