Add support for other backends in get_preferred_device (#132118)

jeejakp12 · pytorchmergebot · commit c184ac0f6b6d · 2024-08-07T07:19:20.000Z
Currenlty get_preferred_device supports only cuda and cpu. Add support for other backends using backend config. Fixes #ISSUE_NUMBER Pull Request resolved: #132118 Approved by: https://github.com/awgu
diff --git a/torch/distributed/_shard/sharded_tensor/api.py b/torch/distributed/_shard/sharded_tensor/api.py
@@ -14,6 +14,7 @@
 import torch
 import torch.distributed as dist
 import torch.distributed._shard.sharding_spec as shard_spec
+from torch._utils import _get_device_module
 from torch.distributed import distributed_c10d, rpc
 from torch.distributed._shard._utils import DEPRECATE_MSG
 from torch.distributed._shard.sharding_spec._internals import (
@@ -370,8 +371,16 @@ def _get_preferred_device(self) -> torch.device:
         Return the preferred device to be used when creating tensors for collectives.
         This method takes into account the associated process group
         """
-        if dist.get_backend(self._process_group) == dist.Backend.NCCL:
+        backend = dist.get_backend(self._process_group)
+        if backend == dist.Backend.NCCL:
             return torch.device(torch.cuda.current_device())
+        else:
+            backend_config = dist.BackendConfig(backend)
+            for device, backend_str in backend_config.get_device_backend_map().items():
+                if backend_str == backend and device != "cpu":
+                    return torch.device(
+                        device, _get_device_module(device).current_device()
+                    )
         return torch.device("cpu")
 
     def gather(  # type: ignore[override]
