first commit distributed module

README.md

@@ -142,7 +142,7 @@ for epoch in range(epochs):
 
 | Development                  | Status      | Feature                                                                |
 | ---------------------------- | ----------- | ---------------------------------------------------------------------- |
-| Operations                   | in progress | <ul><li>[X] GPU Support</li><li>[X] Autograd</li><li>[X] Broadcasting</li></ul>                 |
+| Operations                   | in progress | <ul><li>[X] GPU Support</li><li>[X] Autograd</li><li>[X] Broadcasting</li><li>[] Memory Management</li></ul>                 |
 | Loss                         | in progress | <ul><li>[x] MSE</li><li>[X] Cross Entropy</li></ul>    |
 | Data                         | in progress    | <ul><li>[X] Dataset</li><li>[X] Batch</li><li>[X] Iterator</li></ul>   |
 | Convolutional Neural Network | in progress    | <ul><li>[ ] Conv2d</li><li>[ ] MaxPool2d</li><li>[ ] Dropout</li></ul> |

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);

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);

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)) {

norch/distributed/distributed.py

@@ -0,0 +1,2 @@
+def init_process_group(world_size, rank, backend='nccl'):
+    pass

norch/distributed/run/__main__.py

@@ -0,0 +1,8 @@
+import sys
+from . import run
+
+def main():
+    run.main()
+
+if __name__ == "__main__":
+    main()

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

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')
 

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