Adding torch accelerator to ddp-tutorial-series example

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

Author: dggaytan

distributed/ddp-tutorial-series/multigpu.py

@@ -17,24 +17,20 @@ def ddp_setup(rank, world_size):
         world_size: Total number of processes
     """
     os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = "12453"
+    os.environ["MASTER_PORT"] = "12455"
 
 
+    rank = int(os.environ["LOCAL_RANK"])
     if torch.accelerator.is_available():
         device_type = torch.accelerator.current_accelerator()
-        torch.accelerator.set_device_idx(rank)
-        device: torch.device = torch.device(f"{device_type}:{rank}")
+        device = torch.device(f"{device_type}:{rank}")
         torch.accelerator.device_index(rank)
         print(f"Running on rank {rank} on device {device}")
-        backend = torch.distributed.get_default_backend_for_device(device)
-        torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size, device_id=device)
     else:
         device = torch.device("cpu")
         print(f"Running on device {device}")
-        torch.distributed.init_process_group(backend="gloo", device_id=device)
 
-    # torch.cuda.set_device(rank)
-    # init_process_group(backend="xccl", rank=rank, world_size=world_size)
+    backend = torch.distributed.get_default_backend_for_device(device)
 
 class Trainer:
     def __init__(
@@ -116,5 +112,4 @@ def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_s
     args = parser.parse_args()
 
     world_size = torch.accelerator.device_count()
-    print(world_size)
     mp.spawn(main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size)

distributed/ddp-tutorial-series/multigpu_torchrun.py

@@ -14,17 +14,16 @@ def ddp_setup():
     rank = int(os.environ["LOCAL_RANK"])
     if torch.accelerator.is_available():
         device_type = torch.accelerator.current_accelerator()
-        device: torch.device = torch.device(f"{device_type}:{rank}")
+        device = torch.device(f"{device_type}:{rank}")
         torch.accelerator.device_index(rank)
         print(f"Running on rank {rank} on device {device}")
-        backend = torch.distributed.get_default_backend_for_device(device)
-        torch.distributed.init_process_group(backend=backend)
-        return device_type
     else:
         device = torch.device("cpu")
         print(f"Running on device {device}")
-        torch.distributed.init_process_group(backend="gloo")
-        return device
+                
+    backend = torch.distributed.get_default_backend_for_device(device)
+    torch.distributed.init_process_group(backend=backend, device_id=device)
+    return device
 
 
 class Trainer:

distributed/ddp-tutorial-series/requirements.txt

@@ -1 +1 @@
-torch>=2.7
+torch>=2.7

distributed/ddp-tutorial-series/run_example.sh

@@ -7,4 +7,4 @@
 # example.py
 
 echo "Launching ${1:-example.py} with ${2:-2} gpus"
-torchrun --nnodes=1 --nproc_per_node=${2:-2} --rdzv_id=101 --rdzv_endpoint="localhost:5972" ${1:-example.py}
+torchrun --nnodes=1 --nproc_per_node=${2:-2} --rdzv_id=101 --rdzv_endpoint="localhost:5972" ${1:-example.py}