From 825c03d376a59ec52400010713bcae13976486e4 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Mon, 28 Sep 2026 19:59:26 -0400 Subject: [PATCH 1/5] fix(onnx): move a negative DequantizeLinear axis when transposing the weight transpose_dequantize_linear_weights swapped the axis attribute by matching the literal values 0 and 1. ONNX also allows a negative axis, and -1 / -2 matched neither, so the weight was transposed while the axis kept pointing at the dimension it had before. Checked against NumPy on a [K, N] weight, per-channel scale, for each legal 2D axis. On a non-square weight the converted model no longer runs; on a square one the scale length still matches, nothing complains, and the output is simply wrong. Normalize the axis before swapping it, which covers every legal 2D value. The regression test runs onnxruntime with graph optimization disabled: onnxruntime's optimizer does not treat a negative axis the same as its non-negative equivalent, so an optimized session is not a trustworthy baseline here. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/graph_surgery/dq_transpose.py | 8 +- .../quantization/test_dq_transpose_surgery.py | 87 +++++++++++++++++++ 3 files changed, 91 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f621df5dea..6af10bed29e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -94,6 +94,7 @@ Changelog - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - 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. - 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 the column-major ``DequantizeLinear`` surgery leaving a negative ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/graph_surgery/dq_transpose.py b/modelopt/onnx/graph_surgery/dq_transpose.py index 9ae9029b087..a22e6cfa4b9 100644 --- a/modelopt/onnx/graph_surgery/dq_transpose.py +++ b/modelopt/onnx/graph_surgery/dq_transpose.py @@ -301,14 +301,12 @@ 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. 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) if verbose: logger.debug(f" Updated axis: {old_axis} -> {attr.i}") break diff --git a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py index 68b28a84775..d67808df4d9 100644 --- a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py +++ b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py @@ -172,3 +172,90 @@ 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 _build_dq_matmul_model_with_axis(axis): + """Same graph as above, but with the given (possibly negative) quantization axis.""" + rng = np.random.RandomState(123) + w_int8 = rng.randint(-127, 127, size=(IN_FEATURES, OUT_FEATURES)).astype(np.int8) + channels = IN_FEATURES if axis % 2 == 0 else OUT_FEATURES + scale = rng.rand(channels).astype(np.float32) * 0.01 + 0.001 + + graph = helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", + inputs=["weight_quantized", "weight_scale", "weight_zp"], + outputs=["weight_dequantized"], + name="dq_weight", + axis=axis, + ), + helper.make_node( + "MatMul", + inputs=["input", "weight_dequantized"], + outputs=["output"], + name="matmul_0", + ), + ], + "test_dq_transpose_axis", + inputs=[ + helper.make_tensor_value_info("input", TensorProto.FLOAT, [BATCH, SEQ, IN_FEATURES]) + ], + outputs=[ + helper.make_tensor_value_info("output", TensorProto.FLOAT, [BATCH, SEQ, OUT_FEATURES]) + ], + initializer=[ + numpy_helper.from_array(w_int8, name="weight_quantized"), + numpy_helper.from_array(scale, name="weight_scale"), + numpy_helper.from_array(np.zeros(channels, dtype=np.int8), name="weight_zp"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = 10 + onnx.checker.check_model(model) + return model + + +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_with_axis(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, + ) From ea6f8bbcab74ed0ba3f779e1ebd95ad47d681589 Mon Sep 17 00:00:00 2001 From: SID-6921 <99672439+SID-6921@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:04:34 -0400 Subject: [PATCH 2/5] Consolidate duplicate DQ-transpose test model builders _build_dq_matmul_model() and _build_dq_matmul_model_with_axis(axis) built the same graph, differing only in whether axis was fixed at 0 or parametrized, and in an accidental ir_version drift (9 vs 10). Add an axis parameter (default 0) to the former and drop the latter, so both TestDqTransposeSurgery and test_negative_axis_is_swapped share one builder. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- .../quantization/test_dq_transpose_surgery.py | 60 ++++--------------- 1 file changed, 10 insertions(+), 50 deletions(-) diff --git a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py index d67808df4d9..5b3ecf7d770 100644 --- a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py +++ b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py @@ -32,19 +32,21 @@ SEQ = 4 -def _build_dq_matmul_model(): +def _build_dq_matmul_model(axis: int = 0): """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. """ rng = np.random.RandomState(123) # 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 scale along `axis` -> one value per channel on that axis + channels = IN_FEATURES if axis % 2 == 0 else OUT_FEATURES + scale = rng.rand(channels).astype(np.float32) * 0.01 + 0.001 # Per-channel zero point - zp = np.zeros(IN_FEATURES, dtype=np.int8) + 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") @@ -55,7 +57,7 @@ def _build_dq_matmul_model(): inputs=["weight_quantized", "weight_scale", "weight_zp"], outputs=["weight_dequantized"], name="dq_weight", - axis=0, + axis=axis, ) matmul_node = helper.make_node( @@ -78,7 +80,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 @@ -174,48 +176,6 @@ def test_output_matches(self, models): np.testing.assert_allclose(orig_out, mod_out, atol=1e-5, rtol=1e-5) -def _build_dq_matmul_model_with_axis(axis): - """Same graph as above, but with the given (possibly negative) quantization axis.""" - rng = np.random.RandomState(123) - w_int8 = rng.randint(-127, 127, size=(IN_FEATURES, OUT_FEATURES)).astype(np.int8) - channels = IN_FEATURES if axis % 2 == 0 else OUT_FEATURES - scale = rng.rand(channels).astype(np.float32) * 0.01 + 0.001 - - graph = helper.make_graph( - [ - helper.make_node( - "DequantizeLinear", - inputs=["weight_quantized", "weight_scale", "weight_zp"], - outputs=["weight_dequantized"], - name="dq_weight", - axis=axis, - ), - helper.make_node( - "MatMul", - inputs=["input", "weight_dequantized"], - outputs=["output"], - name="matmul_0", - ), - ], - "test_dq_transpose_axis", - inputs=[ - helper.make_tensor_value_info("input", TensorProto.FLOAT, [BATCH, SEQ, IN_FEATURES]) - ], - outputs=[ - helper.make_tensor_value_info("output", TensorProto.FLOAT, [BATCH, SEQ, OUT_FEATURES]) - ], - initializer=[ - numpy_helper.from_array(w_int8, name="weight_quantized"), - numpy_helper.from_array(scale, name="weight_scale"), - numpy_helper.from_array(np.zeros(channels, dtype=np.int8), name="weight_zp"), - ], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) - model.ir_version = 10 - onnx.checker.check_model(model) - return model - - def _run_unoptimized(model_proto, feeds): """Run with graph optimization off. @@ -234,7 +194,7 @@ def _run_unoptimized(model_proto, 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_with_axis(axis) + 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) From bc652f1b83209cc3ab49557abff206f198253e55 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:51:38 -0400 Subject: [PATCH 3/5] fix(onnx): correct the default (unset) DequantizeLinear axis too Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- modelopt/onnx/graph_surgery/dq_transpose.py | 10 ++++- .../quantization/test_dq_transpose_surgery.py | 40 +++++++++++++++++-- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dq_transpose.py b/modelopt/onnx/graph_surgery/dq_transpose.py index a22e6cfa4b9..4b245c85f0e 100644 --- a/modelopt/onnx/graph_surgery/dq_transpose.py +++ b/modelopt/onnx/graph_surgery/dq_transpose.py @@ -302,14 +302,22 @@ def transpose_dequantize_linear_weights( stats["zero_points_transposed"] += 1 # Update axis attribute (0 -> 1 after transpose). The weight is 2D here and ONNX - # allows a negative axis, so normalize before swapping. + # 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. + axis_attr_found = False for attr in dq_node.attribute: if attr.name == "axis": old_axis = attr.i 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: + dq_node.attribute.append(helper.make_attribute("axis", 0)) + if verbose: + logger.debug(" Added axis: 1 (default) -> 0") # 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 5b3ecf7d770..30724d4f425 100644 --- a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py +++ b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py @@ -32,18 +32,21 @@ SEQ = 4 -def _build_dq_matmul_model(axis: int = 0): +def _build_dq_matmul_model(axis: int | None = 0): """Build minimal model: input -> MatMul(input, DQ(weight, scale, zp)) -> output. 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. + [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. """ 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` -> one value per channel on that axis - channels = IN_FEATURES if axis % 2 == 0 else OUT_FEATURES + 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) @@ -52,12 +55,13 @@ def _build_dq_matmul_model(axis: int = 0): 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=axis, + **dq_kwargs, ) matmul_node = helper.make_node( @@ -219,3 +223,31 @@ def test_negative_axis_is_swapped(axis, expected_axis, tmp_path): 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, + ) From f62f877bbccad9e57d30543b2e63ea5199cc0e1e Mon Sep 17 00:00:00 2001 From: SID-6921 <99672439+SID-6921@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:26:55 -0400 Subject: [PATCH 4/5] Only add a default DQ axis when the scale is per-axis A DQ node with no explicit axis attribute got one appended unconditionally. Per-tensor DQ (scalar scale) ignores axis and, for opset < 13, does not even define the attribute, so adding it there made an otherwise-valid node invalid for its own opset. Only add the default axis when the scale has more than one element, which is only possible at opset 13+ since per-axis quantization is an opset-13 addition. Explicit axis attributes are untouched either way. Added test_per_tensor_dq_axis_is_not_added, confirmed it fails against the prior unconditional-add behavior. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/onnx/graph_surgery/dq_transpose.py | 16 ++++-- .../quantization/test_dq_transpose_surgery.py | 49 ++++++++++++++++--- 3 files changed, 55 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6fe24598bbb..d9e9d6b338f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -194,7 +194,7 @@ Changelog - 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. -- Fix the column-major ``DequantizeLinear`` surgery leaving a negative ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. +- Fix the column-major ``DequantizeLinear`` surgery leaving a negative ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output; the same fix also corrects a DQ node that relies on the ONNX-spec default axis instead of setting it explicitly, without adding an axis attribute to a per-tensor DQ node, which has none. - Fix the ``evaluation`` agent skill's LiveCodeBench recipe scoring correct solutions that use binary standard I/O (``sys.stdin.buffer`` / ``sys.stdout.buffer``) as failures: the recipe now pins a fixed LiveCodeBench judge in the NeMo-Skills container. LiveCodeBench scores change as a result, so re-run earlier baselines before comparing. 0.47.0 (2026-09-23) diff --git a/modelopt/onnx/graph_surgery/dq_transpose.py b/modelopt/onnx/graph_surgery/dq_transpose.py index 4b245c85f0e..369c647d5e4 100644 --- a/modelopt/onnx/graph_surgery/dq_transpose.py +++ b/modelopt/onnx/graph_surgery/dq_transpose.py @@ -304,7 +304,11 @@ def transpose_dequantize_linear_weights( # 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. + # 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": @@ -315,9 +319,13 @@ def transpose_dequantize_linear_weights( logger.debug(f" Updated axis: {old_axis} -> {attr.i}") break if not axis_attr_found: - dq_node.attribute.append(helper.make_attribute("axis", 0)) - if verbose: - logger.debug(" Added axis: 1 (default) -> 0") + 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 30724d4f425..933d4615014 100644 --- a/tests/unit/onnx/quantization/test_dq_transpose_surgery.py +++ b/tests/unit/onnx/quantization/test_dq_transpose_surgery.py @@ -32,12 +32,14 @@ SEQ = 4 -def _build_dq_matmul_model(axis: int | None = 0): +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 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. + 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) @@ -45,11 +47,15 @@ def _build_dq_matmul_model(axis: int | None = 0): # 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` -> 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) + 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") @@ -225,6 +231,35 @@ def test_negative_axis_is_swapped(axis, expected_axis, tmp_path): ) +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. From 4e2d309b9de1b3bb16f28f3788d9d0b450befeeb Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:36:12 -0400 Subject: [PATCH 5/5] Move and trim CHANGELOG entry per CONTRIBUTING guidance Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f51429f96e5..3dba418ad33 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -64,6 +64,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) ^^^^^^^^^^^^^^^^^^^ @@ -258,7 +259,6 @@ Changelog - 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. -- Fix the column-major ``DequantizeLinear`` surgery leaving a negative ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output; the same fix also corrects a DQ node that relies on the ONNX-spec default axis instead of setting it explicitly, without adding an axis attribute to a per-tensor DQ node, which has none. - Fix the ``evaluation`` agent skill's LiveCodeBench recipe scoring correct solutions that use binary standard I/O (``sys.stdin.buffer`` / ``sys.stdout.buffer``) as failures: the recipe now pins a fixed LiveCodeBench judge in the NeMo-Skills container. LiveCodeBench scores change as a result, so re-run earlier baselines before comparing. - Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. - Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``.