From 2821e9d4e14a6a4d6584458259c24422aeb78a72 Mon Sep 17 00:00:00 2001 From: Yi-111-a <153097222+Yi-111-a@users.noreply.github.com> Date: Thu, 1 Oct 2026 09:15:02 +0800 Subject: [PATCH] fix(onnx): support per-tensor scales in qdq_to_dq _convert_weight treated every scale as per-axis: it compared scale.shape[0] against the weight's quantization axis and reshaped the scale along that axis. A per-tensor scale is the ONNX default for QuantizeLinear/DequantizeLinear and has scalar or single-element shape, so a scalar scale raised IndexError on scale.shape[0] and a (1,) scale raised a spurious shape mismatch. Either way the conversion of a per-tensor quantized weight failed. Decide per-tensor vs per-axis before any scale indexing: a scale with a single element is broadcast over the whole weight, and only a multi-element scale is validated against the axis length. The per-axis path and its error are unchanged. Signed-off-by: Yi-111-a <153097222+Yi-111-a@users.noreply.github.com> --- modelopt/onnx/quantization/qdq_utils.py | 12 +++- .../unit/onnx/quantization/test_qdq_utils.py | 64 +++++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py index d00105710355..993a4166604b 100644 --- a/modelopt/onnx/quantization/qdq_utils.py +++ b/modelopt/onnx/quantization/qdq_utils.py @@ -662,13 +662,19 @@ def _convert_weight( axis = axis_map[op_type] - if scale_array.shape and scale_array.shape[0] != weight_shape[axis]: + # A scalar (or single-element) scale is per-tensor quantization: one value broadcast over the + # whole weight, which is the ONNX default for QuantizeLinear/DequantizeLinear. Only a scale + # with one entry per channel is per-axis, and only that shape is checked against the axis. + is_per_tensor = scale_array.size == 1 + if not is_per_tensor and scale_array.shape[0] != weight_shape[axis]: raise ValueError( - f"Scale shape {scale_array.shape} does not match weight shape {weight_shape} along axis {axis}" + f"Scale shape {scale_array.shape} does not match weight shape {weight_shape} " + f"along axis {axis}" ) reshape_dims = [1] * len(weight_shape) - reshape_dims[axis] = scale_array.shape[0] + if not is_per_tensor: + reshape_dims[axis] = scale_array.shape[0] scale_array = scale_array.reshape(*reshape_dims) zp_array = zp_array.reshape(*reshape_dims) diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 95944e1f7ab3..77c2f53e15ad 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -1179,6 +1179,70 @@ def test_quantize_linear_without_inputs_raises_value_error(self): qdq_to_dq(model) +def create_test_model_with_qdq_matmul_weight(scale_shape, zero_point_shape=None): + """Create a QDQ model whose weight feeds a MatMul through a QuantizeLinear/DequantizeLinear pair. + + The scale and zero point take the shapes given, so callers can build per-tensor (scalar or + single-element) and per-axis (one entry per input channel) quantization graphs.""" + weight = np.random.RandomState(3).randn(16, 32).astype(np.float32) + amax = np.abs(weight).max() / 127.0 + scale = np.full(scale_shape, amax, dtype=np.float32) + if zero_point_shape is None: + zero_point = np.zeros(scale_shape, dtype=np.int8) + else: + zero_point = np.full(zero_point_shape, 0, dtype=np.int8) + + graph = helper.make_graph( + nodes=[ + helper.make_node("QuantizeLinear", ["W", "s", "z"], ["wq"], name="q"), + helper.make_node("DequantizeLinear", ["wq", "s", "z"], ["wdq"], name="dq"), + helper.make_node("MatMul", ["x", "wdq"], ["y"], name="mm"), + ], + name="test_graph", + inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 16])], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [4, 32])], + initializer=[ + numpy_helper.from_array(weight, "W"), + numpy_helper.from_array(scale, "s"), + numpy_helper.from_array(zero_point, "z"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = 10 + return model + + +def run_matmul_model(model): + """Run a MatMul model and return its output.""" + session = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + inputs = np.random.RandomState(7).randn(4, 16).astype(np.float32) + return session.run(None, {"x": inputs})[0] + + +class TestQdqToDqPerTensorScale: + """Per-tensor quantized weights must survive qdq_to_dq and stay numerically equivalent.""" + + @pytest.mark.parametrize("scale_shape", [(), (1,)], ids=["scalar", "single_element"]) + def test_per_tensor_scale_matches_qdq_output(self, scale_shape): + model = create_test_model_with_qdq_matmul_weight(scale_shape) + expected = run_matmul_model(model) + + converted = qdq_to_dq(model) + + assert "QuantizeLinear" not in [node.op_type for node in converted.graph.node] + converted_weight = numpy_helper.to_array( + next(init for init in converted.graph.initializer if init.name == "W") + ) + assert converted_weight.dtype == np.int8 + np.testing.assert_allclose(run_matmul_model(converted), expected, atol=1e-5) + + def test_per_axis_scale_still_rejected_when_axis_length_mismatches(self): + model = create_test_model_with_qdq_matmul_weight((8,)) + + with pytest.raises(RuntimeError, match=r"Scale shape \(8,\) does not match weight shape"): + qdq_to_dq(model) + + class TestLegacyEdgeLLMShims: """Smoke tests for the deprecated top-level shims kept for TensorRT-Edge-LLM 0.6.1.