diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 370ab1bf358..f0c8966e27e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,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 b480bb7c24e..d5ed09215fc 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -2220,6 +2220,25 @@ 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_combines_op_and_node_exclusions(simple_model): + model, *_ = simple_model + converted = convert_to_f16( + model, + keep_io_types=False, + op_block_list=["MatMul"], + nodes_to_exclude=[r"^add$"], + ) + + value_types = { + value.name: value.type.tensor_type.elem_type + for value in (*converted.graph.output, *converted.graph.value_info) + } + 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) + + def test_convert_to_f16_refreshes_gathernd_pre_cast_declaration(monkeypatch): def discover_test_plugins_without_trt(self): self.custom_ops = {