diff --git a/docs/usage/index.rst b/docs/usage/index.rst index c4acb764..59c297ea 100644 --- a/docs/usage/index.rst +++ b/docs/usage/index.rst @@ -20,6 +20,14 @@ For detailed tutorials on using these components, see: - :ref:`Advanced Usage Tutorial ` - :ref:`YOLO Tutorials ` +Guides +------ + +.. toctree:: + :maxdepth: 1 + + Strongly-Typed Networks + API Reference ------------- diff --git a/docs/usage/strongly_typed.rst b/docs/usage/strongly_typed.rst new file mode 100644 index 00000000..3d9366e8 --- /dev/null +++ b/docs/usage/strongly_typed.rst @@ -0,0 +1,235 @@ +.. _strongly_typed: + +Strongly-Typed Networks and Blackwell INT8/FP8 +============================================== + +TensorRT 10.x supports two ways to drive engine precision: + +* **Weakly-typed** (the default, and the only form before TRT 10): the builder + takes FP32 ONNX plus precision flags (``fp16``, ``int8``, ``fp8``) and + decides per-layer precision at build time, optionally calibrating INT8 + layers at runtime using a calibration cache or data batcher. +* **Strongly-typed**: precision is entirely determined by the ONNX graph + itself — FP16 casts, explicit input dtypes, and ``QuantizeLinear`` / + ``DequantizeLinear`` (Q/DQ) nodes. The builder performs no calibration and + ignores per-layer precision flags. + +On **Blackwell (SM 10.0+, compute capability ≥ 100)**, mixed INT8+FP8 +precision is only available under strongly-typed mode. Attempting it +weakly-typed raises: + +.. code-block:: text + + Error Code 9: API Usage Error (INT8 and FP8 mixed precision is allowed + only when building network with kSTRONGLY_TYPED mode on Blackwell+ + platforms.) + +Platform / version support matrix +--------------------------------- + ++------------------------------+------------------+------------------+-----------------+ +| Platform | INT8 only | FP8 only | INT8 + FP8 | ++==============================+==================+==================+=================+ +| Ampere / Ada / Hopper | weakly-typed OK | weakly-typed OK | weakly-typed OK | ++------------------------------+------------------+------------------+-----------------+ +| Blackwell + TRT 10.x | weakly-typed OK | weakly-typed OK | **strongly** | +| | | | **-typed only** | ++------------------------------+------------------+------------------+-----------------+ +| Blackwell + future TRT (≥11) | likely strongly | likely strongly | strongly-typed | +| | -typed only | -typed only | only | ++------------------------------+------------------+------------------+-----------------+ + +``trtutils`` detects both dimensions at import time: + +* :py:data:`~trtutils.FLAGS.IS_BLACKWELL` — True on SM 10.0+ GPUs. +* :py:data:`~trtutils.FLAGS.STRONGLY_TYPED_SUPPORTED` — True if the installed + TensorRT exposes ``trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED``. + +When to use ``strongly_typed=True`` +----------------------------------- + +Pass ``strongly_typed=True`` to :py:func:`trtutils.builder.build_engine` when +you are building from a **pre-quantized ONNX** that already contains Q/DQ +nodes (or explicit FP16 / FP8 dtypes). This is required on Blackwell for +mixed INT8+FP8 and recommended any time the precision decisions have been +made ahead of time. + +.. code-block:: python + + from trtutils.builder import build_engine + + build_engine( + "model_qdq.onnx", # pre-quantized with Q/DQ nodes + "model.engine", + strongly_typed=True, # precision comes from the ONNX graph + ) + +The following arguments are **mutually exclusive** with ``strongly_typed=True`` +and raise :py:class:`ValueError` if combined: + +* ``fp16``, ``int8``, ``fp8`` — precision is carried by the graph. +* ``calibration_cache``, ``data_batcher`` — runtime calibration does not + apply; scales must already be in the ONNX. +* ``layer_precision`` — per-layer overrides are ignored. + +``input_tensor_formats`` and ``output_tensor_formats`` may still be used, +but the dtype portion is ignored — only the tensor format (``LINEAR``, +``CHW4``, ``HWC``, …) is applied. A warning is logged in this case. + +Generating a Q/DQ ONNX +---------------------- + +``trtutils`` ships a thin wrapper around ``nvidia-modelopt`` that covers the +common PTQ path end-to-end. Both a Python API and a CLI are provided. + +**Python API** + +.. code-block:: python + + from trtutils.builder import quantize, build_engine, ImageBatcher + + # 1. Produce calibration data (.npy) from an image directory. + batcher = ImageBatcher( + image_dir="calibration_images/", + shape=(224, 224, 3), # HWC + dtype="float32", + batch_size=8, + order="NCHW", + ) + batcher.save_calibration_data("calib.npy") + + # 2. Quantize the ONNX, baking Q/DQ nodes into a new file. + quantize.quantize_onnx( + onnx_path="model.onnx", + output_path="model_qdq.onnx", + calibration_data="calib.npy", + quantize_mode="int8", # "int4" | "int8" | "fp8" + calibration_method="max", # "max" | "entropy" | "percentile" | "mse" + ) + + # 3. Build a strongly-typed engine from the Q/DQ ONNX. + build_engine( + "model_qdq.onnx", + "model.engine", + strongly_typed=True, + ) + +**CLI** + +The same three steps are available through subcommands of ``python -m trtutils``: + +.. code-block:: bash + + # 1. Generate calibration data from a directory of images + python -m trtutils generate_calibration \ + --calibration_dir calibration_images/ \ + --input_shape 224 224 3 \ + --input_dtype float32 \ + --batch_size 8 \ + --data_order NCHW \ + --output calib.npy + + # 2. Quantize the ONNX + python -m trtutils quantize \ + --onnx model.onnx \ + --output model_qdq.onnx \ + --calibration_data calib.npy \ + --quantize_mode int8 \ + --calibration_method max + + # 3. Build a strongly-typed engine + python -m trtutils build \ + --onnx model_qdq.onnx \ + --output model.engine \ + --strongly_typed + +Under the hood ``quantize_onnx`` calls +``modelopt.onnx.quantization.quantize``. For workflows the wrapper does not +cover — QAT, custom quantizer placement, non-image calibration — reach for +the underlying tools directly: + +* `NVIDIA TensorRT Model Optimizer `_ + (``nvidia-modelopt``) — full PTQ/QAT toolkit. +* `ONNX Runtime quantization + `_ + — alternative INT8 PTQ. +* `pytorch_quantization + `_ + (legacy) — pre-modelopt pipeline. + +For FP16-only models, exporting via PyTorch with ``model.half()`` before +``torch.onnx.export(...)`` is sufficient — no Q/DQ needed. + +DLA + strongly-typed (Jetson Orin) +---------------------------------- + +Jetson Orin (SM 87, Ampere) has DLA cores but no FP8 hardware. On +**JetPack 6.1+** (TRT 10.1+) the ``STRONGLY_TYPED`` flag is available, which +makes modelopt-quantized INT8 ONNX the cleanest path for DLA deployment. +Pass ``strongly_typed=True`` to :py:func:`trtutils.builder.build_dla_engine` +or the ``build_dla`` CLI and the builder will: + +* Skip the weakly-typed INT8 calibration step entirely — the Q/DQ nodes in + the ONNX define the scales. +* Still auto-assign DLA-compatible layer chunks to DLA via ``layer_device``. + Per-layer ``layer_precision`` overrides are **not** applied; precision + comes from the graph, so your Q/DQ quantization must be placed on the + DLA-bound subgraph for the INT8 hardware to be exercised. + +**Python API** + +.. code-block:: python + + from trtutils.builder import build_dla_engine + + build_dla_engine( + "model_qdq.onnx", + "model_dla.engine", + dla_core=0, + strongly_typed=True, + ) + +**CLI** + +.. code-block:: bash + + python -m trtutils build_dla \ + --onnx model_qdq.onnx \ + --output model_dla.engine \ + --dla_core 0 \ + --strongly_typed + +Under strongly-typed, ``build_dla_engine`` rejects ``data_batcher``, +``calibration_cache``, and ``fp8`` — those are weakly-typed artifacts. Omit +them; precision lives in the graph. + +.. note:: + + Thor (Blackwell, JetPack 7) drops DLA hardware, so Orin is the last + Jetson generation where this DLA + strongly-typed workflow applies. On + Thor you'll use the GPU path (:py:func:`trtutils.builder.build_engine` + with ``strongly_typed=True``) covered earlier in this page. + +Common errors +------------- + +``"INT8+FP8 mixed precision on Blackwell requires strongly_typed=True …"`` + You passed both ``int8=True`` and ``fp8=True`` on a Blackwell GPU without + ``strongly_typed=True``. Re-export your ONNX with Q/DQ nodes using modelopt + or similar, then pass ``strongly_typed=True`` and drop the ``int8``/``fp8`` + flags. + +``"strongly_typed=True does not support runtime calibration …"`` + You passed ``strongly_typed=True`` together with ``calibration_cache`` or + ``data_batcher``. Strongly-typed networks have no calibration step; scales + must live in the ONNX Q/DQ nodes. Remove the calibration arguments. + +``"strongly_typed=True derives precision from the ONNX graph; remove +fp16/int8/fp8 builder flags."`` + You passed ``strongly_typed=True`` together with precision flags. Drop the + flags — precision is now the ONNX's job. + +``"Installed TensorRT does not support strongly-typed networks"`` + The installed TensorRT version predates the ``STRONGLY_TYPED`` network + creation flag (TRT < 10.1). Upgrade TensorRT, or build weakly-typed on a + non-Blackwell GPU. diff --git a/src/trtutils/__main__.py b/src/trtutils/__main__.py index 4e8f3c7e..067055bc 100644 --- a/src/trtutils/__main__.py +++ b/src/trtutils/__main__.py @@ -336,6 +336,7 @@ def _build(args: SimpleNamespace, *, add_yolo_hook: bool = False) -> None: fp16=args.fp16, fp8=args.fp8, int8=args.int8, + strongly_typed=args.strongly_typed, hooks=hooks, cache=args.cache, verbose=args.verbose, @@ -566,36 +567,39 @@ def _build_dla(args: SimpleNamespace) -> None: If required DLA build parameters are missing. """ - # require calibration data for dla builds - if args.calibration_dir is None: - err_msg = "Calibration directory is required for DLA builds" - raise ValueError(err_msg) - if args.input_shape is None: - err_msg = "Input shape is required for DLA builds" - raise ValueError(err_msg) - if args.input_dtype is None: - err_msg = "Input dtype is required for DLA builds" - raise ValueError(err_msg) + # strongly_typed skips calibration entirely; precision is in the ONNX graph + batcher = None + if not args.strongly_typed: + # require calibration data for weakly-typed dla builds + if args.calibration_dir is None: + err_msg = "Calibration directory is required for DLA builds" + raise ValueError(err_msg) + if args.input_shape is None: + err_msg = "Input shape is required for DLA builds" + raise ValueError(err_msg) + if args.input_dtype is None: + err_msg = "Input dtype is required for DLA builds" + raise ValueError(err_msg) - input_shape = args.input_shape - input_dtype = args.input_dtype + input_shape = args.input_shape + input_dtype = args.input_dtype + + batcher = trtutils.builder.ImageBatcher( + image_dir=args.calibration_dir, + shape=input_shape, + dtype=np.dtype(input_dtype).type, + batch_size=args.batch_size, + order=args.data_order, + max_images=args.max_images, + resize_method=args.resize_method, + input_scale=_to_float_pair(args.input_scale), + verbose=args.verbose, + ) # Set default dla_core to 0 if not provided if args.dla_core is None: args.dla_core = 0 - batcher = trtutils.builder.ImageBatcher( - image_dir=args.calibration_dir, - shape=input_shape, - dtype=np.dtype(input_dtype).type, - batch_size=args.batch_size, - order=args.data_order, - max_images=args.max_images, - resize_method=args.resize_method, - input_scale=_to_float_pair(args.input_scale), - verbose=args.verbose, - ) - shapes = _parse_shapes_arg(getattr(args, "shape", None)) trtutils.builder.build_dla_engine( @@ -615,6 +619,7 @@ def _build_dla(args: SimpleNamespace) -> None: reject_empty_algorithms=args.reject_empty_algorithms, ignore_timing_mismatch=args.ignore_timing_mismatch, fp8=args.fp8, + strongly_typed=args.strongly_typed, cache=args.cache, verbose=args.verbose, ) @@ -1533,6 +1538,16 @@ def _main() -> None: action="store_true", help="Quantize the engine to INT8 precision.", ) + build_device_parser.add_argument( + "--strongly_typed", + action="store_true", + help=( + "Build a strongly-typed network whose precision is determined by " + "the ONNX graph (Q/DQ nodes, explicit dtypes). Required on " + "Blackwell (SM 10.0+) for mixed INT8+FP8. Mutually exclusive with " + "--fp16/--int8/--fp8 and calibration arguments." + ), + ) # main parser parser = argparse.ArgumentParser( @@ -1762,6 +1777,15 @@ def _main() -> None: action="store_true", help="Enable FP8 precision for GPU layers. Requires compute capability >= 8.9.", ) + build_dla_parser.add_argument( + "--strongly_typed", + action="store_true", + help=( + "Build a strongly-typed DLA engine from a pre-quantized ONNX " + "(Q/DQ nodes). Skips calibration and precision flags -- precision " + "is derived from the ONNX graph. Target: Jetson Orin + JetPack 6.1+." + ), + ) build_dla_parser.set_defaults(func=_build_dla) # detect parser diff --git a/src/trtutils/_flags.py b/src/trtutils/_flags.py index 9bb02878..08f071f5 100644 --- a/src/trtutils/_flags.py +++ b/src/trtutils/_flags.py @@ -55,6 +55,11 @@ class Flags: SM_VERSION : int The SM (compute capability) version as an integer. E.g. SM 7.5 -> 75, SM 10.0 -> 100. + IS_BLACKWELL : bool + Whether the detected GPU is Blackwell (SM 10.0) or newer. + STRONGLY_TYPED_SUPPORTED : bool + Whether the installed TensorRT exposes + trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED. SM_ARCH : str The GPU architecture name. E.g. "turing", "blackwell". DEVICE_NAME : str @@ -84,6 +89,7 @@ class Flags: # TensorRT flags TRT_VERSION: tuple[int, int] = (0, 0) TRT_10: bool = False + STRONGLY_TYPED_SUPPORTED: bool = False TRT_HAS_UINT8: bool = False TRT_HAS_INT64: bool = False NEW_CAN_RUN_ON_DLA: bool = False @@ -99,6 +105,7 @@ class Flags: # System flags IS_JETSON: bool = False SM_VERSION: int = 0 + IS_BLACKWELL: bool = False SM_ARCH: str = "unknown" DEVICE_NAME: str = "unknown" HAS_DLA: bool = False @@ -117,6 +124,7 @@ def init_device_flags(self) -> None: _sm = get_compute_capability() self.SM_VERSION = _sm[0] * 10 + _sm[1] + self.IS_BLACKWELL = self.SM_VERSION >= 100 # noqa: PLR2004 self.SM_ARCH = get_sm_arch(*_sm) self.DEVICE_NAME = get_device_name() @@ -162,6 +170,10 @@ def _get_version(package: str) -> tuple[int, int]: FLAGS.EXEC_ASYNC_V1 = hasattr(trt.IExecutionContext, "execute_async") FLAGS.EXEC_V2 = hasattr(trt.IExecutionContext, "execute_v2") FLAGS.EXEC_V1 = hasattr(trt.IExecutionContext, "execute") +FLAGS.STRONGLY_TYPED_SUPPORTED = hasattr( + trt.NetworkDefinitionCreationFlag, + "STRONGLY_TYPED", +) # Set system flags FLAGS.IS_JETSON = Path("/etc/nv_tegra_release").exists() diff --git a/src/trtutils/builder/_build.py b/src/trtutils/builder/_build.py index f1019096..a7870767 100644 --- a/src/trtutils/builder/_build.py +++ b/src/trtutils/builder/_build.py @@ -62,6 +62,7 @@ def build_engine( fp16: bool | None = None, fp8: bool | None = None, int8: bool | None = None, + strongly_typed: bool = False, cache: bool | None = None, verbose: bool | None = None, ) -> None: @@ -185,6 +186,15 @@ def build_engine( Requires compute capability >= 8.9 (Ada Lovelace / Hopper or newer). int8 : bool, optional If True, quantize the engine to INT8 precision. + strongly_typed : bool, optional + If True, build a strongly-typed network whose precision is entirely + determined by the ONNX graph (Q/DQ nodes, explicit dtypes). Required + on Blackwell (SM 10.0+) when mixing INT8 and FP8. Mutually exclusive + with fp16/int8/fp8 builder flags and with calibration_cache, + data_batcher, and layer_precision. Tensor-format dtype overrides are + also ignored when enabled. By default, False. + See ``docs/usage/strongly_typed.rst`` for guidance on generating a + Q/DQ ONNX. cache : bool, optional Whether or not to cache the engine in the trtutils engine cache. If an existing version is found will use that. @@ -255,10 +265,46 @@ def build_engine( trt.DeviceType.GPU if default_device == trt.DeviceType.GPU else trt.DeviceType.DLA ) + # validate strongly_typed vs. incompatible args + if strongly_typed: + if calibration_cache is not None or data_batcher is not None: + err_msg = ( + "strongly_typed=True does not support runtime calibration. " + "Re-export the ONNX with explicit Q/DQ nodes and omit " + "calibration_cache / data_batcher. " + "See docs/usage/strongly_typed.rst." + ) + raise ValueError(err_msg) + if layer_precision is not None: + err_msg = ( + "strongly_typed=True ignores per-layer precision overrides. " + "Bake precision into the ONNX graph instead." + ) + raise ValueError(err_msg) + if fp16 or int8 or fp8: + err_msg = ( + "strongly_typed=True derives precision from the ONNX graph; " + "remove fp16/int8/fp8 builder flags." + ) + raise ValueError(err_msg) + if input_tensor_formats is not None or output_tensor_formats is not None: + LOG.warning( + "strongly_typed=True ignores tensor dtype overrides; " + "only the format portion of input/output_tensor_formats will apply.", + ) + elif FLAGS.IS_BLACKWELL and int8 and fp8: + err_msg = ( + "INT8+FP8 mixed precision on Blackwell requires strongly_typed=True " + "with a pre-quantized ONNX (explicit Q/DQ nodes). " + "See docs/usage/strongly_typed.rst." + ) + raise ValueError(err_msg) + # read the onnx model network, builder, config, _ = read_onnx( onnx, workspace, + strongly_typed=strongly_typed, ) # handle all hooks to start @@ -316,7 +362,8 @@ def build_engine( for idx in range(network.num_inputs): inp = network.get_input(idx) if inp.name == tensor_name: - inp.dtype = tensor_dtype + if not strongly_typed: + inp.dtype = tensor_dtype inp.allowed_formats = 1 << int(tensor_format) found = True break @@ -329,7 +376,8 @@ def build_engine( for idx in range(network.num_outputs): out = network.get_output(idx) if out.name == tensor_name: - out.dtype = tensor_dtype + if not strongly_typed: + out.dtype = tensor_dtype out.allowed_formats = 1 << int(tensor_format) found = True break @@ -337,24 +385,26 @@ def build_engine( err_msg = f"Output tensor '{tensor_name}' not found in network" raise ValueError(err_msg) - # setup the precision sets - if fp16 or fp8 or int8: - # want to enable fp16 for int8, fp8, and fp16 since fp16 may be faster - if not builder.platform_has_fast_fp16: - LOG.warning("Platform does not have native fast FP16.") - config.set_flag(trt.BuilderFlag.FP16) - if fp8: - config.set_flag(trt.BuilderFlag.FP8) - if int8: - if not builder.platform_has_fast_int8: - LOG.warning("Platform does not have native fast INT8.") - config.set_flag(trt.BuilderFlag.INT8) - if calibration_cache is None and data_batcher is None: - err_msg = "Neither calibration cache or data batcher passed during model building, INT8 build will not be accurate." - LOG.warning(err_msg) - config.int8_calibrator = EngineCalibrator(calibration_cache=calibration_cache) - if data_batcher is not None: - config.int8_calibrator.set_batcher(data_batcher) + # setup the precision sets -- skipped under strongly_typed since the + # precision is carried by the ONNX graph itself + if not strongly_typed: + if fp16 or fp8 or int8: + # want to enable fp16 for int8, fp8, and fp16 since fp16 may be faster + if not builder.platform_has_fast_fp16: + LOG.warning("Platform does not have native fast FP16.") + config.set_flag(trt.BuilderFlag.FP16) + if fp8: + config.set_flag(trt.BuilderFlag.FP8) + if int8: + if not builder.platform_has_fast_int8: + LOG.warning("Platform does not have native fast INT8.") + config.set_flag(trt.BuilderFlag.INT8) + if calibration_cache is None and data_batcher is None: + err_msg = "Neither calibration cache or data batcher passed during model building, INT8 build will not be accurate." + LOG.warning(err_msg) + config.int8_calibrator = EngineCalibrator(calibration_cache=calibration_cache) + if data_batcher is not None: + config.int8_calibrator.set_batcher(data_batcher) # assign the default device config.default_device_type = default_device diff --git a/src/trtutils/builder/_dla.py b/src/trtutils/builder/_dla.py index fffb2195..1932a5f5 100644 --- a/src/trtutils/builder/_dla.py +++ b/src/trtutils/builder/_dla.py @@ -114,8 +114,8 @@ def can_run_on_dla( def build_dla_engine( onnx: Path | str, output_path: Path | str, - data_batcher: AbstractBatcher, - dla_core: int, + data_batcher: AbstractBatcher | None = None, + dla_core: int = 0, max_chunks: int = 1, min_layers: int = 20, workspace: float = 4.0, @@ -132,6 +132,7 @@ def build_dla_engine( reject_empty_algorithms: bool = False, ignore_timing_mismatch: bool = False, fp8: bool | None = None, + strongly_typed: bool = False, cache: bool | None = None, verbose: bool | None = None, ) -> None: @@ -208,6 +209,13 @@ def build_dla_engine( If True, enable FP8 precision for GPU layers. Requires compute capability >= 8.9 (Ada Lovelace / Hopper or newer). DLA layers will still use INT8 precision. + strongly_typed : bool, optional + If True, build a strongly-typed network whose precision is determined + by the ONNX graph (Q/DQ nodes). Use with modelopt-quantized ONNX on + Jetson Orin + JetPack 6.1+. The builder still steers DLA-compatible + chunks to DLA via layer_device, but per-layer precision is dictated + by the graph. Mutually exclusive with data_batcher, calibration_cache, + and fp8. By default, False. cache : bool, optional Whether or not to cache the engine in the trtutils engine cache. If an existing version is found will use that. @@ -219,8 +227,36 @@ def build_dla_engine( Whether to print verbose output, by default False """ + # validate strongly_typed vs. incompatible args + if strongly_typed: + if data_batcher is not None: + err_msg = ( + "strongly_typed=True does not support runtime calibration. " + "Re-export the ONNX with Q/DQ nodes and omit data_batcher. " + "See docs/usage/strongly_typed.rst." + ) + raise ValueError(err_msg) + if calibration_cache is not None: + err_msg = ( + "strongly_typed=True does not support calibration_cache. " + "See docs/usage/strongly_typed.rst." + ) + raise ValueError(err_msg) + if fp8: + err_msg = ( + "strongly_typed=True derives precision from the ONNX graph; " + "remove fp8 and bake FP8 into the ONNX instead." + ) + raise ValueError(err_msg) + elif data_batcher is None: + err_msg = ( + "data_batcher is required for weakly-typed DLA builds. " + "Pass an AbstractBatcher or use strongly_typed=True with a Q/DQ ONNX." + ) + raise ValueError(err_msg) + # read the onnx path - network, _, config, _ = read_onnx(onnx) + network, _, config, _ = read_onnx(onnx, strongly_typed=strongly_typed) # check layers for DLA compatibility and use int8 precision full_dla, chunks = can_run_on_dla( @@ -236,29 +272,50 @@ def build_dla_engine( # case where the entire model can run on DLA if full_dla: - build_engine( - onnx, - output_path, - default_device=trt.DeviceType.DLA, - data_batcher=data_batcher, - workspace=workspace, - timing_cache=timing_cache, - calibration_cache=calibration_cache, - dla_core=dla_core, - shapes=shapes, - input_tensor_formats=input_tensor_formats, - output_tensor_formats=output_tensor_formats, - hooks=hooks, - direct_io=direct_io, - prefer_precision_constraints=prefer_precision_constraints, - reject_empty_algorithms=reject_empty_algorithms, - ignore_timing_mismatch=ignore_timing_mismatch, - cache=cache, - fp16=True, - fp8=fp8, - int8=True, - verbose=verbose, - ) + if strongly_typed: + build_engine( + onnx, + output_path, + default_device=trt.DeviceType.DLA, + workspace=workspace, + timing_cache=timing_cache, + dla_core=dla_core, + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + cache=cache, + strongly_typed=True, + verbose=verbose, + ) + else: + build_engine( + onnx, + output_path, + default_device=trt.DeviceType.DLA, + data_batcher=data_batcher, + workspace=workspace, + timing_cache=timing_cache, + calibration_cache=calibration_cache, + dla_core=dla_core, + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + cache=cache, + fp16=True, + fp8=fp8, + int8=True, + verbose=verbose, + ) return # identify if any chunks contain DLA layers @@ -267,27 +324,46 @@ def build_dla_engine( # case where no DLA layers are found if not dla_chunks: LOG.warning("No DLA-compatible layers found. Building GPU-only engine.") - build_engine( - onnx, - output_path, - workspace=workspace, - timing_cache=timing_cache, - calibration_cache=calibration_cache, - data_batcher=data_batcher, - shapes=shapes, - input_tensor_formats=input_tensor_formats, - output_tensor_formats=output_tensor_formats, - hooks=hooks, - direct_io=direct_io, - prefer_precision_constraints=prefer_precision_constraints, - reject_empty_algorithms=reject_empty_algorithms, - ignore_timing_mismatch=ignore_timing_mismatch, - fp16=True, - fp8=fp8, - int8=True, - cache=cache, - verbose=verbose, - ) + if strongly_typed: + build_engine( + onnx, + output_path, + workspace=workspace, + timing_cache=timing_cache, + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + strongly_typed=True, + cache=cache, + verbose=verbose, + ) + else: + build_engine( + onnx, + output_path, + workspace=workspace, + timing_cache=timing_cache, + calibration_cache=calibration_cache, + data_batcher=data_batcher, + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + fp16=True, + fp8=fp8, + int8=True, + cache=cache, + verbose=verbose, + ) return # sort chunks by len and filter by min_layers or until max_chunks is reached @@ -299,7 +375,9 @@ def build_dla_engine( ) # define lists for storing layer assignments - layer_precision: list[tuple[int, trt.DataType | None]] = [] + # under strongly_typed, precision is dictated by the ONNX Q/DQ graph so we + # only build a layer_device map -- layer_precision stays None + layer_precision: list[tuple[int, trt.DataType | None]] | None = None if strongly_typed else [] layer_device: list[tuple[int, trt.DeviceType | None]] = [] # assign default to GPU/FP16 @@ -309,12 +387,13 @@ def build_dla_engine( layer_name: str = layer.name layer_name = layer_name.lower() layer_device.append((idx, trt.DeviceType.GPU)) - # intelligently assign precision level to HALF unless layer - # is Constant, Shuffle, or Tile - if layer.type in exclude_layer_types or "tile" in layer_name: - layer_precision.append((idx, None)) - else: - layer_precision.append((idx, trt.DataType.HALF)) + if layer_precision is not None: + # intelligently assign precision level to HALF unless layer + # is Constant, Shuffle, or Tile + if layer.type in exclude_layer_types or "tile" in layer_name: + layer_precision.append((idx, None)) + else: + layer_precision.append((idx, trt.DataType.HALF)) # iterate over chunks and assign to DLA matched_chunks = 0 @@ -327,45 +406,78 @@ def build_dla_engine( continue for layer_id in range(start, end + 1, 1): - layer_precision[layer_id] = (layer_id, trt.DataType.INT8) layer_device[layer_id] = (layer_id, trt.DeviceType.DLA) + if layer_precision is not None: + layer_precision[layer_id] = (layer_id, trt.DataType.INT8) matched_chunks += 1 # verbose iteration if verbose: - for (idx, device), (_, datatype) in zip(layer_device, layer_precision): - LOG.info( - f"Layer {idx}: {network.get_layer(idx).name}, " - f"{'DLA' if device == trt.DeviceType.DLA else 'GPU'}, " - f"{'INT8' if datatype == trt.DataType.INT8 else 'FP16'}" - ) + if layer_precision is not None: + for (idx, device), (_, datatype) in zip(layer_device, layer_precision): + LOG.info( + f"Layer {idx}: {network.get_layer(idx).name}, " + f"{'DLA' if device == trt.DeviceType.DLA else 'GPU'}, " + f"{'INT8' if datatype == trt.DataType.INT8 else 'FP16'}" + ) + else: + for idx, device in layer_device: + LOG.info( + f"Layer {idx}: {network.get_layer(idx).name}, " + f"{'DLA' if device == trt.DeviceType.DLA else 'GPU'}, " + "precision from ONNX" + ) # build engine with specific layer assignments - build_engine( - onnx, - output_path, - default_device=trt.DeviceType.DLA, # default device DLA - timing_cache=timing_cache, - workspace=workspace, - calibration_cache=calibration_cache, - data_batcher=data_batcher, - layer_precision=layer_precision, - layer_device=layer_device, - dla_core=dla_core, # ensure DLA core is maintained - shapes=shapes, - input_tensor_formats=input_tensor_formats, - output_tensor_formats=output_tensor_formats, - hooks=hooks, - optimization_level=optimization_level, - gpu_fallback=True, # enable GPU fallback to account for input/copy - direct_io=direct_io, - prefer_precision_constraints=prefer_precision_constraints, - reject_empty_algorithms=reject_empty_algorithms, - ignore_timing_mismatch=ignore_timing_mismatch, - fp16=True, - fp8=fp8, - int8=True, - cache=cache, - verbose=verbose, - ) + if strongly_typed: + build_engine( + onnx, + output_path, + default_device=trt.DeviceType.DLA, + timing_cache=timing_cache, + workspace=workspace, + layer_device=layer_device, + dla_core=dla_core, + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + optimization_level=optimization_level, + gpu_fallback=True, + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + strongly_typed=True, + cache=cache, + verbose=verbose, + ) + else: + build_engine( + onnx, + output_path, + default_device=trt.DeviceType.DLA, # default device DLA + timing_cache=timing_cache, + workspace=workspace, + calibration_cache=calibration_cache, + data_batcher=data_batcher, + layer_precision=layer_precision, + layer_device=layer_device, + dla_core=dla_core, # ensure DLA core is maintained + shapes=shapes, + input_tensor_formats=input_tensor_formats, + output_tensor_formats=output_tensor_formats, + hooks=hooks, + optimization_level=optimization_level, + gpu_fallback=True, # enable GPU fallback to account for input/copy + direct_io=direct_io, + prefer_precision_constraints=prefer_precision_constraints, + reject_empty_algorithms=reject_empty_algorithms, + ignore_timing_mismatch=ignore_timing_mismatch, + fp16=True, + fp8=fp8, + int8=True, + cache=cache, + verbose=verbose, + ) diff --git a/src/trtutils/builder/_onnx.py b/src/trtutils/builder/_onnx.py index ecd6990f..963d4b07 100644 --- a/src/trtutils/builder/_onnx.py +++ b/src/trtutils/builder/_onnx.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024 Justin Davis (davisjustin302@gmail.com) +# Copyright (c) 2024-2026 Justin Davis (davisjustin302@gmail.com) # # MIT License # mypy: disable-error-code="import-untyped" @@ -7,6 +7,7 @@ from pathlib import Path from trtutils._config import CONFIG +from trtutils._flags import FLAGS from trtutils._log import LOG from trtutils.compat._libs import trt @@ -14,6 +15,8 @@ def read_onnx( onnx: Path | str, workspace: float = 4.0, + *, + strongly_typed: bool = False, ) -> tuple[ trt.INetworkDefinition, trt.IBuilder, @@ -30,6 +33,11 @@ def read_onnx( workspace : float The size of the workspace in gigabytes. Default is 4.0 GiB. + strongly_typed : bool, optional + If True, create the network with the STRONGLY_TYPED flag so that + precision is determined by the ONNX graph (Q/DQ nodes) rather than + builder flags. Required on Blackwell (SM 10.0+) for mixed INT8+FP8. + By default, False. Returns ------- @@ -45,7 +53,8 @@ def read_onnx( ValueError If the onnx model path does not have .onnx extension RuntimeError - If the ONNX model cannot be parsed + If the ONNX model cannot be parsed, or if strongly_typed is True + but the installed TensorRT does not support strongly-typed networks. """ # load libnvinfer plugins @@ -73,9 +82,16 @@ def read_onnx( config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes) # make network - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH), - ) + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + if strongly_typed: + if not FLAGS.STRONGLY_TYPED_SUPPORTED: + err_msg = ( + "Installed TensorRT does not support strongly-typed networks " + "(NetworkDefinitionCreationFlag.STRONGLY_TYPED not found)." + ) + raise RuntimeError(err_msg) + network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + network = builder.create_network(network_flags) # setup parser parser = trt.OnnxParser(network, LOG) diff --git a/tests/_cpu_stubs.py b/tests/_cpu_stubs.py index 871b71b1..a4abe174 100644 --- a/tests/_cpu_stubs.py +++ b/tests/_cpu_stubs.py @@ -46,6 +46,7 @@ def inject() -> None: trt.IBuilderConfig = type("IBuilderConfig", (), {}) trt.Builder = type("Builder", (), {}) trt.IExecutionContext = type("IExecutionContext", (), {}) + trt.NetworkDefinitionCreationFlag = type("NetworkDefinitionCreationFlag", (), {}) # classes with attrs used as default parameter values trt.TensorFormat = type("TensorFormat", (), {"LINEAR": 0}) trt.DeviceType = type("DeviceType", (), {"GPU": 0, "DLA": 1}) diff --git a/tests/_gpu_fixtures.py b/tests/_gpu_fixtures.py index 217aff64..408b41de 100644 --- a/tests/_gpu_fixtures.py +++ b/tests/_gpu_fixtures.py @@ -52,13 +52,14 @@ class TestImage: @pytest.fixture(autouse=True) -def _skip_jetson_tests(request: pytest.FixtureRequest) -> None: +def skip_jetson_tests(request: pytest.FixtureRequest) -> None: + # underscore-prefixed names are dropped by `from tests._gpu_fixtures import *` if request.node.get_closest_marker("jetson") and not FLAGS.IS_JETSON: pytest.skip("Jetson tests require Jetson hardware") @pytest.fixture(autouse=True) -def _skip_dla_tests(request: pytest.FixtureRequest) -> None: +def skip_dla_tests(request: pytest.FixtureRequest) -> None: if request.node.get_closest_marker("dla") and not FLAGS.HAS_DLA: pytest.skip("DLA tests require DLA hardware") diff --git a/tests/builder/_gpu_fixtures.py b/tests/builder/_gpu_fixtures.py index 25c87e27..45080f1c 100644 --- a/tests/builder/_gpu_fixtures.py +++ b/tests/builder/_gpu_fixtures.py @@ -8,9 +8,13 @@ import tempfile from typing import TYPE_CHECKING +import numpy as np +import onnx import pytest +from onnx import TensorProto, helper, numpy_helper from tests.conftest import DATA_DIR +from trtutils._flags import FLAGS from trtutils.builder._build import build_engine if TYPE_CHECKING: @@ -27,6 +31,64 @@ def onnx_path() -> Path: return ONNX_PATH +@pytest.fixture(scope="session") +def quantized_onnx_path(tmp_path_factory) -> Path: + """ + Minimal Q/DQ ONNX: fp32 input -> Quantize -> Dequantize -> Identity -> fp32 output. + + Scale and zero_point are constant initializers. No weights, no dynamic + shapes. Used for exercising the strongly-typed build path. + """ + if not FLAGS.STRONGLY_TYPED_SUPPORTED: + pytest.skip("Installed TensorRT does not support strongly-typed networks") + + out_dir = tmp_path_factory.mktemp("qdq") + out_path = out_dir / "qdq.onnx" + + shape = [1, 3, 8, 8] + scale = numpy_helper.from_array(np.array(0.01, dtype=np.float32), name="scale") + zero_point = numpy_helper.from_array(np.array(0, dtype=np.int8), name="zero_point") + + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, shape) + output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, shape) + + quantize = helper.make_node( + "QuantizeLinear", + inputs=["input", "scale", "zero_point"], + outputs=["quantized"], + name="quantize", + ) + dequantize = helper.make_node( + "DequantizeLinear", + inputs=["quantized", "scale", "zero_point"], + outputs=["dequantized"], + name="dequantize", + ) + identity = helper.make_node( + "Identity", + inputs=["dequantized"], + outputs=["output"], + name="identity", + ) + + graph = helper.make_graph( + nodes=[quantize, dequantize, identity], + name="qdq_graph", + inputs=[input_tensor], + outputs=[output_tensor], + initializer=[scale, zero_point], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 13)], + producer_name="trtutils-tests", + ) + model.ir_version = 8 + onnx.checker.check_model(model) + onnx.save(model, str(out_path)) + return out_path + + @pytest.fixture(scope="session") def _can_build_engine(onnx_path) -> bool: """Check if TRT can build engines on this hardware (session-cached).""" diff --git a/tests/builder/test_build.py b/tests/builder/test_build.py index 449b1ca3..65f07e21 100644 --- a/tests/builder/test_build.py +++ b/tests/builder/test_build.py @@ -15,18 +15,47 @@ from trtutils.compat._libs import trt from trtutils.core import cache as caching_tools +_BLACKWELL_WEAK_FAIL_REASON = ( + "Weakly-typed INT8/FP8 mixes require strongly_typed=True on Blackwell " + "(SM 10.0+); see test_build_strongly_typed_from_qdq for the replacement path." +) + @pytest.mark.parametrize( ("fp16", "int8", "fp8"), [ pytest.param(False, False, False, id="default"), pytest.param(True, False, False, id="fp16"), - pytest.param(False, True, False, id="int8"), + pytest.param( + False, + True, + False, + id="int8", + marks=pytest.mark.skipif(REAL_FLAGS.IS_BLACKWELL, reason=_BLACKWELL_WEAK_FAIL_REASON), + ), pytest.param(False, False, True, id="fp8"), - pytest.param(True, True, False, id="fp16-int8"), + pytest.param( + True, + True, + False, + id="fp16-int8", + marks=pytest.mark.skipif(REAL_FLAGS.IS_BLACKWELL, reason=_BLACKWELL_WEAK_FAIL_REASON), + ), pytest.param(True, False, True, id="fp16-fp8"), - pytest.param(False, True, True, id="int8-fp8"), - pytest.param(True, True, True, id="all"), + pytest.param( + False, + True, + True, + id="int8-fp8", + marks=pytest.mark.skipif(REAL_FLAGS.IS_BLACKWELL, reason=_BLACKWELL_WEAK_FAIL_REASON), + ), + pytest.param( + True, + True, + True, + id="all", + marks=pytest.mark.skipif(REAL_FLAGS.IS_BLACKWELL, reason=_BLACKWELL_WEAK_FAIL_REASON), + ), ], ) def test_build_precision( @@ -49,6 +78,10 @@ def test_build_precision( assert output_engine_path.stat().st_size > 0 +@pytest.mark.skipif( + REAL_FLAGS.IS_BLACKWELL, + reason="Runtime INT8/FP8 calibration is a weakly-typed concept; see test_build_strongly_typed_from_qdq.", +) @pytest.mark.parametrize( ("precision", "calibration"), [ @@ -347,3 +380,79 @@ def test_build_failure_raises(onnx_path, output_engine_path) -> None: with pytest.raises(RuntimeError, match="Failed to build engine"): build_engine(onnx_path, output_engine_path, optimization_level=1) + + +def test_build_strongly_typed_from_qdq(quantized_onnx_path, output_engine_path) -> None: + """Build from a pre-quantized Q/DQ ONNX with strongly_typed=True produces an engine.""" + build_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + optimization_level=1, + ) + assert output_engine_path.exists() + assert output_engine_path.stat().st_size > 0 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"fp16": True}, id="with-fp16"), + pytest.param({"int8": True}, id="with-int8"), + pytest.param({"fp8": True}, id="with-fp8"), + pytest.param({"layer_precision": [(0, trt.DataType.HALF)]}, id="with-layer-precision"), + ], +) +def test_build_strongly_typed_rejects_weak_args( + quantized_onnx_path, output_engine_path, kwargs +) -> None: + """strongly_typed=True combined with precision or per-layer overrides raises ValueError.""" + with pytest.raises(ValueError, match="strongly_typed"): + build_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + optimization_level=1, + **kwargs, + ) + + +def test_build_strongly_typed_rejects_calibration_args( + quantized_onnx_path, + output_engine_path, + calibration_cache_path, + synthetic_batcher, +) -> None: + """strongly_typed=True with calibration_cache or data_batcher raises ValueError.""" + with pytest.raises(ValueError, match="runtime calibration"): + build_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + calibration_cache=calibration_cache_path, + optimization_level=1, + ) + with pytest.raises(ValueError, match="runtime calibration"): + build_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + data_batcher=synthetic_batcher, + optimization_level=1, + ) + + +@pytest.mark.skipif( + not REAL_FLAGS.IS_BLACKWELL, + reason="Blackwell-only helpful-error check", +) +def test_build_blackwell_int8_fp8_rejects_weakly_typed(onnx_path, output_engine_path) -> None: + """On Blackwell, int8+fp8 without strongly_typed=True raises with a docs pointer.""" + with pytest.raises(ValueError, match=r"strongly_typed=True.*strongly_typed\.rst"): + build_engine( + onnx_path, + output_engine_path, + int8=True, + fp8=True, + optimization_level=1, + ) diff --git a/tests/builder/test_dla.py b/tests/builder/test_dla.py index e4869485..08116524 100644 --- a/tests/builder/test_dla.py +++ b/tests/builder/test_dla.py @@ -270,3 +270,64 @@ def test_layer_assignment_mapping( else: assert layer_precision[i] == (i, trt.DataType.HALF) assert layer_device[i] == (i, trt.DeviceType.GPU) + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"fp8": True}, id="with-fp8"), + pytest.param({"calibration_cache": "some_cache.cache"}, id="with-calibration-cache"), + ], +) +def test_build_dla_strongly_typed_rejects_weak_args( + quantized_onnx_path, output_engine_path, kwargs +) -> None: + """strongly_typed=True combined with fp8 or calibration_cache raises ValueError.""" + with pytest.raises(ValueError, match="strongly_typed"): + build_dla_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + dla_core=0, + optimization_level=1, + **kwargs, + ) + + +def test_build_dla_strongly_typed_rejects_batcher( + quantized_onnx_path, output_engine_path, synthetic_batcher +) -> None: + """strongly_typed=True combined with data_batcher raises ValueError.""" + with pytest.raises(ValueError, match="runtime calibration"): + build_dla_engine( + quantized_onnx_path, + output_engine_path, + data_batcher=synthetic_batcher, + strongly_typed=True, + dla_core=0, + optimization_level=1, + ) + + +def test_build_dla_weakly_typed_requires_batcher(onnx_path, output_engine_path) -> None: + """Weakly-typed DLA build without data_batcher raises ValueError.""" + with pytest.raises(ValueError, match="data_batcher is required"): + build_dla_engine( + onnx_path, + output_engine_path, + dla_core=0, + optimization_level=1, + ) + + +def test_build_dla_strongly_typed_from_qdq(quantized_onnx_path, output_engine_path) -> None: + """build_dla_engine with strongly_typed=True builds from a Q/DQ ONNX without a batcher.""" + build_dla_engine( + quantized_onnx_path, + output_engine_path, + strongly_typed=True, + dla_core=0, + optimization_level=1, + ) + assert output_engine_path.exists() + assert output_engine_path.stat().st_size > 0