diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py index 57ad689179da9..a565e8932fca4 100644 --- a/test/distributed/test_c10d_gloo.py +++ b/test/distributed/test_c10d_gloo.py @@ -55,6 +55,7 @@ run_tests, skip_but_pass_in_sandcastle, skipIfRocmArch, + skipIfRocm, TestCase, ) @@ -395,6 +396,7 @@ def test_broadcast_stress(self): inputs = [torch.tensor([i * self.world_size + self.rank]) for i in range(1000)] self._test_broadcast_stress(inputs) + @skipIfRocm @skip_if_lt_x_gpu(2) @requires_gloo() def test_broadcast_stress_cuda(self): @@ -500,6 +502,7 @@ def test_allreduce_stress(self): inputs = [torch.tensor([i + self.rank]) for i in range(1000)] self._test_allreduce_stress(inputs) + @skipIfRocm @skip_if_lt_x_gpu(2) @requires_gloo() def test_allreduce_stress_cuda(self): @@ -953,6 +956,8 @@ def test_scatter_stress(self): @skip_but_pass_in_sandcastle( "Test is flaky, see https://github.com/pytorch/pytorch/issues/15963" ) + + @skipIfRocm @skip_if_lt_x_gpu(2) @requires_gloo() def test_scatter_stress_cuda(self): @@ -1127,6 +1132,7 @@ def test_gather_stress(self): inputs = [torch.tensor([i + self.rank]) for i in range(1000)] self._test_gather_stress(inputs, lambda t: t.clone()) + @skipIfRocm @skip_if_lt_x_gpu(2) @skipIfRocmArch(MI300_ARCH) @requires_gloo() @@ -1263,6 +1269,7 @@ def test_allgather_stress(self): inputs = [torch.tensor([i + self.rank]) for i in range(1000)] self._test_allgather_stress(inputs, lambda t: t.clone()) + @skipIfRocm @skip_if_lt_x_gpu(2) @requires_gloo() def test_allgather_stress_cuda(self): @@ -1449,6 +1456,7 @@ def test_reduce_stress(self): inputs = [torch.tensor([i + self.rank]) for i in range(1000)] self._test_reduce_stress(inputs) + @skipIfRocm @skip_if_lt_x_gpu(2) @requires_gloo() def test_reduce_stress_cuda(self):