From b152ae2a1d9e19799f443accdd7e65c704fdbf3d Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 23 May 2024 19:15:30 -0300 Subject: [PATCH] first commit distributed module --- README.md | 2 +- norch/__pycache__/__init__.cpython-38.pyc | Bin 365 -> 365 bytes norch/csrc/cuda.cu | 11 +++++- norch/csrc/cuda.h | 2 +- norch/csrc/tensor.cpp | 9 ++++- norch/distributed/distributed.py | 2 ++ norch/distributed/run/__main__.py | 8 +++++ norch/distributed/run/run.py | 39 ++++++++++++++++++++++ norch/tensor.py | 2 ++ train.py | 12 +++++++ 10 files changed, 83 insertions(+), 4 deletions(-) create mode 100644 norch/distributed/distributed.py create mode 100644 norch/distributed/run/__main__.py create mode 100644 norch/distributed/run/run.py create mode 100644 train.py 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/norch/__pycache__/__init__.cpython-38.pyc b/norch/__pycache__/__init__.cpython-38.pyc index 176d505192d006a204e01bbdff5524e7ed201c38..71962d9e9e29bd7e443f82302b507406eb8b5080 100644 GIT binary patch delta 107 zcmaFM^p=S`l$V!_0SHzl_)p|+U^JQ7qvsdJl$RI9nqN?onH$AgT9TPl9L1fNUzD6t zmRX#cpBKd)A77SQ1Y*X=-{OjoPb@9T$S;bI2T3FsrKV(-6vxM}WGG?-np?y+nS(JE E08|zuPyhe` delta 107 zcmaFM^p=S`l$V!_0SHpH{U&lZFd9$n(et~-l$RI9nqN?onR|=1v?Md9IEp(jzbHAQ zEVDQ>KkpWIe0*7I5r`Qde~T+VKC!eUBfkj921+CsrKV(-6vxM}WGDg|T*Nh*gE18V Dng}G^ 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..3619452 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -1,7 +1,7 @@ #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/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/distributed.py b/norch/distributed/distributed.py new file mode 100644 index 0000000..a11a82e --- /dev/null +++ b/norch/distributed/distributed.py @@ -0,0 +1,2 @@ +def init_process_group(world_size, rank, backend='nccl'): + pass \ No newline at end of file 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..c2eaba1 --- /dev/null +++ b/norch/distributed/run/run.py @@ -0,0 +1,39 @@ +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) + +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 + + 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) + + for p in processes: + p.join() + +if __name__ == '__main__': + main() 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/train.py b/train.py new file mode 100644 index 0000000..539cd65 --- /dev/null +++ b/train.py @@ -0,0 +1,12 @@ +import os + +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}") + +if __name__ == "__main__": + main()