已合并
fix DTensor #29015
jizewei创建于 1月4日
fix DTensor #29015
已合并
共 10 个文件变更+201-83
| @@ -10,6 +10,7 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 10 | import torch_npu | 10 | import torch_npu |
| 11 | from torch_npu.testing.testcase import run_tests | 11 | from torch_npu.testing.testcase import run_tests |
| 12 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 12 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 13 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 13 | 14 | ||
| 14 | 15 | ||
| 15 | def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536): | 16 | def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536): |
| @@ -27,7 +28,16 @@ def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536): | |||
| 27 | 28 | ||
| 28 | 29 | ||
| 29 | class TestAttentionOps(DTensorTestBase): | 30 | class TestAttentionOps(DTensorTestBase): |
| 30 | - @skipIfUnsupportMultiNPU(4) | 31 | + @property |
| 32 | + def world_size(self): | ||
| 33 | + device_count = torch.npu.device_count() | ||
| 34 | + device_num = 4 | ||
| 35 | + if device_count > 1: | ||
| 36 | + device_num = min(device_num, device_count) | ||
| 37 | + return device_num | ||
| 38 | + | ||
| 39 | + | ||
| 40 | + | ||
| 31 | 41 | ||
| 32 | 42 | ||
| 33 | "sparse_mode,pre_tokens,next_tokens", | 43 | "sparse_mode,pre_tokens,next_tokens", |
| @@ -79,7 +89,8 @@ class TestAttentionOps(DTensorTestBase): | |||
| 79 | else: | 89 | else: |
| 80 | test_placement_comb([placement], [placement], [placement], [placement]) | 90 | test_placement_comb([placement], [placement], [placement], [placement]) |
| 81 | 91 | ||
| 82 | - @skipIfUnsupportMultiNPU(4) | 92 | + @SupportedDevices(['Ascend910B']) |
| 93 | + | ||
| 83 | 94 | ||
| 84 | 95 | ||
| 85 | "sparse_mode,pre_tokens,next_tokens", | 96 | "sparse_mode,pre_tokens,next_tokens", |
| @@ -138,7 +149,8 @@ class TestAttentionOps(DTensorTestBase): | |||
| 138 | else: | 149 | else: |
| 139 | test_placement_comb([placement], [placement], [placement], [placement]) | 150 | test_placement_comb([placement], [placement], [placement], [placement]) |
| 140 | 151 | ||
| 141 | - @skipIfUnsupportMultiNPU(4) | 152 | + @SupportedDevices(['Ascend910B']) |
| 153 | + | ||
| 142 | 154 | ||
| 143 | def test_npu_fusion_attention_bsnd(self): | 155 | def test_npu_fusion_attention_bsnd(self): |
| 144 | device_mesh = self.build_device_mesh() | 156 | device_mesh = self.build_device_mesh() |
| @@ -174,7 +186,8 @@ class TestAttentionOps(DTensorTestBase): | |||
| 174 | for placement in placements: | 186 | for placement in placements: |
| 175 | test_placement_comb([placement], [placement], [placement]) | 187 | test_placement_comb([placement], [placement], [placement]) |
| 176 | 188 | ||
| 177 | - @skipIfUnsupportMultiNPU(4) | 189 | + @SupportedDevices(['Ascend910B']) |
| 190 | + | ||
| 178 | 191 | ||
| 179 | def test_npu_fusion_attention_bsh(self): | 192 | def test_npu_fusion_attention_bsh(self): |
| 180 | device_mesh = self.build_device_mesh() | 193 | device_mesh = self.build_device_mesh() |
| @@ -210,7 +223,8 @@ class TestAttentionOps(DTensorTestBase): | |||
| 210 | for placement in placements: | 223 | for placement in placements: |
| 211 | test_placement_comb([placement], [placement], [placement]) | 224 | test_placement_comb([placement], [placement], [placement]) |
| 212 | 225 | ||
| 213 | - @skipIfUnsupportMultiNPU(4) | 226 | + @SupportedDevices(['Ascend910B']) |
| 227 | + | ||
| 214 | 228 | ||
| 215 | def test_npu_fusion_attention_tnd(self): | 229 | def test_npu_fusion_attention_tnd(self): |
| 216 | device_mesh = self.build_device_mesh() | 230 | device_mesh = self.build_device_mesh() |
| @@ -5,14 +5,24 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 5 | import torch_npu | 5 | import torch_npu |
| 6 | from torch_npu.testing.testcase import run_tests | 6 | from torch_npu.testing.testcase import run_tests |
| 7 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 7 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 8 | 9 | ||
| 9 | 10 | ||
| 10 | class TestRegisterSharding(DTensorTestBase): | 11 | class TestRegisterSharding(DTensorTestBase): |
| 11 | - @skipIfUnsupportMultiNPU(4) | 12 | + @property |
| 13 | + def world_size(self): | ||
| 14 | + device_count = torch.npu.device_count() | ||
| 15 | + device_num = 4 | ||
| 16 | + if device_count > 1: | ||
| 17 | + device_num = min(device_num, device_count) | ||
| 18 | + return device_num | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 12 | 22 | ||
| 13 | def test_torch_npu_npu_swiglu_replicate(self): | 23 | def test_torch_npu_npu_swiglu_replicate(self): |
| 14 | mesh = self.build_device_mesh() | 24 | mesh = self.build_device_mesh() |
| 15 | - | 25 | + |
| 16 | input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True) | 26 | input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True) |
| 17 | grad_tensor = torch.randn(1024, 512, device="npu") | 27 | grad_tensor = torch.randn(1024, 512, device="npu") |
| 18 | dim = -1 | 28 | dim = -1 |
| @@ -26,8 +36,9 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 26 | output.backward(grad_dtensor) | 36 | output.backward(grad_dtensor) |
| 27 | self.assertEqual(output.full_tensor(), out_tensor) | 37 | self.assertEqual(output.full_tensor(), out_tensor) |
| 28 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) | 38 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) |
| 29 | - | 39 | + |
| 30 | - @skipIfUnsupportMultiNPU(4) | 40 | + @SupportedDevices(['Ascend910B']) |
| 41 | + | ||
| 31 | 42 | ||
| 32 | def test_torch_npu_npu_swiglu_shard0(self): | 43 | def test_torch_npu_npu_swiglu_shard0(self): |
| 33 | mesh = self.build_device_mesh() | 44 | mesh = self.build_device_mesh() |
| @@ -46,14 +57,15 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 46 | self.assertEqual(output.full_tensor(), out_tensor) | 57 | self.assertEqual(output.full_tensor(), out_tensor) |
| 47 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) | 58 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) |
| 48 | 59 | ||
| 49 | - @skipIfUnsupportMultiNPU(4) | 60 | + @SupportedDevices(['Ascend910B']) |
| 61 | + | ||
| 50 | 62 | ||
| 51 | def test_torch_npu_npu_swiglu_shard1(self): | 63 | def test_torch_npu_npu_swiglu_shard1(self): |
| 52 | mesh = self.build_device_mesh() | 64 | mesh = self.build_device_mesh() |
| 53 | 65 | ||
| 54 | input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True) | 66 | input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True) |
| 55 | grad_tensor = torch.randn(1024, 512, device="npu") | 67 | grad_tensor = torch.randn(1024, 512, device="npu") |
| 56 | - dim = -1 | 68 | + dim = -1 |
| 57 | out_tensor = torch_npu.npu_swiglu(input_tensor, dim) | 69 | out_tensor = torch_npu.npu_swiglu(input_tensor, dim) |
| 58 | out_tensor.backward(grad_tensor) | 70 | out_tensor.backward(grad_tensor) |
| 59 | 71 | ||
| @@ -64,8 +76,9 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 64 | output.backward(grad_dtensor) | 76 | output.backward(grad_dtensor) |
| 65 | self.assertEqual(output.full_tensor(), out_tensor) | 77 | self.assertEqual(output.full_tensor(), out_tensor) |
| 66 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) | 78 | self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad) |
| 67 | - | 79 | + |
| 68 | - @skipIfUnsupportMultiNPU(4) | 80 | + @SupportedDevices(['Ascend910B']) |
| 81 | + | ||
| 69 | 82 | ||
| 70 | def test_torch_gather_replicate(self): | 83 | def test_torch_gather_replicate(self): |
| 71 | mesh = self.build_device_mesh() | 84 | mesh = self.build_device_mesh() |
| @@ -88,8 +101,9 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 88 | out.backward(grad) | 101 | out.backward(grad) |
| 89 | self.assertEqual(out.full_tensor(), out_tensor) | 102 | self.assertEqual(out.full_tensor(), out_tensor) |
| 90 | self.assertEqual(x.grad.full_tensor(), input_tensor.grad) | 103 | self.assertEqual(x.grad.full_tensor(), input_tensor.grad) |
| 91 | - | 104 | + |
| 92 | - @skipIfUnsupportMultiNPU(4) | 105 | + @SupportedDevices(['Ascend910B']) |
| 106 | + | ||
| 93 | 107 | ||
| 94 | def test_torch_gather_shard0(self): | 108 | def test_torch_gather_shard0(self): |
| 95 | mesh = self.build_device_mesh() | 109 | mesh = self.build_device_mesh() |
| @@ -105,15 +119,16 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 105 | 119 | ||
| 106 | x = distribute_tensor(input_tensor, mesh, [Shard(0)]) | 120 | x = distribute_tensor(input_tensor, mesh, [Shard(0)]) |
| 107 | index = distribute_tensor(index_tensor, mesh, [Shard(0)]) | 121 | index = distribute_tensor(index_tensor, mesh, [Shard(0)]) |
| 108 | - | 122 | + |
| 109 | out = torch.gather(input=x, dim=dim, index=index) | 123 | out = torch.gather(input=x, dim=dim, index=index) |
| 110 | grad = distribute_tensor(grad_tensor, mesh, out.placements) | 124 | grad = distribute_tensor(grad_tensor, mesh, out.placements) |
| 111 | 125 | ||
| 112 | out.backward(grad) | 126 | out.backward(grad) |
| 113 | self.assertEqual(out.full_tensor(), out_tensor) | 127 | self.assertEqual(out.full_tensor(), out_tensor) |
| 114 | self.assertEqual(x.grad.full_tensor(), input_tensor.grad) | 128 | self.assertEqual(x.grad.full_tensor(), input_tensor.grad) |
| 115 | - | 129 | + |
| 116 | - @skipIfUnsupportMultiNPU(4) | 130 | + @SupportedDevices(['Ascend910B']) |
| 131 | + | ||
| 117 | 132 | ||
| 118 | def test_torch_gather_shard1(self): | 133 | def test_torch_gather_shard1(self): |
| 119 | mesh = self.build_device_mesh() | 134 | mesh = self.build_device_mesh() |
| @@ -129,7 +144,7 @@ class TestRegisterSharding(DTensorTestBase): | |||
| 129 | 144 | ||
| 130 | x = distribute_tensor(input_tensor, mesh, [Shard(1)]) | 145 | x = distribute_tensor(input_tensor, mesh, [Shard(1)]) |
| 131 | index = distribute_tensor(index_tensor, mesh, [Shard(1)]) | 146 | index = distribute_tensor(index_tensor, mesh, [Shard(1)]) |
| 132 | - | 147 | + |
| 133 | out = torch.gather(input=x, dim=dim, index=index) | 148 | out = torch.gather(input=x, dim=dim, index=index) |
| 134 | grad = distribute_tensor(grad_tensor, mesh, out.placements) | 149 | grad = distribute_tensor(grad_tensor, mesh, out.placements) |
| 135 | 150 | ||
| @@ -11,6 +11,7 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 11 | import torch_npu | 11 | import torch_npu |
| 12 | from torch_npu.testing.testcase import run_tests | 12 | from torch_npu.testing.testcase import run_tests |
| 13 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 13 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 14 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 14 | 15 | ||
| 15 | 16 | ||
| 16 | def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str): | 17 | def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str): |
| @@ -32,7 +33,16 @@ def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: | |||
| 32 | 33 | ||
| 33 | 34 | ||
| 34 | class TestMathOps(DTensorTestBase): | 35 | class TestMathOps(DTensorTestBase): |
| 35 | - @skipIfUnsupportMultiNPU(4) | 36 | + @property |
| 37 | + def world_size(self): | ||
| 38 | + device_count = torch.npu.device_count() | ||
| 39 | + device_num = 4 | ||
| 40 | + if device_count > 1: | ||
| 41 | + device_num = min(device_num, device_count) | ||
| 42 | + return device_num | ||
| 43 | + | ||
| 44 | + | ||
| 45 | + | ||
| 36 | 46 | ||
| 37 | def test_npu_rms_norm_forward(self): | 47 | def test_npu_rms_norm_forward(self): |
| 38 | device_mesh = self.build_device_mesh() | 48 | device_mesh = self.build_device_mesh() |
| @@ -50,7 +60,8 @@ class TestMathOps(DTensorTestBase): | |||
| 50 | self.assertEqual(dist_y.full_tensor(), y) | 60 | self.assertEqual(dist_y.full_tensor(), y) |
| 51 | self.assertEqual(dist_gamma.full_tensor(), gamma) | 61 | self.assertEqual(dist_gamma.full_tensor(), gamma) |
| 52 | 62 | ||
| 53 | - @skipIfUnsupportMultiNPU(4) | 63 | + @SupportedDevices(['Ascend910B']) |
| 64 | + | ||
| 54 | 65 | ||
| 55 | def test_npu_rms_norm_backward(self): | 66 | def test_npu_rms_norm_backward(self): |
| 56 | device_mesh = self.build_device_mesh() | 67 | device_mesh = self.build_device_mesh() |
| @@ -85,7 +96,8 @@ class TestMathOps(DTensorTestBase): | |||
| 85 | self.assertEqual(dist_dx.full_tensor(), dx) | 96 | self.assertEqual(dist_dx.full_tensor(), dx) |
| 86 | self.assertEqual(dist_dw.full_tensor(), dw) | 97 | self.assertEqual(dist_dw.full_tensor(), dw) |
| 87 | 98 | ||
| 88 | - @skipIfUnsupportMultiNPU(4) | 99 | + @SupportedDevices(['Ascend910B']) |
| 100 | + | ||
| 89 | 101 | ||
| 90 | def test_npu_add_rms_norm_forward(self): | 102 | def test_npu_add_rms_norm_forward(self): |
| 91 | device_mesh = self.build_device_mesh() | 103 | device_mesh = self.build_device_mesh() |
| @@ -110,7 +122,8 @@ class TestMathOps(DTensorTestBase): | |||
| 110 | for comb in placement_combs: | 122 | for comb in placement_combs: |
| 111 | test_placement_comb([comb[0]], [comb[1]]) | 123 | test_placement_comb([comb[0]], [comb[1]]) |
| 112 | 124 | ||
| 113 | - @skipIfUnsupportMultiNPU(4) | 125 | + @SupportedDevices(['Ascend910B']) |
| 126 | + | ||
| 114 | 127 | ||
| 115 | 128 | ||
| 116 | "rotary_mode,input_layout,sin_cos_layout", | 129 | "rotary_mode,input_layout,sin_cos_layout", |
| @@ -158,7 +171,8 @@ class TestMathOps(DTensorTestBase): | |||
| 158 | else: | 171 | else: |
| 159 | test_placement_comb([placement], [placement], [placement]) | 172 | test_placement_comb([placement], [placement], [placement]) |
| 160 | 173 | ||
| 161 | - @skipIfUnsupportMultiNPU(4) | 174 | + @SupportedDevices(['Ascend910B']) |
| 175 | + | ||
| 162 | 176 | ||
| 163 | 177 | ||
| 164 | "rotary_mode,input_layout,sin_cos_layout", | 178 | "rotary_mode,input_layout,sin_cos_layout", |
| @@ -16,6 +16,14 @@ from torch_npu.testing.common_utils import SupportedDevices | |||
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | class TestAllGatherBaseMmOp(DTensorTestBase): | 18 | class TestAllGatherBaseMmOp(DTensorTestBase): |
| 19 | + | ||
| 20 | + def world_size(self): | ||
| 21 | + device_count = torch.npu.device_count() | ||
| 22 | + device_num = 4 | ||
| 23 | + if device_count > 1: | ||
| 24 | + device_num = min(device_num, device_count) | ||
| 25 | + return device_num | ||
| 26 | + | ||
| 19 | def _get_global_tensor(self, x1_list, x2_list, bias_list=None, x1_scale_list=None, x2_scale_list=None): | 27 | def _get_global_tensor(self, x1_list, x2_list, bias_list=None, x1_scale_list=None, x2_scale_list=None): |
| 20 | # Example: | 28 | # Example: |
| 21 | # x1_rank0 = [[1, 1]], x2_rank0 = [[2], [2]], x1_rank1 = [[3, 3]], x2_rank1 = [[4], [4]] | 29 | # x1_rank0 = [[1, 1]], x2_rank0 = [[2], [2]], x1_rank1 = [[3, 3]], x2_rank1 = [[4], [4]] |
| @@ -38,7 +46,7 @@ class TestAllGatherBaseMmOp(DTensorTestBase): | |||
| 38 | return x1, x2, output, gather_out | 46 | return x1, x2, output, gather_out |
| 39 | 47 | ||
| 40 | 48 | ||
| 41 | - @skipIfUnsupportMultiNPU(4) | 49 | + @skipIfUnsupportMultiNPU(2) |
| 42 | 50 | ||
| 43 | def test_npu_all_gather_base_mm(self): | 51 | def test_npu_all_gather_base_mm(self): |
| 44 | mesh = self.build_device_mesh() | 52 | mesh = self.build_device_mesh() |
| @@ -81,6 +89,14 @@ class TestAllGatherBaseMmOp(DTensorTestBase): | |||
| 81 | 89 | ||
| 82 | 90 | ||
| 83 | class TestMmReduceScatterBaseOp(DTensorTestBase): | 91 | class TestMmReduceScatterBaseOp(DTensorTestBase): |
| 92 | + | ||
| 93 | + def world_size(self): | ||
| 94 | + device_count = torch.npu.device_count() | ||
| 95 | + device_num = 4 | ||
| 96 | + if device_count > 1: | ||
| 97 | + device_num = min(device_num, device_count) | ||
| 98 | + return device_num | ||
| 99 | + | ||
| 84 | def _get_global_tensor(self, x1_list, x2_list, bias_list=None, x1_scale_list=None, x2_scale_list=None): | 100 | def _get_global_tensor(self, x1_list, x2_list, bias_list=None, x1_scale_list=None, x2_scale_list=None): |
| 85 | # Example: | 101 | # Example: |
| 86 | # x1_rank0 = [[1], [1]], x2_rank0 = [[2, 2]], x1_rank1 = [[3], [3]], x2_rank1 = [[4, 4]] | 102 | # x1_rank0 = [[1], [1]], x2_rank0 = [[2, 2]], x1_rank1 = [[3], [3]], x2_rank1 = [[4, 4]] |
| @@ -109,7 +125,7 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 109 | return x1, x2, output | 125 | return x1, x2, output |
| 110 | 126 | ||
| 111 | 127 | ||
| 112 | - @skipIfUnsupportMultiNPU(4) | 128 | + @skipIfUnsupportMultiNPU(2) |
| 113 | 129 | ||
| 114 | def test_npu_mm_reduce_scatter_base(self): | 130 | def test_npu_mm_reduce_scatter_base(self): |
| 115 | with DeterministicGuard(True): | 131 | with DeterministicGuard(True): |
| @@ -151,7 +167,7 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 151 | test_placement_comb([comb[0]], [comb[1]]) | 167 | test_placement_comb([comb[0]], [comb[1]]) |
| 152 | 168 | ||
| 153 | 169 | ||
| 154 | - @skipIfUnsupportMultiNPU(4) | 170 | + @skipIfUnsupportMultiNPU(2) |
| 155 | 171 | ||
| 156 | def test_npu_mm_reduce_scatter_base_bias(self): | 172 | def test_npu_mm_reduce_scatter_base_bias(self): |
| 157 | with DeterministicGuard(True): | 173 | with DeterministicGuard(True): |
| @@ -165,7 +181,7 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 165 | for _ in range(self.world_size): | 181 | for _ in range(self.world_size): |
| 166 | x1_list.append(torch.randn(m, k, dtype=dtype, device="npu")) | 182 | x1_list.append(torch.randn(m, k, dtype=dtype, device="npu")) |
| 167 | x2_list.append(torch.randn(k, n, dtype=dtype, device="npu")) | 183 | x2_list.append(torch.randn(k, n, dtype=dtype, device="npu")) |
| 168 | - bias_list.append(torch.randn(n, dtype=dtype, device="npu")) | 184 | + bias_list.append(torch.zeros(n, dtype=dtype, device="npu")) |
| 169 | 185 | ||
| 170 | global_x1, global_x2, global_output = self._get_global_tensor(x1_list, x2_list, bias_list=bias_list) | 186 | global_x1, global_x2, global_output = self._get_global_tensor(x1_list, x2_list, bias_list=bias_list) |
| 171 | global_bias = torch.cat(bias_list, dim=0) | 187 | global_bias = torch.cat(bias_list, dim=0) |
| @@ -187,7 +203,7 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 187 | dist_output = torch_npu.npu_mm_reduce_scatter_base( | 203 | dist_output = torch_npu.npu_mm_reduce_scatter_base( |
| 188 | dist_x1, dist_x2, hcom_name, self.world_size, bias=dist_bias | 204 | dist_x1, dist_x2, hcom_name, self.world_size, bias=dist_bias |
| 189 | ) | 205 | ) |
| 190 | - self.assertEqual(dist_output.full_tensor(), global_output.to(dtype), atol=0.05, rtol=0.05) | 206 | + self.assertEqual(dist_output.full_tensor(), global_output.to(dtype)) |
| 191 | self.assertEqual(dist_output.to_local(), output) | 207 | self.assertEqual(dist_output.to_local(), output) |
| 192 | 208 | ||
| 193 | placement = [Shard(0), Shard(1), Replicate()] | 209 | placement = [Shard(0), Shard(1), Replicate()] |
| @@ -197,7 +213,7 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 197 | test_placement_comb([comb[0]], [comb[1]], [bias_placement]) | 213 | test_placement_comb([comb[0]], [comb[1]], [bias_placement]) |
| 198 | 214 | ||
| 199 | 215 | ||
| 200 | - @skipIfUnsupportMultiNPU(4) | 216 | + @skipIfUnsupportMultiNPU(2) |
| 201 | 217 | ||
| 202 | def test_npu_mm_reduce_scatter_base_quant(self): | 218 | def test_npu_mm_reduce_scatter_base_quant(self): |
| 203 | with DeterministicGuard(True): | 219 | with DeterministicGuard(True): |
| @@ -252,8 +268,16 @@ class TestMmReduceScatterBaseOp(DTensorTestBase): | |||
| 252 | 268 | ||
| 253 | 269 | ||
| 254 | class TestGroupedMatMulOp(DTensorTestBase): | 270 | class TestGroupedMatMulOp(DTensorTestBase): |
| 271 | + | ||
| 272 | + def world_size(self): | ||
| 273 | + device_count = torch.npu.device_count() | ||
| 274 | + device_num = 4 | ||
| 275 | + if device_count > 1: | ||
| 276 | + device_num = min(device_num, device_count) | ||
| 277 | + return device_num | ||
| 278 | + | ||
| 255 | 279 | ||
| 256 | - @skipIfUnsupportMultiNPU(4) | 280 | + @skipIfUnsupportMultiNPU(2) |
| 257 | 281 | ||
| 258 | 282 | ||
| 259 | 283 | ||
| @@ -293,16 +317,16 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 293 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) | 317 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) |
| 294 | 318 | ||
| 295 | 319 | ||
| 296 | - @skipIfUnsupportMultiNPU(4) | 320 | + @skipIfUnsupportMultiNPU(2) |
| 297 | 321 | ||
| 298 | 322 | ||
| 299 | def test_npu_grouped_matmul_x1w1y1(self, with_bias): | 323 | def test_npu_grouped_matmul_x1w1y1(self, with_bias): |
| 300 | mesh = self.build_device_mesh() | 324 | mesh = self.build_device_mesh() |
| 301 | 325 | ||
| 302 | - x = [torch.randn(112, 64, dtype=torch.float32, device="npu")] | 326 | + x = [torch.randn(8, 8, dtype=torch.float16, device="npu")] |
| 303 | - weight = [torch.randn(4, 64, 16, dtype=torch.float32, device="npu")] | 327 | + weight = [torch.randn(2, 8, 8, dtype=torch.float16, device="npu")] |
| 304 | - bias = [torch.randn(4, 16, dtype=torch.float32, device="npu")] if with_bias else None | 328 | + bias = [torch.randn(2, 8, dtype=torch.float16, device="npu")] if with_bias else None |
| 305 | - group_list = torch.tensor([16, 48, 64, 112], device="npu") | 329 | + group_list = torch.tensor([2, 8], device="npu") |
| 306 | split_item = 3 | 330 | split_item = 3 |
| 307 | group_type = 0 | 331 | group_type = 0 |
| 308 | 332 | ||
| @@ -320,7 +344,7 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 320 | split_item=split_item, group_type=group_type | 344 | split_item=split_item, group_type=group_type |
| 321 | ) | 345 | ) |
| 322 | for dist_y_i, y_i in zip(dist_y, y): | 346 | for dist_y_i, y_i in zip(dist_y, y): |
| 323 | - self.assertEqual(dist_y_i.full_tensor(), y_i) | 347 | + self.assertEqual(dist_y_i.full_tensor(), y_i, atol=0.001, rtol=0.02) |
| 324 | 348 | ||
| 325 | placement = [Shard(0), Shard(1), Replicate()] | 349 | placement = [Shard(0), Shard(1), Replicate()] |
| 326 | placement_combs = itertools.product(placement, placement, [Shard(0), Replicate()]) | 350 | placement_combs = itertools.product(placement, placement, [Shard(0), Replicate()]) |
| @@ -328,7 +352,7 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 328 | test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]]) | 352 | test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]]) |
| 329 | 353 | ||
| 330 | 354 | ||
| 331 | - @skipIfUnsupportMultiNPU(4) | 355 | + @skipIfUnsupportMultiNPU(2) |
| 332 | 356 | ||
| 333 | 357 | ||
| 334 | def test_npu_grouped_matmul_xNwNy1(self, with_bias): | 358 | def test_npu_grouped_matmul_xNwNy1(self, with_bias): |
| @@ -367,7 +391,7 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 367 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) | 391 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) |
| 368 | 392 | ||
| 369 | 393 | ||
| 370 | - @skipIfUnsupportMultiNPU(4) | 394 | + @skipIfUnsupportMultiNPU(2) |
| 371 | 395 | ||
| 372 | 396 | ||
| 373 | 397 | ||
| @@ -410,7 +434,7 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 410 | test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]]) | 434 | test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]]) |
| 411 | 435 | ||
| 412 | 436 | ||
| 413 | - @skipIfUnsupportMultiNPU(4) | 437 | + @skipIfUnsupportMultiNPU(2) |
| 414 | 438 | ||
| 415 | 439 | ||
| 416 | def test_npu_grouped_matmul_x1wNyN(self, with_bias): | 440 | def test_npu_grouped_matmul_x1wNyN(self, with_bias): |
| @@ -444,7 +468,7 @@ class TestGroupedMatMulOp(DTensorTestBase): | |||
| 444 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) | 468 | test_placement_comb([comb[0]], [comb[1]], [comb[2]]) |
| 445 | 469 | ||
| 446 | 470 | ||
| 447 | - @skipIfUnsupportMultiNPU(4) | 471 | + @skipIfUnsupportMultiNPU(2) |
| 448 | 472 | ||
| 449 | def test_npu_grouped_matmul_quant(self): | 473 | def test_npu_grouped_matmul_quant(self): |
| 450 | mesh = self.build_device_mesh() | 474 | mesh = self.build_device_mesh() |
| @@ -7,10 +7,20 @@ from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTe | |||
| 7 | import torch_npu | 7 | import torch_npu |
| 8 | from torch_npu.testing.testcase import run_tests | 8 | from torch_npu.testing.testcase import run_tests |
| 9 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | 9 | from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU |
| 10 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 10 | 11 | ||
| 11 | 12 | ||
| 12 | class TestMoeOps(DTensorTestBase): | 13 | class TestMoeOps(DTensorTestBase): |
| 13 | - @skipIfUnsupportMultiNPU(4) | 14 | + @property |
| 15 | + def world_size(self): | ||
| 16 | + device_count = torch.npu.device_count() | ||
| 17 | + device_num = 4 | ||
| 18 | + if device_count > 1: | ||
| 19 | + device_num = min(device_num, device_count) | ||
| 20 | + return device_num | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 14 | 24 | ||
| 15 | def test_npu_moe_token_permute_forward(self): | 25 | def test_npu_moe_token_permute_forward(self): |
| 16 | device_mesh = self.build_device_mesh() | 26 | device_mesh = self.build_device_mesh() |
| @@ -35,7 +45,8 @@ class TestMoeOps(DTensorTestBase): | |||
| 35 | for comb in placement_combs: | 45 | for comb in placement_combs: |
| 36 | test_placement_comb([comb[0]], [comb[1]]) | 46 | test_placement_comb([comb[0]], [comb[1]]) |
| 37 | 47 | ||
| 38 | - @skipIfUnsupportMultiNPU(4) | 48 | + @SupportedDevices(['Ascend910B']) |
| 49 | + | ||
| 39 | 50 | ||
| 40 | def test_npu_moe_token_permute_backward(self): | 51 | def test_npu_moe_token_permute_backward(self): |
| 41 | device_mesh = self.build_device_mesh() | 52 | device_mesh = self.build_device_mesh() |
| @@ -67,7 +78,8 @@ class TestMoeOps(DTensorTestBase): | |||
| 67 | for comb in placement_combs: | 78 | for comb in placement_combs: |
| 68 | test_placement_comb([comb[0]], [comb[1]]) | 79 | test_placement_comb([comb[0]], [comb[1]]) |
| 69 | 80 | ||
| 70 | - @skipIfUnsupportMultiNPU(4) | 81 | + @SupportedDevices(['Ascend910B']) |
| 82 | + | ||
| 71 | 83 | ||
| 72 | def test_npu_moe_token_permute_clip(self): | 84 | def test_npu_moe_token_permute_clip(self): |
| 73 | device_mesh = self.build_device_mesh() | 85 | device_mesh = self.build_device_mesh() |
| @@ -95,7 +107,8 @@ class TestMoeOps(DTensorTestBase): | |||
| 95 | self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices) | 107 | self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices) |
| 96 | self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad) | 108 | self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad) |
| 97 | 109 | ||
| 98 | - @skipIfUnsupportMultiNPU(4) | 110 | + @SupportedDevices(['Ascend910B']) |
| 111 | + | ||
| 99 | 112 | ||
| 100 | def test_npu_moe_token_unpermute_forward(self): | 113 | def test_npu_moe_token_unpermute_forward(self): |
| 101 | device_mesh = self.build_device_mesh() | 114 | device_mesh = self.build_device_mesh() |
| @@ -120,7 +133,8 @@ class TestMoeOps(DTensorTestBase): | |||
| 120 | for comb in placement_combs: | 133 | for comb in placement_combs: |
| 121 | test_placement_comb([comb[0]], [comb[1]]) | 134 | test_placement_comb([comb[0]], [comb[1]]) |
| 122 | 135 | ||
| 123 | - @skipIfUnsupportMultiNPU(4) | 136 | + @SupportedDevices(['Ascend910B']) |
| 137 | + | ||
| 124 | 138 | ||
| 125 | def test_npu_moe_token_unpermute_backward(self): | 139 | def test_npu_moe_token_unpermute_backward(self): |
| 126 | device_mesh = self.build_device_mesh() | 140 | device_mesh = self.build_device_mesh() |
| @@ -157,7 +171,8 @@ class TestMoeOps(DTensorTestBase): | |||
| 157 | for comb in placement_combs: | 171 | for comb in placement_combs: |
| 158 | test_placement_comb([comb[0]], [comb[1]]) | 172 | test_placement_comb([comb[0]], [comb[1]]) |
| 159 | 173 | ||
| 160 | - @skipIfUnsupportMultiNPU(4) | 174 | + @SupportedDevices(['Ascend910B']) |
| 175 | + | ||
| 161 | 176 | ||
| 162 | def test_npu_moe_token_permute_unpermute(self): | 177 | def test_npu_moe_token_permute_unpermute(self): |
| 163 | device_mesh = self.build_device_mesh() | 178 | device_mesh = self.build_device_mesh() |
| @@ -30,13 +30,13 @@ def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse | |||
| 30 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, | 30 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, |
| 31 | next_tockens=2147483647, inner_precise=0, prefix=None, actual_seq_qlen=None, | 31 | next_tockens=2147483647, inner_precise=0, prefix=None, actual_seq_qlen=None, |
| 32 | actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False, | 32 | actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False, |
| 33 | - softmax_layout=""): | 33 | + softmax_layout="", sink=None): |
| 34 | # func: npu_fusion_attention(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, | 34 | # func: npu_fusion_attention(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, |
| 35 | # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., | 35 | # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., |
| 36 | # float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, | 36 | # float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, |
| 37 | # int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, | 37 | # int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, |
| 38 | # int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, | 38 | # int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, |
| 39 | - # bool sync=False, str softmax_layout="") | 39 | + # bool sync=False, str softmax_layout="", Tensor? sink=None) |
| 40 | # -> (Tensor, Tensor, Tensor, Tensor, int, int, int) | 40 | # -> (Tensor, Tensor, Tensor, Tensor, int, int, int) |
| 41 | strategies = [] | 41 | strategies = [] |
| 42 | 42 | ||
| @@ -58,14 +58,15 @@ def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse | |||
| 58 | None if pse is None else Replicate(), # pse | 58 | None if pse is None else Replicate(), # pse |
| 59 | None if padding_mask is None else Replicate(), # padding_mask | 59 | None if padding_mask is None else Replicate(), # padding_mask |
| 60 | None if atten_mask is None else Replicate(), # atten_mask | 60 | None if atten_mask is None else Replicate(), # atten_mask |
| 61 | - None, None, None, None, None, None, None, None, None, None, None, None # others | 61 | + None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 62 | + None if sink is None else Replicate() # sink | ||
| 62 | ] | 63 | ] |
| 63 | ) | 64 | ) |
| 64 | strategies.append(replicate_strategy) | 65 | strategies.append(replicate_strategy) |
| 65 | 66 | ||
| 66 | # only support sharding for sdpa currently, in which pse and padding_mask are not used | 67 | # only support sharding for sdpa currently, in which pse and padding_mask are not used |
| 67 | # keep_prob < 1.0 may effect different results under sharding | 68 | # keep_prob < 1.0 may effect different results under sharding |
| 68 | - unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen] | 69 | + unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen, sink] |
| 69 | if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: | 70 | if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: |
| 70 | return strategies | 71 | return strategies |
| 71 | 72 | ||
| @@ -97,7 +98,8 @@ def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse | |||
| 97 | None, # pse | 98 | None, # pse |
| 98 | None, # padding_mask | 99 | None, # padding_mask |
| 99 | atten_mask_sharding, # atten_mask | 100 | atten_mask_sharding, # atten_mask |
| 100 | - None, None, None, None, None, None, None, None, None, None, None, None # others | 101 | + None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 102 | + None # sink | ||
| 101 | ] | 103 | ] |
| 102 | ) | 104 | ) |
| 103 | strategies.append(dp_sharding_strategy) | 105 | strategies.append(dp_sharding_strategy) |
| @@ -128,7 +130,8 @@ def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse | |||
| 128 | None, # pse | 130 | None, # pse |
| 129 | None, # padding_mask | 131 | None, # padding_mask |
| 130 | atten_mask_sharding, # atten_mask | 132 | atten_mask_sharding, # atten_mask |
| 131 | - None, None, None, None, None, None, None, None, None, None, None, None # others | 133 | + None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 134 | + None # sink | ||
| 132 | ] | 135 | ] |
| 133 | ) | 136 | ) |
| 134 | strategies.append(tp_sharding_strategy) | 137 | strategies.append(tp_sharding_strategy) |
| @@ -142,7 +145,7 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 142 | attention_in=None, scale_value=1., keep_prob=1., pre_tockens=2147483647, | 145 | attention_in=None, scale_value=1., keep_prob=1., pre_tockens=2147483647, |
| 143 | next_tockens=2147483647, inner_precise=0, seed=0, offset=0, numels=0, | 146 | next_tockens=2147483647, inner_precise=0, seed=0, offset=0, numels=0, |
| 144 | prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, | 147 | prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, |
| 145 | - gen_mask_parallel=True, sync=False, softmax_layout=""): | 148 | + gen_mask_parallel=True, sync=False, softmax_layout="", sink=None): |
| 146 | # npu_fusion_attention_grad(Tensor query, Tensor key, Tensor value, Tensor dy, int head_num, str input_layout, *, | 149 | # npu_fusion_attention_grad(Tensor query, Tensor key, Tensor value, Tensor dy, int head_num, str input_layout, *, |
| 147 | # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, | 150 | # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, |
| 148 | # Tensor? softmax_max=None, Tensor? softmax_sum=None, Tensor? softmax_in=None, | 151 | # Tensor? softmax_max=None, Tensor? softmax_sum=None, Tensor? softmax_in=None, |
| @@ -150,8 +153,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 150 | # int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, | 153 | # int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, |
| 151 | # int seed=0, int offset=0, int numels=0, int[]? prefix=None, | 154 | # int seed=0, int offset=0, int numels=0, int[]? prefix=None, |
| 152 | # int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, | 155 | # int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, |
| 153 | - # bool gen_mask_parallel=True, bool sync=False, str softmax_layout="") | 156 | + # bool gen_mask_parallel=True, bool sync=False, str softmax_layout="", Tensor? sink=None) |
| 154 | - # -> (Tensor, Tensor, Tensor, Tensor) | 157 | + # -> (Tensor, Tensor, Tensor, Tensor, Tensor) |
| 155 | strategies = [] | 158 | strategies = [] |
| 156 | 159 | ||
| 157 | # all replicate strategy | 160 | # all replicate strategy |
| @@ -160,7 +163,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 160 | Replicate(), # grad_query | 163 | Replicate(), # grad_query |
| 161 | Replicate(), # grad_key | 164 | Replicate(), # grad_key |
| 162 | Replicate(), # grad_value | 165 | Replicate(), # grad_value |
| 163 | - Replicate() # grad_pse(reserve, unused now) | 166 | + Replicate(), # grad_pse(reserve, unused now) |
| 167 | + Replicate() # grad_sink | ||
| 164 | ], | 168 | ], |
| 165 | [ | 169 | [ |
| 166 | Replicate(), # query | 170 | Replicate(), # query |
| @@ -176,14 +180,15 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 176 | None if softmax_sum is None else Replicate(), # softmax_sum | 180 | None if softmax_sum is None else Replicate(), # softmax_sum |
| 177 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | 181 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) |
| 178 | None if attention_in is None else Replicate(), # attention_in | 182 | None if attention_in is None else Replicate(), # attention_in |
| 179 | - None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | 183 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 184 | + None if sink is None else Replicate() # sink | ||
| 180 | ] | 185 | ] |
| 181 | ) | 186 | ) |
| 182 | strategies.append(replicate_strategy) | 187 | strategies.append(replicate_strategy) |
| 183 | 188 | ||
| 184 | # only support sharding for sdpa currently, in which pse and padding_mask are not used | 189 | # only support sharding for sdpa currently, in which pse and padding_mask are not used |
| 185 | # keep_prob < 1.0 may effect different results under sharding | 190 | # keep_prob < 1.0 may effect different results under sharding |
| 186 | - unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen] | 191 | + unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen, sink] |
| 187 | if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: | 192 | if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: |
| 188 | return strategies | 193 | return strategies |
| 189 | 194 | ||
| @@ -203,7 +208,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 203 | Shard(batch_dim), # grad_query | 208 | Shard(batch_dim), # grad_query |
| 204 | Shard(batch_dim), # grad_key | 209 | Shard(batch_dim), # grad_key |
| 205 | Shard(batch_dim), # grad_value | 210 | Shard(batch_dim), # grad_value |
| 206 | - Replicate() # grad_pse(reserve, unused now) | 211 | + Replicate(), # grad_pse(reserve, unused now) |
| 212 | + Replicate() # grad_sink(unsupported now) | ||
| 207 | ], | 213 | ], |
| 208 | [ | 214 | [ |
| 209 | Shard(batch_dim), # query | 215 | Shard(batch_dim), # query |
| @@ -219,7 +225,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 219 | Shard(0) if softmax_sum is not None else None, # softmax_sum layout: BNS8 | 225 | Shard(0) if softmax_sum is not None else None, # softmax_sum layout: BNS8 |
| 220 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | 226 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) |
| 221 | Shard(batch_dim) if attention_in is not None else None, # attention_in | 227 | Shard(batch_dim) if attention_in is not None else None, # attention_in |
| 222 | - None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | 228 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 229 | + None # sink | ||
| 223 | ] | 230 | ] |
| 224 | ) | 231 | ) |
| 225 | strategies.append(dp_sharding_strategy) | 232 | strategies.append(dp_sharding_strategy) |
| @@ -238,7 +245,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 238 | Shard(head_dim), # grad_query | 245 | Shard(head_dim), # grad_query |
| 239 | Shard(head_dim), # grad_key | 246 | Shard(head_dim), # grad_key |
| 240 | Shard(head_dim), # grad_value | 247 | Shard(head_dim), # grad_value |
| 241 | - Replicate() # grad_pse(reserve, unused now) | 248 | + Replicate(), # grad_pse(reserve, unused now) |
| 249 | + Replicate() # grad_sink(unsupported now) | ||
| 242 | ], | 250 | ], |
| 243 | [ | 251 | [ |
| 244 | Shard(head_dim), # query | 252 | Shard(head_dim), # query |
| @@ -254,7 +262,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la | |||
| 254 | Shard(1) if softmax_sum is not None else None, # softmax_sum layout: BNS8 | 262 | Shard(1) if softmax_sum is not None else None, # softmax_sum layout: BNS8 |
| 255 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | 263 | None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) |
| 256 | Shard(head_dim) if attention_in is not None else None, # attention_in | 264 | Shard(head_dim) if attention_in is not None else None, # attention_in |
| 257 | - None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | 265 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, # others |
| 266 | + None # sink | ||
| 258 | ] | 267 | ] |
| 259 | ) | 268 | ) |
| 260 | strategies.append(tp_sharding_strategy) | 269 | strategies.append(tp_sharding_strategy) |
| @@ -24,8 +24,7 @@ def get_redistributed_local_args( | |||
| 24 | if output_sharding.needs_redistribute: | 24 | if output_sharding.needs_redistribute: |
| 25 | DTensor._op_dispatcher.redistribute_local_args( | 25 | DTensor._op_dispatcher.redistribute_local_args( |
| 26 | op_info, | 26 | op_info, |
| 27 | - output_sharding.redistribute_schema, | 27 | + output_sharding.redistribute_schema |
| 28 | - output_sharding.use_val_from_redistribute_schema, | ||
| 29 | ) | 28 | ) |
| 30 | local_args = ( | 29 | local_args = ( |
| 31 | pytree.tree_unflatten( | 30 | pytree.tree_unflatten( |
| @@ -1,3 +1,5 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | + | ||
| 1 | import itertools | 3 | import itertools |
| 2 | from typing import Callable, Optional | 4 | from typing import Callable, Optional |
| 3 | 5 | ||
| @@ -8,6 +10,7 @@ from torch.distributed.tensor._op_schema import ( | |||
| 8 | OpStrategy, | 10 | OpStrategy, |
| 9 | PlacementStrategy, | 11 | PlacementStrategy, |
| 10 | PlacementList, | 12 | PlacementList, |
| 13 | + TupleStrategy | ||
| 11 | ) | 14 | ) |
| 12 | from torch.distributed.tensor._ops.utils import ( | 15 | from torch.distributed.tensor._ops.utils import ( |
| 13 | generate_redistribute_costs, | 16 | generate_redistribute_costs, |
| @@ -16,9 +19,9 @@ from torch.distributed.tensor._ops.utils import ( | |||
| 16 | from torch.distributed.tensor.device_mesh import DeviceMesh | 19 | from torch.distributed.tensor.device_mesh import DeviceMesh |
| 17 | 20 | ||
| 18 | try: | 21 | try: |
| 19 | - from torch.utils._cxx_pytree import tree_leaves | 22 | + from torch.utils._cxx_pytree import register_pytree_node, tree_leaves |
| 20 | except ImportError: | 23 | except ImportError: |
| 21 | - from torch.utils._pytree import tree_leaves | 24 | + from torch.utils._pytree import register_pytree_node, tree_leaves |
| 22 | 25 | ||
| 23 | 26 | ||
| 24 | def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]: | 27 | def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]: |
| @@ -99,6 +102,18 @@ def _patched_expand_to_full_mesh_op_strategy( | |||
| 99 | return OpStrategy(all_strategies) | 102 | return OpStrategy(all_strategies) |
| 100 | 103 | ||
| 101 | 104 | ||
| 105 | +def _patched_register_tuple_strategy(): | ||
| 106 | + try: | ||
| 107 | + register_pytree_node( | ||
| 108 | + TupleStrategy, | ||
| 109 | + lambda node: (node.childs, None), | ||
| 110 | + lambda childs, _: TupleStrategy(tuple(childs)), | ||
| 111 | + ) | ||
| 112 | + except ValueError: | ||
| 113 | + # already registered TupleStrategy, skip | ||
| 114 | + pass | ||
| 115 | + | ||
| 116 | + | ||
| 102 | def _apply_dtensor_patch(): | 117 | def _apply_dtensor_patch(): |
| 103 | # adding kwarg inputs handling in register sharding for previous pytorch version | 118 | # adding kwarg inputs handling in register sharding for previous pytorch version |
| 104 | # See pytorch/pytorch/pull/168249 | 119 | # See pytorch/pytorch/pull/168249 |
| @@ -107,5 +122,10 @@ def _apply_dtensor_patch(): | |||
| 107 | OpSchema.kwargs_strategy = property(_patched_kwargs_strategy) | 122 | OpSchema.kwargs_strategy = property(_patched_kwargs_strategy) |
| 108 | torch.distributed.tensor._ops.utils.expand_to_full_mesh_op_strategy = _patched_expand_to_full_mesh_op_strategy | 123 | torch.distributed.tensor._ops.utils.expand_to_full_mesh_op_strategy = _patched_expand_to_full_mesh_op_strategy |
| 109 | 124 | ||
| 125 | + # register TupleStrategy pytree node to support flattening tensor lists for previous pytorch version | ||
| 126 | + # See pytorch/pytorch/pull/158046 | ||
| 127 | + if torch.__version__ < "2.9": | ||
| 128 | + _patched_register_tuple_strategy() | ||
| 129 | + | ||
| 110 | 130 | ||
| 111 | _apply_dtensor_patch() | 131 | _apply_dtensor_patch() |
| @@ -239,7 +239,7 @@ def npu_add_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 239 | output_target_spec = (y_target_spec, rstd_target_spec, x_target_spec) | 239 | output_target_spec = (y_target_spec, rstd_target_spec, x_target_spec) |
| 240 | 240 | ||
| 241 | output_strategy.strategies.append( | 241 | output_strategy.strategies.append( |
| 242 | - OpSpec( | 242 | + PlacementStrategy( |
| 243 | output_specs=output_target_spec, | 243 | output_specs=output_target_spec, |
| 244 | input_specs=op_args_target_specs, | 244 | input_specs=op_args_target_specs, |
| 245 | redistribute_cost=redistribute_costs, | 245 | redistribute_cost=redistribute_costs, |
| @@ -399,7 +399,7 @@ def custom_npu_conv2d_strategy(x, weight, bias, stride, padding, dilation, group | |||
| 399 | ] | 399 | ] |
| 400 | ) | 400 | ) |
| 401 | acceptable_shardings.append(replicate_strategy) | 401 | acceptable_shardings.append(replicate_strategy) |
| 402 | - | 402 | + |
| 403 | # x layout: (N, Ci, Hi, Wi) | 403 | # x layout: (N, Ci, Hi, Wi) |
| 404 | # weight layout: (Co, Ci/groups, Hk, Wk) | 404 | # weight layout: (Co, Ci/groups, Hk, Wk) |
| 405 | # bias layout: (Co) | 405 | # bias layout: (Co) |
| @@ -526,14 +526,14 @@ def custom_cross_entropy_loss_sharding(op_schema: OpSchema): | |||
| 526 | single_mesh_dim_strategies = [] | 526 | single_mesh_dim_strategies = [] |
| 527 | 527 | ||
| 528 | args_schema = op_schema.args_schema | 528 | args_schema = op_schema.args_schema |
| 529 | - | 529 | + |
| 530 | input_strategy = args_schema[0] if len(args_schema) > 0 else None | 530 | input_strategy = args_schema[0] if len(args_schema) > 0 else None |
| 531 | target_strategy = args_schema[1] if len(args_schema) > 1 else None | 531 | target_strategy = args_schema[1] if len(args_schema) > 1 else None |
| 532 | weight_strategy = args_schema[2] if len(args_schema) > 2 else None | 532 | weight_strategy = args_schema[2] if len(args_schema) > 2 else None |
| 533 | reduction = args_schema[3] if len(args_schema) > 3 else 'mean' | 533 | reduction = args_schema[3] if len(args_schema) > 3 else 'mean' |
| 534 | 534 | ||
| 535 | mesh = input_strategy.mesh | 535 | mesh = input_strategy.mesh |
| 536 | - | 536 | + |
| 537 | all_replicate: PlacementList = [ | 537 | all_replicate: PlacementList = [ |
| 538 | Replicate(), # loss | 538 | Replicate(), # loss |
| 539 | Replicate(), # log_prob | 539 | Replicate(), # log_prob |
| @@ -6,6 +6,7 @@ from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta | |||
| 6 | from torch.distributed.tensor._ops.utils import register_op_strategy, expand_to_full_mesh_op_strategy | 6 | from torch.distributed.tensor._ops.utils import register_op_strategy, expand_to_full_mesh_op_strategy |
| 7 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard | 7 | from torch.distributed.tensor import DTensor, Partial, Replicate, Shard |
| 8 | from torch.distributed.tensor._op_schema import ( | 8 | from torch.distributed.tensor._op_schema import ( |
| 9 | + _is_out_variant_op, | ||
| 9 | OpInfo, | 10 | OpInfo, |
| 10 | OpSchema, | 11 | OpSchema, |
| 11 | OpStrategy, | 12 | OpStrategy, |
| @@ -54,8 +55,6 @@ def _handle_tensor_list_in_kwargs(kwargs: Dict[str, object], op_info: OpInfo) -> | |||
| 54 | op_info.schema.kwargs_schema[key] = tuple(new_schema) # list is not hashable for cache | 55 | op_info.schema.kwargs_schema[key] = tuple(new_schema) # list is not hashable for cache |
| 55 | op_info.local_kwargs[key] = new_local_tensors | 56 | op_info.local_kwargs[key] = new_local_tensors |
| 56 | 57 | ||
| 57 | - op_info.schema._recompute_comparison_key() | ||
| 58 | - | ||
| 59 | 58 | ||
| 60 | 59 | ||
| 61 | def custom_matmul_sharding( | 60 | def custom_matmul_sharding( |
| @@ -194,12 +193,14 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 194 | # int? split_item=0, int? group_type=None, int? group_list_type=0, int? act_type=0, | 193 | # int? split_item=0, int? group_type=None, int? group_list_type=0, int? act_type=0, |
| 195 | # int[]? tuning_config=None, int? output_dtype=None, int? x_dtype=None, int? weight_dtype=None, | 194 | # int[]? tuning_config=None, int? output_dtype=None, int? x_dtype=None, int? weight_dtype=None, |
| 196 | # int? scale_dtype=None, int? per_token_scale_dtype=None) -> Tensor[] | 195 | # int? scale_dtype=None, int? per_token_scale_dtype=None) -> Tensor[] |
| 196 | + if op_schema.schema_info is None: | ||
| 197 | + op_schema.schema_info = RuntimeSchemaInfo(needs_pytree=True) # to flatten tensor list in arguments | ||
| 197 | x_src_strategy: TupleStrategy = op_schema.args_schema[0] | 198 | x_src_strategy: TupleStrategy = op_schema.args_schema[0] |
| 198 | - x_num = len(x_src_strategy.children) | 199 | + x_num = len(x_src_strategy.childs) |
| 199 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] | 200 | weight_src_strategy: TupleStrategy = op_schema.args_schema[1] |
| 200 | - weight_num = len(weight_src_strategy.children) | 201 | + weight_num = len(weight_src_strategy.childs) |
| 201 | bias_src_strategy: Optional[Union[TupleStrategy, list]] = op_schema.kwargs_schema.get("bias", []) | 202 | bias_src_strategy: Optional[Union[TupleStrategy, list]] = op_schema.kwargs_schema.get("bias", []) |
| 202 | - bias_num = len(bias_src_strategy.children) if isinstance(bias_src_strategy, TupleStrategy) else len(bias_src_strategy) | 203 | + bias_num = len(bias_src_strategy.childs) if isinstance(bias_src_strategy, TupleStrategy) else len(bias_src_strategy) |
| 203 | group_list_num = 1 if ( | 204 | group_list_num = 1 if ( |
| 204 | op_schema.op == npu.npu_grouped_matmul.default and | 205 | op_schema.op == npu.npu_grouped_matmul.default and |
| 205 | op_schema.kwargs_schema.get("group_list", None) is not None | 206 | op_schema.kwargs_schema.get("group_list", None) is not None |
| @@ -219,7 +220,7 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 219 | ] | 220 | ] |
| 220 | for key in unsupported_arguments: | 221 | for key in unsupported_arguments: |
| 221 | schema = op_schema.kwargs_schema.get(key, None) | 222 | schema = op_schema.kwargs_schema.get(key, None) |
| 222 | - if schema is not None and isinstance(schema, TupleStrategy) and len(schema.children) > 0: | 223 | + if schema is not None and isinstance(schema, TupleStrategy) and len(schema.childs) > 0: |
| 223 | full_mesh_strategies = expand_to_full_mesh_op_strategy( | 224 | full_mesh_strategies = expand_to_full_mesh_op_strategy( |
| 224 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num | 225 | op_schema.get_mesh_from_args(), op_schema, strategies, input_index=y_num |
| 225 | ) | 226 | ) |
| @@ -249,7 +250,7 @@ def npu_grouped_matmul_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 249 | pair_strategies = [] | 250 | pair_strategies = [] |
| 250 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) | 251 | # x: 2-6D, weight: 2D, weight: 1D (equals to weight.shape[1]) |
| 251 | # shard x | 252 | # shard x |
| 252 | - x_ndim = x_src_strategy.children[0].ndim | 253 | + x_ndim = x_src_strategy.childs[0].ndim |
| 253 | for i in range(x_ndim - 1): | 254 | for i in range(x_ndim - 1): |
| 254 | pair_strategies.append([Shard(i), Shard(i), Replicate(), Replicate()]) # y, x, weight, bias | 255 | pair_strategies.append([Shard(i), Shard(i), Replicate(), Replicate()]) # y, x, weight, bias |
| 255 | # shard weight | 256 | # shard weight |
| @@ -342,6 +343,14 @@ def _npu_grouped_matmul_handler( | |||
| 342 | # since upwrap_to_op_info does not process List[DTensor] in kwargs, we need to handle it here | 343 | # since upwrap_to_op_info does not process List[DTensor] in kwargs, we need to handle it here |
| 343 | _handle_tensor_list_in_kwargs(kwargs, op_info) | 344 | _handle_tensor_list_in_kwargs(kwargs, op_info) |
| 344 | 345 | ||
| 346 | + # return type of npu_grouped_matmul is tensor list, which caused output_spec to be None after propagation, and | ||
| 347 | + # v2.9.0 fixed it. We set return_type_tensor to True to avoid patching the entire propagate_op_sharding_non_cached | ||
| 348 | + # function in previous versions. | ||
| 349 | + def _return_type_tensor(): | ||
| 350 | + return True | ||
| 351 | + | ||
| 352 | + op_info.schema.return_type_tensor = _return_type_tensor | ||
| 353 | + | ||
| 345 | # sharding propagation | 354 | # sharding propagation |
| 346 | DTensor._op_dispatcher.sharding_propagator.propagate(op_info) | 355 | DTensor._op_dispatcher.sharding_propagator.propagate(op_info) |
| 347 | output_sharding = op_info.output_sharding | 356 | output_sharding = op_info.output_sharding |
| @@ -583,7 +592,7 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 583 | max_gard_norm_arg_index = 8 | 592 | max_gard_norm_arg_index = 8 |
| 584 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] | 593 | grad_strategy: OpStrategy = op_schema.args_schema[grad_arg_index] |
| 585 | if "out" in op_schema.kwargs_schema.keys(): | 594 | if "out" in op_schema.kwargs_schema.keys(): |
| 586 | - grad_spec: DTensorSpec = op_schema.kwargs_schema["out"].children[0].strategies[0].output_spec | 595 | + grad_spec: DTensorSpec = op_schema.kwargs_schema["out"].childs[0].strategies[0].output_spec |
| 587 | else: | 596 | else: |
| 588 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec | 597 | grad_spec: DTensorSpec = grad_strategy.strategies[0].output_spec |
| 589 | input_target_specs = [] | 598 | input_target_specs = [] |
| @@ -607,7 +616,7 @@ def npu_apply_adam_w_strategy(op_schema: OpSchema) -> OpStrategy: | |||
| 607 | output_spec = [] | 616 | output_spec = [] |
| 608 | for k, values in op_schema.kwargs_schema.items(): | 617 | for k, values in op_schema.kwargs_schema.items(): |
| 609 | if k == 'out': | 618 | if k == 'out': |
| 610 | - for v in values.children: | 619 | + for v in values.childs: |
| 611 | output_spec.append(v.strategies[0].output_spec) | 620 | output_spec.append(v.strategies[0].output_spec) |
| 612 | output_strategy = OpStrategy([ | 621 | output_strategy = OpStrategy([ |
| 613 | PlacementStrategy(output_specs=tuple(output_spec), input_specs=input_target_specs) | 622 | PlacementStrategy(output_specs=tuple(output_spec), input_specs=input_target_specs) |
| @@ -637,8 +646,7 @@ def _npu_apply_adam_w_handler( | |||
| 637 | if output_sharding.needs_redistribute: | 646 | if output_sharding.needs_redistribute: |
| 638 | DTensor._op_dispatcher.redistribute_local_args( | 647 | DTensor._op_dispatcher.redistribute_local_args( |
| 639 | op_info, | 648 | op_info, |
| 640 | - output_sharding.redistribute_schema, | 649 | + output_sharding.redistribute_schema |
| 641 | - output_sharding.use_val_from_redistribute_schema, | ||
| 642 | ) | 650 | ) |
| 643 | local_args = ( | 651 | local_args = ( |
| 644 | pytree.tree_unflatten( | 652 | pytree.tree_unflatten( |
| @@ -650,7 +658,7 @@ def _npu_apply_adam_w_handler( | |||
| 650 | 658 | ||
| 651 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) | 659 | local_results = torch_npu.npu_apply_adam_w(*local_args, **op_info.local_kwargs) |
| 652 | 660 | ||
| 653 | - if op_info.schema.is_out_variant_op(): | 661 | + if _is_out_variant_op(op_call): |
| 654 | output_specs = ( | 662 | output_specs = ( |
| 655 | (output_sharding.output_spec,) | 663 | (output_sharding.output_spec,) |
| 656 | if not isinstance(output_sharding.output_spec, tuple) | 664 | if not isinstance(output_sharding.output_spec, tuple) |