fix(onnx): trace Transpose between DequantizeLinear and its consumer in qdq_to_dq

CHANGELOG.rst

@@ -163,6 +163,7 @@ Changelog
 - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths.
 - Fix YAML config I/O decoding with the locale codepage instead of UTF-8, which made a config containing any non-ASCII byte fail to load on a machine whose locale is not UTF-8 (notably Windows, where the default is cp1252). ``modelopt/recipe/loader.py``, the two ONNX autotune state files, the two transformers config readers, the distill config and the puzzletron profile now pass ``encoding="utf-8"`` explicitly. Only the YAML config paths are covered: these are the files most likely to carry non-ASCII text in comments, model names or paths, and the only ones read inside a user's process.
 - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand.
+- Fix ONNX export of quantized models failing with ``Unsupported op_type for real weight quantization: Transpose`` when a weight reaches its consumer transposed, as in MaxViT's attention blocks. The same export now also quantizes each weight along the axis its Q/DQ pair declares, correcting silently wrong ``dq_only`` output when that axis differed from the consumer's default.
 
 0.47.0 (2026-09-23)
 ^^^^^^^^^^^^^^^^^^^

modelopt/onnx/quantization/qdq_utils.py

@@ -577,13 +577,22 @@ def _get_scale_and_zp(
     return scale, zp
 
 
+# Consumers whose quantized weights qdq_to_dq knows how to convert.
+_REAL_QUANT_WEIGHT_CONSUMERS = {"Conv", "ConvTranspose", "Gemm", "MatMul"}
+
+
+def _quantization_axis(node: onnx.NodeProto) -> int:
+    """Returns a Q/DQ node's quantization axis; ONNX defaults the attribute to 1."""
+    return next((attr.i for attr in node.attribute if attr.name == "axis"), 1)
+
+
 def _get_successive_consumers(
     node: onnx.NodeProto, tensor_consumers: dict[str, list[onnx.NodeProto]]
 ) -> tuple[onnx.NodeProto, onnx.NodeProto]:
     """Get the DequantizeLinear node and its consumer node for a given QuantizeLinear node.
 
     This function validates and retrieves the next two nodes in the quantization chain:
-    QuantizeLinear -> DequantizeLinear -> Operation
+    QuantizeLinear -> DequantizeLinear -> [Cast | Transpose]* -> Operation
 
     Args:
         node: The QuantizeLinear node to find consumers for
@@ -601,10 +610,13 @@ def _get_successive_consumers(
     quantized_node = tensor_consumers.get(dq_node.output[0], [None])[0]
     if not quantized_node:
         raise ValueError(f"No consumer found for {dq_node.name}")
-    if quantized_node.op_type == "Cast":
+
+    while quantized_node.op_type in ("Cast", "Transpose"):
         next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
         if not next_node:
-            raise ValueError(f"No consumer found after Cast for {quantized_node.name}")
+            raise ValueError(
+                f"No consumer found after {quantized_node.op_type} for {quantized_node.name}"
+            )
         quantized_node = next_node
 
     return dq_node, quantized_node
@@ -615,6 +627,8 @@ def _convert_weight(
     scale: onnx.TensorProto,
     zp: onnx.TensorProto,
     quantized_node: onnx.NodeProto,
+    *,
+    dq_axis: int | None = None,
 ) -> np.ndarray:
     """Convert a weight tensor to INT8/FP8 format based on scale and zero point.
 
@@ -623,6 +637,7 @@ def _convert_weight(
         scale: The scale tensor for quantization
         zp: The zero point tensor for quantization
         quantized_node: The operation node that will use the converted weight
+        dq_axis: The axis declared by the source Q/DQ pair, if it declares one
 
     Returns:
         The converted weight tensor as a numpy array
@@ -634,41 +649,33 @@ def _convert_weight(
         - INT8 weights are clipped to [-128, 127]
         - FP8 weights use float8e4m3fn format
     """
-    # Per-op quantization axis mapping (must match ORT config)
     weight_shape = weight_array.shape
     op_type = quantized_node.op_type
 
     # Convert onnx tensors to numpy array
     scale_array = onnx.numpy_helper.to_array(scale)
     zp_array = onnx.numpy_helper.to_array(zp)
 
-    # Dynamically determine transB for Gemm
-    trans_b = 0
-    if op_type == "Gemm":
-        for attr in quantized_node.attribute:
-            if attr.name == "transB":
-                trans_b = attr.i
-                break
-
-    axis_map = {
-        "Conv": 0,
-        "ConvTranspose": 1,
-        "Gemm": 0 if trans_b else 1,
-        "MatMul": 1,
-    }
-
-    if op_type not in axis_map:
+    if op_type not in _REAL_QUANT_WEIGHT_CONSUMERS:
         raise ValueError(f"Unsupported op_type for real weight quantization: {op_type}")
 
-    axis = axis_map[op_type]
-
-    if scale_array.shape 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}"
-        )
-
+    # A single scale covers the whole weight, so it broadcasts and no axis applies to it.
     reshape_dims = [1] * len(weight_shape)
-    reshape_dims[axis] = scale_array.shape[0]
+    if scale_array.size > 1:
+        # Per-axis: the DequantizeLinear left in the graph dequantizes along the axis it
+        # declares, and that axis indexes the stored weight whatever the consumer's layout is.
+        axis = 1 if dq_axis is None else dq_axis
+        if not -len(weight_shape) <= axis < len(weight_shape):
+            raise ValueError(
+                f"Quantization axis {axis} is out of range for weight shape {weight_shape}"
+            )
+        axis %= len(weight_shape)
+        if 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}"
+            )
+        reshape_dims[axis] = scale_array.shape[0]
+
     scale_array = scale_array.reshape(*reshape_dims)
     zp_array = zp_array.reshape(*reshape_dims)
 
@@ -756,8 +763,16 @@ def qdq_to_dq(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
             # Validate Q->DQ->Op pattern and get consumers
             dq_node, quantized_node = _get_successive_consumers(node, tensor_consumers)
 
-            # Convert weight
-            scaled = _convert_weight(weight_array, scale, zp, quantized_node)
+            # Convert weight. The conversion follows the DQ and drops the Q, so a per-axis
+            # pair that disagrees on the axis would silently change the output.
+            dq_axis, q_axis = _quantization_axis(dq_node), _quantization_axis(node)
+            if int(np.prod(scale.dims)) > 1 and (
+                q_axis % weight_array.ndim != dq_axis % weight_array.ndim
+            ):
+                raise ValueError(
+                    f"QuantizeLinear axis {q_axis} disagrees with DequantizeLinear axis {dq_axis}"
+                )
+            scaled = _convert_weight(weight_array, scale, zp, quantized_node, dq_axis=dq_axis)
 
             # Create and update new weight tensor
             if zp.data_type == onnx_dtype_map["Float8"]:

tests/unit/onnx/quantization/test_qdq_utils.py

@@ -1162,6 +1162,203 @@ def test_constant_node_scale_path_still_patched(self):
         assert (scale_arr > 0).all()
 
 
+def create_test_model_with_dq_transpose_matmul(
+    out_features: int,
+    in_features: int,
+    *,
+    perm: tuple[int, ...] | None = (1, 0),
+    q_axis: int = 0,
+    transpose: bool = True,
+    declare_axis: bool = True,
+    single_scale: bool = False,
+):
+    """Create weight[out, in] -> Q -> DQ -> Transpose -> MatMul, per-channel on the out axis.
+
+    This is how a torch ``Linear`` exports when the weight reaches MatMul transposed, as in
+    MaxViT's attention blocks. ``perm=None`` omits the attribute, which ONNX defines as
+    reversing the axes.
+    """
+    rng = np.random.RandomState(7)
+    weight = rng.randn(out_features, in_features).astype(np.float32)
+    if single_scale:
+        # One element: covers the whole weight and broadcasts, whatever the axis says.
+        scale = np.array([np.abs(weight).max() / 127.0], dtype=np.float32)
+        zero_point = np.zeros(1, dtype=np.int8)
+    else:
+        reduced = tuple(axis for axis in range(2) if axis != q_axis)
+        scale = (np.abs(weight).max(axis=reduced) / 127.0).astype(np.float32)
+        zero_point = np.zeros(weight.shape[q_axis], dtype=np.int8)
+
+    # ONNX defaults an omitted `axis` to 1, so a caller omitting it must pass q_axis=1.
+    axis_attr = {"axis": q_axis} if declare_axis else {}
+    nodes = [
+        helper.make_node(
+            "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q", **axis_attr
+        ),
+        helper.make_node(
+            "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq", **axis_attr
+        ),
+    ]
+    matmul_rhs = "w_dq"
+    if transpose:
+        nodes.append(
+            helper.make_node(
+                "Transpose",
+                ["w_dq"],
+                ["w_t"],
+                name="w_t",
+                **({} if perm is None else {"perm": list(perm)}),
+            )
+        )
+        matmul_rhs = "w_t"
+    nodes.append(helper.make_node("MatMul", ["input", matmul_rhs], ["output"], name="matmul"))
+
+    rows, cols = (in_features, out_features) if transpose else (out_features, in_features)
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="dq_transpose_matmul",
+        inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [3, rows])],
+        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [3, cols])],
+        initializer=[
+            numpy_helper.from_array(weight, "weight"),
+            numpy_helper.from_array(scale, "w_scale"),
+            numpy_helper.from_array(zero_point, "w_zp"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+    model.ir_version = 10
+    onnx.checker.check_model(model)
+    return model
+
+
+class TestQdqToDqTranspose:
+    """qdq_to_dq must trace a Transpose and follow the Q/DQ pair's own quantization axis."""
+
+    # Square weights are the dangerous shape: the scale length matches either axis, so a
+    # wrong axis passes validation and silently changes the model's output.
+    @pytest.mark.parametrize(
+        ("out_features", "in_features", "q_axis", "transpose", "declare_axis"),
+        [
+            (96, 32, 0, True, True),
+            (32, 32, 0, True, True),
+            (32, 32, 0, False, True),
+            (32, 32, 1, True, True),
+            (32, 32, 1, True, False),
+        ],
+    )
+    def test_conversion_matches_qdq(
+        self, out_features, in_features, q_axis, transpose, declare_axis
+    ):
+        """The DequantizeLinear left in the graph dequantizes along its own axis."""
+        model = create_test_model_with_dq_transpose_matmul(
+            out_features,
+            in_features,
+            q_axis=q_axis,
+            transpose=transpose,
+            declare_axis=declare_axis,
+        )
+        rows = in_features if transpose else out_features
+        inputs = {"input": np.random.RandomState(11).randn(3, rows).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+
+        assert not [n for n in converted.graph.node if n.op_type == "QuantizeLinear"]
+        weight = next(t for t in converted.graph.initializer if t.name == "weight")
+        assert weight.data_type == TensorProto.INT8
+        assert np.array_equal(run(converted), expected)
+
+    def test_dq_transpose_without_perm_is_traced(self):
+        """ONNX lets Transpose omit `perm`; the MatMul behind it must still be found."""
+        model = create_test_model_with_dq_transpose_matmul(96, 32, perm=None)
+        initializer = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}
+
+        converted = qdq_to_dq(model)
+
+        # onnxruntime cannot execute DQ -> Transpose(no perm) -> MatMul, so check the
+        # converted weight directly. The Q/DQ pair declares axis 0.
+        expected = np.clip(
+            np.round(initializer["weight"] / initializer["w_scale"][:, None]), -128, 127
+        ).astype(np.int8)
+        actual = next(t for t in converted.graph.initializer if t.name == "weight")
+        assert np.array_equal(numpy_helper.to_array(actual), expected)
+
+    def test_one_element_scale_is_not_per_axis(self):
+        """A one-element scale is broadcast over the whole weight, not indexed per axis."""
+        # A single-output-channel Conv: reading the (1,) scale as per-axis would pick the
+        # DQ's default axis 1 and reject a graph that converts fine.
+        weight = np.random.RandomState(3).randn(1, 3, 3, 3).astype(np.float32)
+        graph = helper.make_graph(
+            nodes=[
+                helper.make_node(
+                    "QuantizeLinear", ["weight", "w_scale", "w_zp"], ["w_q"], name="w_q"
+                ),
+                helper.make_node(
+                    "DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["w_dq"], name="w_dq"
+                ),
+                helper.make_node(
+                    "Conv", ["input", "w_dq"], ["output"], name="conv", kernel_shape=[3, 3]
+                ),
+            ],
+            name="one_element_scale",
+            inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 8, 8])],
+            outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1, 6, 6])],
+            initializer=[
+                numpy_helper.from_array(weight, "weight"),
+                numpy_helper.from_array(
+                    np.array([np.abs(weight).max() / 127.0], np.float32), "w_scale"
+                ),
+                numpy_helper.from_array(np.array([0], np.int8), "w_zp"),
+            ],
+        )
+        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+        model.ir_version = 10
+        onnx.checker.check_model(model)
+        inputs = {"input": np.random.RandomState(5).randn(1, 3, 8, 8).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+        assert np.array_equal(run(converted), expected)
+
+    @pytest.mark.parametrize(("out_features", "in_features"), [(1, 32), (32, 1)])
+    def test_one_element_scale_broadcasts_through_transpose(self, out_features, in_features):
+        """A single scale covers the whole weight, so no axis of it has to be length 1."""
+        model = create_test_model_with_dq_transpose_matmul(
+            out_features, in_features, single_scale=True
+        )
+        inputs = {"input": np.random.RandomState(11).randn(3, in_features).astype(np.float32)}
+
+        def run(m):
+            return ort.InferenceSession(
+                m.SerializeToString(), providers=["CPUExecutionProvider"]
+            ).run(None, inputs)[0]
+
+        expected = run(model)
+        converted = qdq_to_dq(onnx.ModelProto.FromString(model.SerializeToString()))
+
+        assert not [n for n in converted.graph.node if n.op_type == "QuantizeLinear"]
+        assert np.array_equal(run(converted), expected)
+
+    def test_disagreeing_qdq_axes_raise(self):
+        """Following the DQ while the Q declares another axis would change the output."""
+        model = create_test_model_with_dq_transpose_matmul(32, 32, q_axis=0, transpose=False)
+        q_node = next(node for node in model.graph.node if node.name == "w_q")
+        next(attr for attr in q_node.attribute if attr.name == "axis").i = 1
+
+        with pytest.raises(RuntimeError, match="disagrees with DequantizeLinear axis"):
+            qdq_to_dq(model)
+
+
 class TestQdqToDqValidation:
     """Regression tests for qdq_to_dq input validation."""