diff --git a/examples/onnx_ptq/_trt_compat.py b/examples/onnx_ptq/_trt_compat.py new file mode 100644 index 00000000000..7ec498090b1 --- /dev/null +++ b/examples/onnx_ptq/_trt_compat.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from pathlib import Path + +import onnx + +from modelopt.onnx.quantization.ort_utils import _check_for_trtexec + +_DYNAMIC_NVFP4_OP = "TRT_FP4DynamicQuantize" +_DYNAMIC_NVFP4_MIN_TRT_VERSION = "11.0" +_DYNAMIC_NVFP4_AUTO_FORMATS = {"nvfp4_awq_lite"} +_DYNAMIC_NVFP4_TRT_ERROR = ( + "Dynamic NVFP4 (W4A4) TensorRT engine builds require TensorRT 11.0 or newer. " + "Upgrade TensorRT, or re-export with `--qformat=fp8` and without `--recipe`. " + "ONNX export without `--trt_build` remains supported." +) + + +def request_needs_dynamic_nvfp4_check( + qformat: str, + auto_quantization_formats: list[str], + *, + recipe_provided: bool, + trt_build: bool, +) -> bool: + if not trt_build or recipe_provided: + return False + return qformat == "nvfp4" or ( + qformat == "auto" + and bool(_DYNAMIC_NVFP4_AUTO_FORMATS.intersection(auto_quantization_formats)) + ) + + +def _nodes_use_dynamic_nvfp4(nodes: Iterable[onnx.NodeProto]) -> bool: + for node in nodes: + if node.op_type == _DYNAMIC_NVFP4_OP: + return True + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + nested_graphs = (attribute.g,) + elif attribute.type == onnx.AttributeProto.GRAPHS: + nested_graphs = attribute.graphs + else: + continue + if any(_nodes_use_dynamic_nvfp4(nested.node) for nested in nested_graphs): + return True + return False + + +def onnx_uses_dynamic_nvfp4(onnx_path: str | Path) -> bool: + model = onnx.load(str(onnx_path), load_external_data=False) + return _nodes_use_dynamic_nvfp4(model.graph.node) or any( + _nodes_use_dynamic_nvfp4(function.node) for function in model.functions + ) + + +def check_dynamic_nvfp4_trt_support() -> None: + try: + _check_for_trtexec(min_version=_DYNAMIC_NVFP4_MIN_TRT_VERSION) + except ImportError as error: + raise ImportError(f"{_DYNAMIC_NVFP4_TRT_ERROR} ({error})") from error diff --git a/examples/onnx_ptq/evaluate.py b/examples/onnx_ptq/evaluate.py index 89d6daca070..c26945a9b10 100644 --- a/examples/onnx_ptq/evaluate.py +++ b/examples/onnx_ptq/evaluate.py @@ -17,6 +17,7 @@ import csv import timm +from _trt_compat import check_dynamic_nvfp4_trt_support, onnx_uses_dynamic_nvfp4 from evaluation import evaluate from modelopt.torch._deploy._runtime import RuntimeRegistry @@ -80,6 +81,12 @@ def main(): ) args = parser.parse_args() + if onnx_uses_dynamic_nvfp4(args.onnx_path): + try: + check_dynamic_nvfp4_trt_support() + except ImportError as error: + parser.error(str(error)) + deployment = { "runtime": "TRT", "precision": args.engine_precision, diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index bf50cfd2c69..e634a9d1832 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -97,7 +97,10 @@ not make MXFP8, NVFP4, INT4_AWQ, or AutoQuantize supported for convolutional arc If the input model is of type image classification, use the following script to evaluate it. The script automatically downloads and uses the [ILSVRC/imagenet-1k](https://huggingface.co/datasets/ILSVRC/imagenet-1k) dataset from Hugging Face. This gated repository requires authentication via Hugging Face access token. See for details. -> *Note: TensorRT 10.11 or later is required to evaluate the MXFP8 or NVFP4 ONNX models.* +> *Note: TensorRT 10.11 or later is required to evaluate MXFP8 ONNX models. Dynamic NVFP4 +> (W4A4) models containing `TRT_FP4DynamicQuantize` require TensorRT 11.0 or later for +> `--trt_build` or evaluation. Export without `--trt_build` remains supported; use +> `--qformat=fp8` without `--recipe` when targeting TensorRT 10.* ```bash python ../onnx_ptq/evaluate.py \ diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 7450f124274..e51a267088c 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -29,6 +29,11 @@ import torch import torch.multiprocessing as mp import torch.nn.functional as F +from _trt_compat import ( + check_dynamic_nvfp4_trt_support, + onnx_uses_dynamic_nvfp4, + request_needs_dynamic_nvfp4_check, +) from datasets import load_dataset from download_example_onnx import export_to_onnx from evaluation import evaluate @@ -592,7 +597,10 @@ def main(): parser.add_argument( "--trt_build", action="store_true", - help="Build a TensorRT engine from the exported ONNX model using trtexec.", + help=( + "Build a TensorRT engine from the exported ONNX model using trtexec. " + "Dynamic NVFP4 engine builds require TensorRT 11.0 or newer." + ), ) parser.add_argument( "--no_pretrained", @@ -616,6 +624,17 @@ def main(): f"Expected a PTQ or AutoQuantize recipe, got {type(recipe).__name__} from {args.recipe}." ) + if request_needs_dynamic_nvfp4_check( + args.qformat, + args.auto_quantization_formats, + recipe_provided=recipe is not None, + trt_build=args.trt_build, + ): + try: + check_dynamic_nvfp4_trt_support() + except ImportError as error: + parser.error(str(error)) + # Create model and move to appropriate device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_kwargs = json.loads(args.model_kwargs) if args.model_kwargs else {} @@ -731,6 +750,11 @@ def main(): print(f"Quantized ONNX model is saved to {args.onnx_save_path}") if args.trt_build: + if recipe is not None and onnx_uses_dynamic_nvfp4(args.onnx_save_path): + try: + check_dynamic_nvfp4_trt_support() + except ImportError as error: + parser.error(str(error)) build_trt_engine(args.onnx_save_path) diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index d2748942972..d567ff7b6e1 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -21,10 +21,19 @@ from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from modelopt.recipe import load_recipe +from modelopt.torch.quantization.backends.utils import fp4_compatible # TODO: Add int4_awq once the INT4 exporter supports non-MatMul/Gemm consumer patterns # (e.g., DQ -> Reshape -> Slice in small ViT / SwinTransformer ONNX graphs). -_QFORMATS = ["fp8", "int8", "mxfp8", "nvfp4", "auto"] +_REQUIRES_FP4 = pytest.mark.skipif(not fp4_compatible(), reason="FP4 is not supported on this GPU") + +_QFORMATS = [ + "fp8", + "int8", + "mxfp8", + pytest.param("nvfp4", marks=_REQUIRES_FP4), + pytest.param("auto", marks=[_REQUIRES_FP4, pytest.mark.timeout(600)]), +] _RESNET_RECIPE_QFORMATS = {"fp8", "int8"} _MODELS = { @@ -88,9 +97,15 @@ def test_torch_onnx(tmp_path, model_key, qformat): calibration_data_size="1", num_score_steps="1", ) - cmd_parts.extend(["--no_pretrained", "--trt_build"]) + cmd_parts.append("--no_pretrained") + if qformat not in {"nvfp4", "auto"}: + cmd_parts.append("--trt_build") run_example_command(cmd_parts, "torch_onnx") + if qformat == "nvfp4": + op_types = {node.op_type for node in onnx.load(onnx_save_path).graph.node} + assert "TRT_FP4DynamicQuantize" in op_types + if model_key == "resnet50" and qformat in _RESNET_RECIPE_QFORMATS: _assert_residual_inputs_are_quantized(onnx_save_path) diff --git a/tests/unit/examples/test_nvfp4_trt_compat.py b/tests/unit/examples/test_nvfp4_trt_compat.py new file mode 100644 index 00000000000..452d7fc3fbd --- /dev/null +++ b/tests/unit/examples/test_nvfp4_trt_compat.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import onnx +import pytest + +from examples.onnx_ptq import _trt_compat + + +@pytest.mark.parametrize( + ("trt_build", "recipe_provided", "qformat", "auto_formats", "expected"), + [ + (True, False, "nvfp4", [], True), + (True, False, "auto", ["nvfp4_awq_lite", "fp8"], True), + (True, True, "nvfp4", [], False), + (True, False, "fp8", [], False), + (False, False, "nvfp4", [], False), + ], +) +def test_request_needs_dynamic_nvfp4_check( + trt_build, recipe_provided, qformat, auto_formats, expected +): + assert ( + _trt_compat.request_needs_dynamic_nvfp4_check( + qformat, + auto_formats, + recipe_provided=recipe_provided, + trt_build=trt_build, + ) + is expected + ) + + +@pytest.mark.parametrize( + ("placement", "expected"), + [ + ("top_level", True), + ("nested_graph", True), + ("local_function", True), + ("absent", False), + ], +) +def test_onnx_uses_dynamic_nvfp4(tmp_path, placement, expected): + op_type = "Identity" if placement == "absent" else "TRT_FP4DynamicQuantize" + node = onnx.helper.make_node(op_type, ["input"], ["output"]) + function = None + if placement == "nested_graph": + subgraph = onnx.helper.make_graph([node], "subgraph", [], []) + node = onnx.helper.make_node("Container", [], [], body=subgraph) + elif placement == "local_function": + function = onnx.helper.make_function( + "local", + "DynamicQuantize", + ["input"], + ["output"], + [node], + opset_imports=[onnx.helper.make_opsetid("", 20)], + ) + node = onnx.helper.make_node("DynamicQuantize", ["input"], ["output"], domain="local") + graph = onnx.helper.make_graph([node], "graph", [], []) + model = onnx.helper.make_model(graph) + if function is not None: + model.functions.append(function) + model.opset_import.append(onnx.helper.make_opsetid("local", 1)) + path = tmp_path / "model.onnx" + onnx.save(model, path) + + assert _trt_compat.onnx_uses_dynamic_nvfp4(path) is expected + + +def test_check_dynamic_nvfp4_trt_support_reports_action(monkeypatch): + def reject_trt10(*, min_version): + assert min_version == "11.0" + raise ImportError("`trtexec` version must be >= 11.0, found 10.16") + + monkeypatch.setattr(_trt_compat, "_check_for_trtexec", reject_trt10) + + with pytest.raises(ImportError) as error: + _trt_compat.check_dynamic_nvfp4_trt_support() + + message = str(error.value) + assert "TensorRT 11.0 or newer" in message + assert "--qformat=fp8" in message + assert "without `--trt_build`" in message diff --git a/tests/unit/onnx/quantization/test_ort_utils.py b/tests/unit/onnx/quantization/test_ort_utils.py index 010fc7409de..f9d1abadc28 100644 --- a/tests/unit/onnx/quantization/test_ort_utils.py +++ b/tests/unit/onnx/quantization/test_ort_utils.py @@ -287,3 +287,26 @@ def test_prepare_ep_list_rejects_unrecognized_trt_rtx_backend(): def test_configure_ort_rejects_unrecognized_trt_rtx_backend(): with pytest.raises(ValueError, match="trt_rtx_backend must be 'legacy' or 'abi'"): ort_utils.configure_ort([], [], calibration_eps=["cpu"], trt_rtx_backend="invalid") + + +@pytest.mark.parametrize( + ("banner", "accepted"), + [ + ("&&&& FAILED TensorRT.trtexec [TensorRT v101601]", False), + ("&&&& PASSED TensorRT.trtexec [TensorRT v110000] [b114]", True), + ], +) +def test_check_for_trtexec_compact_version_banner(monkeypatch, banner, accepted): + path = "/usr/bin/trtexec" + monkeypatch.setattr(ort_utils.shutil, "which", lambda _: path) + monkeypatch.setattr( + ort_utils, + "_run_trtexec", + lambda **_: types.SimpleNamespace(stdout=banner, stderr=""), + ) + + if accepted: + assert ort_utils._check_for_trtexec(min_version="11.0") == path + else: + with pytest.raises(ImportError, match=r">= 11\.0, found 10\.16"): + ort_utils._check_for_trtexec(min_version="11.0")