multigpu training support v0.0.5

README.md

@@ -1,13 +1,13 @@
 # PyNorch
-Recreating PyTorch from scratch (C/C++, CUDA and Python, with GPU support and automatic differentiation!)
+Recreating PyTorch from scratch (C/C++, CUDA and Python, with multi-GPU support and automatic differentiation!)
 
 Project details explanations can also be found on [medium](https://towardsdatascience.com/recreating-pytorch-from-scratch-with-gpu-support-and-automatic-differentiation-8f565122a3cc).
 
 # 1 - About
-**PyNorch** is a deep learning framework constructed using C/C++, CUDA and Python. This is a personal project with educational purpose only! `Norch` means **NOT** PyTorch, and we have **NO** claims to rivaling the already established PyTorch. The main objective of **PyNorch** was to give a brief understanding of how a deep learning framework works internally. It implements the Tensor object, GPU support and an automatic differentiation system. 
+**PyNorch** is a deep learning framework constructed using C/C++, CUDA and Python. This is a personal project with educational purpose only! `Norch` means **NOT** PyTorch, and we have **NO** claims to rivaling the already established PyTorch. The main objective of **PyNorch** was to give a brief understanding of how a deep learning framework works internally. It implements the Tensor object, multi-GPU support and an automatic differentiation system. 
 
 # 2 - Installation
-Install this package from PyPi (you can test on Colab!)
+Install this package from PyPi (you can test on Colab! Also tested on AWS g4dn.12xlarge instance with image ami-061debf863768593d)
 
 ```css
 $ pip install norch
@@ -60,8 +60,10 @@ class MyModel(nn.Module):
         return out
 ```
 
-### 3.3 - Example training
+### 3.3 - Example single GPU training
 ```python
+# examples/train_singlegpu.py
+
 import norch
 from norch.utils.data.dataloader import DataLoader
 from norch.norchvision import transforms as T
@@ -135,6 +137,107 @@ for epoch in range(epochs):
 
 ```
 
+### 3.4 - Example multi-GPU training
+First create a file .py as the example below
+
+```python
+# examples/train_multigpu.py
+
+import os
+import norch
+import norch.distributed as dist
+import norch.distributed
+import norch.nn as nn
+import norch.optim as optim
+from norch.nn.parallel import DistributedDataParallel
+from norch.utils.data.distributed import DistributedSampler
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+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
+)
+
+BATCH_SIZE = 32
+device = local_rank
+epochs = 10
+
+transform = T.Compose(
+    [
+        T.ToTensor(),
+        T.Reshape([-1, 784, 1])
+    ]
+)
+
+target_transform = T.Compose(
+    [
+        T.ToTensor()
+    ]
+)
+
+train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank)
+train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler)
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.fc1 = nn.Linear(784, 30)
+        self.sigmoid1 = nn.Sigmoid()
+        self.fc2 = nn.Linear(30, 10)
+        self.sigmoid2 = nn.Sigmoid()
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.sigmoid1(out)
+        out = self.fc2(out)
+        out = self.sigmoid2(out)
+        
+        return out
+
+model = MyModel().to(device)
+model = DistributedDataParallel(model)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(model.parameters(), lr=0.01)
+loss_list = []
+
+print(f"Starting training on Rank {rank}/{world_size}\n\n")
+
+for epoch in range(epochs):    
+    for idx, batch in enumerate(train_loader):
+
+        inputs, target = batch
+
+        inputs = inputs.to(device)
+        target = target.to(device)
+
+        outputs = model(inputs)
+        
+        loss = criterion(outputs, target)
+        
+        optimizer.zero_grad()
+        
+        loss.backward()
+
+        optimizer.step()
+    
+    if rank == 0:
+        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+        loss_list.append(loss[0])
+```
+Then you can run using
+
+```css
+$ python3 -m norch.distributed.run --nproc_per_node 4 examples/train_multigpu.py
+```
+
+
 
 # 4 - Progress
 
@@ -144,4 +247,4 @@ for epoch in range(epochs):
 | 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> |
-| Distributed                  | in progress | <ul><li>[ ] All-reduce</li><li>[X] DistributedSampler</li><li>[ ] DistributedDataParallel</li></ul>                 |
+| Distributed                  | in progress | <ul><li>[X] All-reduce</li><li>[X] Broadcast</li><li>[X] DistributedSampler</li><li>[X] DistributedDataParallel</li></ul>                 |

examples/train_multigpu.py

@@ -0,0 +1,93 @@
+import os
+import norch
+import norch.distributed as dist
+import norch.distributed
+import norch.nn as nn
+import norch.optim as optim
+from norch.nn.parallel import DistributedDataParallel
+from norch.utils.data.distributed import DistributedSampler
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+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
+    )
+
+    BATCH_SIZE = 32
+    device = local_rank
+    epochs = 10
+
+    transform = T.Compose(
+        [
+            T.ToTensor(),
+            T.Reshape([-1, 784, 1])
+        ]
+    )
+
+    target_transform = T.Compose(
+        [
+            T.ToTensor()
+        ]
+    )
+
+    train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+    distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank)
+    train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler)
+
+    class MyModel(nn.Module):
+        def __init__(self):
+            super(MyModel, self).__init__()
+            self.fc1 = nn.Linear(784, 30)
+            self.sigmoid1 = nn.Sigmoid()
+            self.fc2 = nn.Linear(30, 10)
+            self.sigmoid2 = nn.Sigmoid()
+
+        def forward(self, x):
+            out = self.fc1(x)
+            out = self.sigmoid1(out)
+            out = self.fc2(out)
+            out = self.sigmoid2(out)
+            
+            return out
+
+    model = MyModel().to(device)
+    model = DistributedDataParallel(model)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+    loss_list = []
+
+    print(f"Starting training on Rank {rank}/{world_size}\n\n")
+
+    for epoch in range(epochs):    
+        for idx, batch in enumerate(train_loader):
+
+            inputs, target = batch
+
+            inputs = inputs.to(device)
+            target = target.to(device)
+
+            outputs = model(inputs)
+            
+            loss = criterion(outputs, target)
+            
+            optimizer.zero_grad()
+            
+            loss.backward()
+
+            optimizer.step()
+        
+        if rank == 0:
+            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+            loss_list.append(loss[0])
+
+if __name__ == "__main__":
+    main()
+

