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 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/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()