已合并
Add handling logic for y2Out shape. #3760
huawuyi创建于 2025年12月11日
Add handling logic for y2Out shape. #3760
已合并
从已删除 :master合入到Ascend/op-pluginmaster
共 3 个文件变更+114-40
| @@ -10553,6 +10553,7 @@ scale2: Tensor类型，第二路量化scale输出，shape同scale1，数据类 | |||
| 10553 | 当outputMaskOptional不为空时，参数smoothScale1Optional有值时，则outputMaskOptional[0]必须为True。参数smoothScale2Optional有值时，则outputMaskOptional[1]必须为True。 | 10553 | 当outputMaskOptional不为空时，参数smoothScale1Optional有值时，则outputMaskOptional[0]必须为True。参数smoothScale2Optional有值时，则outputMaskOptional[1]必须为True。 |
| 10554 | 当outputMaskOptional不为空时，outputMaskOptional[0]与outputMaskOptional[1]不能同时为False。 | 10554 | 当outputMaskOptional不为空时，outputMaskOptional[0]与outputMaskOptional[1]不能同时为False。 |
| 10555 | 当outputMaskOptional为空时，参数smoothScale2Optional有值时，参数smoothScale1Optional也必须有值。 | 10555 | 当outputMaskOptional为空时，参数smoothScale2Optional有值时，参数smoothScale1Optional也必须有值。 |
| 10556 | +如果y2Out为有效输出时，shape需要与y1Out保持一致；如果y2Out为无效输出时，shape为[1]。 | ||
| 10556 | 10557 | ||
| 10557 | 支持的型号 | 10558 | 支持的型号 |
| 10558 | Atlas A3训练系列产品/Atlas A3推理系列产品 | 10559 | Atlas A3训练系列产品/Atlas A3推理系列产品 |
| @@ -6308,7 +6308,7 @@ custom: | |||
| 6308 | size: x1 | 6308 | size: x1 |
| 6309 | dtype: at::kChar | 6309 | dtype: at::kChar |
| 6310 | y2: | 6310 | y2: |
| 6311 | - size: x1 | 6311 | + size: 'output_mask[1] ? x1.sizes() : at::IntArrayRef{}' |
| 6312 | dtype: at::kChar | 6312 | dtype: at::kChar |
| 6313 | x_out: | 6313 | x_out: |
| 6314 | size: x1 | 6314 | size: x1 |
| @@ -6317,7 +6317,7 @@ custom: | |||
| 6317 | size: reduce_lastdim_output_size(x1) | 6317 | size: reduce_lastdim_output_size(x1) |
| 6318 | dtype: at::kFloat | 6318 | dtype: at::kFloat |
| 6319 | scale2: | 6319 | scale2: |
| 6320 | - size: reduce_lastdim_output_size(x1) | 6320 | + size: 'output_mask[1] ? reduce_lastdim_output_size(x1) : c10::SmallVector<int64_t, op_infer::SIZE>{}' |
| 6321 | dtype: at::kFloat | 6321 | dtype: at::kFloat |
| 6322 | exec: aclnnAddRmsNormDynamicQuantV2, x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon, output_mask, y1, y2, x_out, scale1, scale2 | 6322 | exec: aclnnAddRmsNormDynamicQuantV2, x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon, output_mask, y1, y2, x_out, scale1, scale2 |
| 6323 | 6323 | ||
| @@ -54,7 +54,8 @@ class TestAddRmsNormDynamicQuant(TestCase): | |||
| 54 | smooth_scale1: torch.Tensor = None, | 54 | smooth_scale1: torch.Tensor = None, |
| 55 | smooth_scale2: torch.Tensor = None, | 55 | smooth_scale2: torch.Tensor = None, |
| 56 | beta: torch.Tensor = None, | 56 | beta: torch.Tensor = None, |
| 57 | - epsilon: float = 1e-6, | 57 | + epsilon: float = 1e-6, |
| 58 | + output_mask=None, | ||
| 58 | ): | 59 | ): |
| 59 | assert x1.shape == x2.shape | 60 | assert x1.shape == x2.shape |
| 60 | last_dim = x1.shape[-1] | 61 | last_dim = x1.shape[-1] |
| @@ -82,20 +83,43 @@ class TestAddRmsNormDynamicQuant(TestCase): | |||
| 82 | def row_max_abs(t: torch.Tensor): | 83 | def row_max_abs(t: torch.Tensor): |
| 83 | return t.abs().amax(dim=-1, keepdim=True) | 84 | return t.abs().amax(dim=-1, keepdim=True) |
| 84 | 85 | ||
| 85 | - input1 = y if smooth_scale1 is None else y * smooth_scale1.to(torch.float32) | 86 | + no_mask = (output_mask is None) |
| 86 | - scale1 = row_max_abs(input1) / 127.0 | 87 | + mask0 = True if no_mask else bool(output_mask[0]) |
| 87 | - scale1 = torch.where(scale1 > 0, scale1, torch.ones_like(scale1)) | 88 | + mask1 = None if no_mask else bool(output_mask[1]) |
| 88 | - q1 = torch.round(input1 / scale1).to(torch.int32) | ||
| 89 | - q1 = torch.clamp(q1, -128, 127).to(torch.int8) | ||
| 90 | 89 | ||
| 91 | - input2 = y if smooth_scale2 is None else y * smooth_scale2.to(torch.float32) | 90 | + if mask0: |
| 92 | - scale2 = row_max_abs(input2) / 127.0 | 91 | + input1 = y if smooth_scale1 is None else y * smooth_scale1.to(torch.float32) |
| 93 | - scale2 = torch.where(scale2 > 0, scale2, torch.ones_like(scale2)) | 92 | + scale1 = row_max_abs(input1) / 127.0 |
| 94 | - q2 = torch.round(input2 / scale2).to(torch.int32) | 93 | + scale1 = torch.where(scale1 > 0, scale1, torch.ones_like(scale1)) |
| 95 | - q2 = torch.clamp(q2, -128, 127).to(torch.int8) | 94 | + q1 = torch.round(input1 / scale1).to(torch.int32) |
| 95 | + q1 = torch.clamp(q1, -128, 127).to(torch.int8) | ||
| 96 | + s1_out = scale1.squeeze(-1).to(torch.float32).contiguous() | ||
| 97 | + else: | ||
| 98 | + q1 = None | ||
| 99 | + s1_out = None | ||
| 96 | 100 | ||
| 97 | - s1_out = scale1.squeeze(-1).to(torch.float32).contiguous() | 101 | + compute_branch2 = False |
| 98 | - s2_out = scale2.squeeze(-1).to(torch.float32).contiguous() | 102 | + if no_mask: |
| 103 | + if (smooth_scale1 is not None) and (smooth_scale2 is not None): | ||
| 104 | + compute_branch2 = True | ||
| 105 | + else: | ||
| 106 | + compute_branch2 = False | ||
| 107 | + else: | ||
| 108 | + if mask1: | ||
| 109 | + compute_branch2 = True | ||
| 110 | + else: | ||
| 111 | + compute_branch2 = False | ||
| 112 | + | ||
| 113 | + if compute_branch2: | ||
| 114 | + input2 = y if smooth_scale2 is None else y * smooth_scale2.to(torch.float32) | ||
| 115 | + scale2 = row_max_abs(input2) / 127.0 | ||
| 116 | + scale2 = torch.where(scale2 > 0, scale2, torch.ones_like(scale2)) | ||
| 117 | + q2 = torch.round(input2 / scale2).to(torch.int32) | ||
| 118 | + q2 = torch.clamp(q2, -128, 127).to(torch.int8) | ||
| 119 | + s2_out = scale2.squeeze(-1).to(torch.float32).contiguous() | ||
| 120 | + else: | ||
| 121 | + q2 = None | ||
| 122 | + s2_out = None | ||
| 99 | 123 | ||
| 100 | x_out = x.to(x1.dtype) | 124 | x_out = x.to(x1.dtype) |
| 101 | return q1, q2, x_out, s1_out, s2_out | 125 | return q1, q2, x_out, s1_out, s2_out |
| @@ -136,45 +160,74 @@ class TestAddRmsNormDynamicQuant(TestCase): | |||
| 136 | ) | 160 | ) |
| 137 | 161 | ||
| 138 | y1_cpu, y2_cpu, x_out_cpu, s1_cpu, s2_cpu = self.npu_add_rms_norm_dynamic_quant_golden( | 162 | y1_cpu, y2_cpu, x_out_cpu, s1_cpu, s2_cpu = self.npu_add_rms_norm_dynamic_quant_golden( |
| 139 | - x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon | 163 | + x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon, output_mask |
| 140 | ) | 164 | ) |
| 141 | 165 | ||
| 142 | - | ||
| 143 | - self.assertEqual(y1_npu.dtype, torch.int8) | ||
| 144 | - self.assertEqual(y2_npu.dtype, torch.int8) | ||
| 145 | self.assertEqual(x_out_npu.dtype, x1.dtype) | 166 | self.assertEqual(x_out_npu.dtype, x1.dtype) |
| 146 | - self.assertEqual(s1_npu.dtype, torch.float32) | ||
| 147 | - self.assertEqual(s2_npu.dtype, torch.float32) | ||
| 148 | - | ||
| 149 | - self.assertEqual(tuple(y1_npu.shape), tuple(x1.shape)) | ||
| 150 | - self.assertEqual(tuple(y2_npu.shape), tuple(x1.shape)) | ||
| 151 | self.assertEqual(tuple(x_out_npu.shape), tuple(x1.shape)) | 167 | self.assertEqual(tuple(x_out_npu.shape), tuple(x1.shape)) |
| 152 | - self.assertEqual(tuple(s1_npu.shape), tuple(x1.shape[:-1])) | ||
| 153 | - self.assertEqual(tuple(s2_npu.shape), tuple(x1.shape[:-1])) | ||
| 154 | 168 | ||
| 155 | - y1_diff = (y1_npu.cpu().to(torch.int16) - y1_cpu.cpu().to(torch.int16)).abs() | 169 | + if output_mask[0]: |
| 156 | - y2_diff = (y2_npu.cpu().to(torch.int16) - y2_cpu.cpu().to(torch.int16)).abs() | 170 | + self.assertEqual(y1_npu.dtype, torch.int8) |
| 157 | - self.assertTrue(int(y1_diff.max()) <= 1, f"max |y1_npu - y1_ref| = {int(y1_diff.max())} > 1") | 171 | + self.assertEqual(s1_npu.dtype, torch.float32) |
| 158 | - self.assertTrue(int(y2_diff.max()) <= 1, f"max |y2_npu - y2_ref| = {int(y2_diff.max())} > 1") | 172 | + self.assertEqual(tuple(y1_npu.shape), tuple(x1.shape)) |
| 173 | + self.assertEqual(tuple(s1_npu.shape), tuple(x1.shape[:-1])) | ||
| 174 | + else: | ||
| 175 | + self.assertTrue(isinstance(y1_npu, torch.Tensor)) | ||
| 176 | + self.assertTrue(y1_npu.numel() == 1, f"Expected y1 numel=1 when mask[0]=False, got {y1_npu.numel()}") | ||
| 177 | + self.assertTrue(isinstance(s1_npu, torch.Tensor)) | ||
| 178 | + self.assertTrue(s1_npu.numel() == 1, f"Expected s1 numel=1 when mask[0]=False, got {s1_npu.numel()}") | ||
| 179 | + | ||
| 180 | + if output_mask[1]: | ||
| 181 | + self.assertEqual(y2_npu.dtype, torch.int8) | ||
| 182 | + self.assertEqual(s2_npu.dtype, torch.float32) | ||
| 183 | + self.assertEqual(tuple(y2_npu.shape), tuple(x1.shape)) | ||
| 184 | + self.assertEqual(tuple(s2_npu.shape), tuple(x1.shape[:-1])) | ||
| 185 | + else: | ||
| 186 | + self.assertTrue(isinstance(y2_npu, torch.Tensor)) | ||
| 187 | + self.assertTrue(y2_npu.numel() == 1, f"Expected y2 numel=1 when mask[1]=False, got {y2_npu.numel()}") | ||
| 188 | + self.assertTrue(isinstance(s2_npu, torch.Tensor)) | ||
| 189 | + self.assertTrue(s2_npu.numel() == 1, f"Expected s2 numel=1 when mask[1]=False, got {s2_npu.numel()}") | ||
| 159 | 190 | ||
| 160 | benchmark = math.pow(2, -7) | 191 | benchmark = math.pow(2, -7) |
| 192 | + | ||
| 161 | x_out_npu_flat = x_out_npu.reshape(-1).cpu().to(torch.float32) | 193 | x_out_npu_flat = x_out_npu.reshape(-1).cpu().to(torch.float32) |
| 162 | x_out_cpu_flat = x_out_cpu.reshape(-1).cpu().to(torch.float32) | 194 | x_out_cpu_flat = x_out_cpu.reshape(-1).cpu().to(torch.float32) |
| 163 | self.assertTrue(self.compare(x_out_cpu_flat, x_out_npu_flat, benchmark)) | 195 | self.assertTrue(self.compare(x_out_cpu_flat, x_out_npu_flat, benchmark)) |
| 164 | 196 | ||
| 165 | - s1_npu_flat = s1_npu.reshape(-1).cpu().to(torch.float32) | 197 | + if output_mask[0]: |
| 166 | - s1_cpu_flat = s1_cpu.reshape(-1).cpu().to(torch.float32) | 198 | + y1_diff = (y1_npu.cpu().to(torch.int16) - y1_cpu.cpu().to(torch.int16)).abs() |
| 167 | - s2_npu_flat = s2_npu.reshape(-1).cpu().to(torch.float32) | 199 | + self.assertTrue(int(y1_diff.max()) <= 1, f"max |y1_npu - y1_ref| = {int(y1_diff.max())} > 1") |
| 168 | - s2_cpu_flat = s2_cpu.reshape(-1).cpu().to(torch.float32) | ||
| 169 | - self.assertTrue(self.compare(s1_cpu_flat, s1_npu_flat, benchmark)) | ||
| 170 | - self.assertTrue(self.compare(s2_cpu_flat, s2_npu_flat, benchmark)) | ||
| 171 | 200 | ||
| 172 | - for t in [x_out_npu, s1_npu, s2_npu]: | 201 | + s1_npu_flat = s1_npu.reshape(-1).cpu().to(torch.float32) |
| 202 | + s1_cpu_flat = s1_cpu.reshape(-1).cpu().to(torch.float32) | ||
| 203 | + self.assertTrue(self.compare(s1_cpu_flat, s1_npu_flat, benchmark)) | ||
| 204 | + | ||
| 205 | + if output_mask[1]: | ||
| 206 | + y2_diff = (y2_npu.cpu().to(torch.int16) - y2_cpu.cpu().to(torch.int16)).abs() | ||
| 207 | + self.assertTrue(int(y2_diff.max()) <= 1, f"max |y2_npu - y2_ref| = {int(y2_diff.max())} > 1") | ||
| 208 | + | ||
| 209 | + s2_npu_flat = s2_npu.reshape(-1).cpu().to(torch.float32) | ||
| 210 | + s2_cpu_flat = s2_cpu.reshape(-1).cpu().to(torch.float32) | ||
| 211 | + self.assertTrue(self.compare(s2_cpu_flat, s2_npu_flat, benchmark)) | ||
| 212 | + | ||
| 213 | + for t in [x_out_npu]: | ||
| 173 | tt = t.float() | 214 | tt = t.float() |
| 174 | self.assertFalse(torch.isnan(tt).any().item()) | 215 | self.assertFalse(torch.isnan(tt).any().item()) |
| 175 | self.assertFalse(torch.isinf(tt).any().item()) | 216 | self.assertFalse(torch.isinf(tt).any().item()) |
| 176 | 217 | ||
| 177 | - @unittest.skip("skip until CANN is updated to support aclnnAddRmsNormDynamicQuantV2") | 218 | + if output_mask[0]: |
| 219 | + for t in [s1_npu]: | ||
| 220 | + tt = t.float() | ||
| 221 | + self.assertFalse(torch.isnan(tt).any().item()) | ||
| 222 | + self.assertFalse(torch.isinf(tt).any().item()) | ||
| 223 | + | ||
| 224 | + if output_mask[1]: | ||
| 225 | + for t in [s2_npu]: | ||
| 226 | + tt = t.float() | ||
| 227 | + self.assertFalse(torch.isnan(tt).any().item()) | ||
| 228 | + self.assertFalse(torch.isinf(tt).any().item()) | ||
| 229 | + | ||
| 230 | + | ||
| 178 | 231 | ||
| 179 | def test_forward_various_shapes(self): | 232 | def test_forward_various_shapes(self): |
| 180 | shape_list = [ | 233 | shape_list = [ |
| @@ -198,7 +251,7 @@ class TestAddRmsNormDynamicQuant(TestCase): | |||
| 198 | 251 | ||
| 199 | self._run_and_check(x1, x2, gamma, smooth_scale1, smooth_scale2, beta) | 252 | self._run_and_check(x1, x2, gamma, smooth_scale1, smooth_scale2, beta) |
| 200 | 253 | ||
| 201 | - @unittest.skip("skip until CANN is updated to support aclnnAddRmsNormDynamicQuantV2") | 254 | + |
| 202 | 255 | ||
| 203 | def test_forward_various_shapes_bf16(self): | 256 | def test_forward_various_shapes_bf16(self): |
| 204 | shape_list = [ | 257 | shape_list = [ |
| @@ -221,6 +274,26 @@ class TestAddRmsNormDynamicQuant(TestCase): | |||
| 221 | smooth_scale2 = torch.ones(last_dim, dtype=torch.bfloat16, device='npu') | 274 | smooth_scale2 = torch.ones(last_dim, dtype=torch.bfloat16, device='npu') |
| 222 | 275 | ||
| 223 | self._run_and_check(x1, x2, gamma, smooth_scale1, smooth_scale2, beta) | 276 | self._run_and_check(x1, x2, gamma, smooth_scale1, smooth_scale2, beta) |
| 224 | - | 277 | + |
| 278 | + | ||
| 279 | + | ||
| 280 | + def test_forward_output_mask_true_false_fp16(self): | ||
| 281 | + x_shape = [2, 3, 32] | ||
| 282 | + last_dim = x_shape[-1] | ||
| 283 | + | ||
| 284 | + x1 = torch.randn(x_shape, dtype=torch.float16, device='npu') | ||
| 285 | + x2 = torch.randn(x_shape, dtype=torch.float16, device='npu') | ||
| 286 | + gamma = torch.ones(last_dim, dtype=torch.float16, device='npu') | ||
| 287 | + beta = torch.zeros(last_dim, dtype=torch.float16, device='npu') | ||
| 288 | + smooth_scale1 = torch.ones(last_dim, dtype=torch.float16, device='npu') | ||
| 289 | + smooth_scale2 = None | ||
| 290 | + | ||
| 291 | + self._run_and_check( | ||
| 292 | + x1, x2, gamma, | ||
| 293 | + smooth_scale1=smooth_scale1, | ||
| 294 | + smooth_scale2=smooth_scale2, | ||
| 295 | + beta=beta, | ||
| 296 | + output_mask=[True, False], | ||
| 297 | + ) | ||
| 225 | if __name__ == "__main__": | 298 | if __name__ == "__main__": |
| 226 | run_tests() | 299 | run_tests() |