examples/train_singlegpu.py

@@ -0,0 +1,71 @@
+import norch
+import norch.nn as nn
+import norch.optim as optim
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+def main():
+
+    BATCH_SIZE = 32
+    device = "cuda"
+    epochs = 10
+
+    transform = T.Compose(
+        [
+            T.ToTensor(),
+            T.Reshape([-1, 784, 1])
+        ]
+    )
+
+    target_transform = T.Compose(
+        [
+            T.ToTensor()
+        ]
+    )
+
+    train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+    train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE)
+
+    class MyModel(nn.Module):
+        def __init__(self):
+            super(MyModel, self).__init__()
+            self.fc1 = nn.Linear(784, 30)
+            self.sigmoid1 = nn.Sigmoid()
+            self.fc2 = nn.Linear(30, 10)
+            self.sigmoid2 = nn.Sigmoid()
+
+        def forward(self, x):
+            out = self.fc1(x)
+            out = self.sigmoid1(out)
+            out = self.fc2(out)
+            out = self.sigmoid2(out)
+            
+            return out
+
+    model = MyModel().to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+    loss_list = []
+
+    for epoch in range(epochs):    
+        for idx, batch in enumerate(train_loader):
+
+            inputs, target = batch
+
+            inputs = inputs.to(device)
+            target = target.to(device)
+
+            outputs = model(inputs)
+            loss = criterion(outputs, target)
+            
+            optimizer.zero_grad()
+            
+            loss.backward()
+
+            optimizer.step()
+
+
+if __name__ == "__main__":
+    main()
+

norch/__init__.py

@@ -5,6 +5,6 @@ from .utils import *
 from .norchvision import *
 from .utils import *
 
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 __author__ = 'Lucas de Lima Nogueira'
 __credits__ = 'Lucas de Lima Nogueira'

norch/csrc/distributed.cpp

@@ -58,19 +58,6 @@ void allreduce_sum_tensor(Tensor* tensor) {
     cudaStreamDestroy(stream);
 }
 
-void allreduce_mean_tensor(Tensor* tensor) {
-    cudaStream_t stream;
-    cudaStreamCreate(&stream);
-
-    // Perform NCCL AllReduce operation to calculate the mean of all tensors across all processes
-    NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
-
-    tensor_div_scalar_cuda(tensor, world_size, tensor->data);
-
-    cudaStreamSynchronize(stream);
-    cudaStreamDestroy(stream);
-}
-
 void end_process_group() {
     MPI_CHECK(MPI_Finalize());
     NCCL_CHECK(ncclCommDestroy(nccl_comm));

norch/distributed/distributed.py

@@ -28,10 +28,3 @@ def allreduce_sum_tensor(tensor):
     Tensor._C.allreduce_sum_tensor.restype = None
     
     Tensor._C.allreduce_sum_tensor(tensor.tensor)
-
-def allreduce_mean_tensor(tensor):
-
-    Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
-    Tensor._C.allreduce_mean_tensor.restype = None
-    
-    Tensor._C.allreduce_mean_tensor(tensor.tensor)

setup.py

@@ -76,7 +76,7 @@ class CustomInstall(install):
 
 setuptools.setup(
     name="norch",
-    version="0.0.4",
+    version="0.0.5",
     author="Lucas de Lima",
     author_email="nogueiralucasdelima@gmail.com",
     description="A deep learning framework",