Star844
1306
代码
代码
Issues1403
Pull Requests2257
流水线
项目讨论
Wiki
分析
Star844
1306
已合并
fix DTensor UT #29013
jizewei创建于 1月4日
fix DTensor UT #29013
已合并
jizewei创建于 1月4日
从jizewei/pytorchv2.9.0_fix_dtensor_ut合入到Ascend/pytorchv2.9.0
共 7 个文件变更+105-62
Mtest/distributed/tensor/test_attention_ops.py+11-5
@@ -10,6 +10,7 @@ import torch_npu
10from torch_npu.testing.testcase import run_tests10from torch_npu.testing.testcase import run_tests
11from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase11from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
12from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU12from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
13+from torch_npu.testing.common_utils import SupportedDevices
13 14 
14 15 
15def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536):16def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536):
@@ -27,7 +28,8 @@ def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536):
27 28 
28 29 
29class TestAttentionOps(NPUDTensorTestBase):30class TestAttentionOps(NPUDTensorTestBase):
30- @skipIfUnsupportMultiNPU(4)31+ @SupportedDevices(['Ascend910B'])
32+ @skipIfUnsupportMultiNPU(2)
31 @with_comms33 @with_comms
32 @parametrize(34 @parametrize(
33 "sparse_mode,pre_tokens,next_tokens",35 "sparse_mode,pre_tokens,next_tokens",
@@ -79,7 +81,8 @@ class TestAttentionOps(NPUDTensorTestBase):
79 else:81 else:
80 test_placement_comb([placement], [placement], [placement], [placement])82 test_placement_comb([placement], [placement], [placement], [placement])
81 83 
82- @skipIfUnsupportMultiNPU(4)84+ @SupportedDevices(['Ascend910B'])
85+ @skipIfUnsupportMultiNPU(2)
83 @with_comms86 @with_comms
84 @parametrize(87 @parametrize(
85 "sparse_mode,pre_tokens,next_tokens",88 "sparse_mode,pre_tokens,next_tokens",
@@ -138,7 +141,8 @@ class TestAttentionOps(NPUDTensorTestBase):
138 else:141 else:
139 test_placement_comb([placement], [placement], [placement], [placement])142 test_placement_comb([placement], [placement], [placement], [placement])
140 143 
141- @skipIfUnsupportMultiNPU(4)144+ @SupportedDevices(['Ascend910B'])
145+ @skipIfUnsupportMultiNPU(2)
142 @with_comms146 @with_comms
143 def test_npu_fusion_attention_bsnd(self):147 def test_npu_fusion_attention_bsnd(self):
144 device_mesh = self.build_device_mesh()148 device_mesh = self.build_device_mesh()
@@ -174,7 +178,8 @@ class TestAttentionOps(NPUDTensorTestBase):
174 for placement in placements:178 for placement in placements:
175 test_placement_comb([placement], [placement], [placement])179 test_placement_comb([placement], [placement], [placement])
176 180 
177- @skipIfUnsupportMultiNPU(4)181+ @SupportedDevices(['Ascend910B'])
182+ @skipIfUnsupportMultiNPU(2)
178 @with_comms183 @with_comms
179 def test_npu_fusion_attention_bsh(self):184 def test_npu_fusion_attention_bsh(self):
180 device_mesh = self.build_device_mesh()185 device_mesh = self.build_device_mesh()
@@ -210,7 +215,8 @@ class TestAttentionOps(NPUDTensorTestBase):
210 for placement in placements:215 for placement in placements:
211 test_placement_comb([placement], [placement], [placement])216 test_placement_comb([placement], [placement], [placement])
212 217 
213- @skipIfUnsupportMultiNPU(4)218+ @SupportedDevices(['Ascend910B'])
219+ @skipIfUnsupportMultiNPU(2)
214 @with_comms220 @with_comms
215 def test_npu_fusion_attention_tnd(self):221 def test_npu_fusion_attention_tnd(self):
216 device_mesh = self.build_device_mesh()222 device_mesh = self.build_device_mesh()
Mtest/distributed/tensor/test_gather_swiglu.py+19-12
@@ -1,18 +1,20 @@
1import torch1import torch
2from torch.distributed._tensor import distribute_tensor, Replicate, Shard2from torch.distributed._tensor import distribute_tensor, Replicate, Shard
3-from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
4 3 
5import torch_npu4import torch_npu
6from torch_npu.testing.testcase import run_tests5from torch_npu.testing.testcase import run_tests
6+from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
7from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU7from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
8+from torch_npu.testing.common_utils import SupportedDevices
8 9 
9 10 
10class TestRegisterSharding(NPUDTensorTestBase):11class TestRegisterSharding(NPUDTensorTestBase):
11- @skipIfUnsupportMultiNPU(4)12+ @SupportedDevices(['Ascend910B'])
13+ @skipIfUnsupportMultiNPU(2)
12 @with_comms14 @with_comms
13 def test_torch_npu_npu_swiglu(self):15 def test_torch_npu_npu_swiglu(self):
14 mesh = self.build_device_mesh()16 mesh = self.build_device_mesh()
15- 17+ 
16 input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True)18 input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True)
17 grad_tensor = torch.randn(1024, 512, device="npu")19 grad_tensor = torch.randn(1024, 512, device="npu")
18 dim = -120 dim = -1
@@ -26,8 +28,9 @@ class TestRegisterSharding(NPUDTensorTestBase):
26 output.backward(grad_dtensor)28 output.backward(grad_dtensor)
27 self.assertEqual(output.full_tensor(), out_tensor)29 self.assertEqual(output.full_tensor(), out_tensor)
28 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)30 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)
29- 31+ 
30- @skipIfUnsupportMultiNPU(4)32+ @SupportedDevices(['Ascend910B'])
33+ @skipIfUnsupportMultiNPU(2)
31 @with_comms34 @with_comms
32 def test_torch_npu_npu_swiglu_shard0(self):35 def test_torch_npu_npu_swiglu_shard0(self):
33 mesh = self.build_device_mesh()36 mesh = self.build_device_mesh()
@@ -46,14 +49,15 @@ class TestRegisterSharding(NPUDTensorTestBase):
46 self.assertEqual(output.full_tensor(), out_tensor)49 self.assertEqual(output.full_tensor(), out_tensor)
47 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)50 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)
48 51 
49- @skipIfUnsupportMultiNPU(4)52+ @SupportedDevices(['Ascend910B'])
53+ @skipIfUnsupportMultiNPU(2)
50 @with_comms54 @with_comms
51 def test_torch_npu_npu_swiglu_shard1(self):55 def test_torch_npu_npu_swiglu_shard1(self):
52 mesh = self.build_device_mesh()56 mesh = self.build_device_mesh()
53 57 
54 input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True)58 input_tensor = torch.randn(1024, 1024, device="npu", requires_grad=True)
55 grad_tensor = torch.randn(1024, 512, device="npu")59 grad_tensor = torch.randn(1024, 512, device="npu")
56- dim = -1 60+ dim = -1
57 out_tensor = torch_npu.npu_swiglu(input_tensor, dim)61 out_tensor = torch_npu.npu_swiglu(input_tensor, dim)
58 out_tensor.backward(grad_tensor)62 out_tensor.backward(grad_tensor)
59 63 
@@ -65,7 +69,8 @@ class TestRegisterSharding(NPUDTensorTestBase):
65 self.assertEqual(output.full_tensor(), out_tensor)69 self.assertEqual(output.full_tensor(), out_tensor)
66 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)70 self.assertEqual(input_dtensor.grad.full_tensor(), input_tensor.grad)
67 71 
68- @skipIfUnsupportMultiNPU(4)72+ @SupportedDevices(['Ascend910B'])
73+ @skipIfUnsupportMultiNPU(2)
69 @with_comms74 @with_comms
70 def test_torch_gather(self):75 def test_torch_gather(self):
71 mesh = self.build_device_mesh()76 mesh = self.build_device_mesh()
@@ -89,7 +94,8 @@ class TestRegisterSharding(NPUDTensorTestBase):
89 self.assertEqual(out.full_tensor(), out_tensor)94 self.assertEqual(out.full_tensor(), out_tensor)
90 self.assertEqual(x.grad.full_tensor(), input_tensor.grad)95 self.assertEqual(x.grad.full_tensor(), input_tensor.grad)
91 96 
92- @skipIfUnsupportMultiNPU(4)97+ @SupportedDevices(['Ascend910B'])
98+ @skipIfUnsupportMultiNPU(2)
93 @with_comms99 @with_comms
94 def test_torch_gather_shard0(self):100 def test_torch_gather_shard0(self):
95 mesh = self.build_device_mesh()101 mesh = self.build_device_mesh()
@@ -105,7 +111,7 @@ class TestRegisterSharding(NPUDTensorTestBase):
105 111 
106 x = distribute_tensor(input_tensor, mesh, [Shard(0)])112 x = distribute_tensor(input_tensor, mesh, [Shard(0)])
107 index = distribute_tensor(index_tensor, mesh, [Shard(0)])113 index = distribute_tensor(index_tensor, mesh, [Shard(0)])
108- 114+ 
109 out = torch.gather(input=x, dim=dim, index=index)115 out = torch.gather(input=x, dim=dim, index=index)
110 grad = distribute_tensor(grad_tensor, mesh, out.placements)116 grad = distribute_tensor(grad_tensor, mesh, out.placements)
111 117 
@@ -113,7 +119,8 @@ class TestRegisterSharding(NPUDTensorTestBase):
113 self.assertEqual(out.full_tensor(), out_tensor)119 self.assertEqual(out.full_tensor(), out_tensor)
114 self.assertEqual(x.grad.full_tensor(), input_tensor.grad)120 self.assertEqual(x.grad.full_tensor(), input_tensor.grad)
115 121 
116- @skipIfUnsupportMultiNPU(4)122+ @SupportedDevices(['Ascend910B'])
123+ @skipIfUnsupportMultiNPU(2)
117 @with_comms124 @with_comms
118 def test_torch_gather_shard1(self):125 def test_torch_gather_shard1(self):
119 mesh = self.build_device_mesh()126 mesh = self.build_device_mesh()
@@ -129,7 +136,7 @@ class TestRegisterSharding(NPUDTensorTestBase):
129 136 
130 x = distribute_tensor(input_tensor, mesh, [Shard(1)])137 x = distribute_tensor(input_tensor, mesh, [Shard(1)])
131 index = distribute_tensor(index_tensor, mesh, [Shard(1)])138 index = distribute_tensor(index_tensor, mesh, [Shard(1)])
132- 139+ 
133 out = torch.gather(input=x, dim=dim, index=index)140 out = torch.gather(input=x, dim=dim, index=index)
134 grad = distribute_tensor(grad_tensor, mesh, out.placements)141 grad = distribute_tensor(grad_tensor, mesh, out.placements)
135 142 
Mtest/distributed/tensor/test_math_ops.py+11-6
@@ -11,7 +11,7 @@ from torch.testing._internal.common_utils import (
11import torch_npu11import torch_npu
12from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase12from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
13from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU13from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
14-from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase14+from torch_npu.testing.common_utils import SupportedDevices
15 15 
16 16 
17def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str):17def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str):
@@ -33,7 +33,8 @@ def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension:
33 33 
34 34 
35class TestMathOps(NPUDTensorTestBase):35class TestMathOps(NPUDTensorTestBase):
36- @skipIfUnsupportMultiNPU(4)36+ @SupportedDevices(['Ascend910B'])
37+ @skipIfUnsupportMultiNPU(2)
37 @with_comms38 @with_comms
38 def test_npu_rms_norm_forward(self):39 def test_npu_rms_norm_forward(self):
39 device_mesh = self.build_device_mesh()40 device_mesh = self.build_device_mesh()
@@ -51,7 +52,8 @@ class TestMathOps(NPUDTensorTestBase):
51 self.assertEqual(dist_y.full_tensor(), y)52 self.assertEqual(dist_y.full_tensor(), y)
52 self.assertEqual(dist_gamma.full_tensor(), gamma)53 self.assertEqual(dist_gamma.full_tensor(), gamma)
53 54 
54- @skipIfUnsupportMultiNPU(4)55+ @SupportedDevices(['Ascend910B'])
56+ @skipIfUnsupportMultiNPU(2)
55 @with_comms57 @with_comms
56 def test_npu_rms_norm_backward(self):58 def test_npu_rms_norm_backward(self):
57 device_mesh = self.build_device_mesh()59 device_mesh = self.build_device_mesh()
@@ -86,7 +88,8 @@ class TestMathOps(NPUDTensorTestBase):
86 self.assertEqual(dist_dx.full_tensor(), dx)88 self.assertEqual(dist_dx.full_tensor(), dx)
87 self.assertEqual(dist_dw.full_tensor(), dw)89 self.assertEqual(dist_dw.full_tensor(), dw)
88 90 
89- @skipIfUnsupportMultiNPU(4)91+ @SupportedDevices(['Ascend910B'])
92+ @skipIfUnsupportMultiNPU(2)
90 @with_comms93 @with_comms
91 def test_npu_add_rms_norm_forward(self):94 def test_npu_add_rms_norm_forward(self):
92 device_mesh = self.build_device_mesh()95 device_mesh = self.build_device_mesh()
@@ -111,7 +114,8 @@ class TestMathOps(NPUDTensorTestBase):
111 for comb in placement_combs:114 for comb in placement_combs:
112 test_placement_comb([comb[0]], [comb[1]])115 test_placement_comb([comb[0]], [comb[1]])
113 116 
114- @skipIfUnsupportMultiNPU(4)117+ @SupportedDevices(['Ascend910B'])
118+ @skipIfUnsupportMultiNPU(2)
115 @with_comms119 @with_comms
116 @parametrize(120 @parametrize(
117 "rotary_mode,input_layout,sin_cos_layout",121 "rotary_mode,input_layout,sin_cos_layout",
@@ -159,7 +163,8 @@ class TestMathOps(NPUDTensorTestBase):
159 else:163 else:
160 test_placement_comb([placement], [placement], [placement])164 test_placement_comb([placement], [placement], [placement])
161 165 
162- @skipIfUnsupportMultiNPU(4)166+ @SupportedDevices(['Ascend910B'])
167+ @skipIfUnsupportMultiNPU(2)
163 @with_comms168 @with_comms
164 @parametrize(169 @parametrize(
165 "rotary_mode,input_layout,sin_cos_layout",170 "rotary_mode,input_layout,sin_cos_layout",
Mtest/distributed/tensor/test_matrix_ops.py+17-17
@@ -38,7 +38,7 @@ class TestAllGatherBaseMmOp(NPUDTensorTestBase):
38 return x1, x2, output, gather_out38 return x1, x2, output, gather_out
39 39 
40 @SupportedDevices(['Ascend910B'])40 @SupportedDevices(['Ascend910B'])
41- @skipIfUnsupportMultiNPU(4)41+ @skipIfUnsupportMultiNPU(2)
42 @with_comms42 @with_comms
43 def test_npu_all_gather_base_mm(self):43 def test_npu_all_gather_base_mm(self):
44 mesh = self.build_device_mesh()44 mesh = self.build_device_mesh()
@@ -109,7 +109,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
109 return x1, x2, output109 return x1, x2, output
110 110 
111 @SupportedDevices(['Ascend910B'])111 @SupportedDevices(['Ascend910B'])
112- @skipIfUnsupportMultiNPU(4)112+ @skipIfUnsupportMultiNPU(2)
113 @with_comms113 @with_comms
114 def test_npu_mm_reduce_scatter_base(self):114 def test_npu_mm_reduce_scatter_base(self):
115 with DeterministicGuard(True):115 with DeterministicGuard(True):
@@ -151,7 +151,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
151 test_placement_comb([comb[0]], [comb[1]])151 test_placement_comb([comb[0]], [comb[1]])
152 152 
153 @SupportedDevices(['Ascend910B'])153 @SupportedDevices(['Ascend910B'])
154- @skipIfUnsupportMultiNPU(4)154+ @skipIfUnsupportMultiNPU(2)
155 @with_comms155 @with_comms
156 def test_npu_mm_reduce_scatter_base_bias(self):156 def test_npu_mm_reduce_scatter_base_bias(self):
157 with DeterministicGuard(True):157 with DeterministicGuard(True):
@@ -165,7 +165,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
165 for _ in range(self.world_size):165 for _ in range(self.world_size):
166 x1_list.append(torch.randn(m, k, dtype=dtype, device="npu"))166 x1_list.append(torch.randn(m, k, dtype=dtype, device="npu"))
167 x2_list.append(torch.randn(k, n, dtype=dtype, device="npu"))167 x2_list.append(torch.randn(k, n, dtype=dtype, device="npu"))
168- bias_list.append(torch.randn(n, dtype=dtype, device="npu"))168+ bias_list.append(torch.zeros(n, dtype=dtype, device="npu"))
169 169 
170 global_x1, global_x2, global_output = self._get_global_tensor(x1_list, x2_list, bias_list=bias_list)170 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)171 global_bias = torch.cat(bias_list, dim=0)
@@ -187,7 +187,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
187 dist_output = torch_npu.npu_mm_reduce_scatter_base(187 dist_output = torch_npu.npu_mm_reduce_scatter_base(
188 dist_x1, dist_x2, hcom_name, self.world_size, bias=dist_bias188 dist_x1, dist_x2, hcom_name, self.world_size, bias=dist_bias
189 )189 )
190- self.assertEqual(dist_output.full_tensor(), global_output.to(dtype), atol=0.05, rtol=0.05)190+ self.assertEqual(dist_output.full_tensor(), global_output.to(dtype))
191 self.assertEqual(dist_output.to_local(), output)191 self.assertEqual(dist_output.to_local(), output)
192 192 
193 placement = [Shard(0), Shard(1), Replicate()]193 placement = [Shard(0), Shard(1), Replicate()]
@@ -197,7 +197,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
197 test_placement_comb([comb[0]], [comb[1]], [bias_placement])197 test_placement_comb([comb[0]], [comb[1]], [bias_placement])
198 198 
199 @SupportedDevices(['Ascend910B'])199 @SupportedDevices(['Ascend910B'])
200- @skipIfUnsupportMultiNPU(4)200+ @skipIfUnsupportMultiNPU(2)
201 @with_comms201 @with_comms
202 def test_npu_mm_reduce_scatter_base_quant(self):202 def test_npu_mm_reduce_scatter_base_quant(self):
203 with DeterministicGuard(True):203 with DeterministicGuard(True):
@@ -253,7 +253,7 @@ class TestMmReduceScatterBaseOp(NPUDTensorTestBase):
253 253 
254class TestGroupedMatMulOp(NPUDTensorTestBase):254class TestGroupedMatMulOp(NPUDTensorTestBase):
255 @SupportedDevices(['Ascend910B'])255 @SupportedDevices(['Ascend910B'])
256- @skipIfUnsupportMultiNPU(4)256+ @skipIfUnsupportMultiNPU(2)
257 @with_comms257 @with_comms
258 @parametrize("x_ndim", [2, 3])258 @parametrize("x_ndim", [2, 3])
259 @parametrize("with_bias", [True, False])259 @parametrize("with_bias", [True, False])
@@ -293,16 +293,16 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
293 test_placement_comb([comb[0]], [comb[1]], [comb[2]])293 test_placement_comb([comb[0]], [comb[1]], [comb[2]])
294 294 
295 @SupportedDevices(['Ascend910B'])295 @SupportedDevices(['Ascend910B'])
296- @skipIfUnsupportMultiNPU(4)296+ @skipIfUnsupportMultiNPU(2)
297 @with_comms297 @with_comms
298 @parametrize("with_bias", [True, False])298 @parametrize("with_bias", [True, False])
299 def test_npu_grouped_matmul_x1w1y1(self, with_bias):299 def test_npu_grouped_matmul_x1w1y1(self, with_bias):
300 mesh = self.build_device_mesh()300 mesh = self.build_device_mesh()
301 301 
302- x = [torch.randn(112, 64, dtype=torch.float32, device="npu")]302+ x = [torch.randn(8, 8, dtype=torch.float16, device="npu")]
303- weight = [torch.randn(4, 64, 16, dtype=torch.float32, device="npu")]303+ 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 None304+ 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")305+ group_list = torch.tensor([2, 8], device="npu")
306 split_item = 3306 split_item = 3
307 group_type = 0307 group_type = 0
308 308 
@@ -320,7 +320,7 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
320 split_item=split_item, group_type=group_type320 split_item=split_item, group_type=group_type
321 )321 )
322 for dist_y_i, y_i in zip(dist_y, y):322 for dist_y_i, y_i in zip(dist_y, y):
323- self.assertEqual(dist_y_i.full_tensor(), y_i)323+ self.assertEqual(dist_y_i.full_tensor(), y_i, atol=0.001, rtol=0.02)
324 324 
325 placement = [Shard(0), Shard(1), Replicate()]325 placement = [Shard(0), Shard(1), Replicate()]
326 placement_combs = itertools.product(placement, placement, [Shard(0), Replicate()])326 placement_combs = itertools.product(placement, placement, [Shard(0), Replicate()])
@@ -328,7 +328,7 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
328 test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]])328 test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]])
329 329 
330 @SupportedDevices(['Ascend910B'])330 @SupportedDevices(['Ascend910B'])
331- @skipIfUnsupportMultiNPU(4)331+ @skipIfUnsupportMultiNPU(2)
332 @with_comms332 @with_comms
333 @parametrize("with_bias", [True, False])333 @parametrize("with_bias", [True, False])
334 def test_npu_grouped_matmul_xNwNy1(self, with_bias):334 def test_npu_grouped_matmul_xNwNy1(self, with_bias):
@@ -367,7 +367,7 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
367 test_placement_comb([comb[0]], [comb[1]], [comb[2]])367 test_placement_comb([comb[0]], [comb[1]], [comb[2]])
368 368 
369 @SupportedDevices(['Ascend910B'])369 @SupportedDevices(['Ascend910B'])
370- @skipIfUnsupportMultiNPU(4)370+ @skipIfUnsupportMultiNPU(2)
371 @with_comms371 @with_comms
372 @parametrize("with_bias", [True, False])372 @parametrize("with_bias", [True, False])
373 @parametrize("group_type", [None, 0])373 @parametrize("group_type", [None, 0])
@@ -410,7 +410,7 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
410 test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]])410 test_placement_comb([comb[0]], [comb[1]], [comb[2]], [comb[2]])
411 411 
412 @SupportedDevices(['Ascend910B'])412 @SupportedDevices(['Ascend910B'])
413- @skipIfUnsupportMultiNPU(4)413+ @skipIfUnsupportMultiNPU(2)
414 @with_comms414 @with_comms
415 @parametrize("with_bias", [True, False])415 @parametrize("with_bias", [True, False])
416 def test_npu_grouped_matmul_x1wNyN(self, with_bias):416 def test_npu_grouped_matmul_x1wNyN(self, with_bias):
@@ -444,7 +444,7 @@ class TestGroupedMatMulOp(NPUDTensorTestBase):
444 test_placement_comb([comb[0]], [comb[1]], [comb[2]])444 test_placement_comb([comb[0]], [comb[1]], [comb[2]])
445 445 
446 @SupportedDevices(['Ascend910B'])446 @SupportedDevices(['Ascend910B'])
447- @skipIfUnsupportMultiNPU(4)447+ @skipIfUnsupportMultiNPU(2)
448 @with_comms448 @with_comms
449 def test_npu_grouped_matmul_quant(self):449 def test_npu_grouped_matmul_quant(self):
450 mesh = self.build_device_mesh()450 mesh = self.build_device_mesh()
Mtest/distributed/tensor/test_moe_ops.py+13-6
@@ -7,10 +7,12 @@ import torch_npu
7from torch_npu.testing.testcase import run_tests7from torch_npu.testing.testcase import run_tests
8from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase8from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
9from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU9from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
10+from torch_npu.testing.common_utils import SupportedDevices
10 11 
11 12 
12class TestMoeOps(NPUDTensorTestBase):13class TestMoeOps(NPUDTensorTestBase):
13- @skipIfUnsupportMultiNPU(4)14+ @SupportedDevices(['Ascend910B'])
15+ @skipIfUnsupportMultiNPU(2)
14 @with_comms16 @with_comms
15 def test_npu_moe_token_permute_forward(self):17 def test_npu_moe_token_permute_forward(self):
16 device_mesh = self.build_device_mesh()18 device_mesh = self.build_device_mesh()
@@ -35,7 +37,8 @@ class TestMoeOps(NPUDTensorTestBase):
35 for comb in placement_combs:37 for comb in placement_combs:
36 test_placement_comb([comb[0]], [comb[1]])38 test_placement_comb([comb[0]], [comb[1]])
37 39 
38- @skipIfUnsupportMultiNPU(4)40+ @SupportedDevices(['Ascend910B'])
41+ @skipIfUnsupportMultiNPU(2)
39 @with_comms42 @with_comms
40 def test_npu_moe_token_permute_backward(self):43 def test_npu_moe_token_permute_backward(self):
41 device_mesh = self.build_device_mesh()44 device_mesh = self.build_device_mesh()
@@ -67,7 +70,8 @@ class TestMoeOps(NPUDTensorTestBase):
67 for comb in placement_combs:70 for comb in placement_combs:
68 test_placement_comb([comb[0]], [comb[1]])71 test_placement_comb([comb[0]], [comb[1]])
69 72 
70- @skipIfUnsupportMultiNPU(4)73+ @SupportedDevices(['Ascend910B'])
74+ @skipIfUnsupportMultiNPU(2)
71 @with_comms75 @with_comms
72 def test_npu_moe_token_permute_clip(self):76 def test_npu_moe_token_permute_clip(self):
73 device_mesh = self.build_device_mesh()77 device_mesh = self.build_device_mesh()
@@ -95,7 +99,8 @@ class TestMoeOps(NPUDTensorTestBase):
95 self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)99 self.assertEqual(dist_sorted_indices.full_tensor(), sorted_indices)
96 self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad)100 self.assertEqual(dist_tokens.grad.full_tensor(), tokens.grad)
97 101 
98- @skipIfUnsupportMultiNPU(4)102+ @SupportedDevices(['Ascend910B'])
103+ @skipIfUnsupportMultiNPU(2)
99 @with_comms104 @with_comms
100 def test_npu_moe_token_unpermute_forward(self):105 def test_npu_moe_token_unpermute_forward(self):
101 device_mesh = self.build_device_mesh()106 device_mesh = self.build_device_mesh()
@@ -120,7 +125,8 @@ class TestMoeOps(NPUDTensorTestBase):
120 for comb in placement_combs:125 for comb in placement_combs:
121 test_placement_comb([comb[0]], [comb[1]])126 test_placement_comb([comb[0]], [comb[1]])
122 127 
123- @skipIfUnsupportMultiNPU(4)128+ @SupportedDevices(['Ascend910B'])
129+ @skipIfUnsupportMultiNPU(2)
124 @with_comms130 @with_comms
125 def test_npu_moe_token_unpermute_backward(self):131 def test_npu_moe_token_unpermute_backward(self):
126 device_mesh = self.build_device_mesh()132 device_mesh = self.build_device_mesh()
@@ -157,7 +163,8 @@ class TestMoeOps(NPUDTensorTestBase):
157 for comb in placement_combs:163 for comb in placement_combs:
158 test_placement_comb([comb[0]], [comb[1]])164 test_placement_comb([comb[0]], [comb[1]])
159 165 
160- @skipIfUnsupportMultiNPU(4)166+ @SupportedDevices(['Ascend910B'])
167+ @skipIfUnsupportMultiNPU(2)
161 @with_comms168 @with_comms
162 def test_npu_moe_token_permute_unpermute(self):169 def test_npu_moe_token_permute_unpermute(self):
163 device_mesh = self.build_device_mesh()170 device_mesh = self.build_device_mesh()
Mtorch_npu/distributed/tensor/_attention.py+25-16
@@ -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(), # pse58 None if pse is None else Replicate(), # pse
59 None if padding_mask is None else Replicate(), # padding_mask59 None if padding_mask is None else Replicate(), # padding_mask
60 None if atten_mask is None else Replicate(), # atten_mask60 None if atten_mask is None else Replicate(), # atten_mask
61- None, None, None, None, None, None, None, None, None, None, None, None # others61+ 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 used67 # 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 sharding68 # 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 strategies71 return strategies
71 72 
@@ -97,7 +98,8 @@ def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse
97 None, # pse98 None, # pse
98 None, # padding_mask99 None, # padding_mask
99 atten_mask_sharding, # atten_mask100 atten_mask_sharding, # atten_mask
100- None, None, None, None, None, None, None, None, None, None, None, None # others101+ 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, # pse130 None, # pse
129 None, # padding_mask131 None, # padding_mask
130 atten_mask_sharding, # atten_mask132 atten_mask_sharding, # atten_mask
131- None, None, None, None, None, None, None, None, None, None, None, None # others133+ 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 strategy160 # all replicate strategy
@@ -160,7 +163,8 @@ def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_la
160 Replicate(), # grad_query163 Replicate(), # grad_query
161 Replicate(), # grad_key164 Replicate(), # grad_key
162 Replicate(), # grad_value165 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(), # query170 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_sum180 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_in182 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 # others183+ 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 used189 # 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 sharding190 # 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 strategies193 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_query208 Shard(batch_dim), # grad_query
204 Shard(batch_dim), # grad_key209 Shard(batch_dim), # grad_key
205 Shard(batch_dim), # grad_value210 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), # query215 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: BNS8225 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_in227 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 # others228+ 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_query245 Shard(head_dim), # grad_query
239 Shard(head_dim), # grad_key246 Shard(head_dim), # grad_key
240 Shard(head_dim), # grad_value247 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), # query252 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: BNS8262 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_in264 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 # others265+ 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)
Mtorch_npu/testing/_internal/common_dtensor.py+9-0
@@ -1,3 +1,4 @@
1+from torch._utils import _get_device_module
1from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase2from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
2 3 
3 4 
@@ -5,3 +6,11 @@ class NPUDTensorTestBase(DTensorTestBase):
5 @property6 @property
6 def device_type(self):7 def device_type(self):
7 return "npu"8 return "npu"
9+ 
10+ @property
11+ def world_size(self):
12+ device_count = _get_device_module(self.device_type).device_count()
13+ device_num = 4
14+ if device_count > 1:
15+ device_num = min(device_num, device_count)
16+ return device_num