已开启
fix: 修复非连续输入触发npu_weight_quant_batchmatmul的EZ1001报错(#184) #261
niwang66创建于 17 天前
fix: 修复非连续输入触发npu_weight_quant_batchmatmul的EZ1001报错(#184) #261
已开启
共 2 个文件变更+122-2
| @@ -204,8 +204,10 @@ class NpuWeightQuantizedLinear(nn.Module): | |||
| 204 | 204 | ||
| 205 | ori_shape = inputs.shape | 205 | ori_shape = inputs.shape |
| 206 | inputs = inputs.to(self.quantized_weight.device) | 206 | inputs = inputs.to(self.quantized_weight.device) |
| 207 | - # input shape reshape to 2d for npu op | 207 | + # input shape reshape to 2d for npu op; npu_weight_quant_batchmatmul |
| 208 | - inputs = inputs.reshape(-1, inputs.shape[-1]) | 208 | + # requires contiguous input, otherwise aclnn reports EZ1001 |
| 209 | + # ("only support x tensor is contiguous or transpose last two dims") | ||
| 210 | + inputs = inputs.reshape(-1, inputs.shape[-1]).contiguous() | ||
| 209 | if self.scale_factor is not None: | 211 | if self.scale_factor is not None: |
| 210 | inputs = torch.mul(inputs, self.scale_factor) | 212 | inputs = torch.mul(inputs, self.scale_factor) |
| 211 | 213 | ||
| @@ -0,0 +1,118 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# -*- coding: UTF-8 -*- | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 4 | +# | ||
| 5 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 6 | +# you may not use this file except in compliance with the License. | ||
| 7 | +# You may obtain a copy of the License at | ||
| 8 | +# | ||
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 10 | + | ||
| 11 | +# Unless required by applicable law or agreed to in writing, software | ||
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | +# ---------------------------------------------------------------------------- | ||
| 17 | +"""Unit tests for NpuWeightQuantizedLinear (torch_npu mocked). | ||
| 18 | + | ||
| 19 | +Regression guard for the aclnn EZ1001 bug: npu_weight_quant_batchmatmul | ||
| 20 | +requires the reshaped 2D input to be contiguous. A 3D input sliced along the | ||
| 21 | +last (K) dim stays non-contiguous after reshape(-1, K), so forward() must | ||
| 22 | +call .contiguous() explicitly. | ||
| 23 | +""" | ||
| 24 | + | ||
| 25 | +import sys | ||
| 26 | +import types | ||
| 27 | +import unittest | ||
| 28 | +from unittest.mock import MagicMock | ||
| 29 | + | ||
| 30 | +import torch | ||
| 31 | +import torch.nn as nn | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def _install_fake_torch_npu(): | ||
| 35 | + """Register a fake torch_npu recording the x handed to the batchmatmul op.""" | ||
| 36 | + mod = types.ModuleType("torch_npu") | ||
| 37 | + mod.hifloat8 = "hifloat8_enum" | ||
| 38 | + calls = [] | ||
| 39 | + | ||
| 40 | + def fake_batchmatmul(x, weight, scale, **kwargs): | ||
| 41 | + calls.append(x) | ||
| 42 | + return torch.zeros(*x.shape[:-1], weight.shape[-1], dtype=x.dtype) | ||
| 43 | + | ||
| 44 | + mod.npu_weight_quant_batchmatmul = fake_batchmatmul | ||
| 45 | + mod.npu_weight_quant_batchmatmul._schemas = {"": MagicMock()} | ||
| 46 | + mod.npu_weight_quant_batchmatmul._schemas[""].arguments = [ | ||
| 47 | + MagicMock(name="weight_dtype") | ||
| 48 | + ] | ||
| 49 | + sys.modules["torch_npu"] = mod | ||
| 50 | + return mod, calls | ||
| 51 | + | ||
| 52 | + | ||
| 53 | +class TestNpuWeightQuantizedLinear(unittest.TestCase): | ||
| 54 | + """forward() must pass a contiguous 2D tensor to the npu op.""" | ||
| 55 | + | ||
| 56 | + def setUp(self): | ||
| 57 | + self._saved_npu = getattr(torch.Tensor, "npu", None) | ||
| 58 | + self._saved_torch_npu = sys.modules.get("torch_npu") | ||
| 59 | + torch.Tensor.npu = lambda self: self | ||
| 60 | + _, self.calls = _install_fake_torch_npu() | ||
| 61 | + | ||
| 62 | + def tearDown(self): | ||
| 63 | + if self._saved_npu is not None: | ||
| 64 | + torch.Tensor.npu = self._saved_npu | ||
| 65 | + elif hasattr(torch.Tensor, "npu"): | ||
| 66 | + delattr(torch.Tensor, "npu") | ||
| 67 | + if self._saved_torch_npu is not None: | ||
| 68 | + sys.modules["torch_npu"] = self._saved_torch_npu | ||
| 69 | + else: | ||
| 70 | + sys.modules.pop("torch_npu", None) | ||
| 71 | + | ||
| 72 | + def _make_module(self, k=256, n=128, group=64): | ||
| 73 | + qm = MagicMock(spec=nn.Linear) | ||
| 74 | + linear = nn.Linear(k, n, bias=False, dtype=torch.float32) | ||
| 75 | + qm.weight = linear.weight | ||
| 76 | + qm.bias = None | ||
| 77 | + qm.wts_type = "int8" | ||
| 78 | + qm.group_size = group | ||
| 79 | + qm.scale_w = torch.rand(n, k // group, 1) * 1e-3 + 1e-4 | ||
| 80 | + qm.offset_w = None | ||
| 81 | + qm.scale = None | ||
| 82 | + | ||
| 83 | + from amct_pytorch.classic.deploy_op.weight_npu_quant_module import ( | ||
| 84 | + NpuWeightQuantizedLinear, | ||
| 85 | + ) | ||
| 86 | + | ||
| 87 | + return NpuWeightQuantizedLinear(qm) | ||
| 88 | + | ||
| 89 | + def test_forward_input_to_op_is_contiguous(self): | ||
| 90 | + """Sliced 3D input stays non-contiguous after reshape; op needs contig.""" | ||
| 91 | + mod = self._make_module() | ||
| 92 | + # [B, S, 2K] sliced to [..., K]: reshape(-1, K) cannot be a view here | ||
| 93 | + # and the result of reshape on this layout is still non-contiguous. | ||
| 94 | + x = torch.randn(2, 16, 512)[:, :, :256] | ||
| 95 | + self.assertFalse( | ||
| 96 | + x.reshape(-1, x.shape[-1]).is_contiguous(), | ||
| 97 | + "precondition: reshape alone must not fix contiguity", | ||
| 98 | + ) | ||
| 99 | + | ||
| 100 | + out = mod(x) | ||
| 101 | + self.assertEqual(out.shape, (2, 16, 128)) | ||
| 102 | + self.assertEqual(len(self.calls), 1) | ||
| 103 | + self.assertTrue( | ||
| 104 | + self.calls[0].is_contiguous(), | ||
| 105 | + "npu_weight_quant_batchmatmul got a non-contiguous x (aclnn EZ1001)", | ||
| 106 | + ) | ||
| 107 | + self.assertEqual(self.calls[0].shape, (32, 256)) | ||
| 108 | + | ||
| 109 | + def test_forward_output_shape_restored(self): | ||
| 110 | + """Output is reshaped back to the original leading dims.""" | ||
| 111 | + mod = self._make_module() | ||
| 112 | + out = mod(torch.randn(3, 5, 256)) | ||
| 113 | + self.assertEqual(out.shape, (3, 5, 128)) | ||
| 114 | + self.assertTrue(self.calls[0].is_contiguous()) | ||
| 115 | + | ||
| 116 | + | ||
| 117 | +if __name__ == "__main__": | ||
| 118 | + unittest.main() | ||