Merge pull request #5 from daisyden/tdey/newskiptest-12ranks

[XPU] Added a new skip test to detect if the world size is set to non-power of 2

Author: newtdms

test/distributed/fsdp/test_fsdp_tp_integration.py

@@ -25,7 +25,11 @@
     parallelize_module,
     RowwiseParallel,
 )
-from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
+from torch.testing._internal.common_distributed import (
+    skip_if_lt_x_gpu,
+    skip_if_not_powerof2_worldsize_xpu,
+
+)
 from torch.testing._internal.common_fsdp import FSDPTest
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -230,6 +234,7 @@ def _get_grads_as_flattened(
         return torch.cat(all_grads_per_param).contiguous()
 
     @skip_if_lt_x_gpu(4)
+    @skip_if_not_powerof2_worldsize_xpu()
     def test_fsdp_tp_integration(self):
         self.run_subtests(
             {

torch/testing/_internal/common_distributed.py

@@ -85,6 +85,7 @@ class TestSkip(NamedTuple):
     "importerror": TestSkip(88, "Test skipped due to missing import"),
     "no_accelerator": TestSkip(89, "accelerator is not available."),
     "not-support-multithread": TestSkip(90, "backend not support multithread."),
+    "power-of-two": TestSkip(91, "world size needs to be power of two on xpu."),
 }
 
 
@@ -210,6 +211,23 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def skip_if_not_powerof2_worldsize_xpu():
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            if TEST_XPU:
+                x = torch.xpu.device_count()
+                print(f"skip_if_not_powerof2_worldsize_xpu x: {x}")
+                if (x & (x - 1)) == 0:
+                    return func(*args, **kwargs)
+                sys.exit(TEST_SKIPS[f"power-of-two"].exit_code)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 # This decorator helps avoiding initializing cuda while testing other backends
 def nccl_skip_if_lt_x_gpu(backend, x):
     def decorator(func):