Add accelerator API to RPC distributed examples:

- ddp_rpc
- parameter_server
- rnn

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

Author: jafraustro

distributed/rpc/ddp_rpc/main.py

@@ -27,12 +27,13 @@ class HybridModel(torch.nn.Module):
     def __init__(self, remote_emb_module, device):
         super(HybridModel, self).__init__()
         self.remote_emb_module = remote_emb_module
-        self.fc = DDP(torch.nn.Linear(16, 8).cuda(device), device_ids=[device])
+        self.fc = DDP(torch.nn.Linear(16, 8).to(device), device_ids=[device])
         self.device = device
 
     def forward(self, indices, offsets):
+        device = torch.accelerator.current_accelerator()
         emb_lookup = self.remote_emb_module.forward(indices, offsets)
-        return self.fc(emb_lookup.cuda(self.device))
+        return self.fc(emb_lookup.to(self.device))
 
 
 def _run_trainer(remote_emb_module, rank):
@@ -83,7 +84,7 @@ def get_next_batch(rank):
                 batch_size += 1
 
             offsets_tensor = torch.LongTensor(offsets)
-            target = torch.LongTensor(batch_size).random_(8).cuda(rank)
+            target = torch.LongTensor(batch_size).random_(8).to(rank)
             yield indices, offsets_tensor, target
 
     # Train for 100 epochs
@@ -145,9 +146,15 @@ def run_worker(rank, world_size):
         for fut in futs:
             fut.wait()
     elif rank <= 1:
+        if torch.accelerator.is_available():
+            acc = torch.accelerator.current_accelerator()
+            device = torch.device(acc)
+        else:
+            device = torch.device("cpu")
+        backend = torch.distributed.get_default_backend_for_device(device)
         # Initialize process group for Distributed DataParallel on trainers.
         dist.init_process_group(
-            backend="gloo", rank=rank, world_size=2, init_method="tcp://localhost:29500"
+            backend=backend, rank=rank, world_size=2, init_method="tcp://localhost:29500"
         )
 
         # Initialize RPC.

distributed/rpc/ddp_rpc/requirements.txt

@@ -1 +1 @@
-torch>=1.6.0
+torch>=2.7.1

distributed/rpc/parameter_server/rpc_parameter_server.py

@@ -20,15 +20,19 @@ def __init__(self, num_gpus=0):
         super(Net, self).__init__()
         print(f"Using {num_gpus} GPUs to train")
         self.num_gpus = num_gpus
-        device = torch.device(
-            "cuda:0" if torch.cuda.is_available() and self.num_gpus > 0 else "cpu")
+        if torch.accelerator.is_available() and self.num_gpus > 0:
+            acc = torch.accelerator.current_accelerator()
+            device = torch.device(f'{acc}:0')
+        else:
+            device = torch.device("cpu")
         print(f"Putting first 2 convs on {str(device)}")
-        # Put conv layers on the first cuda device
+        # Put conv layers on the first accelerator device
         self.conv1 = nn.Conv2d(1, 32, 3, 1).to(device)
         self.conv2 = nn.Conv2d(32, 64, 3, 1).to(device)
-        # Put rest of the network on the 2nd cuda device, if there is one
-        if "cuda" in str(device) and num_gpus > 1:
-            device = torch.device("cuda:1")
+        # Put rest of the network on the 2nd accelerator device, if there is one
+        if torch.accelerator.is_available() and self.num_gpus > 0:
+            acc = torch.accelerator.current_accelerator()
+            device = torch.device(f'{acc}:1')
 
         print(f"Putting rest of layers on {str(device)}")
         self.dropout1 = nn.Dropout2d(0.25).to(device)
@@ -72,21 +76,22 @@ def call_method(method, rref, *args, **kwargs):
 # <foo_instance>.bar(arg1, arg2) on the remote node and getting the result
 # back.
 
-
 def remote_method(method, rref, *args, **kwargs):
     args = [method, rref] + list(args)
     return rpc.rpc_sync(rref.owner(), call_method, args=args, kwargs=kwargs)
 
-
 # --------- Parameter Server --------------------
 class ParameterServer(nn.Module):
     def __init__(self, num_gpus=0):
         super().__init__()
         model = Net(num_gpus=num_gpus)
         self.model = model
-        self.input_device = torch.device(
-            "cuda:0" if torch.cuda.is_available() and num_gpus > 0 else "cpu")
-
+        if torch.accelerator.is_available() and num_gpus > 0:
+            acc = torch.accelerator.current_accelerator()
+            self.input_device = torch.device(f'{acc}:0')
+        else:
+            self.input_device = torch.device("cpu")
+            
     def forward(self, inp):
         inp = inp.to(self.input_device)
         out = self.model(inp)
@@ -113,11 +118,9 @@ def get_param_rrefs(self):
         param_rrefs = [rpc.RRef(param) for param in self.model.parameters()]
         return param_rrefs
 
-
 param_server = None
 global_lock = Lock()
 
-
 def get_parameter_server(num_gpus=0):
     global param_server
     # Ensure that we get only one handle to the ParameterServer.
@@ -197,8 +200,11 @@ def get_accuracy(test_loader, model):
     model.eval()
     correct_sum = 0
     # Use GPU to evaluate if possible
-    device = torch.device("cuda:0" if model.num_gpus > 0
-        and torch.cuda.is_available() else "cpu")
+    if torch.accelerator.is_available() and model.num_gpus > 0:
+        acc = torch.accelerator.current_accelerator()
+        device = torch.device(f'{acc}:0')
+    else:
+        device = torch.device("cpu")
     with torch.no_grad():
         for i, (data, target) in enumerate(test_loader):
             out = model(data)

distributed/rpc/rnn/rnn.py

@@ -43,13 +43,15 @@ def __init__(self, ntoken, ninp, dropout):
         super(EmbeddingTable, self).__init__()
         self.drop = nn.Dropout(dropout)
         self.encoder = nn.Embedding(ntoken, ninp)
-        if torch.cuda.is_available():
-            self.encoder = self.encoder.cuda()
+        if torch.accelerator.is_available():
+            device = torch.accelerator.current_accelerator()
+            self.encoder = self.encoder.to(device)
         nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
 
     def forward(self, input):
-        if torch.cuda.is_available():
-            input = input.cuda()
+        if torch.accelerator.is_available():
+            device = torch.accelerator.current_accelerator()
+            input = input.to(device)
         return self.drop(self.encoder(input)).cpu()