fix(onnx): move a negative DequantizeLinear axis when transposing the weight

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)
 ^^^^^^^^^^^^^^^^^^^

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]

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,
+    )