implemented mpi for distributed run

build/Makefile

@@ -1,10 +1,19 @@
 # Compiler
 CC = g++
 NVCC = nvcc
+MPI_CC = mpiCC
 
 # Compiler flags
 CFLAGS = -Wall -Wextra -std=c++11
 NVCCFLAGS = -std=c++11
+MPICC_FLAGS = -Wall -Wextra -std=c++11
+
+# CUDA flags and libraries
+CUDAFLAGS = -arch=sm_75
+CUDALIBS = -I/usr/local/cuda/include -L/usr/local/cuda/lib64 -lcudart -lcuda
+
+# MPI flags and libraries
+MPILIBS = -lmpi_cxx
 
 # Directories
 SRCDIR = ../norch/csrc
@@ -12,18 +21,17 @@ BUILDDIR = ../build
 TARGET = ../norch/libtensor.so
 
 # Files
-SRCS = $(wildcard $(SRCDIR)/*.cpp)
+SRCS := $(filter-out $(SRCDIR)/distributed.cpp, $(wildcard $(SRCDIR)/*.cpp))
 CU_SRCS = $(wildcard $(SRCDIR)/*.cu)
 OBJS = $(patsubst $(SRCDIR)/%.cpp, $(BUILDDIR)/%.o, $(SRCS))
 CU_OBJS = $(patsubst $(SRCDIR)/%.cu, $(BUILDDIR)/%.cu.o, $(CU_SRCS))
+MPI_SRCS := $(SRCDIR)/distributed.cpp
+MPI_OBJS := $(BUILDDIR)/distributed.o
 
-# CUDA flags and libraries
-CUDAFLAGS = -arch=sm_75
-CUDALIBS = -I/usr/local/cuda/include -L/usr/local/cuda/lib64 -lcudart -lcuda
 
 # Rule to build the target
-$(TARGET): $(OBJS) $(CU_OBJS)
-	$(NVCC) --shared -o $(TARGET) $(OBJS) $(CU_OBJS) $(CUDALIBS)
+$(TARGET): $(OBJS) $(MPI_OBJS) $(CU_OBJS)
+	$(NVCC) --shared -o $(TARGET) $(OBJS) $(MPI_OBJS) $(CU_OBJS) $(CUDALIBS) $(MPILIBS)
 
 # Rule to compile C++ source files
 $(BUILDDIR)/%.o: $(SRCDIR)/%.cpp
@@ -33,6 +41,10 @@ $(BUILDDIR)/%.o: $(SRCDIR)/%.cpp
 $(BUILDDIR)/%.cu.o: $(SRCDIR)/%.cu
 	$(NVCC) $(NVCCFLAGS) $(CUDAFLAGS) -Xcompiler -fPIC -c $< -o $@
 
+# Rule to compile distributed.cpp with mpiCC
+$(BUILDDIR)/distributed.o: $(SRCDIR)/distributed.cpp
+	$(MPI_CC) $(MPICC_FLAGS) -fPIC -c $< -o $@
+
 # Clean rule
 clean:
 	rm -f $(BUILDDIR)/*.o $(TARGET)

norch/csrc/cuda.h

@@ -1,6 +1,7 @@
 #ifndef CUDA_KERNEL_H_
 #define CUDA_KERNEL_H_
 
+
     __host__ void cpu_to_cuda(Tensor* tensor, int device_id);
     __host__ void cuda_to_cpu(Tensor* tensor);
     

norch/csrc/distributed.cpp

@@ -0,0 +1,27 @@
+//#include <nccl.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include "distributed.h"
+#include <mpi.h>
+
+extern "C" {
+  void init_process_group(int rank, int world_size) {
+
+      MPI_Init(NULL, NULL);
+      
+      int id;
+
+      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);
+}

norch/csrc/distributed.h

@@ -0,0 +1,26 @@
+#ifndef DISTRIBUTED_H
+#define DISTRIBUTED_H
+
+#define MPICHECK(cmd) do {                          \
+  int e = cmd;                                      \
+  if( e != MPI_SUCCESS ) {                          \
+    printf("Failed: MPI error %s:%d '%d'\n",        \
+        __FILE__,__LINE__, e);   \
+    exit(EXIT_FAILURE);                             \
+  }                                                 \
+} while(0)
+
+/*#define NCCLCHECK(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)*/
+
+extern "C" {
+  void init_process_group(int rank, int world_size);
+}
+
+#endif /* DISTRIBUTED_H */

norch/distributed/__init__.py

@@ -0,0 +1 @@
+from .distributed import *

norch/distributed/distributed.py

@@ -1,2 +1,11 @@
-def init_process_group(world_size, rank, backend='nccl'):
-    pass
+import os
+import ctypes
+from norch import Tensor
+
+def init_process_group(rank, world_size, backend='nccl'):
+
+    Tensor._C.init_process_group.argtypes = [ctypes.c_int, ctypes.c_int]
+    Tensor._C.init_process_group.restype = None
+
+    Tensor._C.init_process_group(rank, world_size)
+

norch/distributed/run/run.py

@@ -1,14 +1,7 @@
 import sys
 import os
 import argparse
-import multiprocessing
-
-def worker(rank, nnodes, world_size, script_name, script_args):
-    os.environ['LOCAL_RANK'] = str(rank // nnodes)
-    os.environ['RANK'] = str(rank)
-    os.environ['WORLD_SIZE'] = str(world_size)
-    script_args = [sys.executable, script_name] + script_args
-    os.execvp(script_args[0], script_args)
+import subprocess
 
 def main():
     parser = argparse.ArgumentParser(description="Distributed runner")
@@ -26,14 +19,13 @@ def main():
 
     world_size = nproc_per_node * nnodes
 
-    processes = []
-    for rank in range(nproc_per_node):
-        p = multiprocessing.Process(target=worker, args=(rank, nnodes, world_size, script_name, script_args))
-        p.start()
-        processes.append(p)
+    mpiexec_command = [
+        'mpiexec',
+        '-n', str(world_size),
+        sys.executable, script_name
+    ] + script_args
 
-    for p in processes:
-        p.join()
+    subprocess.run(mpiexec_command)
 
 if __name__ == '__main__':
     main()

setup.py

@@ -9,9 +9,17 @@ apt_dependencies = [
     'nvidia-cuda-toolkit'
 ]
 
+apt_get_dependencies = [
+    'mpi',
+    'libopenmpi-dev'
+]
+
 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])
+
 class CustomInstall(install):
     def run(self):
         subprocess.call(['make', '-C', 'build'])

train.py

@@ -1,12 +1,13 @@
 import os
+import norch.distributed as dist
 
 def main():
-    local_rank = int(os.getenv('LOCAL_RANK', -1))
-    rank = int(os.getenv('RANK', -1))
-    world_size = int(os.getenv('WORLD_SIZE', -1))
 
-    # Your main script logic here
-    print(f"Hello from process {local_rank}, {rank} out of {world_size}")
+    local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', -1))
+    rank = int(os.getenv('OMPI_COMM_WORLD_RANK', -1))
+    world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE', -1))
+
+    dist.init_process_group(rank, world_size)
 
 if __name__ == "__main__":
     main()