diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..e69de29 diff --git a/README.md b/README.md index fad9f4a..adac45f 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ for epoch in range(epochs): | Development | Status | Feature | | ---------------------------- | ----------- | ---------------------------------------------------------------------- | -| Operations | in progress | | +| Operations | in progress | | | Loss | in progress | | | Data | in progress | | | Convolutional Neural Network | in progress | | diff --git a/build/Makefile b/build/Makefile index 615abf7..5c36c07 100644 --- a/build/Makefile +++ b/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) diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 4d96c07..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 new file mode 100644 index 0000000..61ad45c Binary files /dev/null and b/build/distributed.o differ diff --git a/build/tensor.o b/build/tensor.o index ee245d0..4246623 100644 Binary files a/build/tensor.o and b/build/tensor.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/__pycache__/__init__.cpython-38.pyc b/norch/__pycache__/__init__.cpython-38.pyc index 176d505..71962d9 100644 Binary files a/norch/__pycache__/__init__.cpython-38.pyc and b/norch/__pycache__/__init__.cpython-38.pyc differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 882aabd..2381a0f 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 87ba247..e39c20c 100644 --- a/norch/csrc/cuda.cu +++ b/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); diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index daeef28..702dea1 100644 --- a/norch/csrc/cuda.h +++ b/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); diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp new file mode 100644 index 0000000..8124ec0 --- /dev/null +++ b/norch/csrc/distributed.cpp @@ -0,0 +1,65 @@ +#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 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)); +} + +} diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h new file mode 100644 index 0000000..1040fc0 --- /dev/null +++ b/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 */ diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index d5bfa71..725b9f5 100644 --- a/norch/csrc/tensor.cpp +++ b/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)) { diff --git a/norch/distributed/__init__.py b/norch/distributed/__init__.py new file mode 100644 index 0000000..a979451 --- /dev/null +++ b/norch/distributed/__init__.py @@ -0,0 +1 @@ +from .distributed import * \ No newline at end of file diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py new file mode 100644 index 0000000..d033ede --- /dev/null +++ b/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) diff --git a/norch/distributed/run/__main__.py b/norch/distributed/run/__main__.py new file mode 100644 index 0000000..a16682e --- /dev/null +++ b/norch/distributed/run/__main__.py @@ -0,0 +1,8 @@ +import sys +from . import run + +def main(): + run.main() + +if __name__ == "__main__": + main() diff --git a/norch/distributed/run/run.py b/norch/distributed/run/run.py new file mode 100644 index 0000000..f40cb63 --- /dev/null +++ b/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() diff --git a/norch/libtensor.so b/norch/libtensor.so index 7ce5299..b70187e 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index 6861f14..2e80e72 100644 --- a/norch/tensor.py +++ b/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') diff --git a/setup.py b/setup.py index 36bff91..8774eb8 100644 --- a/setup.py +++ b/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']) diff --git a/train.py b/train.py new file mode 100644 index 0000000..8f07a41 --- /dev/null +++ b/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()