🐛 Add barrier

Author: ZihaoZhou

mnist_ddp.py

@@ -138,7 +138,19 @@ def main():
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
     ])
+
+    if args.distributed:
+        if torch.distributed.get_rank() != 0:
+            # might be downloading mnist data, let rank 0 download first
+            torch.distributed.barrier()
+
     train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
+
+    if args.distributed:
+        if torch.distributed.get_rank() == 0:
+            # mnist data is downloaded, indicate other ranks can proceed
+            torch.distributed.barrier()
+
     val_dataset = datasets.MNIST('./data', train=False, transform=transform)
     if args.distributed:
         train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
@@ -171,6 +183,9 @@ def main():
             # Only run validation on GPU 0 process, for simplicity, so we do not run validation on multi gpu.
             if dist.get_rank() == 0:
                 test(model_without_ddp, device, test_loader)
+                torch.distributed.barrier()
+            else:
+                torch.distributed.barrier()
         else:
             test(model, device, test_loader)
         scheduler.step()
@@ -183,7 +198,10 @@ def main():
         else:
             torch.save(model.state_dict(), f"mnist_cnn_.pt")
 
-    return dist.get_rank(), total_time
+    if args.distributed:
+        return dist.get_rank(), total_time
+    else:
+        return 0, total_time
 
 
 if __name__ == '__main__':

mnist_ds.py

@@ -85,7 +85,17 @@ def main():
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
     ])
-    dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform)
+
+    if torch.distributed.get_rank() != 0:
+        # might be downloading mnist data, let rank 0 download first
+        torch.distributed.barrier()
+
+    dataset1 = datasets.MNIST('./data', train=True, download=True, transform=transform)
+
+    if torch.distributed.get_rank() == 0:
+        # mnist data is downloaded, indicate other ranks can proceed
+        torch.distributed.barrier()
+
     dataset2 = datasets.MNIST('../data', train=False, transform=transform)
     test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
 

mnist_hf.py

@@ -24,7 +24,7 @@ def train(args, model, device, train_loader, optimizer, epoch):
         accelerator.backward(loss)
         optimizer.step()
         if batch_idx % args.log_interval == 0:
-            if local_rank == 0:
+            if accelerator.is_main_process:
                 print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                     epoch, AcceleratorState().num_processes * batch_idx * len(data), len(train_loader.dataset),
                            100. * batch_idx / len(train_loader), loss.item()))
@@ -46,7 +46,7 @@ def test(model, device, test_loader):
 
     test_loss /= len(test_loader.dataset)
 
-    if local_rank == 0:
+    if accelerator.is_main_process:
         print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
             test_loss, correct, len(test_loader.dataset),
             100. * correct / len(test_loader.dataset)))
@@ -100,8 +100,18 @@ def main():
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
     ])
-    dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform)
-    dataset2 = datasets.MNIST('../data', train=False, transform=transform)
+
+    if not accelerator.is_main_process:
+        # might be downloading mnist data, let rank 0 download first
+        accelerator.wait_for_everyone()
+
+    dataset1 = datasets.MNIST('./data', train=True, download=True, transform=transform)
+
+    if accelerator.is_main_process:
+        # mnist data is downloaded, indicate other ranks can proceed
+        accelerator.wait_for_everyone()
+
+    dataset2 = datasets.MNIST('./data', train=False, transform=transform)
     train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
     test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
 

mnist_hvd.py

@@ -122,9 +122,18 @@ def main():
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
         ])
+
+    if hvd.rank() != 0:
+        # might be downloading mnist data, let rank 0 download first
+        hvd.barrier()
+
     # train_dataset = datasets.MNIST('data-%d' % hvd.rank(), train=True, download=True, transform=transform)
     train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
-    
+
+    if hvd.rank() == 0:
+        # mnist data is downloaded, indicate other ranks can proceed
+        hvd.barrier()
+
     # Horovod: use DistributedSampler to partition the training data.
     train_sampler = dist.DistributedSampler(train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
     train_loader = torch.utils.data.DataLoader(