diff --git a/build/Makefile b/build/Makefile index 615abf7..b1da9bd 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 +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) diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 4d96c07..e7e5316 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..ba9cac7 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/__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.h b/norch/csrc/cuda.h index 3619452..702dea1 100644 --- a/norch/csrc/cuda.h +++ b/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); diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp new file mode 100644 index 0000000..c0d85ee --- /dev/null +++ b/norch/csrc/distributed.cpp @@ -0,0 +1,27 @@ +//#include +#include +#include +#include "distributed.h" +#include + +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); +} \ No newline at end of file diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h new file mode 100644 index 0000000..a1bf4e7 --- /dev/null +++ b/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 */ diff --git a/norch/csrc/distributed.o b/norch/csrc/distributed.o new file mode 100755 index 0000000..144f330 Binary files /dev/null and b/norch/csrc/distributed.o differ 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 index a11a82e..ae22777 100644 --- a/norch/distributed/distributed.py +++ b/norch/distributed/distributed.py @@ -1,2 +1,11 @@ -def init_process_group(world_size, rank, backend='nccl'): - pass \ No newline at end of file +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) + diff --git a/norch/distributed/run/run.py b/norch/distributed/run/run.py index c2eaba1..f40cb63 100644 --- a/norch/distributed/run/run.py +++ b/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() diff --git a/norch/libtensor.so b/norch/libtensor.so index 7ce5299..4db92c0 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/setup.py b/setup.py index 36bff91..4df3c72 100644 --- a/setup.py +++ b/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']) diff --git a/train.py b/train.py index 539cd65..c51a652 100644 --- a/train.py +++ b/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()