Adding torch accelerator to ddp-tutorial-series example

Signed-off-by: dggaytan <[email protected]>

Author: dggaytan

distributed/ddp-tutorial-series/multigpu.py

@@ -25,8 +25,7 @@ def ddp_setup(rank, world_size):
         torch.accelerator.set_device_index(rank)
         print(f"Running on rank {rank} on device {device}")
     else:
-        device = torch.device("cpu")
-        print(f"Running on device {device}")
+        print(f"Multi-GPU environment not detected")
 
     backend = torch.distributed.get_default_backend_for_device(device)
     init_process_group(backend=backend, rank=rank, world_size=world_size)

distributed/ddp-tutorial-series/multigpu_torchrun.py

@@ -17,12 +17,11 @@ def ddp_setup():
         torch.accelerator.set_device_index(rank)
         print(f"Running on rank {rank} on device {device}")
     else:
-        device = torch.device("cpu")
-        print(f"Running on device {device}")
-                
-    backend = torch.distributed.get_default_backend_for_device(device)
-    torch.distributed.init_process_group(backend=backend, device_id=device)
-    return device
+        print(f"Multi-GPU environment not detected")
+
+    backend = torch.distributed.get_default_backend_for_device(rank)
+    torch.distributed.init_process_group(backend=backend, rank=rank, device_id=rank)
+   
 
 
 class Trainer:
@@ -33,7 +32,6 @@ def __init__(
         optimizer: torch.optim.Optimizer,
         save_every: int,
         snapshot_path: str,
-        device: torch.device, 
     ) -> None:
         self.gpu_id = int(os.environ["LOCAL_RANK"])
         self.model = model.to(self.gpu_id)
@@ -42,15 +40,14 @@ def __init__(
         self.save_every = save_every
         self.epochs_run = 0
         self.snapshot_path = snapshot_path
-        self.device = device
         if os.path.exists(snapshot_path):
             print("Loading snapshot")
             self._load_snapshot(snapshot_path)
 
         self.model = DDP(self.model, device_ids=[self.gpu_id])
 
     def _load_snapshot(self, snapshot_path):
-        loc = str(self.device)
+        loc = str(torch.accelerator.current_accelerator())
         snapshot = torch.load(snapshot_path, map_location=loc)
         self.model.load_state_dict(snapshot["MODEL_STATE"])
         self.epochs_run = snapshot["EPOCHS_RUN"]
@@ -105,10 +102,10 @@ def prepare_dataloader(dataset: Dataset, batch_size: int):
 
 
 def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"):
-    device = ddp_setup()
+    ddp_setup()
     dataset, model, optimizer = load_train_objs()
     train_data = prepare_dataloader(dataset, batch_size)
-    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path, device)
+    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)
     trainer.train(total_epochs)
     destroy_process_group()
 

distributed/ddp-tutorial-series/multinode.py

@@ -17,12 +17,11 @@ def ddp_setup():
         torch.accelerator.set_device_index(rank)
         print(f"Running on rank {rank} on device {device}")
     else:
-        device = torch.device("cpu")
-        print(f"Running on device {device}")
-                
-    backend = torch.distributed.get_default_backend_for_device(device)
-    torch.distributed.init_process_group(backend=backend, device_id=device)
-    return device
+        print(f"Multi-GPU environment not detected")
+
+    backend = torch.distributed.get_default_backend_for_device(rank)
+    torch.distributed.init_process_group(backend=backend, rank=rank, device_id=rank)
+
 
 class Trainer:
     def __init__(
@@ -32,7 +31,6 @@ def __init__(
         optimizer: torch.optim.Optimizer,
         save_every: int,
         snapshot_path: str,
-        device: torch.device,
     ) -> None:
         self.local_rank = int(os.environ["LOCAL_RANK"])
         self.global_rank = int(os.environ["RANK"])
@@ -42,15 +40,14 @@ def __init__(
         self.save_every = save_every
         self.epochs_run = 0
         self.snapshot_path = snapshot_path
-        self.device = device
         if os.path.exists(snapshot_path):
             print("Loading snapshot")
             self._load_snapshot(snapshot_path)
 
         self.model = DDP(self.model, device_ids=[self.local_rank])
 
     def _load_snapshot(self, snapshot_path):
-        loc = str(self.device)
+        loc = str(torch.accelerator.current_accelerator())
         snapshot = torch.load(snapshot_path, map_location=loc)
         self.model.load_state_dict(snapshot["MODEL_STATE"])
         self.epochs_run = snapshot["EPOCHS_RUN"]
@@ -105,10 +102,10 @@ def prepare_dataloader(dataset: Dataset, batch_size: int):
 
 
 def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"):
-    device = ddp_setup()
+    ddp_setup()
     dataset, model, optimizer = load_train_objs()
     train_data = prepare_dataloader(dataset, batch_size)
-    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path, device)
+    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)
     trainer.train(total_epochs)
     destroy_process_group()