update

Author: Chao1Han

torch/testing/_internal/common_fsdp.py

@@ -1215,7 +1215,7 @@ def _run(cls, rank, test_name, file_name, pipe, **kwargs):
         device_ids = None
         device_id = self.rank % DEVICE_COUNT
         if TEST_CUDA or TEST_XPU:
-            torch.accelerator.set_device_idx(device_id)
+            torch.accelerator.set_device_index(device_id)
         device_ids = [device_id]
 
         # Execute barrier prior to running test to ensure that every process

torch/testing/_internal/distributed/_tensor/common_dtensor.py

@@ -35,6 +35,7 @@
 from torch.testing._internal.common_utils import (
     TEST_HPU,
     TEST_CUDA,
+    TEST_XPU,
 )
 from torch.testing._internal.common_distributed import (
     MultiProcessTestCase,
@@ -55,6 +56,10 @@
     DEVICE_TYPE = "hpu"
     PG_BACKEND = "hccl"
     DEVICE_COUNT = _get_device_module("hpu").device_count()
+elif TEST_HPU:
+    DEVICE_TYPE = "xpu"
+    PG_BACKEND = "xccl"
+    DEVICE_COUNT = _get_device_module("xpu").device_count()
 else:
     DEVICE_TYPE = "cpu"
     PG_BACKEND = "gloo"
@@ -328,6 +333,8 @@ def backend(self) -> str:
             backend = "nccl"
         elif TEST_HPU:
             backend = "hccl"
+        elif TEST_XPU:
+            backend = "xccl"
         else:
             backend = "gloo"
         return backend
@@ -345,7 +352,7 @@ def init_pg(self, eager_init) -> None:
         device_id = None
         if "nccl" or "xccl" in self.backend:
             # set device for nccl pg for collectives
-            torch.accelerator.set_device_idx(self.rank)
+            torch.accelerator.set_device_index(self.rank)
             # we only need to set device_id for nccl backend with eager init
             device_id = torch.device(f"{self.device_type}:{self.rank}") if eager_init else None
         # For nccl backend, bind the device to the process if device_id is not None