diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..e69de29 diff --git a/build/Makefile b/build/Makefile index b1da9bd..5c36c07 100644 --- a/build/Makefile +++ b/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 diff --git a/build/cuda.cu.o b/build/cuda.cu.o index e7e5316..01ab4b3 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/distributed.o b/build/distributed.o index ba9cac7..61ad45c 100644 Binary files a/build/distributed.o and b/build/distributed.o differ diff --git a/norch/__init__.py b/norch/__init__.py index 7bc1d53..c2a3122 100644 --- a/norch/__init__.py +++ b/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 * diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp index c0d85ee..666c435 100644 --- a/norch/csrc/distributed.cpp +++ b/norch/csrc/distributed.cpp @@ -1,27 +1,69 @@ -//#include +#include #include #include #include "distributed.h" +#include "tensor.h" +#include "cuda.h" #include +#include + +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); -} \ No newline at end of file +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)); +} + +} diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h index a1bf4e7..40da266 100644 --- a/norch/csrc/distributed.h +++ b/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); diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py index ae22777..3d5574e 100644 --- a/norch/distributed/distributed.py +++ b/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) diff --git a/norch/libtensor.so b/norch/libtensor.so index 4db92c0..b70187e 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/setup.py b/setup.py index 4df3c72..8774eb8 100644 --- a/setup.py +++ b/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']) diff --git a/train.py b/train.py index c51a652..6cea436 100644 --- a/train.py +++ b/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()