Enable split_group API when TorchComms is used as a backend for TorchTitan on XPU (pytorch#178236)

When TP>1 is enabled for TorchTitan models using TorchComms backend, the execution hangs on XPU (new_comms creation times out). Investigation shows that split_group API is being enabled only for cuda device for the same model configuration. This PR makes the calls generic using the accelerator API, and makes TP>1 cases functional on XPU.
Pull Request resolved: pytorch#178236
Approved by: https://github.com/d4l3k, https://github.com/guangyey

Co-authored-by: Yu, Guangye <guangye.yu@intel.com>

Author: 2 people

torch/distributed/device_mesh.py

@@ -549,10 +549,12 @@ def _init_one_process_group(
                     getattr(default_group, "bound_device_id", None) is not None
                     or dist_config.use_torchcomms
                 )
-                and torch.cuda.is_available()
+                and torch.accelerator.is_available()
                 and (
                     backend is None
-                    or default_group._get_backend(torch.device("cuda")).name()
+                    or default_group._get_backend(
+                        torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+                    ).name()
                     == backend
                 )
             ):

torch/distributed/distributed_c10d.py

@@ -5453,10 +5453,18 @@ def split_group(
         )
 
     parent_group_rank = parent_global_to_group_ranks[global_rank]
-    parent_backend = parent_pg._get_backend(torch.device("cuda"))
+
+    if torch.accelerator.is_available():
+        parent_backend = parent_pg._get_backend(
+            torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+        )
+    else:
+        raise RuntimeError(
+            "No backend for the parent process group or its backend does not support splitting"
+        )
 
     # if the parent backend does not support splitting, raise error
-    # currently this API only support NCCL backend
+    # currently this API only support NCCL and XCCL backend
     if (
         not parent_backend or not parent_backend.supports_splitting
     ) and not _use_torchcomms_enabled():
@@ -5522,7 +5530,16 @@ def split_group(
 
     global_ranks_in_my_group = [parent_group_to_global_ranks[rank] for rank in my_group]
     split_pg.bound_device_id = device_id  # type: ignore[union-attr]
-    split_backend_class = split_pg._get_backend(torch.device("cuda"))
+
+    if torch.accelerator.is_available():
+        split_backend_class = split_pg._get_backend(
+            torch.accelerator.current_accelerator()  # pyrefly: ignore[bad-argument-type]
+        )
+    else:
+        raise RuntimeError(
+            "No backend for the parent process group or its backend does not support splitting"
+        )
+
     if not _use_torchcomms_enabled():
         split_backend_class._set_sequence_number_for_group()
     if split_pg.group_name != group_name: