Model-Optimizer/examples/onnx_ptq

ONNX Post-training quantization (PTQ)

This ONNX PTQ Toolkit provides a comprehensive suite of tools designed to optimize ONNX (Open Neural Network Exchange) models through quantization. Our toolkit is aimed at developers looking to enhance performance, reduce model size, and accelerate inference times without compromising the accuracy of their neural networks when deployed with TensorRT.

Quantization is an effective model optimization technique that compresses your models. Quantization with Model Optimizer can compress model size by 2x-4x, speeding up inference while preserving model quality.

Model Optimizer enables highly performant quantization formats including NVFP4, FP8, INT8, INT4 and supports advanced algorithms such as AWQ and Double Quantization with easy-to-use Python APIs.

Section	Description	Link	Docs
Pre-Requisites	Required & optional packages to use this technique	Link
Getting Started	Learn how to optimize your models using PTQ to reduce precision and improve inference efficiency	Link	docs
PyTorch to ONNX	Example scripts demonstrating how to quantize with PyTorch and then convert to ONNX	Link
Advanced Features	Examples demonstrating use advanced ONNX quantization features	Link
Resources	Extra links to relevant resources	Link

Pre-Requisites

Docker

Please use the TensorRT docker image (e.g., nvcr.io/nvidia/tensorrt:26.02-py3) or visit our installation docs for more information.

Note: If you are using onnxruntime-gpu, we recommend using nvcr.io/nvidia/tensorrt:25.06-py3 as it is built with CUDA 12, which is required by the stable onnxruntime-gpu package.

PETR and FAR3D containers

PETR and FAR3D share two targets from one Dockerfile. The evaluator target contains the legacy OpenMMLab stack used for data preparation, ONNX export, calibration, and final accuracy evaluation. The modelopt target uses the PyTorch 26.07 container for Model Optimizer, ONNX Runtime CUDA, and TensorRT 11.1 engine builds. Neither target creates a virtual environment.

From the Model Optimizer repository root:

docker build --target evaluator -f examples/onnx_ptq/Dockerfile -t modelopt-onnx-evaluator .
docker build --target modelopt -f examples/onnx_ptq/Dockerfile -t modelopt-onnx-trt11 .

Mount the same workspace into both containers to hand off ONNX models, calibration batches, and TensorRT engines:

docker run --rm -it --gpus=all --ipc=host \
  --user "$(id -u):$(id -g)" -e HOME=/tmp \
  -e USER="$(id -un)" -e LOGNAME="$(id -un)" \
  -v /path/to/workspace:/workspace \
  modelopt-onnx-evaluator

docker run --rm -it --gpus=all --ipc=host \
  --user "$(id -u):$(id -g)" -e HOME=/tmp \
  -e USER="$(id -un)" -e LOGNAME="$(id -un)" \
  -v /path/to/workspace:/workspace \
  modelopt-onnx-trt11

TensorRT engines must be built and evaluated with TensorRT 11.1.0.106 on the same GPU architecture. See the PETR and FAR3D guides for their source and dataset mounts.

Set the following environment variables inside the TensorRT docker.

export CUDNN_LIB_DIR=/usr/lib/x86_64-linux-gnu/
export LD_LIBRARY_PATH="${CUDNN_LIB_DIR}:${LD_LIBRARY_PATH}"

Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies.

Local Installation

Install Model Optimizer with onnx dependencies using pip from PyPI and install the requirements for the example:

pip install -U nvidia-modelopt[onnx]
pip install -r requirements.txt

For TensorRT Compiler framework workloads:

Install the latest TensorRT from here.

Getting Started

Prepare the example model

Most of the examples in this doc use vit_base_patch16_224.onnx as the input model. The model can be downloaded with the following script:

python download_example_onnx.py \
    --timm_model_name=vit_base_patch16_224 \
    --onnx_save_path=vit_base_patch16_224.onnx \
    --fp16 # <Optional, if the desired output ONNX precision is FP16>

