From f09885abece3c9667d02b9048e6dcaadf261bbc2 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Fri, 31 May 2024 17:36:10 -0300 Subject: [PATCH 1/4] test cloud multigpu --- train.py | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index 29d8bc4..22a04c5 100644 --- a/train.py +++ b/train.py @@ -1,6 +1,7 @@ import os import norch import norch.distributed as dist +import norch.distributed def main(): @@ -10,7 +11,7 @@ def main(): dist.init_process_group(rank, world_size) - tensor = norch.Tensor([1,1,1]) + tensor = norch.Tensor([1,1,1]).to(rank) tensor = (rank + 1) * tensor print(f"BEFORE on rank {rank}: {tensor} \n\n") @@ -27,6 +28,93 @@ def main(): print(f"AFTER BROADCAST on rank {rank}: {tensor} \n\n") +def main2(): + import norch + import norch.nn as nn + import norch.optim as optim + from norch.utils.data.dataloader import DataLoader + from norch.nn.parallel import DistributedDataParallel + from norch.norchvision import transforms as T + import numpy as np + import matplotlib.pyplot as plt + 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 = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=world_size, rank=local_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"Local rank: {local_rank}") + print(f"World size: {world_size}") + + 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() + print(f"AFTER rank {local_rank}: {model.module.fc1.weight.grad}") + + + optimizer.step() + + print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}') + loss_list.append(loss[0]) + if __name__ == "__main__": main() From c8eb01ca92214fc2f11eb49ab63cbc18d7dd35ff Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Fri, 31 May 2024 17:37:09 -0300 Subject: [PATCH 2/4] test cloud multigpu --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 22a04c5..ff71c50 100644 --- a/train.py +++ b/train.py @@ -111,9 +111,9 @@ def main2(): optimizer.step() + break - print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}') - loss_list.append(loss[0]) + break if __name__ == "__main__": main() From 9735007805c6140bfeb6fdde29818eb4148ee0ec Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Fri, 31 May 2024 17:42:15 -0300 Subject: [PATCH 3/4] test cloud multigpu --- train.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index ff71c50..8abb8cd 100644 --- a/train.py +++ b/train.py @@ -34,6 +34,7 @@ def main2(): import norch.optim as optim from norch.utils.data.dataloader import DataLoader from norch.nn.parallel import DistributedDataParallel + from norch.utils.data.distributed import DistributedSampler from norch.norchvision import transforms as T import numpy as np import matplotlib.pyplot as plt @@ -64,7 +65,7 @@ def main2(): ) train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform) - distributed_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=world_size, rank=local_rank) + distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=local_rank) train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler) class MyModel(nn.Module): @@ -107,7 +108,8 @@ def main2(): optimizer.zero_grad() loss.backward() - print(f"AFTER rank {local_rank}: {model.module.fc1.weight.grad}") + print(f"AFTER rank {local_rank}: {model.module.fc2.bias.grad}") + print("\n\n") optimizer.step() @@ -116,5 +118,5 @@ def main2(): break if __name__ == "__main__": - main() + main2() From ef9a74b2573570cf6227f8e46ee11b1310057d3d Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Fri, 31 May 2024 17:43:22 -0300 Subject: [PATCH 4/4] test cloud multigpu --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 8abb8cd..96cbc8d 100644 --- a/train.py +++ b/train.py @@ -118,5 +118,5 @@ def main2(): break if __name__ == "__main__": - main2() + main()