已合并
[Test] Sync DeviceMesh tests with MeshLayout changes #43811
MengJiayi创建于 8月5日
[Test] Sync DeviceMesh tests with MeshLayout changes #43811
已合并
共 1 个文件变更+14-22
| @@ -1,16 +1,12 @@ | |||
| 1 | # Copyright (c) Meta Platforms, Inc. and affiliates | 1 | # Copyright (c) Meta Platforms, Inc. and affiliates |
| 2 | # Owner(s): ["oncall: distributed"] | 2 | # Owner(s): ["oncall: distributed"] |
| 3 | import os | 3 | import os |
| 4 | -import unittest | ||
| 5 | -from datetime import timedelta | ||
| 6 | from functools import wraps | 4 | from functools import wraps |
| 7 | from typing import Tuple, Dict, Any | 5 | from typing import Tuple, Dict, Any |
| 8 | 6 | ||
| 9 | import torch | 7 | import torch |
| 10 | import torch.distributed as dist | 8 | import torch.distributed as dist |
| 11 | import torch.distributed._functional_collectives as funcol | 9 | import torch.distributed._functional_collectives as funcol |
| 12 | -from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 13 | -from torch.distributed._mesh_layout import _MeshLayout as _Layout | ||
| 14 | from torch.distributed.device_mesh import _mesh_resources, DeviceMesh, init_device_mesh | 10 | from torch.distributed.device_mesh import _mesh_resources, DeviceMesh, init_device_mesh |
| 15 | from torch.distributed.distributed_c10d import ( | 11 | from torch.distributed.distributed_c10d import ( |
| 16 | _get_default_group, | 12 | _get_default_group, |
| @@ -30,11 +26,14 @@ from torch.distributed.tensor._collective_utils import ( | |||
| 30 | ) | 26 | ) |
| 31 | from torch.distributed.tensor.placement_types import _Partial, Shard | 27 | from torch.distributed.tensor.placement_types import _Partial, Shard |
| 32 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 28 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 33 | -from torch.testing._internal.distributed.fake_pg import FakeProcessGroup, FakeStore | 29 | +from torch.testing._internal.distributed.fake_pg import FakeStore |
| 34 | from torch.utils._typing_utils import not_none | 30 | from torch.utils._typing_utils import not_none |
| 35 | 31 | ||
| 36 | -import torch_npu | 32 | +from torch_npu.testing.common_distributed import ( |
| 37 | -from torch_npu.testing.common_distributed import with_comms, init_pg, skipIfUnsupportMultiNPU, TEST_SKIPS | 33 | + TEST_SKIPS, |
| 34 | + init_pg, | ||
| 35 | + skipIfUnsupportMultiNPU, | ||
| 36 | +) | ||
| 38 | from torch_npu.testing.testcase import run_tests | 37 | from torch_npu.testing.testcase import run_tests |
| 39 | 38 | ||
| 40 | 39 | ||
| @@ -217,19 +216,12 @@ class DeviceMeshTest(NPUDTensorTestBase): | |||
| 217 | ) | 216 | ) |
| 218 | 217 | ||
| 219 | 218 | ||
| 220 | -#DeviceMeshTest with resetting world_size to 4. | 219 | +# DeviceMeshTest with resetting world_size to 4. |
| 221 | class DeviceMeshTestF(NPUDTensorTestBase): | 220 | class DeviceMeshTestF(NPUDTensorTestBase): |
| 222 | 221 | ||
| 223 | def world_size(self): | 222 | def world_size(self): |
| 224 | return 4 | 223 | return 4 |
| 225 | 224 | ||
| 226 | - | ||
| 227 | - | ||
| 228 | - def test_assert_invalid_mesh_tensor(self): | ||
| 229 | - mesh = torch.arange(self.world_size).to(self.rank) | ||
| 230 | - with self.assertRaises(ValueError): | ||
| 231 | - device_mesh = DeviceMesh(self.device_type, mesh) | ||
| 232 | - | ||
| 233 | 225 | ||
| 234 | 226 | ||
| 235 | def test_get_local_rank(self): | 227 | def test_get_local_rank(self): |
| @@ -341,7 +333,7 @@ class DeviceMeshTestNDim(NPUDTensorTestBase): | |||
| 341 | self.assertEqual(ep_mesh, another_mesh) | 333 | self.assertEqual(ep_mesh, another_mesh) |
| 342 | 334 | ||
| 343 | 335 | ||
| 344 | -#DeviceMeshTestNDim with resetting world_size to 8. | 336 | +# DeviceMeshTestNDim with resetting world_size to 8. |
| 345 | class DeviceMeshTestNDimE(NPUDTensorTestBase): | 337 | class DeviceMeshTestNDimE(NPUDTensorTestBase): |
| 346 | 338 | ||
| 347 | def world_size(self): | 339 | def world_size(self): |
| @@ -654,7 +646,7 @@ class TestDeviceMeshGetItem(NPUDTensorTestBase): | |||
| 654 | self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d) | 646 | self.assertEqual(_mesh_resources.get_root_mesh(dp_cp_mesh), mesh_4d) |
| 655 | 647 | ||
| 656 | 648 | ||
| 657 | -#TestDeviceMeshGetItem with resetting world_size to 8. | 649 | +# TestDeviceMeshGetItem with resetting world_size to 8. |
| 658 | class TestDeviceMeshGetItemE(NPUDTensorTestBase): | 650 | class TestDeviceMeshGetItemE(NPUDTensorTestBase): |
| 659 | 651 | ||
| 660 | def world_size(self): | 652 | def world_size(self): |
| @@ -802,7 +794,7 @@ class TestDeviceMeshGetItemE(NPUDTensorTestBase): | |||
| 802 | flatten_mesh_layout = root_mesh._flatten_mapping["dp_cp"]._layout | 794 | flatten_mesh_layout = root_mesh._flatten_mapping["dp_cp"]._layout |
| 803 | self.assertEqual(flatten_mesh_layout, flattened_dp_cp_mesh._layout) | 795 | self.assertEqual(flatten_mesh_layout, flattened_dp_cp_mesh._layout) |
| 804 | self.assertEqual( | 796 | self.assertEqual( |
| 805 | - flattened_dp_cp_mesh._layout.global_ranks(8), | 797 | + flattened_dp_cp_mesh._layout.collapse().global_ranks(8), |
| 806 | [[0, 2, 4, 6], [1, 3, 5, 7]], | 798 | [[0, 2, 4, 6], [1, 3, 5, 7]], |
| 807 | ) | 799 | ) |
| 808 | 800 | ||
| @@ -822,13 +814,13 @@ class TestDeviceMeshGetItemE(NPUDTensorTestBase): | |||
| 822 | flatten_mesh_root_layout = root_mesh._flatten_mapping["dp_tp"]._layout | 814 | flatten_mesh_root_layout = root_mesh._flatten_mapping["dp_tp"]._layout |
| 823 | self.assertEqual(flatten_mesh_root_layout, flattened_dp_tp_mesh._layout) | 815 | self.assertEqual(flatten_mesh_root_layout, flattened_dp_tp_mesh._layout) |
| 824 | self.assertEqual( | 816 | self.assertEqual( |
| 825 | - flattened_dp_tp_mesh._layout.global_ranks(8), | 817 | + flattened_dp_tp_mesh._layout.collapse().global_ranks(8), |
| 826 | [[0, 1, 4, 5], [2, 3, 6, 7]], | 818 | [[0, 1, 4, 5], [2, 3, 6, 7]], |
| 827 | ) | 819 | ) |
| 828 | with self.assertRaisesRegex( | 820 | with self.assertRaisesRegex( |
| 829 | - NotImplementedError, | 821 | + KeyError, "Mesh dim indices should be in ascending order" |
| 830 | - "Currently, this only allows slicing out a contiguous flattened dim", | ||
| 831 | ): | 822 | ): |
| 823 | + # dp_tp is partly "above" and partly "below" cp | ||
| 832 | mesh_3d["dp_tp", "cp"] | 824 | mesh_3d["dp_tp", "cp"] |
| 833 | 825 | ||
| 834 | # Test flatten with a flattened mesh_dim_name | 826 | # Test flatten with a flattened mesh_dim_name |
| @@ -1193,4 +1185,4 @@ class DeviceMeshCollectiveTest(NPUDTensorTestBase): | |||
| 1193 | 1185 | ||
| 1194 | 1186 | ||
| 1195 | if __name__ == "__main__": | 1187 | if __name__ == "__main__": |
| 1196 | - run_tests() | 1188 | + run_tests() |