Prepare calibration data

Calibration data is a representative subset of your training or validation dataset used during quantization to determine the optimal scale factors for converting floating-point values to lower precision formats (INT8, FP8, INT4). This data helps maintain model accuracy after quantization by analyzing the distribution of activations throughout the network.

First, prepare some calibration data. TensorRT recommends calibration data size to be at least 500 for CNN and ViT models. The following command picks up 500 images from the tiny-imagenet dataset and converts them to a numpy-format calibration array. Reduce the calibration data size for resource constrained environments.

python image_prep.py \
    --calibration_data_size=500 \
    --output_path=calib.npy \
    --fp16 # <Optional, if the input ONNX is in FP16 precision>

For Int4 quantization, it is recommended to set --calibration_data_size=64.

Quantize ONNX Model to FP8, INT8 or INT4

The model can be quantized as an FP8, INT8 or INT4 model using either the CLI or Python API. For FP8 and INT8 quantization, you have a choice between max and entropy calibration algorithms. For INT4 quantization, awq_clip or rtn_dq algorithms can be chosen.

For NVFP4 and MXFP8 ONNX, see the PyTorch to ONNX example.

Minimum opset requirements: int8 (13+), fp8 (21+), int4 (21+). ModelOpt will automatically upgrade lower opset versions to meet these requirements.

Option 1: Command-line interface

python -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch16_224.onnx \
    --quantize_mode=<fp8|int8|int4> \
    --calibration_data=calib.npy \
    --calibration_method=<max|entropy|awq_clip|rtn_dq> \
    --output_path=vit_base_patch16_224.quant.onnx

Option 2: Python API

from modelopt.onnx.quantization import quantize

quantize(
    onnx_path="vit_base_patch16_224.onnx",
    quantize_mode="int8",       # fp8, int8, int4 etc.
    calibration_data="calib.npy",
    calibration_method="max",   # max, entropy, awq_clip, rtn_dq etc.
    output_path="vit_base_patch16_224.quant.onnx",
)

Evaluate the quantized ONNX model

The evaluation script automatically downloads and uses the ILSVRC/imagenet-1k dataset from Hugging Face. This gated repository requires authentication via Hugging Face access token. See https://huggingface.co/docs/hub/en/security-tokens for details. The quantized ONNX ViT model can be evaluated on the ImageNet dataset as follows:

python evaluate.py \
    --onnx_path=<path to classification model> \
    --imagenet_path=<HF dataset card or local path to the ImageNet dataset> \
    --engine_precision=stronglyTyped \
    --model_name=vit_base_patch16_224

This script converts the quantized ONNX model to a TensorRT engine and does the evaluation with that engine. Finally, the evaluation result will be reported as follows:

The top1 accuracy of the model is <accuracy score between 0-100%>
The top5 accuracy of the model is <accuracy score between 0-100%>
Inference latency of the model is <X> ms

FAR3D 3D object detection

The FAR3D example exports and quantizes the FAR3D ONNX image encoder, builds TensorRT engines, and evaluates 3D object detection mAP on the Argoverse 2 validation set.

PETR 3D object detection

The PETR example exports and quantizes the PETRv1 and PETRv2 ONNX backbones, builds TensorRT engines, and evaluates 3D object detection mAP on the nuScenes validation set.

Advanced Features

Per node calibration of ONNX models

Per node calibration is a memory optimization feature designed to reduce memory consumption during quantization of large ONNX models. Instead of running inference over the entire network at once, this feature processes the model node-by-node, which can significantly reduce peak memory usage and prevent out-of-memory (OOM) errors.

How it works

When per node calibration is enabled, the quantization process:

