Star186
854
代码
代码
Issues179
Pull Requests335
项目讨论
Wiki
分析
Star186
854
已合并
[fix] A5 代际 Dropout 反向随机数算子数值对齐 PyTorch/GPU #5737
yucaopanmu创建于 16 天前
[fix] A5 代际 Dropout 反向随机数算子数值对齐 PyTorch/GPU #5737
已合并
yucaopanmu创建于 16 天前
从yucaopanmu/op-pluginfix/a5-native-dropout-backward合入到Ascend/op-pluginmaster
共 2 个文件变更+137-6
Mop_plugin/ops/opapi/NativeDropoutKernelNpuOpApi.cpp+31-0
@@ -153,6 +153,37 @@ std::tuple<at::Tensor, at::Tensor> native_dropout(const at::Tensor& input, doubl
153}153}
154 154 
155at::Tensor native_dropout_backward(const at::Tensor& grad_output, const at::Tensor& mask, double scale) {155at::Tensor native_dropout_backward(const at::Tensor& grad_output, const at::Tensor& mask, double scale) {
156+ // On A5, to align the backward precision with GPU(H20), the original scale is passed
157+ // to aclnnDropoutV3Grad directly. The kernel computes gradX = gradY * mask * scale with
158+ // a pure multiplication chain, avoiding the extra float computation of restoring the
159+ // scale factor (p = 1 - 1 / scale) inside the operator.
160+ if (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950 &&
161+ check_aclnn_kernel_available("aclnnDropoutV3Grad")) {
162+ TORCH_CHECK(
163+ scale == NUMBER_ZERO || scale >= NUMBER_ONE,
164+ "native_dropout_backward scale has to be 0 or greater than or equal to 1, but got ",
165+ scale,
166+ OPS_ERROR(ErrCode::VALUE));
167+ 
C
Cchengpeng2513 天前

缺少 aclnn接口是否可用的校验；不可用则回退原有的方案

likedislike
168+ if (mask.numel() == 0) {
169+ return at_npu::native::OpPreparation::apply_tensor_without_format(mask.sizes(), grad_output.options());
170+ }
171+ // Branch on scale: scale == 1 (p == 0, no dropout) returns grad_output itself;
172+ // scale == 0 or scale == inf (p == 1, all dropped) returns zeros;
173+ // otherwise aclnnDropoutV3Grad receives the original scale.
174+ double p = (scale == 0.0) ? 1 : (1 - 1 / scale);
175+ if (p == 0) {
176+ return grad_output.clone();
177+ }
178+ if (p == 1) {
179+ at::TensorOptions options = grad_output.options();
180+ return at::zeros(grad_output.sizes(), options);
181+ }
182+ at::Tensor result = at_npu::native::OpPreparation::apply_tensor_without_format(grad_output);
183+ EXEC_NPU_CMD(aclnnDropoutV3Grad, grad_output, mask, scale, result);
184+ return result;
185+ }
186+ 
156 DO_COMPATIBILITY(aclnnDropoutDoMask, acl_op::native_dropout_backward(grad_output, mask, scale));187 DO_COMPATIBILITY(aclnnDropoutDoMask, acl_op::native_dropout_backward(grad_output, mask, scale));
157 TORCH_CHECK(188 TORCH_CHECK(
158 scale == NUMBER_ZERO || scale >= NUMBER_ONE,189 scale == NUMBER_ZERO || scale >= NUMBER_ONE,
Mtest/test_base_ops/test_native_dropout_backward.py+106-6
@@ -1,24 +1,124 @@
1import torch1import torch
2import torch_npu2import torch_npu
3from torch_npu.testing.testcase import TestCase, run_tests3from torch_npu.testing.testcase import TestCase, run_tests
4+from torch_npu.testing.common_utils import SupportedDevices
4import random5import random
5import numpy as np6import numpy as np
6 7 
8+ 
7class TestDropout(TestCase):9class TestDropout(TestCase):
10+ @SupportedDevices(['Ascend910A', 'Ascend910B'])
11+ def test_native_dropout_backward_fp32(self):
12+ torch.manual_seed(0)
13+ self._test_native_dropout_backward(torch.float32, 2)
14+ 
15+ @SupportedDevices(['Ascend910A', 'Ascend910B'])
16+ def test_native_dropout_backward_fp16(self):
17+ torch.manual_seed(0)
18+ self._test_native_dropout_backward(torch.float16, 3)
19+ 
20+ @SupportedDevices(['Ascend950'])
21+ def test_native_dropout_backward_scale_zero(self, device="npu"):
22+ """A5 branch: scale == 0 (p == 1, all dropped) returns zeros."""
23+ torch.manual_seed(0)
24+ grad_output = torch.arange(32, dtype=torch.float32).reshape(32)
25+ mask = torch.zeros(32, dtype=torch.bool)
26+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, mask, 0.0)
27+ output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), mask.npu(), 0.0)
28+ self.assertRtolEqual(output_cpu.numpy(), output_npu.cpu().numpy(), 0.0001)
29+ 
30+ @SupportedDevices(['Ascend950'])
31+ def test_native_dropout_backward_scale_one(self, device="npu"):
32+ """A5 branch: scale == 1 (p == 0, no dropout) returns grad_output itself."""
33+ torch.manual_seed(0)
34+ grad_output = torch.arange(2 * 4 * 32, dtype=torch.float32).reshape(2, 4, 32)
35+ mask = torch.ones(grad_output.shape, dtype=torch.bool)
36+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, mask, 1.0)
37+ output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), mask.npu(), 1.0)
38+ self.assertEqual(output_npu.shape, grad_output.shape)
39+ self.assertRtolEqual(output_cpu.numpy(), output_npu.cpu().numpy(), 0.0001)
40+ 
41+ @SupportedDevices(['Ascend950'])
42+ def test_native_dropout_backward_scale_gt_one_fp32(self, device="npu"):
43+ """A5 kernel path: gradX = gradY * mask * scale (pure multiply chain, 0xAA mask pattern)."""
44+ torch.manual_seed(0)
45+ shape = (4, 32)
46+ grad_output = torch.arange(4 * 32, dtype=torch.float32).reshape(shape)
47+ packed = self._packed_bit_mask(shape, 0xAA)
48+ bits = self._expand_bit_mask(packed, shape)
49+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, bits, 2.0)
50+ output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), packed.npu(), 2.0)
51+ self.assertRtolEqual(output_cpu.numpy(), output_npu.cpu().numpy(), 0.0001)
52+ 
53+ @SupportedDevices(['Ascend950'])
54+ def test_native_dropout_backward_scale_gt_one_fp16(self, device="npu"):
55+ """A5 kernel path fp16: golden computed in fp32 then cast to fp16 (0x55 mask pattern)."""
56+ torch.manual_seed(0)
57+ shape = (4, 32)
58+ grad_output = torch.arange(4 * 32, dtype=torch.float32).reshape(shape)
59+ packed = self._packed_bit_mask(shape, 0x55)
60+ bits = self._expand_bit_mask(packed, shape)
61+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, bits, 3.0).to(torch.float16)
62+ grad_fp16 = grad_output.to(torch.float16)
63+ output_npu = torch.ops.aten.native_dropout_backward(grad_fp16.npu(), packed.npu(), 3.0)
64+ self.assertRtolEqual(output_cpu.numpy(), output_npu.cpu().numpy(), 0.001)
65+ 
66+ @SupportedDevices(['Ascend950'])
67+ def test_native_dropout_backward_scale_gt_one_bf16(self, device="npu"):
68+ """A5 kernel path bf16: golden computed in fp32 then cast to bf16 (0x55 mask pattern)."""
69+ torch.manual_seed(0)
70+ shape = (4, 32)
71+ grad_output = torch.arange(4 * 32, dtype=torch.float32).reshape(shape)
72+ packed = self._packed_bit_mask(shape, 0x55)
73+ bits = self._expand_bit_mask(packed, shape)
74+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, bits, 3.0).to(torch.bfloat16)
75+ grad_bf16 = grad_output.to(torch.bfloat16)
76+ output_npu = torch.ops.aten.native_dropout_backward(grad_bf16.npu(), packed.npu(), 3.0)
77+ self.assertRtolEqual(output_cpu.float().numpy(), output_npu.cpu().float().numpy(), 0.004)
78+ 
79+ @SupportedDevices(['Ascend950'])
80+ def test_native_dropout_backward_empty_mask(self, device="npu"):
81+ """A5 branch: empty mask returns an empty result shaped like mask.sizes()."""
82+ torch.manual_seed(0)
83+ grad_output = torch.tensor(1.2, dtype=torch.float32)
84+ mask = torch.zeros(0, dtype=torch.uint8)
85+ output_cpu = torch.ops.aten.native_dropout_backward(grad_output, mask, 2.0)
86+ output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), mask.npu(), 2.0)
87+ self.assertRtolEqual(output_cpu.numpy(), output_npu.cpu().numpy(), 0.0001)
88+ 
89+ @SupportedDevices(['Ascend950'])
90+ def test_neg_scale_range(self, device="npu"):
91+ """scale must be 0 or >= 1 (TORCH_CHECK on the A5 branch)."""
92+ torch.manual_seed(0)
93+ grad_output = torch.randn((4, 32), dtype=torch.float32).npu()
94+ mask = self._packed_bit_mask((4, 32), 0xAA).npu()
95+ with self.assertRaisesRegex(RuntimeError, "scale has to be 0"):
96+ torch.ops.aten.native_dropout_backward(grad_output, mask, 0.5)
97+ 
8 def _test_native_dropout_backward(self, dtype, p):98 def _test_native_dropout_backward(self, dtype, p):
9- grad_output = torch.tensor(1.2,dtype=dtype)99+ grad_output = torch.tensor(1.2, dtype=dtype)
10- b = np.random.randint(0,100,size=(0)).astype(np.uint8)100+ b = np.random.randint(0, 100, size=(0)).astype(np.uint8)
11 mask = torch.tensor(b).to(torch.uint8)101 mask = torch.tensor(b).to(torch.uint8)
12 output_cpu = torch.ops.aten.native_dropout_backward(grad_output, mask, p)102 output_cpu = torch.ops.aten.native_dropout_backward(grad_output, mask, p)
13 output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), mask.npu(), p)103 output_npu = torch.ops.aten.native_dropout_backward(grad_output.npu(), mask.npu(), p)
14 self.assertEqual(output_cpu, output_npu)104 self.assertEqual(output_cpu, output_npu)
15 105 
106+ def _packed_bit_mask(self, grad_shape, pattern):
107+ """Build a UINT8 bit mask with align(numel(grad), 128) / 8 elements, LSB-first per byte."""
108+ numel = 1
109+ for s in grad_shape:
110+ numel *= s
111+ packed_len = (numel + 127) // 128 * 16
112+ return torch.tensor([pattern] * packed_len, dtype=torch.uint8)
16 113 
17- def test_native_dropout_backward_fp32(self):114+ def _expand_bit_mask(self, packed_mask, grad_shape):
18- self._test_native_dropout_backward(torch.float32, 2)115+ """Expand a packed UINT8 bit mask to a bool mask of grad_shape (LSB-first, CPU golden side)."""
116+ numel = 1
117+ for s in grad_shape:
118+ numel *= s
119+ bits = (packed_mask.unsqueeze(1) >> torch.arange(8, dtype=torch.uint8)) & 1
120+ return bits.bool().reshape(-1)[:numel].reshape(grad_shape)
19 121 
20- def test_native_dropout_backward_fp16(self):
21- self._test_native_dropout_backward(torch.float16, 3)
22 122 
23if __name__ == '__main__':123if __name__ == '__main__':
24 run_tests()124 run_tests()