已合并
【API一致性任务】test: add Partial placement test for compute_local_shape_and_global_offset #45657
m0_53076416创建于 9月1日
【API一致性任务】test: add Partial placement test for compute_local_shape_and_global_offset #45657
已合并
共 1 个文件变更+17-8
| @@ -1,17 +1,13 @@ | |||
| 1 | import itertools | 1 | import itertools |
| 2 | import torch | 2 | import torch |
| 3 | from torch.distributed._tensor import distribute_tensor | 3 | from torch.distributed._tensor import distribute_tensor |
| 4 | -from torch.distributed._tensor._utils import ( | 4 | +from torch.distributed._tensor._utils import compute_local_shape_and_global_offset |
| 5 | - compute_local_shape, | ||
| 6 | - compute_local_shape_and_global_offset, | ||
| 7 | -) | ||
| 8 | from torch.distributed._tensor.device_mesh import DeviceMesh | 5 | from torch.distributed._tensor.device_mesh import DeviceMesh |
| 9 | -from torch.distributed._tensor.placement_types import Replicate, Shard | 6 | +from torch.distributed._tensor.placement_types import Replicate, Shard, Partial |
| 10 | 7 | ||
| 11 | from torch.testing._internal.common_utils import run_tests | 8 | from torch.testing._internal.common_utils import run_tests |
| 12 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 9 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 13 | 10 | ||
| 14 | -import torch_npu | ||
| 15 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 11 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 16 | 12 | ||
| 17 | 13 | ||
| @@ -31,7 +27,7 @@ class UtilTest(DTensorTestBase): | |||
| 31 | 27 | ||
| 32 | # replicate, shard | 28 | # replicate, shard |
| 33 | placements2 = [Replicate(), Shard(0)] | 29 | placements2 = [Replicate(), Shard(0)] |
| 34 | - local_size2 = compute_local_shape(size, mesh, placements2) | 30 | + local_size2, _ = compute_local_shape_and_global_offset(size, mesh, placements2) |
| 35 | if rank_coordinates[1] < 1: | 31 | if rank_coordinates[1] < 1: |
| 36 | self.assertEqual(local_size2, torch.Size([4, 7])) | 32 | self.assertEqual(local_size2, torch.Size([4, 7])) |
| 37 | else: | 33 | else: |
| @@ -39,7 +35,7 @@ class UtilTest(DTensorTestBase): | |||
| 39 | 35 | ||
| 40 | # shard, shard | 36 | # shard, shard |
| 41 | placements3 = [Shard(0), Shard(1)] | 37 | placements3 = [Shard(0), Shard(1)] |
| 42 | - local_size3 = compute_local_shape(size, mesh, placements3) | 38 | + local_size3, _ = compute_local_shape_and_global_offset(size, mesh, placements3) |
| 43 | # first dim | 39 | # first dim |
| 44 | if rank_coordinates[0] < 3: | 40 | if rank_coordinates[0] < 3: |
| 45 | self.assertEqual(local_size3[0], 2) | 41 | self.assertEqual(local_size3[0], 2) |
| @@ -110,6 +106,19 @@ class UtilTest(DTensorTestBase): | |||
| 110 | global_tensor[dim0_start:dim0_end, dim1_start:dim1_end], | 106 | global_tensor[dim0_start:dim0_end, dim1_start:dim1_end], |
| 111 | ) | 107 | ) |
| 112 | 108 | ||
| 109 | + @skipIfUnsupportMultiNPU(4) | ||
| 110 | + @with_comms | ||
| 111 | + def test_compute_local_shape_and_global_offset_partial(self): | ||
| 112 | + # Verify local size and global offset for Partial placement. | ||
| 113 | + placements = [Partial()] | ||
| 114 | + mesh_tensor = torch.arange(self.world_size) | ||
| 115 | + device_mesh = DeviceMesh(self.device_type, mesh_tensor) | ||
| 116 | + global_tensor = torch.arange(64).view(8, 8) | ||
| 117 | + global_shape = global_tensor.size() | ||
| 118 | + local_size, global_offset = compute_local_shape_and_global_offset(global_shape, device_mesh, placements) | ||
| 119 | + self.assertEqual(local_size, tuple(global_shape)) | ||
| 120 | + self.assertEqual(global_offset, tuple([0] * len(global_shape))) | ||
| 121 | + | ||
| 113 | 122 | ||
| 114 | if __name__ == "__main__": | 123 | if __name__ == "__main__": |
| 115 | run_tests() | 124 | run_tests() |