# Owner(s): ["oncall: distributed"] from unittest import mock import torch import torch.distributed as c10d from torch.distributed.collective_utils import ( _check_rng_sync, _check_rng_sync_internal, _summarize_ranks, all_gather, broadcast, ) from torch.distributed.device_mesh import init_device_mesh from torch.testing import FileCheck from torch.testing._internal.common_distributed import MultiProcessTestCase from torch.testing._internal.common_utils import ( instantiate_parametrized_tests, parametrize, run_tests, TestCase, ) from torch.testing._internal.distributed.fake_pg import FakeStore class TestCollectiveUtils(MultiProcessTestCase): def setUp(self): super().setUp() self._spawn_processes() def tearDown(self) -> None: super().tearDown() def opts(self, threads=2): opts = c10d.ProcessGroupGloo._Options() opts._timeout = 50.0 opts._threads = threads return opts def test_broadcast_result(self) -> None: """ Basic unit test for broadcast using a process group of default world size. """ store = c10d.FileStore(self.file_name, self.world_size) c10d.init_process_group( backend="gloo", store=store, rank=self.rank, world_size=self.world_size ) pg = c10d.new_group(pg_options=self.opts()) func = mock.MagicMock() func.return_value = pg.rank() res = broadcast(data_or_fn=func, rank=0, pg=pg) assert res == 0, f"Expect res to be 0 (got {res})" if pg.rank() == 0: func.assert_called_once() else: func.assert_not_called() func.reset_mock() res = broadcast(data_or_fn=func, rank=1, pg=pg) assert res == 1, f"Expect res to be 1 (got {res})" if pg.rank() == 1: func.assert_called_once() else: func.assert_not_called() def test_broadcast_result_no_pg(self) -> None: """ Ensure broadcast has no dependency on torch.distributed when run in single process. """ func = mock.MagicMock() broadcast(data_or_fn=func, rank=0) func.assert_called_once() def test_broadcast_result_raises_exceptions_from_func( self, ) -> None: """ Ensure broadcast exception is propagated properly. """ # no process group func = mock.MagicMock() exc = Exception("test exception") func.side_effect = exc expected_exception = "test exception" with self.assertRaisesRegex(Exception, expected_exception): broadcast(data_or_fn=func, rank=0) def test_all_gather_result(self) -> None: """ Basic unit test for all_gather using a process group of default world size. """ store = c10d.FileStore(self.file_name, self.world_size) c10d.init_process_group( backend="gloo", store=store, rank=self.rank, world_size=self.world_size ) pg = c10d.new_group(pg_options=self.opts()) func = mock.MagicMock() func.return_value = pg.rank() res = all_gather(data_or_fn=func, pg=pg) func.assert_called_once() assert res == list(range(self.world_size)), ( f"Expect res to be list of 0 through {self.world_size} (got {res})" ) def test_all_gather_result_no_pg(self) -> None: """ Ensure all_gather has no dependency on torch.distributed when run in single process. """ func = mock.MagicMock() all_gather(data_or_fn=func) func.assert_called_once() def test_all_gather_result_raises_exceptions_from_func( self, ) -> None: """ Ensure all_gather exception is propagated properly. """ # no process group func = mock.MagicMock() exc = Exception("test exception") func.side_effect = exc expected_exception = "test exception" with self.assertRaisesRegex(Exception, expected_exception): all_gather(data_or_fn=func) @parametrize("device", ["cpu", "cuda"]) def test_check_rng_sync( self, device, ) -> None: if device == "cuda" and not torch.cuda.is_available(): self.skipTest("Cuda is not available") store = c10d.FileStore(self.file_name, self.world_size) c10d.init_process_group( backend="gloo", store=store, rank=self.rank, world_size=self.world_size ) group = torch.distributed.distributed_c10d._get_default_group() generator = torch.Generator(device=device) generator.manual_seed(123) value_ranks, _ = _check_rng_sync_internal(generator, group) self.assertEqual(len(value_ranks), 1, value_ranks) for actual, expected in zip(value_ranks.values(), [{0, 1, 2, 3}]): self.assertEqual(actual, expected, actual) if torch.distributed.get_rank() == 1: torch.randn((10,), device=device, generator=generator) value_ranks, _ = _check_rng_sync_internal(generator, group) self.assertEqual(len(value_ranks), 2, value_ranks) for actual, expected in zip(value_ranks.values(), [{0, 2, 3}, {1}]): self.assertEqual(actual, expected, actual) if torch.distributed.get_rank() == 0: generator.manual_seed(456) value_ranks, _ = _check_rng_sync_internal(generator, group) self.assertEqual(len(value_ranks), 3, value_ranks) for actual, expected in zip(value_ranks.values(), [{0}, {1}, {2, 3}]): self.assertEqual(actual, expected, actual) log_str = _check_rng_sync(generator, group) FileCheck().check("Generator desync detected").check("Ranks").check("0").check( "1" ).check("2:4").run(log_str) class TestUtils(TestCase): def setUp(self): super().setUp() if not c10d.is_initialized(): self.rank = 0 self.world_size = 4096 store = FakeStore() c10d.init_process_group( backend="fake", world_size=self.world_size, rank=self.rank, store=store, ) def tearDown(self): c10d.destroy_process_group() def test_summarize_ranks(self): mesh_dim_names = ("pp", "dp", "tp") mesh = init_device_mesh("cpu", (8, 64, 8), mesh_dim_names=mesh_dim_names) ranks_lists = {name: mesh[name].mesh.tolist() for name in mesh_dim_names} summaries = { name: _summarize_ranks(ranks_lists[name]) for name in mesh_dim_names } self.assertEqual(summaries["pp"], "0:4096:512") self.assertEqual(summaries["dp"], "0:512:8") self.assertEqual(summaries["tp"], "0:8") self.assertEqual( _summarize_ranks([1, 2, 3, 6, 7, 8, 10, 12, 14, 16]), "1:4,6:9,10:18:2", ) self.assertEqual( _summarize_ranks([1]), "1", ) instantiate_parametrized_tests(TestCollectiveUtils) if __name__ == "__main__": run_tests()