[Distributed tests] Add skip for odd world_size condition

As per pytorch#74995, the tests
needs to be skipped for odd WORLD_SIZE

Signed-off-by: Jagadish Krishnamoorthy <[email protected]>

Fixes pytorch#74995

Pull Request resolved: pytorch#76136
Approved by: https://github.com/kumpera, https://github.com/wayi1

Author: jaglinux

torch/testing/_internal/common_distributed.py

@@ -47,6 +47,7 @@ class TestSkip(NamedTuple):
         72, "Skipped because distributed backend is not available."
     ),
     "small_worldsize": TestSkip(73, "Skipped due to small world size."),
+    "odd_worldsize": TestSkip(87, "Skipped due to odd world size."),
     "no_cuda": TestSkip(74, "CUDA is not available."),
     "multi-gpu-1": TestSkip(75, "Need at least 1 CUDA device"),
     "multi-gpu-2": TestSkip(77, "Need at least 2 CUDA devices"),
@@ -108,6 +109,15 @@ def wrapper(*args, **kwargs):
 
     return wrapper
 
+def skip_if_odd_worldsize(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if (os.environ["BACKEND"] != "mpi") and int(os.environ["WORLD_SIZE"]) % 2 == 1:
+            sys.exit(TEST_SKIPS["odd_worldsize"].exit_code)
+
+        return func(*args, **kwargs)
+
+    return wrapper
 
 def require_n_gpus_for_nccl_backend(n, backend):
     def decorator(func):

torch/testing/_internal/distributed/distributed_test.py

@@ -49,6 +49,7 @@
     simple_sparse_reduce_tests,
     skip_if_rocm,
     skip_if_small_worldsize,
+    skip_if_odd_worldsize,
     skip_if_lt_x_gpu,
     nccl_skip_if_lt_x_gpu,
     skip_if_no_gpu,
@@ -4871,6 +4872,7 @@ def _create_hierarchical_model_averager(self):
             )
 
         @skip_if_lt_x_gpu(4)
+        @skip_if_odd_worldsize
         @sandcastle_skip_if(
             BACKEND not in DistTestCases.backend_feature["ddp"],
             f"The {BACKEND} backend does not support DistributedDataParallel"
@@ -4883,6 +4885,7 @@ def test_post_localSGD_optimizer_parity_with_hierarchical_sgd(self):
             )
 
         @skip_if_lt_x_gpu(4)
+        @skip_if_odd_worldsize
         @sandcastle_skip_if(
             BACKEND not in DistTestCases.backend_feature["ddp"],
             f"The {BACKEND} backend does not support DistributedDataParallel"