diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index 6977ace6360..b4e399e30dc 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -64,7 +64,7 @@ jobs: # nvcr.io/nvidia/tensorrt:26.05-py3 ships cuDNN 9.22 with no preinstalled torch, and # torch 2.14 pins cuDNN 9.24: mixing them fails with CUDNN_SUBLIBRARY_LOADING_FAILED. - if [[ "${{ inputs.docker_image }}" == *"/tensorrt:"* ]]; then + if [[ "${{ inputs.docker_image }}" == *"/tensorrt:26.05-"* ]]; then echo "torch<2.14" > /tmp/pip-constraints.txt export PIP_CONSTRAINT=/tmp/pip-constraints.txt fi @@ -79,6 +79,22 @@ jobs: fi find examples/${{ inputs.example }} -name "requirements.txt" | while read req_file; do python -m pip install -r "$req_file" || exit 1; done + + if [[ "${{ inputs.example }}" == "torch_onnx" ]]; then + # Prefer the CUDA libraries bundled with PyPI torch to avoid mixing cuDNN sublibraries. + torch_lib_path=$(python - <<'PY' + from pathlib import Path + + import torch + + root = Path(torch.__file__).resolve().parent.parent / "nvidia" + paths = sorted(str(path) for path in root.glob("*/lib") if path.is_dir()) + assert paths + print(":".join(paths)) + PY + ) + echo "LD_LIBRARY_PATH=${torch_lib_path}:${LD_LIBRARY_PATH}" >> "$GITHUB_ENV" + fi - name: Run tests id: run_tests continue-on-error: ${{ inputs.allow_failure }} diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 68d9a0c1e54..491c2938aa1 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -168,18 +168,26 @@ jobs: strategy: fail-fast: false matrix: - example: [diffusers, torch_onnx, torch_trt] + include: + - example: diffusers + docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3" + timeout_minutes: 45 + - example: torch_onnx + docker_image: "nvcr.io/nvidia/tensorrt:26.06-py3" + timeout_minutes: 90 + - example: torch_trt + docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3" + timeout_minutes: 45 uses: ./.github/workflows/_example_tests_runner.yml permissions: contents: read secrets: inherit with: - # Pinned to 26.05 (TensorRT 10): torch-tensorrt is capped at <2.13 (== 2.12.1), - # which needs libnvinfer.so.10; newer tensorrt containers drop it. Bump only once - # a torch-tensorrt build for the newer TensorRT is available. - docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3" + # torch_trt stays on TensorRT 10 because torch-tensorrt needs libnvinfer.so.10; + # torch_onnx uses TensorRT 11 for W4A4 NVFP4 engine coverage. + docker_image: ${{ matrix.docker_image }} example: ${{ matrix.example }} - timeout_minutes: 45 + timeout_minutes: ${{ matrix.timeout_minutes }} pip_install_extras: "[onnx,hf,dev-test]" runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }} diff --git a/examples/onnx_ptq/_trt_compat.py b/examples/onnx_ptq/_trt_compat.py new file mode 100644 index 00000000000..615674f02c5 --- /dev/null +++ b/examples/onnx_ptq/_trt_compat.py @@ -0,0 +1,58 @@ +# 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 pathlib import Path + +import onnx + +from modelopt.onnx.quantization.ort_utils import _check_for_trtexec +from modelopt.onnx.utils import has_node_op_type + +DYNAMIC_NVFP4_OP = "TRT_FP4DynamicQuantize" +DYNAMIC_NVFP4_MIN_TRT_VERSION = "11.0" +W4A16_NVFP4_RECIPE = "w4a16_nvfp4" + +_DYNAMIC_NVFP4_FORMATS = {"nvfp4", "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 " + "`--quantize_mode=nvfp4 --recipe=w4a16_nvfp4` to use the weight-only NVFP4 " + "recipe on TensorRT 10.16." +) + + +def request_uses_dynamic_nvfp4( + quantize_mode: str, recipe: str | None, auto_quantization_formats: list[str] +) -> bool: + """Return whether the requested quantization can emit dynamic NVFP4 activations.""" + if quantize_mode == "nvfp4": + return recipe != W4A16_NVFP4_RECIPE + return quantize_mode == "auto" and bool( + _DYNAMIC_NVFP4_FORMATS.intersection(auto_quantization_formats) + ) + + +def onnx_uses_dynamic_nvfp4(onnx_path: str | Path) -> bool: + """Return whether an ONNX graph contains dynamic NVFP4 activation quantization.""" + model = onnx.load(str(onnx_path), load_external_data=False) + return has_node_op_type(model.graph, DYNAMIC_NVFP4_OP) + + +def check_dynamic_nvfp4_trt_support() -> None: + """Require a TensorRT release that reliably compiles dynamic NVFP4 graphs.""" + try: + _check_for_trtexec(min_version=DYNAMIC_NVFP4_MIN_TRT_VERSION) + except ImportError as e: + raise ImportError(f"{_DYNAMIC_NVFP4_TRT_ERROR} ({e})") from e diff --git a/examples/onnx_ptq/evaluate.py b/examples/onnx_ptq/evaluate.py index 89d6daca070..9eef1448010 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 @@ -25,7 +26,13 @@ def main(): - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser( + description=( + "Dynamic NVFP4 (W4A4) models require TensorRT 11.0 or newer. " + "For TensorRT 10.16, re-export with " + "--quantize_mode=nvfp4 --recipe=w4a16_nvfp4." + ) + ) parser.add_argument( "--onnx_path", type=str, @@ -80,6 +87,12 @@ def main(): ) args = parser.parse_args() + if onnx_uses_dynamic_nvfp4(args.onnx_path): + try: + check_dynamic_nvfp4_trt_support() + except ImportError as e: + parser.error(str(e)) + deployment = { "runtime": "TRT", "precision": args.engine_precision, diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index f479bab3ae1..b9bb5af791e 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -74,6 +74,25 @@ Quantization configs are loaded from the YAML preset recipes under `--recipe=` to use a different recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`). +### TensorRT Compatibility + +Building a dynamic W4A4 NVFP4 vision model requires TensorRT 11.0 or later. +TensorRT 10.16 users can explicitly select the validated weight-only W4A16 +NVFP4 recipe instead: + +```bash +python torch_quant_to_onnx.py \ + --timm_model_name=vit_small_patch16_224 \ + --quantize_mode=nvfp4 \ + --recipe=w4a16_nvfp4 \ + --onnx_save_path=vit_small_patch16_224.w4a16_nvfp4.onnx \ + --trt_build +``` + +The fallback changes the quantization behavior: Linear weights use NVFP4 while +their activations remain in higher precision. The existing FP8 Conv2d override +still applies. + ### Conv2d Quantization Override TensorRT only supports FP8 and INT8 for convolution operations. When quantizing models with Conv2d layers (like SwinTransformer), the script automatically applies the following overrides: @@ -88,7 +107,8 @@ TensorRT only supports FP8 and INT8 for convolution operations. When quantizing 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. W4A4 +> NVFP4 vision models require TensorRT 11.0 or later.* ```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 e0ffc75a294..7c357c3fe4e 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -29,6 +29,7 @@ import torch import torch.multiprocessing as mp import torch.nn.functional as F +from _trt_compat import check_dynamic_nvfp4_trt_support, request_uses_dynamic_nvfp4 from datasets import load_dataset from download_example_onnx import export_to_onnx from evaluation import evaluate @@ -525,7 +526,11 @@ 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 (W4A4) builds require TensorRT 11.0 or newer. " + "For TensorRT 10.16, use --quantize_mode=nvfp4 --recipe=w4a16_nvfp4." + ), ) parser.add_argument( "--no_pretrained", @@ -546,6 +551,15 @@ def main(): "--recipe is not supported with --quantize_mode=auto; " "use --auto_quantization_formats instead." ) + if args.trt_build and request_uses_dynamic_nvfp4( + args.quantize_mode, + args.recipe, + args.auto_quantization_formats, + ): + try: + check_dynamic_nvfp4_trt_support() + except ImportError as e: + parser.error(str(e)) # Create model and move to appropriate device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..a13152e2041 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -105,7 +105,7 @@ def _add_initializer(initializer): if initializer.name not in initializer_indices: graph.initializer.append(initializer) - def _add_input_value_info(graph, tensor_proto): + def _add_initializer_value_info(graph, tensor_proto): assert tensor_proto.name not in graph_inputs, ( f"{tensor_proto.name} already in graph inputs." ) @@ -116,7 +116,7 @@ def _add_input_value_info(graph, tensor_proto): value_info = onnx.helper.make_tensor_value_info( tensor_proto.name, tensor_proto.data_type, tensor_proto.dims ) - graph.input.append(value_info) + graph.value_info.append(value_info) # Remove the original node from the graph graph.node.remove(node) @@ -148,9 +148,9 @@ def _add_input_value_info(graph, tensor_proto): ) # Add ValueInfo for the initializers if not present - _add_input_value_info(graph, w_f4_proto) - _add_input_value_info(graph, sw_f32_per_tensor_proto) - _add_input_value_info(graph, sw_f8_per_block_proto) + _add_initializer_value_info(graph, w_f4_proto) + _add_initializer_value_info(graph, sw_f32_per_tensor_proto) + _add_initializer_value_info(graph, sw_f8_per_block_proto) # Add the initializers to the graph _add_initializer(w_f4_proto) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..98517859166 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -21,6 +21,7 @@ import tempfile import uuid from collections import defaultdict +from collections.abc import Collection from typing import Any import numpy as np @@ -128,6 +129,21 @@ def get_node_names(model: onnx.ModelProto) -> list[str]: return [node.name for node in model.graph.node] +def has_node_op_type(graph: onnx.GraphProto, op_type: str) -> bool: + """Return whether a graph or any nested subgraph contains an operator type.""" + for node in graph.node: + if node.op_type == op_type: + return True + for attr in node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + if has_node_op_type(attr.g, op_type): + return True + elif attr.type == onnx.AttributeProto.GRAPHS: + if any(has_node_op_type(subgraph, op_type) for subgraph in attr.graphs): + return True + return False + + def _get_tensor_shape(tensor: onnx.ValueInfoProto) -> list[int]: """This function returns the shape of the input onnx tensor. @@ -1859,18 +1875,25 @@ def remove_node_training_mode(onnx_model: onnx.ModelProto, node_op_type: str) -> return onnx_model -def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) -> onnx.ModelProto: - """Change FP16-to-FP32 Cast nodes whose entire fanout feeds target ops to cast to FP16 instead. +def change_casts_to_fp16( + model: onnx.ModelProto, + target_op_types: list[str], + source_types: Collection[int] | None = None, +) -> onnx.ModelProto: + """Retarget eligible Cast nodes whose entire fanout feeds target ops to FP16. Args: model: The ONNX model to modify. target_op_types: List of op types to check for. Cast nodes feeding exclusively into these will be changed from FP32 to FP16. + source_types: Source element types eligible for retargeting. Defaults to FP16. Returns: The modified ONNX model with Cast nodes updated. """ type_map = _build_tensor_type_map(model) + if source_types is None: + source_types = {onnx.TensorProto.FLOAT16} # Build a map of tensor name -> consumer nodes tensor_to_consumers: dict[str, list[onnx.NodeProto]] = {} @@ -1879,17 +1902,16 @@ def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) -> if inp: tensor_to_consumers.setdefault(inp, []).append(node) - # Find Cast nodes that feed into target ops and change FP16->FP32 to FP16->FP16 + # Find Cast nodes that feed into target ops and change their destination to FP16 for node in model.graph.node: if node.op_type != "Cast": continue - # Only retarget FP16->FP32 casts; leave other casts (e.g. FP64->FP32) alone cast_to = get_cast_to_type(node) if cast_to != onnx.TensorProto.FLOAT: continue source_type = type_map.get(node.input[0]) - if source_type != onnx.TensorProto.FLOAT16: + if source_type not in source_types: continue # Only change when ALL consumers are target ops to avoid breaking non-target branches diff --git a/modelopt/torch/_deploy/_runtime/trt_client.py b/modelopt/torch/_deploy/_runtime/trt_client.py index a9c300eca5f..2d18a791e19 100644 --- a/modelopt/torch/_deploy/_runtime/trt_client.py +++ b/modelopt/torch/_deploy/_runtime/trt_client.py @@ -180,10 +180,6 @@ def initialize_input_output_tensors(self, engine): tensor_name, output_tensors[idx - len(input_tensors)].data_ptr(), ) - assert self.execution_context.all_shape_inputs_specified, ( - "Not all shape inputs are specified." - ) - # Set selected profile idx self.execution_context.set_optimization_profile_async(0, self.stream.cuda_stream) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 01fb754bbae..62e801f8d91 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -35,6 +35,7 @@ from torch.nn.parallel import DataParallel, DistributedDataParallel from modelopt.onnx.autocast.convert import convert_to_f16 +from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.export import ( FP8QuantExporter, INT4QuantExporter, @@ -52,9 +53,12 @@ fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, + get_min_opset_for_precisions, get_node_names, get_output_names, get_output_shapes, + get_qdq_precisions, + has_node_op_type, infer_shapes, remove_node_training_mode, remove_redundant_casts, @@ -347,14 +351,21 @@ def is_int4_quantized(model: nn.Module) -> bool: return False +def _is_enabled_fp4_quantizer(quantizer: Any) -> bool: + block_sizes = getattr(quantizer, "block_sizes", None) + return bool( + getattr(quantizer, "is_enabled", False) + and block_sizes + and block_sizes.get("scale_bits", None) == (4, 3) + ) + + def is_fp4_quantized(model: nn.Module) -> bool: """Check if the model is quantized in NVFP4 mode.""" for _, module in model.named_modules(): - if ( - hasattr(module, "input_quantizer") - and module.input_quantizer.block_sizes - and module.input_quantizer.block_sizes.get("scale_bits", None) == (4, 3) - ): + input_is_fp4 = _is_enabled_fp4_quantizer(getattr(module, "input_quantizer", None)) + weight_is_fp4 = _is_enabled_fp4_quantizer(getattr(module, "weight_quantizer", None)) + if input_is_fp4 or weight_is_fp4: return True return False @@ -635,8 +646,10 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = qdq_to_dq(onnx_opt_graph) if weights_dtype in ["fp16", "bf16"]: + has_dynamic_fp4 = has_node_op_type(onnx_opt_graph.graph, "TRT_FP4DynamicQuantize") if ( - is_int4_quantized(model) + has_dynamic_fp4 + or is_int4_quantized(model) or is_mxfp8_quantized(model) or is_fp8_quantized(model) or is_int8_quantized(model) @@ -651,12 +664,37 @@ def get_onnx_bytes_and_metadata( ) # Change FP32 cast nodes feeding into Concat/Add to FP16 op_list = ["Concat", "Add", "Sqrt", "LayerNormalization", "Clip", "Mul", "Exp"] - onnx_opt_graph = change_casts_to_fp16(onnx_opt_graph, op_list) + source_types = None + if has_dynamic_fp4: + op_list.extend(["Pow", "ReduceSum", "MatMul", "Unsqueeze"]) + source_types = { + onnx.TensorProto.BOOL, + onnx.TensorProto.INT8, + onnx.TensorProto.INT16, + onnx.TensorProto.INT32, + onnx.TensorProto.INT64, + onnx.TensorProto.UINT8, + onnx.TensorProto.UINT16, + onnx.TensorProto.UINT32, + onnx.TensorProto.UINT64, + onnx.TensorProto.FLOAT16, + } + onnx_opt_graph = change_casts_to_fp16( + onnx_opt_graph, op_list, source_types=source_types + ) # Remove Cast(FP32->FP16) nodes after DQ by setting DQ output to FP16 directly onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) # Remove Cast(FP16->FP32) feeding Q/DQ scales so DQ stays FP16 for downstream # MatMul/Add layers under strongly-typed TRT parsing. onnx_opt_graph = fold_qdq_scale_fp16_to_fp32_casts(onnx_opt_graph) + if has_dynamic_fp4: + qdq_precisions = get_qdq_precisions(onnx_opt_graph) + qdq_precisions.add("float4_e2m1fn") + min_opset = get_min_opset_for_precisions(qdq_precisions) + sanitizer = GraphSanitizer(onnx_opt_graph, min_opset=min_opset) + sanitizer.find_custom_nodes() + sanitizer.convert_opset() + onnx_opt_graph = sanitizer.model else: onnx_opt_graph = convert_to_f16( onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False diff --git a/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py index 1394d6a2400..fc7915bc739 100644 --- a/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py @@ -22,6 +22,8 @@ create_tiny_llama_seq_cls_dir, ) +from modelopt.torch.quantization.backends.utils import fp4_compatible + # Tiny stand-ins for the target architectures: a plain encoder for the embedding # path and a sequence-classification model (with a `score` head, kept unquantized # by the recipe) for the reranking path. @@ -43,10 +45,13 @@ @pytest.mark.parametrize( ("model_kind", "recipe", "expected_op"), [ - ( + pytest.param( "embedding", "huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj", "TRT_FP4DynamicQuantize", + marks=pytest.mark.skipif( + not fp4_compatible(), reason="FP4 is not supported on this GPU" + ), ), ( "reranking", 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 fe99cae9bc3..315160ed374 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -18,9 +18,19 @@ import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command +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). -_QUANT_MODES = ["fp8", "int8", "mxfp8", "nvfp4", "auto"] +_REQUIRES_FP4 = pytest.mark.skipif(not fp4_compatible(), reason="FP4 is not supported on this GPU") + +_QUANT_MODES = [ + "fp8", + "int8", + "mxfp8", + pytest.param("nvfp4", marks=_REQUIRES_FP4), + pytest.param("auto", marks=[_REQUIRES_FP4, pytest.mark.timeout(600)]), +] _MODELS = { "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'), @@ -49,6 +59,46 @@ def test_torch_onnx(model_key, quantize_mode): run_example_command(cmd_parts, "torch_onnx") +def _run_vit_small_nvfp4(tmp_path, recipe=None): + onnx_save_path = tmp_path / f"vit_small_patch16_224.{recipe or 'nvfp4'}.onnx" + cmd_parts = extend_cmd_parts( + ["python", "torch_quant_to_onnx.py"], + timm_model_name="vit_small_patch16_224", + quantize_mode="nvfp4", + recipe=recipe, + onnx_save_path=onnx_save_path, + calibration_data_size="1", + ) + cmd_parts.extend(["--no_pretrained", "--trt_build"]) + run_example_command(cmd_parts, "torch_onnx") + return onnx.load(onnx_save_path) + + +@_REQUIRES_FP4 +@pytest.mark.timeout(600) +def test_vit_small_nvfp4_trt_build(tmp_path): + model = _run_vit_small_nvfp4(tmp_path) + + assert any(node.op_type == "TRT_FP4DynamicQuantize" for node in model.graph.node) + + +@_REQUIRES_FP4 +@pytest.mark.timeout(600) +def test_vit_small_w4a16_nvfp4_trt_build(tmp_path): + model = _run_vit_small_nvfp4(tmp_path, recipe="w4a16_nvfp4") + initializer_types = { + initializer.name: initializer.data_type for initializer in model.graph.initializer + } + + assert not any(node.op_type == "TRT_FP4DynamicQuantize" for node in model.graph.node) + assert any( + node.op_type == "DequantizeLinear" + and node.input + and initializer_types.get(node.input[0]) == onnx.TensorProto.FLOAT4E2M1 + for node in model.graph.node + ) + + def test_torch_onnx_recipe_flag(tmp_path): timm_model_name, model_kwargs = _MODELS["vit_tiny"] onnx_save_path = tmp_path / "vit_tiny.recipe.onnx" 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..11483873890 --- /dev/null +++ b/tests/unit/examples/test_nvfp4_trt_compat.py @@ -0,0 +1,207 @@ +# 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 importlib.util +import sys +import types +from pathlib import Path +from unittest import mock + +import onnx +import pytest + +from examples.onnx_ptq import _trt_compat + +REPO_ROOT = Path(__file__).parents[3] + + +@pytest.mark.parametrize( + ("quantize_mode", "recipe", "auto_formats", "expected"), + [ + ("nvfp4", None, [], True), + ("nvfp4", "w4a16_nvfp4", [], False), + ("nvfp4", "custom_recipe", [], True), + ("auto", None, ["nvfp4_awq_lite", "fp8"], True), + ("auto", None, ["mxfp8", "fp8"], False), + ("fp8", None, [], False), + ], +) +def test_request_uses_dynamic_nvfp4(quantize_mode, recipe, auto_formats, expected): + assert _trt_compat.request_uses_dynamic_nvfp4(quantize_mode, recipe, auto_formats) is expected + + +@pytest.mark.parametrize("op_type", ["TRT_FP4DynamicQuantize", "QuantizeLinear"]) +def test_onnx_uses_dynamic_nvfp4(tmp_path, op_type): + graph = onnx.helper.make_graph( + [onnx.helper.make_node(op_type, ["input"], ["output"])], + "test_graph", + [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + ) + path = tmp_path / "model.onnx" + onnx.save(onnx.helper.make_model(graph), path) + + assert _trt_compat.onnx_uses_dynamic_nvfp4(path) is (op_type == "TRT_FP4DynamicQuantize") + + +@pytest.mark.parametrize("attribute_type", [onnx.AttributeProto.GRAPH, onnx.AttributeProto.GRAPHS]) +def test_onnx_uses_dynamic_nvfp4_in_subgraph(tmp_path, attribute_type): + subgraph = onnx.helper.make_graph( + [onnx.helper.make_node("TRT_FP4DynamicQuantize", ["input"], ["output"])], + "subgraph", + [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + ) + container = onnx.helper.make_node("Container", [], []) + attribute = container.attribute.add() + attribute.name = "subgraphs" + attribute.type = attribute_type + if attribute_type == onnx.AttributeProto.GRAPH: + attribute.g.CopyFrom(subgraph) + else: + attribute.graphs.append(subgraph) + graph = onnx.helper.make_graph([container], "test_graph", [], []) + path = tmp_path / "model.onnx" + onnx.save(onnx.helper.make_model(graph), path) + + assert _trt_compat.onnx_uses_dynamic_nvfp4(path) + + +def test_check_dynamic_nvfp4_trt_support_uses_minimum_version(monkeypatch): + check = mock.Mock() + monkeypatch.setattr(_trt_compat, "_check_for_trtexec", check) + + _trt_compat.check_dynamic_nvfp4_trt_support() + + check.assert_called_once_with(min_version="11.0") + + +def test_check_dynamic_nvfp4_trt_support_reports_fallback(monkeypatch): + def reject_trt10(*, min_version): + raise ImportError(f"trtexec version must be >= {min_version}, found 10.16") + + monkeypatch.setattr(_trt_compat, "_check_for_trtexec", reject_trt10) + + with pytest.raises(ImportError, match="--recipe=w4a16_nvfp4"): + _trt_compat.check_dynamic_nvfp4_trt_support() + + +def _module(name, **attributes): + module = types.ModuleType(name) + for key, value in attributes.items(): + setattr(module, key, value) + return module + + +def _load_example(monkeypatch, name, path): + monkeypatch.syspath_prepend(str(path.parent)) + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + +def test_torch_build_rejects_before_model_creation(monkeypatch, capsys, tmp_path): + create_model = mock.Mock() + monkeypatch.setitem(sys.modules, "timm", _module("timm", create_model=create_model)) + monkeypatch.setitem(sys.modules, "datasets", _module("datasets", load_dataset=mock.Mock())) + monkeypatch.setitem( + sys.modules, + "download_example_onnx", + _module("download_example_onnx", export_to_onnx=mock.Mock()), + ) + monkeypatch.setitem( + sys.modules, + "evaluation", + _module("evaluation", evaluate=mock.Mock()), + ) + module = _load_example( + monkeypatch, + "test_torch_quant_to_onnx_entrypoint", + REPO_ROOT / "examples/torch_onnx/torch_quant_to_onnx.py", + ) + + def reject_trt10(): + raise ImportError("dynamic NVFP4 requires TensorRT 11.0") + + monkeypatch.setattr(module, "check_dynamic_nvfp4_trt_support", reject_trt10) + monkeypatch.setattr( + sys, + "argv", + [ + "torch_quant_to_onnx.py", + "--quantize_mode=nvfp4", + f"--onnx_save_path={tmp_path / 'model.onnx'}", + "--trt_build", + ], + ) + + with pytest.raises(SystemExit, match="2"): + module.main() + + create_model.assert_not_called() + assert "dynamic NVFP4 requires TensorRT 11.0" in capsys.readouterr().err + + +def test_evaluate_rejects_before_runtime_creation(monkeypatch, capsys): + runtime_get = mock.Mock() + monkeypatch.setitem(sys.modules, "timm", _module("timm", create_model=mock.Mock())) + monkeypatch.setitem( + sys.modules, + "evaluation", + _module("evaluation", evaluate=mock.Mock()), + ) + monkeypatch.setitem( + sys.modules, + "modelopt.torch._deploy._runtime", + _module( + "modelopt.torch._deploy._runtime", + RuntimeRegistry=types.SimpleNamespace(get=runtime_get), + ), + ) + monkeypatch.setitem( + sys.modules, + "modelopt.torch._deploy.device_model", + _module("modelopt.torch._deploy.device_model", DeviceModel=mock.Mock()), + ) + monkeypatch.setitem( + sys.modules, + "modelopt.torch._deploy.utils", + _module("modelopt.torch._deploy.utils", OnnxBytes=mock.Mock()), + ) + module = _load_example( + monkeypatch, + "test_onnx_ptq_evaluate_entrypoint", + REPO_ROOT / "examples/onnx_ptq/evaluate.py", + ) + monkeypatch.setattr(module, "onnx_uses_dynamic_nvfp4", lambda _: True) + + def reject_trt10(): + raise ImportError("dynamic NVFP4 requires TensorRT 11.0") + + monkeypatch.setattr(module, "check_dynamic_nvfp4_trt_support", reject_trt10) + monkeypatch.setattr( + sys, + "argv", + ["evaluate.py", "--onnx_path=model.onnx", "--model_name=vit_small_patch16_224"], + ) + + with pytest.raises(SystemExit, match="2"): + module.main() + + runtime_get.assert_not_called() + assert "dynamic NVFP4 requires TensorRT 11.0" in capsys.readouterr().err diff --git a/tests/unit/onnx/quantization/test_ort_utils.py b/tests/unit/onnx/quantization/test_ort_utils.py index ee4f81f5645..66ef8ec1166 100644 --- a/tests/unit/onnx/quantization/test_ort_utils.py +++ b/tests/unit/onnx/quantization/test_ort_utils.py @@ -17,6 +17,8 @@ import sys import types +import pytest + from modelopt.onnx.quantization import ort_utils from modelopt.onnx.quantization.ort_utils import create_input_shapes_profile @@ -25,6 +27,33 @@ def _raise_trt_unavailable(): raise RuntimeError("trt unavailable") +def test_check_for_trtexec_rejects_version_below_minimum(monkeypatch): + monkeypatch.setattr(ort_utils.shutil, "which", lambda _: "/usr/bin/trtexec") + monkeypatch.setattr( + ort_utils, + "_run_trtexec", + lambda **_: types.SimpleNamespace( + stdout="&&&& FAILED TensorRT.trtexec [TensorRT v101601]", stderr="" + ), + ) + + with pytest.raises(ImportError, match=r"version must be >= 11\.0, found 10\.16"): + ort_utils._check_for_trtexec("11.0") + + +def test_check_for_trtexec_accepts_version_at_minimum(monkeypatch): + monkeypatch.setattr(ort_utils.shutil, "which", lambda _: "/usr/bin/trtexec") + monkeypatch.setattr( + ort_utils, + "_run_trtexec", + lambda **_: types.SimpleNamespace( + stdout="&&&& PASSED TensorRT.trtexec [TensorRT v110000] [b114]", stderr="" + ), + ) + + assert ort_utils._check_for_trtexec("11.0") == "/usr/bin/trtexec" + + def test_create_input_shapes_profile_forwards_trust_remote_code(monkeypatch): calls = [] diff --git a/tests/unit/onnx/test_fold_casts.py b/tests/unit/onnx/test_fold_casts.py index 59a434d1206..414d676360f 100644 --- a/tests/unit/onnx/test_fold_casts.py +++ b/tests/unit/onnx/test_fold_casts.py @@ -19,7 +19,12 @@ import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.utils import fold_dq_fp32_to_fp16_casts, fold_q_fp16_to_fp32_casts +from modelopt.onnx.utils import ( + change_casts_to_fp16, + fold_dq_fp32_to_fp16_casts, + fold_q_fp16_to_fp32_casts, + get_cast_to_type, +) def _dq_cast_model(opset): @@ -71,6 +76,69 @@ def _cast_q_model(opset): ) +def test_change_casts_to_fp16_retargets_only_allowed_sources_and_fanout(): + nodes = [ + helper.make_node("Cast", ["x"], ["matmul_in"], "matmul_cast", to=TensorProto.FLOAT), + helper.make_node("MatMul", ["matmul_in", "weight"], ["matmul_out"], "matmul"), + helper.make_node("Cast", ["x"], ["unsqueeze_in"], "unsqueeze_cast", to=TensorProto.FLOAT), + helper.make_node("Unsqueeze", ["unsqueeze_in", "axes"], ["unsqueeze_out"], "unsqueeze"), + helper.make_node("Cast", ["mask"], ["mask_float"], "mask_cast", to=TensorProto.FLOAT), + helper.make_node("ReduceSum", ["mask_float"], ["mask_sum"], "mask_sum"), + helper.make_node("Mul", ["mask_float", "half"], ["masked"], "mask_mul"), + helper.make_node("Cast", ["x"], ["mixed"], "mixed_cast", to=TensorProto.FLOAT), + helper.make_node("Mul", ["mixed", "half"], ["mixed_mul"], "mixed_mul"), + helper.make_node("Div", ["mixed", "float"], ["mixed_div"], "mixed_div"), + helper.make_node("Cast", ["x"], ["div_in"], "div_cast", to=TensorProto.FLOAT), + helper.make_node("Div", ["div_in", "float"], ["div_out"], "div"), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "g", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 1]), + helper.make_tensor_value_info("mask", TensorProto.INT64, [1, 1]), + ], + [], + initializer=[ + numpy_helper.from_array(np.ones((1, 1), dtype=np.float16), "weight"), + numpy_helper.from_array(np.array([0], dtype=np.int64), "axes"), + numpy_helper.from_array(np.array(1, dtype=np.float16), "half"), + numpy_helper.from_array(np.array(1, dtype=np.float32), "float"), + ], + ) + ) + + target_ops = ["MatMul", "Unsqueeze", "ReduceSum", "Mul"] + change_casts_to_fp16(model, target_ops, source_types=set()) + assert all( + get_cast_to_type(node) == TensorProto.FLOAT + for node in model.graph.node + if node.op_type == "Cast" + ) + + change_casts_to_fp16(model, target_ops) + mask_cast = next(node for node in model.graph.node if node.name == "mask_cast") + assert get_cast_to_type(mask_cast) == TensorProto.FLOAT + + change_casts_to_fp16( + model, + target_ops, + source_types={TensorProto.FLOAT16, TensorProto.INT64}, + ) + cast_types = { + node.name: get_cast_to_type(node) for node in model.graph.node if node.op_type == "Cast" + } + + assert cast_types == { + "matmul_cast": TensorProto.FLOAT16, + "unsqueeze_cast": TensorProto.FLOAT16, + "mask_cast": TensorProto.FLOAT16, + "mixed_cast": TensorProto.FLOAT, + "div_cast": TensorProto.FLOAT, + } + + @pytest.mark.parametrize( ("fold_fn", "build_model", "scale_name"), [ diff --git a/tests/unit/torch/deploy/_runtime/test_trt_client.py b/tests/unit/torch/deploy/_runtime/test_trt_client.py new file mode 100644 index 00000000000..9c2a40b05fc --- /dev/null +++ b/tests/unit/torch/deploy/_runtime/test_trt_client.py @@ -0,0 +1,84 @@ +# 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 sys +from types import ModuleType, SimpleNamespace +from unittest import mock + +import torch + +tensorrt = ModuleType("tensorrt") +tensorrt.Logger = mock.Mock() +tensorrt.TensorIOMode = SimpleNamespace(INPUT=object()) +tensorrt.tensorrt = mock.Mock() +tensorrt.__version__ = "11.0" +sys.modules["tensorrt"] = tensorrt + +from modelopt.torch._deploy._runtime import trt_client + + +class _ExecutionContext: + def __init__(self): + self.infer_shapes_called = False + + def set_tensor_address(self, *_): + return True + + def set_input_shape(self, *_): + return True + + def set_optimization_profile_async(self, *_): + return True + + def infer_shapes(self): + self.infer_shapes_called = True + return [] + + +class _Engine: + num_io_tensors = 1 + + def get_tensor_name(self, _): + return "input" + + def get_tensor_profile_shape(self, *_): + return ((1,), (1,), (1,)) + + def get_tensor_dtype(self, _): + return None + + def get_tensor_mode(self, _): + return trt_client.trt.TensorIOMode.INPUT + + +def test_initialize_io_tensors_without_deprecated_shape_property(monkeypatch): + monkeypatch.setattr( + trt_client.trt, "TensorIOMode", SimpleNamespace(INPUT=object()), raising=False + ) + context = _ExecutionContext() + session = trt_client.TRTLocalClient.TRTSession.__new__(trt_client.TRTLocalClient.TRTSession) + session.execution_context = context + session.stream = SimpleNamespace(cuda_stream=0) + session.io_shapes = {} + tensor = mock.Mock() + tensor.data_ptr.return_value = 1 + monkeypatch.setattr(trt_client, "convert_trt_dtype_to_torch", lambda _: torch.float32) + monkeypatch.setattr(trt_client.torch, "empty", lambda *_args, **_kwargs: tensor) + + inputs, outputs = session.initialize_input_output_tensors(_Engine()) + + assert inputs == [tensor] + assert outputs == [] + assert context.infer_shapes_called diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 399085ad699..a466d314b02 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -31,7 +31,9 @@ generate_onnx_input, get_onnx_bytes_and_metadata, ) -from modelopt.torch._deploy.utils.torch_onnx import _to_expected_onnx_type +from modelopt.torch._deploy.utils.torch_onnx import _to_expected_onnx_type, is_fp4_quantized +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import standardize_model_args, unflatten_tree deploy_benchmark_all = get_deploy_models() @@ -167,6 +169,26 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): return torch.add(x, y) - x +def test_is_fp4_quantized_detects_enabled_input_or_weight_quantizer(): + model = nn.Module() + config = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + model.input_quantizer = TensorQuantizer(config) + model.weight_quantizer = TensorQuantizer(config) + + model.input_quantizer.disable() + + assert is_fp4_quantized(model) + + model.weight_quantizer.disable() + assert not is_fp4_quantized(model) + + model.input_quantizer.enable() + assert is_fp4_quantized(model) + + @pytest.mark.parametrize( ("model", "n_args", "batch_size"), [ diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ce2ef626d63..421a3bcf416 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -32,9 +32,11 @@ import modelopt.torch.quantization as mtq import modelopt.torch.quantization.tensor_quant as tensor_quant from modelopt.onnx import utils +from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.export import NVFP4QuantExporter from modelopt.onnx.export.nvfp4_exporter import _encode_nvfp4_block_scale from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata from modelopt.torch.quantization.qtensor import NVFP4QTensor from modelopt.torch.quantization.utils import is_quantized_linear @@ -98,13 +100,89 @@ def cpu_dynamic_block_quantize(inputs, *args): buffer.seek(0) exported_model = onnx.load_model_from_string(buffer.read()) + original_inputs = [value.name for value in exported_model.graph.input] assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) converted_model = NVFP4QuantExporter.process_model(exported_model) assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node) + assert [value.name for value in converted_model.graph.input] == original_inputs + assert {value.name for value in converted_model.graph.input}.isdisjoint( + initializer.name for initializer in converted_model.graph.initializer + ) onnx.checker.check_model(converted_model) +def test_nvfp4_weight_only_deploy_export(monkeypatch): + sample_input = SimpleLinear().get_input() + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) + model = mtq.quantize( + SimpleLinear().eval(), + mtq.NVFP4_DEFAULT_CFG, + forward_loop=lambda calibrated_model: calibrated_model(sample_input), + ) + + for module in model.modules(): + if isinstance(module, torch.nn.Linear): + module.input_quantizer.disable() + module.weight_quantizer._onnx_quantizer_type = "static" + + onnx_bytes, _ = get_onnx_bytes_and_metadata(model, sample_input, weights_dtype="fp16") + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + exported_model = onnx.load_model_from_string(next(iter(onnx_bytes_obj.onnx_model.values()))) + initializer_types = { + initializer.name: initializer.data_type for initializer in exported_model.graph.initializer + } + + assert not any( + node.op_type in {"TRT_FP4QDQ", "TRT_FP4DynamicQuantize"} + for node in exported_model.graph.node + ) + assert any( + node.op_type == "DequantizeLinear" + and node.input + and initializer_types.get(node.input[0]) == onnx.TensorProto.FLOAT4E2M1 + for node in exported_model.graph.node + ) + onnx.checker.check_model(exported_model) + + +def test_nvfp4_dynamic_deploy_export(monkeypatch): + sample_input = SimpleLinear().get_input() + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) + + def find_custom_nodes_without_tensorrt(sanitizer): + sanitizer.custom_ops = { + node.op_type + for node in sanitizer.model.graph.node + if node.op_type == "TRT_FP4DynamicQuantize" + } + + monkeypatch.setattr(GraphSanitizer, "find_custom_nodes", find_custom_nodes_without_tensorrt) + model = mtq.quantize( + SimpleLinear().eval(), + mtq.NVFP4_DEFAULT_CFG, + forward_loop=lambda calibrated_model: calibrated_model(sample_input), + ) + + onnx_bytes, _ = get_onnx_bytes_and_metadata(model, sample_input, weights_dtype="fp16") + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + exported_model = onnx.load_model_from_string(next(iter(onnx_bytes_obj.onnx_model.values()))) + initializer_types = { + initializer.name: initializer.data_type for initializer in exported_model.graph.initializer + } + + assert any(node.op_type == "TRT_FP4DynamicQuantize" for node in exported_model.graph.node) + assert any( + node.op_type == "DequantizeLinear" + and node.input + and initializer_types.get(node.input[0]) == onnx.TensorProto.FLOAT4E2M1 + for node in exported_model.graph.node + ) + assert utils.get_opset_version(exported_model) >= 23 + assert {value.name for value in exported_model.graph.input}.isdisjoint(initializer_types) + onnx.checker.check_model(exported_model) + + @pytest.mark.parametrize( ("convert", "deprecated"), [