Merge pull request #70 from lucasdelimanogueira/tmp

Tmp

README.md

@@ -142,7 +142,7 @@ for epoch in range(epochs):
 
 | Development                  | Status      | Feature                                                                |
 | ---------------------------- | ----------- | ---------------------------------------------------------------------- |
-| Operations                   | in progress | <ul><li>[X] GPU Support</li><li>[X] Autograd</li><li>[X] Broadcasting</li></ul>                 |
+| Operations                   | in progress | <ul><li>[X] GPU Support</li><li>[X] Autograd</li><li>[X] Broadcasting</li><li>[] Memory Management</li></ul>                 |
 | Loss                         | in progress | <ul><li>[x] MSE</li><li>[X] Cross Entropy</li></ul>    |
 | Data                         | in progress    | <ul><li>[X] Dataset</li><li>[X] Batch</li><li>[X] Iterator</li></ul>   |
 | Convolutional Neural Network | in progress    | <ul><li>[ ] Conv2d</li><li>[ ] MaxPool2d</li><li>[ ] Dropout</li></ul> |

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
+DISTRIBUTEDLIBS = -lmpi_cxx -lnccl
 
 # 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) $(DISTRIBUTEDLIBS)
 
 # 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/__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/cuda.cu

@@ -7,8 +7,17 @@
 #define THREADS_PER_BLOCK 128
 #define TILE_SIZE 32
 
-__host__ void cpu_to_cuda(Tensor* tensor) {
+__host__ void cpu_to_cuda(Tensor* tensor, int device_id) {
+
+    int deviceCount;
+    cudaGetDeviceCount(&deviceCount);
+    if (device_id >= deviceCount) {
+        fprintf(stderr, "Could not send tensor to device %d, only %d devices available\n", device_id, deviceCount);
+            exit(1);
+    }
     
+    cudaSetDevice(device_id); 
+
     float* data_tmp;
     cudaMalloc((void **)&data_tmp, tensor->size * sizeof(float));
     cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);

norch/csrc/cuda.h

@@ -1,7 +1,8 @@
 #ifndef CUDA_KERNEL_H_
 #define CUDA_KERNEL_H_
 
-    __host__ void cpu_to_cuda(Tensor* tensor);
+
+    __host__ void cpu_to_cuda(Tensor* tensor, int device_id);
     __host__ void cuda_to_cpu(Tensor* tensor);
     
     __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);

norch/csrc/distributed.cpp

@@ -0,0 +1,65 @@
+#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 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));
+
+}
+
+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_sum_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));
+
+    cudaStreamSynchronize(stream);
+    cudaStreamDestroy(stream);
+}
+
+void end_process_group() {
+    MPI_CHECK(MPI_Finalize());
+    NCCL_CHECK(ncclCommDestroy(nccl_comm));
+}
+
+}

norch/csrc/distributed.h

@@ -0,0 +1,28 @@
+#ifndef DISTRIBUTED_H
+#define DISTRIBUTED_H
+
+#define MPI_CHECK(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 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)
+
+extern "C" {
+  void init_process_group(int rank, int world_size);
+  void broadcast_tensor(Tensor* tensor);
+  void allreduce_sum_tensor(Tensor* tensor);
+}
+
+#endif /* DISTRIBUTED_H */

norch/csrc/tensor.cpp

@@ -64,8 +64,15 @@ extern "C" {
     }
 
     void to_device(Tensor* tensor, char* target_device) {
+        int device_id = 0;
+        char* endptr;
+        long num = strtol(target_device, &endptr, 10);
+        if (*endptr == '\0') {
+            device_id = (int)num;
+        }
+
         if ((strcmp(target_device, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
-            cpu_to_cuda(tensor);
+            cpu_to_cuda(tensor, device_id);
         }
 
         else if ((strcmp(target_device, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {

norch/distributed/__init__.py

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

norch/distributed/distributed.py

@@ -0,0 +1,24 @@
+import os
+import ctypes
+from norch import Tensor, CTensor
+
+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)
+
+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_sum_tensor(tensor):
+
+    Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
+    Tensor._C.allreduce_mean_tensor.restype = None
+    
+    Tensor._C.allreduce_mean_tensor(tensor)

norch/distributed/run/__main__.py

@@ -0,0 +1,8 @@
+import sys
+from . import run
+
+def main():
+    run.main()
+
+if __name__ == "__main__":
+    main()

norch/distributed/run/run.py

@@ -0,0 +1,31 @@
+import sys
+import os
+import argparse
+import subprocess
+
+def main():
+    parser = argparse.ArgumentParser(description="Distributed runner")
+    parser.add_argument('--nproc_per_node', type=int, required=True, help='Number of processes')
+    parser.add_argument('--nnodes', type=int, required=False, default=1, help='Number of nodes')
+    parser.add_argument('script', type=str, help='The script to run')
+    parser.add_argument('script_args', nargs=argparse.REMAINDER, help='Arguments for the script')
+
+    args = parser.parse_args()
+
+    nproc_per_node = args.nproc_per_node
+    nnodes = args.nnodes
+    script_name = args.script
+    script_args = args.script_args
+
+    world_size = nproc_per_node * nnodes
+
+    mpiexec_command = [
+        'mpiexec',
+        '-n', str(world_size),
+        sys.executable, script_name
+    ] + script_args
+
+    subprocess.run(mpiexec_command)
+
+if __name__ == '__main__':
+    main()

norch/tensor.py

@@ -190,6 +190,8 @@ class Tensor:
         return self.reshape(new_shape)
 
     def to(self, device):
+        device = str(device)
+        
         self.device = device
         self.device_ctype = self.device.encode('utf-8')
 

setup.py

@@ -5,13 +5,36 @@ import subprocess
 with open("README.md", "r", encoding = "utf-8") as fh:
     long_description = fh.read()
 
+
+
 apt_dependencies = [
     'nvidia-cuda-toolkit'
 ]
 
+apt_get_dependencies = [
+    'mpi',
+    '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

@@ -0,0 +1,31 @@
+import os
+import norch
+import norch.distributed as dist
+
+def main():
+
+    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)
+
+    tensor = norch.Tensor([1,1,1])
+    tensor = (rank + 1) * tensor
+    print(f"BEFORE on rank {rank}: {tensor} \n\n")
+
+    dist.allreduce_sum_tensor(tensor)
+
+    print(f"AFTER ALLREDUCE on rank {rank}: {tensor} \n\n")
+
+    print("###############\n\n\n")
+
+    tensor = tensor * 10
+    print(f"BEFORE BROADCAST on rank {rank}: {tensor} \n\n")
+
+    dist.broadcast_tensor(tensor)
+
+    print(f"AFTER BROADCAST on rank {rank}: {tensor} \n\n")
+
+if __name__ == "__main__":
+    main()