From 43ff93f226488a838f8a9eb33611036500f372d5 Mon Sep 17 00:00:00 2001 From: Jai Prajapati Date: Fri, 21 Aug 2026 18:14:44 -0700 Subject: [PATCH 1/3] Add QDQ node exclusion regression test Signed-off-by: Jai Prajapati --- .../onnx/autocast/test_precisionconverter.py | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index b480bb7c24e..d5ac636bbd2 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy + import numpy as np import onnx import pytest @@ -2220,6 +2222,112 @@ def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary(): onnx.checker.check_model(converted, full_check=True) +def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): + graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) + initializers = [ + numpy_helper.from_array(np.full(4, 0.25, dtype=np.float32), name="q_scale"), + numpy_helper.from_array(np.zeros(4, dtype=np.uint8), name="q_zero_point"), + numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="pow_exponent"), + numpy_helper.from_array(np.array([1], dtype=np.int64), name="reduce_axes"), + numpy_helper.from_array(np.array(1e-6, dtype=np.float32), name="epsilon"), + numpy_helper.from_array(np.ones(4, dtype=np.float32), name="gamma"), + ] + nodes = [ + helper.make_node( + "QuantizeLinear", + ["X", "q_scale", "q_zero_point"], + ["X_quantized"], + name="input/QuantizeLinear", + axis=1, + ), + helper.make_node( + "DequantizeLinear", + ["X_quantized", "q_scale", "q_zero_point"], + ["X_dequantized"], + name="input/DequantizeLinear", + axis=1, + ), + helper.make_node( + "Pow", ["X_dequantized", "pow_exponent"], ["pow_out"], name="/rms/Pow" + ), + helper.make_node( + "ReduceMean", + ["pow_out", "reduce_axes"], + ["mean_out"], + name="/rms/ReduceMean", + keepdims=1, + ), + helper.make_node("Add", ["mean_out", "epsilon"], ["add_out"], name="/rms/Add"), + helper.make_node("Sqrt", ["add_out"], ["sqrt_out"], name="/rms/Sqrt"), + helper.make_node( + "Div", ["X_dequantized", "sqrt_out"], ["div_out"], name="/rms/Div" + ), + helper.make_node("Mul", ["div_out", "gamma"], ["Y"], name="/rms/Mul"), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "quantized_rmsnorm", + [graph_input], + [graph_output], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 19)], + ir_version=10, + ) + onnx.checker.check_model(model, full_check=True) + + qdq_before = { + node.name: node.SerializeToString() + for node in model.graph.node + if node.op_type in {"QuantizeLinear", "DequantizeLinear"} + } + qdq_initializers_before = { + initializer.name: initializer.SerializeToString() + for initializer in model.graph.initializer + if initializer.name in {"q_scale", "q_zero_point"} + } + opsets_before = [(opset.domain, opset.version) for opset in model.opset_import] + + converted = convert_to_f16( + deepcopy(model), + keep_io_types=True, + op_block_list=["QuantizeLinear", "DequantizeLinear"], + nodes_to_exclude=[r"^/rms/(Pow|ReduceMean|Add|Sqrt|Div)$"], + trt_plugins=[], + opset=19, + ) + + qdq_after = { + node.name: node.SerializeToString() + for node in converted.graph.node + if node.name in qdq_before + } + qdq_initializers_after = { + initializer.name: initializer.SerializeToString() + for initializer in converted.graph.initializer + if initializer.name in qdq_initializers_before + } + assert qdq_after == qdq_before + assert qdq_initializers_after == qdq_initializers_before + assert [(opset.domain, opset.version) for opset in converted.opset_import] == opsets_before + + value_types = { + value.name: value.type.tensor_type.elem_type + for value in ( + *converted.graph.input, + *converted.graph.output, + *converted.graph.value_info, + ) + } + assert value_types["X_quantized"] == TensorProto.UINT8 + assert value_types["X_dequantized"] == TensorProto.FLOAT + for output_name in ["pow_out", "mean_out", "add_out", "sqrt_out", "div_out"]: + assert value_types[output_name] == TensorProto.FLOAT + onnx.checker.check_model(converted, full_check=True) + + def test_convert_to_f16_refreshes_gathernd_pre_cast_declaration(monkeypatch): def discover_test_plugins_without_trt(self): self.custom_ops = { From 7c4a338d9b0e3d6a33a0b095fbcb0262171f3285 Mon Sep 17 00:00:00 2001 From: Jai Prajapati Date: Fri, 21 Aug 2026 18:18:34 -0700 Subject: [PATCH 2/3] Support node exclusions in QDQ autocast Signed-off-by: Jai Prajapati --- CHANGELOG.rst | 1 + modelopt/onnx/autocast/convert.py | 18 +++++++++++++++--- .../onnx/autocast/test_precisionconverter.py | 14 ++++++-------- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7be270b3edc..77a8fffe1b8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -23,6 +23,7 @@ Changelog *Misc* +- Add ``nodes_to_exclude`` regex support to the Q/DQ-aware ONNX ``convert_to_f16`` API, matching ``convert_to_mixed_precision`` node-name exclusion semantics while composing with the existing operation and tensor block lists. - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. - Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). A tracked run records the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries, with every command-line argument as a searchable param; failed runs are recorded with their traceback. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. - Add ``--mlflow `` to ``examples/vllm_serve/vllm_serve_fakequant.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a fake-quant serve records what it quantized and an evaluation of that endpoint can be traced back to a recipe. A tracked run uploads the launcher command, the resolved ``RECIPE_PATH`` (or the merged ``QUANT_CFG``/``KV_QUANT_CFG`` when presets are used), the worker log and the quantizer summary; the experiment defaults to ``$USER/vllm_serve_fakequant/-`` and can be overridden with ``--mlflow-experiment`` / ``--mlflow-run-name``. diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index d9bc8d68c99..65cabe86974 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -32,7 +32,11 @@ import modelopt.onnx.utils as onnx_utils from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.autocast.logging_config import logger -from modelopt.onnx.autocast.nodeclassifier import NodeClassifier, NodeRuleBase +from modelopt.onnx.autocast.nodeclassifier import ( + DisabledNodeNameRegexRule, + NodeClassifier, + NodeRuleBase, +) from modelopt.onnx.autocast.precisionconverter import PrecisionConverter from modelopt.onnx.autocast.referencerunner import ReferenceRunner from modelopt.onnx.utils import get_min_opset_for_precisions, get_qdq_precisions @@ -221,6 +225,7 @@ def convert_to_f16( trt_plugins: list[str] | None = [], use_standalone_type_inference: bool = False, opset: int | None = None, + nodes_to_exclude: list[str] | None = None, ) -> onnx.ModelProto: """Convert model to mixed precision, using PrecisionConverter. @@ -240,6 +245,7 @@ def convert_to_f16( (22 for bf16, 19 for fp16) and Q/DQ node requirements. The opset may be automatically increased if Q/DQ nodes in the model require a higher version (e.g., FP8 requires 19, INT4 requires 21, NVFP4 requires 23). + nodes_to_exclude: List of regex patterns to match node names that should remain in FP32. """ assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16" original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types) @@ -303,9 +309,15 @@ def convert_to_f16( use_standalone_type_inference=use_standalone_type_inference, original_network_io_metadata=original_network_io_metadata, ) - high_precision_nodes = [node.name for node in model.graph.node if node.op_type in op_block_list] + node_name_rule = DisabledNodeNameRegexRule(nodes_to_exclude or []) + high_precision_nodes = [ + node.name + for node in model.graph.node + if node.op_type in op_block_list or node_name_rule.check(node) + ] + high_precision_node_set = set(high_precision_nodes) low_precision_nodes = [ - node.name for node in model.graph.node if node.op_type not in op_block_list + node.name for node in model.graph.node if node.name not in high_precision_node_set ] model_mod = precision_converter.convert(high_precision_nodes, low_precision_nodes) return model_mod diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index d5ac636bbd2..6629b65bdf3 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -2228,6 +2228,8 @@ def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): initializers = [ numpy_helper.from_array(np.full(4, 0.25, dtype=np.float32), name="q_scale"), numpy_helper.from_array(np.zeros(4, dtype=np.uint8), name="q_zero_point"), + numpy_helper.from_array(np.full(4, 0.25, dtype=np.float32), name="dq_scale"), + numpy_helper.from_array(np.zeros(4, dtype=np.uint8), name="dq_zero_point"), numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="pow_exponent"), numpy_helper.from_array(np.array([1], dtype=np.int64), name="reduce_axes"), numpy_helper.from_array(np.array(1e-6, dtype=np.float32), name="epsilon"), @@ -2243,14 +2245,12 @@ def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): ), helper.make_node( "DequantizeLinear", - ["X_quantized", "q_scale", "q_zero_point"], + ["X_quantized", "dq_scale", "dq_zero_point"], ["X_dequantized"], name="input/DequantizeLinear", axis=1, ), - helper.make_node( - "Pow", ["X_dequantized", "pow_exponent"], ["pow_out"], name="/rms/Pow" - ), + helper.make_node("Pow", ["X_dequantized", "pow_exponent"], ["pow_out"], name="/rms/Pow"), helper.make_node( "ReduceMean", ["pow_out", "reduce_axes"], @@ -2260,9 +2260,7 @@ def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): ), helper.make_node("Add", ["mean_out", "epsilon"], ["add_out"], name="/rms/Add"), helper.make_node("Sqrt", ["add_out"], ["sqrt_out"], name="/rms/Sqrt"), - helper.make_node( - "Div", ["X_dequantized", "sqrt_out"], ["div_out"], name="/rms/Div" - ), + helper.make_node("Div", ["X_dequantized", "sqrt_out"], ["div_out"], name="/rms/Div"), helper.make_node("Mul", ["div_out", "gamma"], ["Y"], name="/rms/Mul"), ] model = helper.make_model( @@ -2286,7 +2284,7 @@ def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): qdq_initializers_before = { initializer.name: initializer.SerializeToString() for initializer in model.graph.initializer - if initializer.name in {"q_scale", "q_zero_point"} + if initializer.name in {"q_scale", "q_zero_point", "dq_scale", "dq_zero_point"} } opsets_before = [(opset.domain, opset.version) for opset in model.opset_import] From 646ed9bb504d07b2a6a4947adfcdc4d4dfe0253e Mon Sep 17 00:00:00 2001 From: Jai Prajapati Date: Mon, 24 Aug 2026 14:58:02 -0700 Subject: [PATCH 3/3] Simplify node exclusion regression test Signed-off-by: Jai Prajapati --- .../onnx/autocast/test_precisionconverter.py | 107 ++---------------- 1 file changed, 10 insertions(+), 97 deletions(-) diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index 6629b65bdf3..d5ed09215fc 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from copy import deepcopy - import numpy as np import onnx import pytest @@ -2222,107 +2220,22 @@ def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary(): onnx.checker.check_model(converted, full_check=True) -def test_convert_to_f16_excludes_named_rmsnorm_nodes_without_changing_qdq(): - graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) - graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) - initializers = [ - numpy_helper.from_array(np.full(4, 0.25, dtype=np.float32), name="q_scale"), - numpy_helper.from_array(np.zeros(4, dtype=np.uint8), name="q_zero_point"), - numpy_helper.from_array(np.full(4, 0.25, dtype=np.float32), name="dq_scale"), - numpy_helper.from_array(np.zeros(4, dtype=np.uint8), name="dq_zero_point"), - numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="pow_exponent"), - numpy_helper.from_array(np.array([1], dtype=np.int64), name="reduce_axes"), - numpy_helper.from_array(np.array(1e-6, dtype=np.float32), name="epsilon"), - numpy_helper.from_array(np.ones(4, dtype=np.float32), name="gamma"), - ] - nodes = [ - helper.make_node( - "QuantizeLinear", - ["X", "q_scale", "q_zero_point"], - ["X_quantized"], - name="input/QuantizeLinear", - axis=1, - ), - helper.make_node( - "DequantizeLinear", - ["X_quantized", "dq_scale", "dq_zero_point"], - ["X_dequantized"], - name="input/DequantizeLinear", - axis=1, - ), - helper.make_node("Pow", ["X_dequantized", "pow_exponent"], ["pow_out"], name="/rms/Pow"), - helper.make_node( - "ReduceMean", - ["pow_out", "reduce_axes"], - ["mean_out"], - name="/rms/ReduceMean", - keepdims=1, - ), - helper.make_node("Add", ["mean_out", "epsilon"], ["add_out"], name="/rms/Add"), - helper.make_node("Sqrt", ["add_out"], ["sqrt_out"], name="/rms/Sqrt"), - helper.make_node("Div", ["X_dequantized", "sqrt_out"], ["div_out"], name="/rms/Div"), - helper.make_node("Mul", ["div_out", "gamma"], ["Y"], name="/rms/Mul"), - ] - model = helper.make_model( - helper.make_graph( - nodes, - "quantized_rmsnorm", - [graph_input], - [graph_output], - initializer=initializers, - ), - opset_imports=[helper.make_opsetid("", 19)], - ir_version=10, - ) - onnx.checker.check_model(model, full_check=True) - - qdq_before = { - node.name: node.SerializeToString() - for node in model.graph.node - if node.op_type in {"QuantizeLinear", "DequantizeLinear"} - } - qdq_initializers_before = { - initializer.name: initializer.SerializeToString() - for initializer in model.graph.initializer - if initializer.name in {"q_scale", "q_zero_point", "dq_scale", "dq_zero_point"} - } - opsets_before = [(opset.domain, opset.version) for opset in model.opset_import] - +def test_convert_to_f16_combines_op_and_node_exclusions(simple_model): + model, *_ = simple_model converted = convert_to_f16( - deepcopy(model), - keep_io_types=True, - op_block_list=["QuantizeLinear", "DequantizeLinear"], - nodes_to_exclude=[r"^/rms/(Pow|ReduceMean|Add|Sqrt|Div)$"], - trt_plugins=[], - opset=19, + model, + keep_io_types=False, + op_block_list=["MatMul"], + nodes_to_exclude=[r"^add$"], ) - qdq_after = { - node.name: node.SerializeToString() - for node in converted.graph.node - if node.name in qdq_before - } - qdq_initializers_after = { - initializer.name: initializer.SerializeToString() - for initializer in converted.graph.initializer - if initializer.name in qdq_initializers_before - } - assert qdq_after == qdq_before - assert qdq_initializers_after == qdq_initializers_before - assert [(opset.domain, opset.version) for opset in converted.opset_import] == opsets_before - value_types = { value.name: value.type.tensor_type.elem_type - for value in ( - *converted.graph.input, - *converted.graph.output, - *converted.graph.value_info, - ) + for value in (*converted.graph.output, *converted.graph.value_info) } - assert value_types["X_quantized"] == TensorProto.UINT8 - assert value_types["X_dequantized"] == TensorProto.FLOAT - for output_name in ["pow_out", "mean_out", "add_out", "sqrt_out", "div_out"]: - assert value_types[output_name] == TensorProto.FLOAT + assert value_types["gemm_output"] == TensorProto.FLOAT + assert value_types["add_output"] == TensorProto.FLOAT + assert value_types["Y"] == TensorProto.FLOAT16 onnx.checker.check_model(converted, full_check=True)