1. Decomposes the model: Splits the original ONNX model into multiple single-node sub-models
2. Manages dependencies: Tracks input/output dependencies between nodes to ensure correct execution order
3. Processes sequentially: Runs calibration on each node individually using a topological processing order
4. Manages memory: Automatically cleans up intermediate results and manages reference counting to minimize memory usage
5. Aggregates results: Combines calibration data from all nodes to produce the final quantized model

When to use per node calibration

Per node calibration is particularly beneficial for:

• Large models that cause OOM errors during standard calibration
• Memory-constrained environments where GPU memory is limited
• Models with complex architectures that have high intermediate memory requirements

Usage

To enable per node calibration, add the --calibrate_per_node flag to your quantization command:

python -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch16_224.onnx \
    --quantize_mode=<int8/fp8> \
    --calibration_data=calib.npy \
    --calibrate_per_node \
    --output_path=vit_base_patch16_224.quant.onnx

Note

: Per node calibration is not available for INT4 quantization methods (awq_clip, rtn_dq)

Quantize an ONNX model with custom op

This feature requires TensorRT 10+ and ORT>=1.20. For proper usage, please make sure that the paths to libcudnn*.so and TensorRT lib/ are in the LD_LIBRARY_PATH env variable and that the tensorrt python package is installed.

A self-contained example is provided in the custom_op_plugin/ subfolder, based on leimao/TensorRT-Custom-Plugin-Example. Please see the steps below.

Step 1: Build the TensorRT plugin and create the sample ONNX model.

  1.1. Compile the TensorRT plugin:

cmake -S custom_op_plugin/plugin -B /tmp/plugin_build
cmake --build /tmp/plugin_build --config Release --parallel

This generates /tmp/plugin_build/libidentity_conv_plugin.so.

  1.2. Create the ONNX model with a custom IdentityConv operator:

python custom_op_plugin/create_identity_neural_network.py \
    --output_path=/tmp/identity_neural_network.onnx

Step 2: Quantize the ONNX model using the compiled plugin.

python -m modelopt.onnx.quantization \
    --onnx_path=/tmp/identity_neural_network.onnx \
    --trt_plugins=/tmp/plugin_build/libidentity_conv_plugin.so

Step 3: Deploy the quantized model with TensorRT.

trtexec --onnx=/tmp/identity_neural_network.quant.onnx \
    --staticPlugins=/tmp/plugin_build/libidentity_conv_plugin.so

Optimize Q/DQ node placement with Autotune

This feature automates Q/DQ (Quantize/Dequantize) node placement optimization for ONNX models using TensorRT performance measurements.

To access this feature in the ONNX quantization workflow, simply add --autotune in your CLI:

python -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch16_224.onnx \
    --quantize_mode=<fp8|int8|int4> \
    --calibration_data=calib.npy \
    --calibration_method=<max|entropy|awq_clip|rtn_dq> \
    --output_path=vit_base_patch16_224.quant.onnx \
    --autotune=<quick,default,extensive>

For more fine-tuned Autotune flags, please refer to the API guide and the Autotune guide.

Resources

• 📅 Roadmap
• 📖 Documentation
• 🎯 Benchmarks
• 💡 Release Notes
• 🐛 File a bug
• ✨ File a Feature Request

Technical Resources

There are many quantization schemes supported in the example scripts:

1. The FP8 format is available on the Hopper and Ada GPUs with CUDA compute capability greater than or equal to 8.9.

2. The INT4 AWQ is an INT4 weight only quantization and calibration method. INT4 AWQ is particularly effective for low batch inference where inference latency is dominated by weight loading time rather than the computation time itself. For low batch inference, INT4 AWQ could give lower latency than FP8/INT8 and lower accuracy degradation than INT8.

3. The NVFP4 is one of the new FP4 formats supported by NVIDIA Blackwell GPU and demonstrates good accuracy compared with other 4-bit alternatives. NVFP4 can be applied to both model weights as well as activations, providing the potential for both a significant increase in math throughput and reductions in memory footprint and memory bandwidth usage compared to the FP8 data format on Blackwell.