diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c45d8025386..e3b258fd84b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -65,6 +65,7 @@ Changelog - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. +- Fix the column-major ``DequantizeLinear`` surgery leaving a negative or default ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. Re-run the surgery on affected models. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/graph_surgery/dq_transpose.py b/modelopt/onnx/graph_surgery/dq_transpose.py index 9ae9029b087..369c647d5e4 100644 --- a/modelopt/onnx/graph_surgery/dq_transpose.py +++ b/modelopt/onnx/graph_surgery/dq_transpose.py @@ -301,17 +301,31 @@ def transpose_dequantize_linear_weights( break stats["zero_points_transposed"] += 1 - # Update axis attribute (0 -> 1 after transpose) + # Update axis attribute (0 -> 1 after transpose). The weight is 2D here and ONNX + # allows a negative axis, so normalize before swapping. The "axis" attribute is + # optional and defaults to 1 per the ONNX spec, so a DQ node that never set it + # explicitly needs the same correction, not a silent no-op -- but only when the + # scale is per-axis (more than one element). A per-tensor DQ ignores axis entirely, + # and opset < 13 does not define the attribute at all, so adding one there would + # make the node invalid for that opset; scale_size > 1 only happens at opset >= 13, + # since per-axis quantization is an opset-13 addition. + axis_attr_found = False for attr in dq_node.attribute: if attr.name == "axis": old_axis = attr.i - if old_axis == 0: - attr.i = 1 - elif old_axis == 1: - attr.i = 0 + attr.i = 1 - (old_axis % 2) + axis_attr_found = True if verbose: logger.debug(f" Updated axis: {old_axis} -> {attr.i}") break + if not axis_attr_found: + scale_size = int(np.prod(scale_init.dims)) + if scale_size > 1: + dq_node.attribute.append(helper.make_attribute("axis", 0)) + if verbose: + logger.debug(" Added axis: 1 (default) -> 0") + elif verbose: + logger.debug(" Per-tensor DQ with no axis attribute; leaving as-is") # Create intermediate output for DQ (transposed shape) dq_output_name = dq_node.output[0] diff --git a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py index 68b28a84775..933d4615014 100644 --- a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py +++ b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py @@ -32,30 +32,42 @@ SEQ = 4 -def _build_dq_matmul_model(): +def _build_dq_matmul_model(axis: int | None = 0, per_tensor: bool = False): """Build minimal model: input -> MatMul(input, DQ(weight, scale, zp)) -> output. - Uses per-channel INT8 quantization on axis 0 of the weight [IN, OUT]. + Uses per-channel INT8 quantization on the given (possibly negative) axis of the weight + [IN, OUT]. Defaults to axis 0, matching the ONNX export this surgery targets. Pass + axis=None to omit the attribute entirely and rely on the ONNX-spec default of 1. Pass + per_tensor=True for a scalar scale/zero-point instead, as in opset < 13 DQ nodes + (which do not have an axis attribute at all, since per-axis quantization is opset 13+). """ rng = np.random.RandomState(123) + effective_axis = 1 if axis is None else axis + # Quantized weight [IN_FEATURES, OUT_FEATURES] int8 w_int8 = rng.randint(-127, 127, size=(IN_FEATURES, OUT_FEATURES)).astype(np.int8) - # Per-channel scale along axis 0 -> shape [IN_FEATURES] (one per row) - scale = rng.rand(IN_FEATURES).astype(np.float32) * 0.01 + 0.001 - # Per-channel zero point - zp = np.zeros(IN_FEATURES, dtype=np.int8) + if per_tensor: + scale = np.array(rng.rand() * 0.01 + 0.001, dtype=np.float32) + zp = np.array(0, dtype=np.int8) + else: + # Per-channel scale along `axis` -> one value per channel on that axis + channels = IN_FEATURES if effective_axis % 2 == 0 else OUT_FEATURES + scale = rng.rand(channels).astype(np.float32) * 0.01 + 0.001 + # Per-channel zero point + zp = np.zeros(channels, dtype=np.int8) w_init = numpy_helper.from_array(w_int8, name="weight_quantized") s_init = numpy_helper.from_array(scale, name="weight_scale") zp_init = numpy_helper.from_array(zp, name="weight_zp") + dq_kwargs = {} if axis is None else {"axis": axis} dq_node = helper.make_node( "DequantizeLinear", inputs=["weight_quantized", "weight_scale", "weight_zp"], outputs=["weight_dequantized"], name="dq_weight", - axis=0, + **dq_kwargs, ) matmul_node = helper.make_node( @@ -78,7 +90,7 @@ def _build_dq_matmul_model(): ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) - model.ir_version = 9 + model.ir_version = 10 onnx.checker.check_model(model) return model @@ -172,3 +184,105 @@ def test_output_matches(self, models): print(f" Mean abs diff: {diff.mean():.6f}") np.testing.assert_allclose(orig_out, mod_out, atol=1e-5, rtol=1e-5) + + +def _run_unoptimized(model_proto, feeds): + """Run with graph optimization off. + + onnxruntime's optimizer does not treat a negative DequantizeLinear axis the same as its + non-negative equivalent, so leaving optimization on would compare against a wrong baseline. + """ + opts = ort.SessionOptions() + opts.log_severity_level = 3 + opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + sess = ort.InferenceSession( + model_proto.SerializeToString(), opts, providers=["CPUExecutionProvider"] + ) + return sess.run([o.name for o in sess.get_outputs()], feeds) + + +@pytest.mark.parametrize(("axis", "expected_axis"), [(0, 1), (1, 0), (-1, 0), (-2, 1)]) +def test_negative_axis_is_swapped(axis, expected_axis, tmp_path): + """ONNX allows a negative axis; it still has to move when the weight is transposed.""" + original = _build_dq_matmul_model(axis) + orig_path = str(tmp_path / "original.onnx") + out_path = str(tmp_path / "transposed.onnx") + onnx.save(original, orig_path) + + modified = transpose_dequantize_linear_weights( + model_path=orig_path, output_path=out_path, use_external_data=False, verbose=False + ) + + axis_val = next( + attr.i + for node in modified.graph.node + if node.op_type == "DequantizeLinear" + for attr in node.attribute + if attr.name == "axis" + ) + assert axis_val == expected_axis, f"axis {axis} should become {expected_axis}, got {axis_val}" + + feeds = {"input": np.random.RandomState(999).randn(BATCH, SEQ, IN_FEATURES).astype(np.float32)} + np.testing.assert_allclose( + _run_unoptimized(modified, feeds)[0], + _run_unoptimized(original, feeds)[0], + atol=1e-5, + rtol=1e-5, + ) + + +def test_per_tensor_dq_axis_is_not_added(tmp_path): + """A per-tensor DQ node (scalar scale) must not gain an axis attribute. + + Per-tensor quantization predates the axis attribute (opset < 13) and ignores axis + entirely, so adding one unconditionally -- as a DQ node with an explicit per-axis + scale needs -- would make an opset < 13 node invalid for its own opset. + """ + original = _build_dq_matmul_model(axis=None, per_tensor=True) + orig_path = str(tmp_path / "original.onnx") + out_path = str(tmp_path / "transposed.onnx") + onnx.save(original, orig_path) + + modified = transpose_dequantize_linear_weights( + model_path=orig_path, output_path=out_path, use_external_data=False, verbose=False + ) + + dq_node = next(n for n in modified.graph.node if n.op_type == "DequantizeLinear") + axis_attrs = [attr.i for attr in dq_node.attribute if attr.name == "axis"] + assert axis_attrs == [], f"expected no axis attribute on a per-tensor DQ, got {axis_attrs}" + + feeds = {"input": np.random.RandomState(999).randn(BATCH, SEQ, IN_FEATURES).astype(np.float32)} + np.testing.assert_allclose( + _run_unoptimized(modified, feeds)[0], + _run_unoptimized(original, feeds)[0], + atol=1e-5, + rtol=1e-5, + ) + + +def test_default_axis_is_corrected(tmp_path): + """A DQ node that never set "axis" relies on the ONNX-spec default of 1. + + The surgery must correct that implicit axis after transposing the weight, the same + as it would for an explicit axis=1, instead of silently leaving it unset. + """ + original = _build_dq_matmul_model(axis=None) + orig_path = str(tmp_path / "original.onnx") + out_path = str(tmp_path / "transposed.onnx") + onnx.save(original, orig_path) + + modified = transpose_dequantize_linear_weights( + model_path=orig_path, output_path=out_path, use_external_data=False, verbose=False + ) + + dq_node = next(n for n in modified.graph.node if n.op_type == "DequantizeLinear") + axis_attrs = [attr.i for attr in dq_node.attribute if attr.name == "axis"] + assert axis_attrs == [0], f"expected a single axis=0 attribute, got {axis_attrs}" + + feeds = {"input": np.random.RandomState(999).randn(BATCH, SEQ, IN_FEATURES).astype(np.float32)} + np.testing.assert_allclose( + _run_unoptimized(modified, feeds)[0], + _run_unoptimized(original, feeds)[0], + atol=1e-5, + rtol=1e-5, + )