Refactor DDP example to use Accelerator API

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

Author: jafraustro

distributed/ddp/README.md

@@ -74,11 +74,10 @@ def spmd_main():
     local_world_size = int(env_dict['LOCAL_WORLD_SIZE'])
 
     print(f"[{os.getpid()}] Initializing process group with: {env_dict}")
-    dist.init_process_group(backend="nccl")
-    print(
-        f"[{os.getpid()}] world_size = {dist.get_world_size()}, "
-        + f"rank = {dist.get_rank()}, backend={dist.get_backend()}"
-    )
+    acc = torch.accelerator.current_accelerator()
+    vendor_backend = torch.distributed.get_default_backend_for_device(acc)
+    torch.accelerator.set_device_index(rank)
+    torch.distributed.init_process_group(backend=vendor_backend)
 
     demo_basic(rank)
 

distributed/ddp/example.py

@@ -72,7 +72,10 @@ def main():
         dist.init_process_group(backend="gloo", init_method=init_method, rank=int(env_dict["RANK"]), world_size=int(env_dict["WORLD_SIZE"]))
     else:
         print(f"[{os.getpid()}] Initializing process group with: {env_dict}")  
-        dist.init_process_group(backend="nccl")
+        acc = torch.accelerator.current_accelerator()
+        backend = torch.distributed.get_default_backend_for_device(acc)
+        torch.accelerator.set_device_index(rank)
+        dist.init_process_group(backend=backend)
 
     print(
         f"[{os.getpid()}]: world_size = {dist.get_world_size()}, "

distributed/ddp/requirements.txt

@@ -1 +1 @@
-torch
+torch>=2.8