support optim cases

Author: Chao1Han

torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py

@@ -233,12 +233,12 @@ def hook_with_zero_step(
     ddp_ref = weakref.ref(ddp)
 
     # NOTE: Gloo may hang with this overlapping approach, so we require
-    # NCCL/HCCL backend for now; see https://github.com/pytorch/pytorch/issues/62300
+    # NCCL/HCCL/XCCL backend for now; see https://github.com/pytorch/pytorch/issues/62300
     pg = dist.get_backend(ddp_ref().process_group)  # type: ignore[union-attr]
-    if (pg != dist.Backend.NCCL) and (pg != "hccl"):
+    if (pg != dist.Backend.NCCL) and (pg != "hccl") and (pg != "xccl"):
         raise RuntimeError(
             "Overlapping DDP with ZeRO using this approach currently requires "
-            "NCCL/HCCL backend to avoid hangs"
+            "NCCL/HCCL/XCCL backend to avoid hangs"
         )
 
     if shard_buckets:
@@ -395,12 +395,12 @@ def hook_with_zero_step_interleaved(
     ddp_ref = weakref.ref(ddp)
 
     # NOTE: Gloo may hang with this overlapping approach, so we require
-    # NCCL/HCCL backend for now; see https://github.com/pytorch/pytorch/issues/62300
+    # NCCL/HCCL/XCCL backend for now; see https://github.com/pytorch/pytorch/issues/62300
     pg = dist.get_backend(ddp_ref().process_group)  # type: ignore[union-attr]
-    if (pg != dist.Backend.NCCL) and (pg != "hccl"):
+    if (pg != dist.Backend.NCCL) and (pg != "hccl") and (pg != "xccl"):
         raise RuntimeError(
             "Overlapping DDP with ZeRO using this approach currently requires "
-            "NCCL/HCCL backend to avoid hangs"
+            "NCCL/HCCL/XCCL backend to avoid hangs"
         )
 
     if shard_buckets: