Model Demo: QAT Vision Transformer (ViT)

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/vit/QAT-ViT.ipynb.

PyTorch QAT -> ONNX Runtime Pipeline (ViT)

In this tutorial we use the PyTorch 2 Export (pt2e) library to perform quantization-aware training (QAT) on ViT-B/16, and export it such that it can be run through ONNX Runtime.

Notes:

• Training is expected to be done with a GPU, which can be mounted to the docker container with the option --gpus all or --gpus device=0 (replace 0 with whichever device you'd like to use)
  • NOTE: make sure you install torch and torchvision with the required CUDA version
    • In ../../requirements_gpu.txt we've included the index for cu124 (i.e. CUDA 12.4, also works for 12.6 on our end), but depending on training setup this may need to change
  • Original training results were produced using the following hardware specs
    • 13th Gen Intel Core i9-13900K
    • NVIDIA GeForce RTX 3090
  • To see the graph we have trained and processed in our own environment, see the included ./vit-quadric-pretrained.onnx
    • This version does not yet have the custom-ops, jupyter notebook has a bug when ORT is run on it, preventing the outputs from correctly displaying
      • The results of the custom-op replaced graph has the same outputs
    • To use it, run step (0) then skip to (7) and go from there, making sure to uncomment the necessary line in (8) to correctly specify using the pretrained quantized file
• pt2e is still in prototype phase (as of 03/19/2025), breaking changes may occur later
• Other PyTorch quantization libraries are available, but we currently make no guarantee the process will work as intended using them
• Each network requires slightly different post-processing, so not all networks may be supported yet with the post-processing steps currently implemented
• Images, labels, and models used in this notebook conform to the ImageNet-1K standard (i.e. use 224x224 image resolution) and have subjects represented in the 1000 ImageNet classes; any different input size or class labels will require user-supplied datasets

High Level Overview

1. Set up training helper functions
2. Define our quantization configuration
3. Prepare the training dataset
4. Load the model from the torchvision library
5. Perform QAT using the pt2e library
6. ONNX export, post-processing, custom op replacement
7. Prepare the validation dataset
8. Validate and compare against FP32 model
9. Lower the graph in CGC
10. Demo

NOTE: This demo is currently only supported on the ISS simulator.

0. Imports and setup

We install required package versions here. You may need to restart the kernel after this step to use them after running pip install.

%pip install -r ../../requirements_gpu.txt

Requirement already satisfied: torch==2.5.0 in /usr/local/lib/python3.10/dist-packages (from -r ../../requirements_gpu.txt (line 1)) (2.5.0)
Requirement already satisfied: torchvision==0.20.0 in /usr/local/lib/python3.10/dist-packages (from -r ../../requirements_gpu.txt (line 2)) (0.20.0)
Requirement already satisfied: s3fs in /usr/local/lib/python3.10/dist-packages (from -r ../../requirements_gpu.txt (line 3)) (2026.6.0)
Requirement already satisfied: onnxscript==0.2.0 in /usr/local/lib/python3.10/dist-packages (from -r ../../requirements_gpu.txt (line 4)) (0.2.0)
Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (1.13.1)
Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.127)
Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (2.21.5)
Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.5.8)
Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (11.2.1.3)
Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (2.8.5)
Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (10.3.5.147)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (3.16.1)
Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.127)
Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.3.1.170)
Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.127)
Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.127)
Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (12.4.127)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (2026.6.0)
Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (3.1.0)
Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (4.16.0)
Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (3.1.6)
Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (9.1.0.70)
Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /usr/local/lib/python3.10/dist-packages (from torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (11.6.1.9)
Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision==0.20.0->-r ../../requirements_gpu.txt (line 2)) (12.3.0)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torchvision==0.20.0->-r ../../requirements_gpu.txt (line 2)) (1.24.4)
Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from onnxscript==0.2.0->-r ../../requirements_gpu.txt (line 4)) (26.2)
Requirement already satisfied: ml_dtypes in /usr/local/lib/python3.10/dist-packages (from onnxscript==0.2.0->-r ../../requirements_gpu.txt (line 4)) (0.3.2)
Requirement already satisfied: onnx>=1.16 in /usr/local/lib/python3.10/dist-packages (from onnxscript==0.2.0->-r ../../requirements_gpu.txt (line 4)) (1.16.2)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (1.3.0)
Requirement already satisfied: aiobotocore<4.0.0,>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from s3fs->-r ../../requirements_gpu.txt (line 3)) (3.8.0)
Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0 in /usr/local/lib/python3.10/dist-packages (from s3fs->-r ../../requirements_gpu.txt (line 3)) (3.14.1)
Requirement already satisfied: botocore<1.43.47,>=1.43.3 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.43.46)
Requirement already satisfied: wrapt<3.0.0,>=1.10.10 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (2.2.2)
Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (2.9.0.post0)
Requirement already satisfied: multidict<7.0.0,>=6.0.0 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (6.7.1)
Requirement already satisfied: aioitertools<1.0.0,>=0.5.1 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (0.13.0)
Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.1.0)
Requirement already satisfied: async-timeout<6.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (5.0.1)
Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (0.5.2)
Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (26.1.0)
Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.24.2)
Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (2.7.1)
Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.4.0)
Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.8.0)
Requirement already satisfied: protobuf>=3.20.2 in /usr/local/lib/python3.10/dist-packages (from onnx>=1.16->onnxscript==0.2.0->-r ../../requirements_gpu.txt (line 4)) (4.25.3)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch==2.5.0->-r ../../requirements_gpu.txt (line 1)) (3.0.3)
Requirement already satisfied: urllib3!=2.2.0,<3,>=1.25.4 in /usr/local/lib/python3.10/dist-packages (from botocore<1.43.47,>=1.43.3->aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.26.20)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil<3.0.0,>=2.1->aiobotocore<4.0.0,>=2.19.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (1.17.0)
Requirement already satisfied: idna>=2.0 in /usr/local/lib/python3.10/dist-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1,>=3.9.0->s3fs->-r ../../requirements_gpu.txt (line 3)) (3.18)
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
Note: you may need to restart the kernel to use updated packages.

import os
import sys
import time
import copy
import logging
import itertools
import warnings
import tempfile
import random
import operator
from pathlib import Path
from dataclasses import dataclass
from datasets import load_from_disk
from typing import Optional
import numpy as np
from PIL import Image
import json

import onnx

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
from torch.utils.data import DataLoader, Subset
from torch._export import capture_pre_autograd_graph
from torch.export import export
from torch.ao.quantization.quantize_pt2e import (
    prepare_qat_pt2e,
    convert_pt2e,
)
from torch.ao.quantization.quantizer import (
    Quantizer,
    QuantizationSpec,
    FixedQParamsQuantizationSpec,
    QuantizationAnnotation,
)
from torch.ao.quantization.fake_quantize import FakeQuantize
from torch.ao.quantization.quantizer import QuantizationSpec
from torch.ao.quantization.observer import MinMaxObserver
from torch.fx.passes.utils.source_matcher_utils import get_source_partitions

import torchvision
from torchvision.datasets import ImageFolder
from torchvision.models import (
    vit_b_16,
    ViT_B_16_Weights,
)
from torchvision.transforms import (
    RandomResizedCrop,
    RandomHorizontalFlip,
    CenterCrop,
    Compose,
    Normalize,
    Resize,
    ToTensor,
)

import sdk_cli.lib.qat_processor as processor

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from tvm.relay.backend.contrib.epu.util import logger as tvm_logger

from tvm.contrib.epu.chimera_job.constants import DEFAULT_ONNX_OPSET
from sdk_cli.utils.dataloaders import CalibrationDataLoader
from sdk_cli.utils.datasets import ImageNet_Mini_Quadric, QuadricCalibration
from sdk_cli.utils.datasets.ImageNet import (
    IMAGENET_1K_NORMALIZATION_PARAMETERS,
    OPTIMIZED_IMAGENET_1K_LABELS,
)
from sdk_cli.utils.model_helpers import ClassifyResult
from sdk_cli.utils.performance_trackers import ClassifierPerformanceTracker
from sdk_cli.lib.quantize import (
    QuantizationExperiment,
    QuantizedONNXModel,
    run_quantization_experiment,
    compute_tensor_ranges,
)
from sdk_cli.lib.inference import InferenceEngine, batch_inference, single_inference
from sdk_cli.visualizers.layouter import ClassifierLayouter

from vit_helpers import create_vit_custom_op_model

warnings.filterwarnings(action="default", module=r"torch.ao.quantization")
sys.tracebacklimit = 0

## NOTE: below are manually set for reproducibility
torch.manual_seed(191009)
np.random.seed(2147483648)
random.seed(2147483648)
torch.use_deterministic_algorithms(True)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

torch.set_default_device("cuda")

1. Set up training helper functions

These will assist us in creating our training loop, as per the pt2e documentation page listed above.

class AverageMeter(object):
    def __init__(self, name, fmt=":f"):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        return f"{name} {val:f} ({avg:f})"


def accuracy(output, target, topk=(1,)):
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def train_one_epoch(model, criterion, optimizer, data_loader, device, ntrain_batches):
    top1 = AverageMeter("Acc@1")
    top5 = AverageMeter("Acc@5")
    avgloss = AverageMeter("Loss")

    cnt = 0
    for batch in data_loader:
        start_time = time.time()
        image, target = batch.values()
        print(".", end="")
        cnt += 1
        image = image.to(device)
        target = target.to(device)
        output = model(image)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        top1.update(acc1[0], image.size(0))
        top5.update(acc5[0], image.size(0))
        avgloss.update(loss, image.size(0))
        if cnt >= ntrain_batches:
            print("Loss", avgloss.avg)
            print(
                "Training: * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}".format(top1=top1, top5=top5)
            )
            return

    print(
        "Full imagenet train set:  * Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}".format(
            top1=top1, top5=top5
        )
    )


def evaluate(model, data_loader, device, neval_batches):
    top1 = AverageMeter("Acc@1", ":6.2f")
    top5 = AverageMeter("Acc@5", ":6.2f")
    cnt = 0
    with torch.no_grad():
        for image, target in data_loader:
            image = image.to(device)
            target = target.to(device)
            output = model(image)
            cnt += 1
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            top1.update(acc1[0], image.size(0))
            top5.update(acc5[0], image.size(0))
            if cnt >= neval_batches:
                return top1, top5
    print("")

    return top1, top5

2. Define our quantization configuration

The pt2e quantization library is extremely flexible. It allows us to define the method of quantization for each tensor of each individual operation. Here we will define a simple quantization configuration, but you may choose to experiment with the options available.

@dataclass(eq=True, frozen=True)
class QuantizationConfig:
    input_activation: Optional[QuantizationSpec]
    output_activation: Optional[QuantizationSpec]
    weight: Optional[QuantizationSpec]
    bias: Optional[QuantizationSpec]
    is_qat: bool = False


quadric_weight_fake_quant = FakeQuantize.with_args(
    observer=MinMaxObserver,
    quant_min=-128,
    quant_max=127,
    dtype=torch.qint8,
    qscheme=torch.per_tensor_symmetric,
)

quadric_activation_fake_quant = FakeQuantize.with_args(
    observer=MinMaxObserver,
    quant_min=-128,
    quant_max=127,
    dtype=torch.qint8,
    qscheme=torch.per_tensor_symmetric,
)


class QuadricQuantizer(Quantizer):
    def __init__(self):
        super().__init__()
        self.global_config: QuantizationConfig = None

    def set_global_config(self, quant_config: QuantizationConfig):
        self.global_config = quant_config
        return self

    def get_default_config(self):
        activation_spec = QuantizationSpec(
            observer_or_fake_quant_ctr=quadric_activation_fake_quant,
            dtype=torch.int8,
            quant_min=-128,
            quant_max=127,
            qscheme=torch.per_tensor_symmetric,
        )
        weight_spec = QuantizationSpec(
            observer_or_fake_quant_ctr=quadric_weight_fake_quant,
            dtype=torch.int8,
            quant_min=-128,
            quant_max=127,
            qscheme=torch.per_tensor_symmetric,  # NOTE: currently only supporting symmetric weight quantization
        )
        quant_config = QuantizationConfig(
            input_activation=activation_spec,
            output_activation=activation_spec,
            weight=weight_spec,
            bias=None,
            is_qat=True,
        )
        return quant_config

    def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
        self._annotate_concat(model)
        self._annotate_conv2d(model)
        self._annotate_gemm(model)
        self._annotate_attention_linear(model)
        self._annotate_attention_dot_product(model)
        return model

    def _annotate_concat(self, model: torch.fx.GraphModule):
        concat_partitions = get_source_partitions(
            model.graph, [torch.cat, torch.concat, torch.concatenate]
        )
        concat_partitions = list(itertools.chain(*concat_partitions.values()))
        for partition in concat_partitions:
            concat_node = partition.output_nodes[0]
            input_qspec_map = {
                concat_node.args[0][0]: self.global_config.input_activation,
                concat_node.args[0][1]: self.global_config.input_activation,
            }
            concat_node.meta["quantization_annotation"] = QuantizationAnnotation(
                input_qspec_map=input_qspec_map,
                output_qspec=self.global_config.output_activation,
                _annotated=True,
            )

    def _annotate_conv2d(self, model: torch.fx.GraphModule):
        conv_partitions = get_source_partitions(model.graph, [nn.Conv2d, F.conv2d])
        conv_partitions = list(itertools.chain(*conv_partitions.values()))
        for partition in conv_partitions:
            conv_node = partition.output_nodes[0]
            input_qspec_map = {
                conv_node.args[0]: self.global_config.input_activation,
                conv_node.args[1]: self.global_config.weight,
            }
            conv_node.meta["quantization_annotation"] = QuantizationAnnotation(
                input_qspec_map=input_qspec_map,
                output_qspec=self.global_config.output_activation,
                _annotated=True,
            )

    def _annotate_gemm(self, model: torch.fx.GraphModule):
        gemm_partitions = get_source_partitions(
            model.graph, [torch.matmul, torch.mm, nn.Linear, F.linear]
        )
        gemm_partitions = list(itertools.chain(*gemm_partitions.values()))
        for partition in gemm_partitions:
            gemm_node = partition.output_nodes[0]
            input_qspec_map = {
                gemm_node.args[0]: self.global_config.input_activation,
                gemm_node.args[1]: self.global_config.input_activation,
            }
            gemm_node.meta["quantization_annotation"] = QuantizationAnnotation(
                input_qspec_map=input_qspec_map,
                output_qspec=self.global_config.output_activation,
                _annotated=True,
            )

    def _annotate_attention_linear(self, model: torch.fx.GraphModule):
        for n in model.graph.nodes:
            if n.op != "call_function" or n.target not in [
                torch.ops.aten.linear.default,
            ]:
                continue

            fc_node = n
            if (
                fc_node.meta.get("quantization_annotation", None)
                and fc_node.meta["quantization_annotation"]._annotated
            ):
                continue
            input_qspec_map = {
                fc_node.args[0]: self.global_config.input_activation,
                fc_node.args[1]: self.global_config.input_activation,
            }
            fc_node.meta["quantization_annotation"] = QuantizationAnnotation(
                input_qspec_map=input_qspec_map,
                output_qspec=self.global_config.output_activation,
                _annotated=True,
            )

    def _annotate_attention_dot_product(self, model: torch.fx.GraphModule):
        for n in model.graph.nodes:
            if n.op != "call_function" or n.target not in [
                torch.ops.aten.scaled_dot_product_attention.default,
            ]:
                continue

            dp_node = n
            if (
                dp_node.meta.get("quantization_annotation", None)
                and dp_node.meta["quantization_annotation"]._annotated
            ):
                continue
            input_qspec_map = {
                dp_node.args[0]: self.global_config.input_activation,
                dp_node.args[1]: self.global_config.input_activation,
                dp_node.args[2]: self.global_config.input_activation,
            }
            dp_node.meta["quantization_annotation"] = QuantizationAnnotation(
                input_qspec_map=input_qspec_map,
                output_qspec=self.global_config.output_activation,
                _annotated=True,
            )

    def validate(self, model):
        pass

    @classmethod
    def get_supported_operators(cls):
        return []

3. Prepare the training dataset

We use images from ImageNet-1K to perform the training and validation. All input tensors in this dataset are of dimension NCHW format [3, 224, 224] and need to be float-normalized with a mean, sigma of [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]. For training, we include the transformations RandomResizedCrop and RandomHorizontalFlip for robustness.

normalize = Normalize(
    IMAGENET_1K_NORMALIZATION_PARAMETERS.channel_means,
    IMAGENET_1K_NORMALIZATION_PARAMETERS.channel_standard_deviations,
)
train_transforms = Compose(
    [
        RandomResizedCrop(224),
        RandomHorizontalFlip(),
        ToTensor(),
        normalize,
    ]
)
test_transforms = Compose(
    [
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        normalize,
    ]
)


def apply_transforms(imgs):
    imgs["image"] = [train_transforms(img.convert("RGB")) for img in imgs["image"]]
    return imgs


train_set = load_from_disk(
    "s3://sdk-cli-datasets/imagenet_1k_val_train.hf/validation",
    storage_options={"anon": True},
)
train_set.set_transform(apply_transforms)
test_set = ImageFolder(
    root="../../common/validation/imagenet-mini-quadric/val", transform=test_transforms
)

train_batch_size = 1
test_batch_size = 1
generator_train = torch.Generator(device="cuda")
generator_train.manual_seed(67280421310721)
generator_test = torch.Generator(device="cuda")
generator_test.manual_seed(672804213107421)
sampler_train = torch.utils.data.RandomSampler(train_set, generator=generator_train)
sampler_test = torch.utils.data.SequentialSampler(test_set)
data_loader_train = torch.utils.data.DataLoader(
    train_set,
    batch_size=train_batch_size,
    generator=generator_train,
    sampler=sampler_train,
)
data_loader_test = torch.utils.data.DataLoader(
    test_set,
    batch_size=test_batch_size,
    generator=generator_test,
    sampler=sampler_test,
)

/usr/local/lib/python3.10/dist-packages/datasets/table.py:1421: FutureWarning: promote has been superseded by promote_options='default'.
  table = cls._concat_blocks(blocks, axis=0)

4. Load vit_b_16 from torchvision

The original graph goes through a 2-step process using torch.export.export_for_training and the previously defined QuadricQuantizer in order to be ready for QAT.

float_model = vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)

example_input = torch.rand(1, 3, 224, 224)
exported_model = torch.export.export_for_training(float_model, (example_input,)).module()

quantizer = QuadricQuantizer()
quantizer.set_global_config(quantizer.get_default_config())
prepared_model = prepare_qat_pt2e(exported_model, quantizer)
for n in prepared_model.graph.nodes:
    if n.target == torch.ops.aten._native_batch_norm_legit.default:
        n.target = torch.ops.aten.cudnn_batch_norm.default
_ = prepared_model.recompile()

/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning: 'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning: 'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()

5. Perform QAT using the pt2e library

We follow the documentation given here. We start by defining some hyperparameters here. You can alter these parameters as you like.

We didn't need them to get within 1% of FP32, but here are other things to consider implementing if you would like to try for greater accuracy:

• More robust training image transformations
• Knowledge distillation
• Learning rate scheduling

num_epochs = 54
num_train_batches = 32
num_eval_batches = 256
num_observer_update_epochs = 1000
num_batch_norm_update_epochs = 0
num_epochs_between_evals = 25

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(prepared_model.parameters(), lr=2e-6, weight_decay=1e-5)

num_observer_update_flag = True
num_batch_norm_update_flag = True

for epoch in range(num_epochs):
    if epoch + 1 >= num_observer_update_epochs and num_observer_update_flag:
        print("Disabling observer for subseq epochs, epoch = ", epoch)
        prepared_model.apply(torch.ao.quantization.disable_observer)
        num_observer_update_flag = False

    if epoch + 1 >= num_batch_norm_update_epochs and num_batch_norm_update_flag:
        print("Freezing BN for subseq epochs, epoch = ", epoch)
        for n in prepared_model.graph.nodes:
            if n.target in [
                torch.ops.aten._native_batch_norm_legit.default,
                torch.ops.aten.cudnn_batch_norm.default,
            ]:
                new_args = list(n.args)
                new_args[5] = False
                n.args = tuple(new_args)
        prepared_model.recompile()
        num_batch_norm_update_flag = False

    train_one_epoch(
        prepared_model,
        criterion,
        optimizer,
        data_loader_train,
        "cuda",
        num_train_batches,
    )
    print("^ Epoch: %d" % (epoch + 1))

    if (epoch + 1) % num_epochs_between_evals == 0:
        prepared_model_copy = copy.deepcopy(prepared_model)
        torch.set_default_device("cpu")
        prepared_model_copy = prepared_model_copy.to("cpu")
        prepared_model_copy.recompile()

        quantized_model = convert_pt2e(prepared_model_copy)
        torch.set_default_device("cuda")
        quantized_model = quantized_model.to("cuda")
        quantized_model.recompile()
        torch.ao.quantization.move_exported_model_to_eval(quantized_model)

        top1, top5 = evaluate(
            quantized_model,
            data_loader_test,
            device="cuda",
            neval_batches=num_eval_batches,
        )
        print(
            "Epoch %d: Evaluation accuracy on %d images, %2.2f"
            % (epoch + 1, num_eval_batches * test_batch_size, top1.avg)
        )

Freezing BN for subseq epochs, epoch =  0
................................Loss tensor(1.2219, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 93.750
^ Epoch: 1
................................Loss tensor(0.7634, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 100.000
^ Epoch: 2
................................Loss tensor(1.0186, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 96.875
^ Epoch: 3
................................Loss tensor(0.7405, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 96.875
^ Epoch: 4
................................Loss tensor(1.0904, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 90.625
^ Epoch: 5
................................Loss tensor(1.2480, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 90.625
^ Epoch: 6
................................Loss tensor(1.5587, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 78.125
^ Epoch: 7
................................Loss tensor(0.5967, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 87.500 Acc@5 93.750
^ Epoch: 8
................................Loss tensor(0.4473, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 90.625 Acc@5 100.000
^ Epoch: 9
................................Loss tensor(0.4586, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 84.375 Acc@5 100.000
^ Epoch: 10
................................Loss tensor(1.4666, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 84.375
^ Epoch: 11
................................Loss tensor(0.7329, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 100.000
^ Epoch: 12
................................Loss tensor(1.9456, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 62.500 Acc@5 78.125
^ Epoch: 13
................................Loss tensor(1.8045, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 78.125
^ Epoch: 14
................................Loss tensor(1.5782, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 62.500 Acc@5 84.375
^ Epoch: 15
................................Loss tensor(1.1447, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 93.750
^ Epoch: 16
................................Loss tensor(0.8192, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 84.375 Acc@5 93.750
^ Epoch: 17
................................Loss tensor(1.4559, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 81.250
^ Epoch: 18
................................Loss tensor(1.0137, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 96.875
^ Epoch: 19
................................Loss tensor(1.2055, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 90.625
^ Epoch: 20
................................Loss tensor(1.1886, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 90.625
^ Epoch: 21
................................Loss tensor(0.5833, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 87.500 Acc@5 96.875
^ Epoch: 22
................................Loss tensor(1.2215, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 84.375
^ Epoch: 23
................................Loss tensor(1.1601, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 84.375
^ Epoch: 24
................................Loss tensor(1.2698, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 93.750
^ Epoch: 25


W0718 14:16:18.338000 602 torch/_export/__init__.py:64] +============================+
W0718 14:16:18.338000 602 torch/_export/__init__.py:65] |     !!!   WARNING   !!!    |
W0718 14:16:18.339000 602 torch/_export/__init__.py:66] +============================+
W0718 14:16:18.339000 602 torch/_export/__init__.py:67] capture_pre_autograd_graph() is deprecated and doesn't provide any function guarantee moving forward.
W0718 14:16:18.339000 602 torch/_export/__init__.py:68] Please switch to use torch.export.export_for_training instead.


Epoch 25: Evaluation accuracy on 256 images, 78.12
................................Loss tensor(0.9063, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 96.875
^ Epoch: 26
................................Loss tensor(0.9574, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 90.625
^ Epoch: 27
................................Loss tensor(1.3191, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 90.625
^ Epoch: 28
................................Loss tensor(1.4324, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 62.500 Acc@5 93.750
^ Epoch: 29
................................Loss tensor(2.0833, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 59.375 Acc@5 75.000
^ Epoch: 30
................................Loss tensor(0.6915, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 96.875
^ Epoch: 31
................................Loss tensor(1.2225, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 87.500
^ Epoch: 32
................................Loss tensor(1.7903, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 62.500 Acc@5 84.375
^ Epoch: 33
................................Loss tensor(0.7729, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 84.375 Acc@5 90.625
^ Epoch: 34
................................Loss tensor(1.2715, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 65.625 Acc@5 90.625
^ Epoch: 35
................................Loss tensor(0.5766, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 93.750 Acc@5 96.875
^ Epoch: 36
................................Loss tensor(1.5359, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 56.250 Acc@5 90.625
^ Epoch: 37
................................Loss tensor(0.9247, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 93.750
^ Epoch: 38
................................Loss tensor(0.9876, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 78.125 Acc@5 90.625
^ Epoch: 39
................................Loss tensor(1.3272, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 90.625
^ Epoch: 40
................................Loss tensor(0.8833, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 90.625
^ Epoch: 41
................................Loss tensor(1.0470, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 87.500
^ Epoch: 42
................................Loss tensor(0.8576, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 100.000
^ Epoch: 43
................................Loss tensor(1.0788, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 90.625
^ Epoch: 44
................................Loss tensor(1.6147, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 68.750 Acc@5 87.500
^ Epoch: 45
................................Loss tensor(1.2469, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 71.875 Acc@5 87.500
^ Epoch: 46
................................Loss tensor(0.7781, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 96.875
^ Epoch: 47
................................Loss tensor(0.6198, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 84.375 Acc@5 93.750
^ Epoch: 48
................................Loss tensor(1.0182, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 90.625
^ Epoch: 49
................................Loss tensor(0.6933, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 96.875
^ Epoch: 50
Epoch 50: Evaluation accuracy on 256 images, 77.34
................................Loss tensor(0.8369, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 84.375 Acc@5 90.625
^ Epoch: 51
................................Loss tensor(0.7830, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 81.250 Acc@5 96.875
^ Epoch: 52
................................Loss tensor(1.2071, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 75.000 Acc@5 90.625
^ Epoch: 53
................................Loss tensor(0.4963, device='cuda:0', grad_fn=<DivBackward0>)
Training: * Acc@1 90.625 Acc@5 96.875
^ Epoch: 54

6. ONNX export, post-processing, custom op replacement

We use torch.onnx.dynamo_export() to export the graph from PyTorch to ONNX and use Quadric's QAT post-processor in order to clean up the graph and format it nicely. We start here with the conversion from PyTorch to ONNX.

<div class="alert alert-block alert-warning"> <b>Note:</b> Because torch.onnx.dynamo_export is a work in progress, there will be many warnings generated for now. There are no adverse effects. These warnings are applicable to other areas of PyTorch but not the library we use here for QAT, so they can be ignored. But they cannot be suppressed as they are thrown by the underlying C++. </div>

## Workaround for a bug in convert_pt2e where the function would fail
## to execute due to improperly handled tensors when GPUs are used
torch.set_default_device("cpu")
prepared_model = prepared_model.to("cpu")
prepared_model.recompile()
quantized_model = convert_pt2e(prepared_model)
torch.set_default_device("cuda")
quantized_model = quantized_model.to("cuda")
quantized_model.recompile()

torch.ao.quantization.move_exported_model_to_eval(quantized_model)
top1, top5 = evaluate(
    quantized_model, data_loader_test, device="cuda", neval_batches=num_eval_batches
)
print(
    "Final evaluation accuracy on %d images, (Top-1): %2.2f"
    % (num_eval_batches * test_batch_size, top1.avg)
)
print(
    "Final evaluation accuracy on %d images, (Top-5): %2.2f"
    % (num_eval_batches * test_batch_size, top5.avg)
)

onnx_program = torch.onnx.dynamo_export(quantized_model, example_input)
onnx_program.save("vit-unprocessed.onnx")

Final evaluation accuracy on 256 images, (Top-1): 78.52
Final evaluation accuracy on 256 images, (Top-5): 95.70


/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/_exporter_legacy.py:116: UserWarning: torch.onnx.dynamo_export only implements opset version 18 for now. If you need to use a different opset version, please register them with register_custom_op.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/passes/readability.py:52: UserWarning: Attempted to insert a get_attr Node with no underlying reference in the owning GraphModule! Call GraphModule.add_submodule to add the necessary submodule, GraphModule.add_parameter to add the necessary Parameter, or nn.Module.register_buffer to add the necessary buffer
  new_node = self.module.graph.get_attr(normalized_name)
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1586: UserWarning: Node _frozen_param0 target _frozen_param0 _frozen_param0 of  does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1586: UserWarning: Node _frozen_param1 target _frozen_param1 _frozen_param1 of  does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1586: UserWarning: Node _frozen_param2 target _frozen_param2 _frozen_param2 of  does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1586: UserWarning: Node _frozen_param3 target _frozen_param3 _frozen_param3 of  does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1586: UserWarning: Node _frozen_param4 target _frozen_param4 _frozen_param4 of  does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py:1593: UserWarning: Additional 46 warnings suppressed about get_attr references
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/onnxfunction_dispatcher.py:503: FutureWarning: 'onnxscript.values.TracedOnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  self.param_schema = self.onnxfunction.param_schemas()


Applied 37 of general pattern rewrite rules.

The next step uses Quadric's PyTorch graph post-processor. We avoid quantizing residual adds here for accuracy preservation

Supported ops/patterns:

• Conv (+ BatchNormalization + activation)
• Gemm
• MatMul
• Concat
• LayerNormalization
• GlobalAveragePool
• Add
• Softmax
• aten_bernoulli_p PyTorch dropout layers
• _aten_gelu_approximate_none PyTorch exported GELU

Unsupported ops/patterns:

• ATen operators from PyTorch aside from aten_bernoulli_p and _aten_gelu_approximate_none

torch.set_default_device("cpu")

output_path = "./vit-processed.onnx"
qat_processor = processor.QATProcessor()
qat_processor.process(
    "./vit-unprocessed.onnx",
    q_add_flag=False,
    ptq_sym_act=True,
    strip_onnx=False,
    output_path=output_path,
)

2026-07-18 14:18 - DEBUG - epu - qat_processor - Preparing calibration data
/usr/local/lib/python3.10/dist-packages/datasets/table.py:1421: FutureWarning: promote has been superseded by promote_options='default'.
  table = cls._concat_blocks(blocks, axis=0)
2026-07-18 14:18 - DEBUG - epu - qat_processor - Fixing PyTorch exported Unsqueeze axes arg
2026-07-18 14:18 - DEBUG - epu - qat_processor - Removing aten.bernoulli ops
2026-07-18 14:18 - DEBUG - epu - qat_processor - Removing NOP aten.pad ops
2026-07-18 14:18 - DEBUG - epu - qat_processor - Removing aten.as_strided ops
2026-07-18 14:18 - DEBUG - epu - qat_processor - Swapping aten.gelu ops for ONNX equivalent
2026-07-18 14:18 - DEBUG - epu - qat_processor - Swapping aten.roll ops for ONNX equivalent
2026-07-18 14:18 - DEBUG - epu - qat_processor - Swapping ReduceMean for GlobalAveragePool
2026-07-18 14:18 - DEBUG - epu - qat_processor - Removing Dropout ops with rate=0
2026-07-18 14:18 - DEBUG - epu - qat_processor - Removing redundant QDQs
2026-07-18 14:18 - DEBUG - epu - qat_processor - Folding Div into BatchNormalization
2026-07-18 14:18 - DEBUG - epu - qat_processor - Folding BatchNormalization into Conv nodes with biases
2026-07-18 14:18 - DEBUG - epu - qat_processor - Folding BatchNormalization into Conv nodes without biases
2026-07-18 14:18 - DEBUG - epu - qat_processor - Quantizing Swin MHA blocks
2026-07-18 14:18 - DEBUG - epu - qat_processor - Quantizing FC layers
WARNING:root:Please use QuantFormat.QDQ for activation type QInt8 and weight type QInt8. Or it will lead to bad performance on x64.
WARNING:root:Please check if the model is already quantized. Note you don't need to quantize a QAT model. OnnxRuntime support to run QAT model directly.
2026-07-18 14:23 - DEBUG - epu - qat_processor - Quantizing QKV MatMul ops
WARNING:root:Please use QuantFormat.QDQ for activation type QInt8 and weight type QInt8. Or it will lead to bad performance on x64.
WARNING:root:Please check if the model is already quantized. Note you don't need to quantize a QAT model. OnnxRuntime support to run QAT model directly.
2026-07-18 14:27 - DEBUG - epu - qat_processor - Changing to Opset 17
2026-07-18 14:27 - DEBUG - epu - qat_processor - Performing ORT optimizations
2026-07-18 14:27 - DEBUG - epu - qat_processor - Unfolding QLinearSoftmax to avoid inaccurate LUT
2026-07-18 14:27 - DEBUG - epu - qat_processor - Fusing QLinearConv
2026-07-18 14:27 - DEBUG - epu - qat_processor - Fusing QGemm
2026-07-18 14:27 - DEBUG - epu - qat_processor - Folding DQ into MatMul ops where ORT does not
2026-07-18 14:27 - DEBUG - epu - qat_processor - Moving DQ outside of Roll patterns
2026-07-18 14:27 - DEBUG - epu - qat_processor - Moving from uint8 to int8

7. Prepare the validation dataset

Now that the QAT model has been saved to disk, we will compare it to the original pretrained FP32 model from PyTorch. Here we use a set of 3500 images in our validation of the model.

torch.set_default_device("cpu")

MAX_NUM_SAMPLES_FOR_ACCURACY_CALCULATION = 3500
graph_input_name = "l_x_"

normalize = Normalize(
    IMAGENET_1K_NORMALIZATION_PARAMETERS.channel_means,
    IMAGENET_1K_NORMALIZATION_PARAMETERS.channel_standard_deviations,
)
test_transforms = Compose(
    [
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        normalize,
    ]
)

dataset = ImageNet_Mini_Quadric.Dataset(transform=test_transforms)
subset_of_dataset = Subset(dataset, range(MAX_NUM_SAMPLES_FOR_ACCURACY_CALCULATION))
calibration_dataloader = CalibrationDataLoader(
    DataLoader(subset_of_dataset, batch_size=1, shuffle=True), [graph_input_name]
)

8. Validate and compare against FP32 model

Now we want to make sure that quantization was done succesfully by checking the performance of the original model against the quantized model. On our hardware, we see FP32 is at 79.97% and INT8 is at 79.2%. This can be verified by uncommenting the necessary line below to switch to the pretrained vit-quadric-pretrained.onnx.

Depending on the task, you may choose different metrics like spearman correlation, f1-score, AUC, ROC, L2 norm, IoU etc.

pytorch_model = vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)
fp32_path = f"./{pytorch_model.__class__.__name__}_float32.onnx"
int8_path = "./vit-processed.onnx"

## Uncomment the line below if using the pretrained graph
## int8_path = "./vit-quadric-pretrained.onnx"

onnx_model_path = Path(fp32_path)
quantized_onnx_model = QuantizedONNXModel(int8_path, None)
input_tensor_shape = (1, 3, 224, 224)
example_input = torch.randn(*input_tensor_shape, requires_grad=True)
torch.onnx.export(
    pytorch_model,
    example_input,
    str(onnx_model_path),
    export_params=True,
    do_constant_folding=True,
    opset_version=DEFAULT_ONNX_OPSET,
    input_names=[graph_input_name],
    output_names=["dequantize_per_tensor_200"],
)

## Evaluation
quantization_experiment = QuantizationExperiment(
    floating_point_onnx_model_path=onnx_model_path,
    quantized_onnx_model=quantized_onnx_model,
    performance_tracker=ClassifierPerformanceTracker(),
)
run_quantization_experiment(
    quantization_experiment,
    calibration_dataloader,
    export_path=None,
    max_num_samples=MAX_NUM_SAMPLES_FOR_ACCURACY_CALCULATION,
)

3500: FP3279.97% <> INT879.20%: 100%|██████████| 3500/3500 [16:27<00:00,  3.54it/s]


Used 3500 image samples for model accuracy comparison.
Original FP32 model accuracy (Top-1): 79.97%
Quantized INT8 model accuracy (Top-1): 79.20%
Change in Top-1 model accuracy due to quantization: 0.77%

Since the user can choose to train the ONNX graph themselves or use a pretrained version, the name of the ONNX file can change. This is needed in the directive file vit-quadric-pretrained_directive.json. This will set the ONNX file name to be the same as the one chosen by the user.

More information about the directives file is mentioned in Step 10.

def modify_json_file(json_file_name, onnx_file_name):
    json_data = {}
    with open(json_file_name, "r") as json_file:
        json_data = json.load(json_file)
        if "model_info" not in json_data or "onnx_file" not in json_data["model_info"]:
            print(f"Malformed json_file {json_file_name}")
            return
        json_data["model_info"]["onnx_file"] = onnx_file_name

    with open(json_file_name, "w") as json_file:
        json.dump(json_data, json_file, indent=2)


directives_file = f"{os.getcwd()}/vit-quadric-pretrained_directives.json"
modify_json_file(directives_file, int8_path)

9. Lower the graph in CGC

We first replace the appropriate patch-creation subgraph with a QuadricCustomOp node, create the tensor-ranges file, then lower the graph.

Previously, we were replacing attention mechanism with a custom-op. However, we have added support in our compiler to capture attention mechanism subgraph an wrap it in a general attention op.

custom_op_path = "./vit-custom-op.onnx"
custom_op_model = create_vit_custom_op_model(int8_path)
onnx.save(custom_op_model, custom_op_path)

warnings.filterwarnings(action="ignore")
tvm_logger.setLevel(logging.WARNING)

calibration_dataset = ImageNet_Mini_Quadric.Dataset(transform=test_transforms)
trange_path = f"{custom_op_path}.tranges"

_ = compute_tensor_ranges(
    onnx_model_path=custom_op_path,
    calibration_dataloader=CalibrationDataLoader(calibration_dataset, [graph_input_name]),
    tensor_ranges_export_path=trange_path,
)

2026-07-18 15:01 - INFO - sdk - quantize - Saved computed tensor ranges to ./vit-custom-op.onnx.tranges.

General Attention Op

If you run the below ChimeraJob, it will extract the following subgraph into a general attention Op.

Subgraph to be replaced with General Attention Op

The compiler creates a C++ file called attention_stubs.hpp. This file contains functions that the user must fill out. We have inserted static_asserts into the body of each such functions so that the compiler can gracefully halt the compilation and the user can rebuild after adding the appropriate body. Here is an example of one such function:

/* @brief Used to perform compute of subgraph of ops following V-projection.
 * Relevant internal data layout visualization at /tmp/sdk-cli/examples/models/vit/ccl_build/vit_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/attention_visuals/epu_column_width_pattern.svg
 *
 * @tparam patchHeight Rows of PE array that are active
 * @tparam patchWidth Columns of PE array that are active
 * @tparam requantize_0_input_scale_frac_bits Fractional bits used for requantize_0_input_scale
 * @tparam requantize_0_output_scale_frac_bits Fractional bits used for requantize_0_output_scale
 * @tparam dequantize_0_input_scale_frac_bits Fractional bits used for dequantize_0_input_scale
 * @tparam quantize_0_output_scale_frac_bits Fractional bits used for quantize_0_output_scale
 * @param head Current head index
 * @param NDArrayV I/O NDArray
 * @param requantize_0_input_scale Arg input_scale of requantize #0 (post-order)
 * @param requantize_0_input_zero_point Arg input_zero_point of requantize #0 (post-order)
 * @param requantize_0_output_scale Arg output_scale of requantize #0 (post-order)
 * @param requantize_0_output_zero_point Arg output_zero_point of requantize #0 (post-order)
 * @param dequantize_0_input_scale Arg input_scale of dequantize #0 (post-order)
 * @param dequantize_0_input_zero_point Arg input_zero_point of dequantize #0 (post-order)
 * @param quantize_0_output_scale Arg output_scale of quantize #0 (post-order)
 * @param quantize_0_output_zero_point Arg output_zero_point of quantize #0 (post-order)
 */
template <std::int32_t patchHeight,
          std::int32_t patchWidth,
          std::int32_t postVFuncId,
          typename NDArrayTypeV,
          typename OcmAllocatorType,
          FracRepType requantize_0_input_scale_frac_bits,
          FracRepType requantize_0_output_scale_frac_bits,
          FracRepType dequantize_0_input_scale_frac_bits,
          FracRepType quantize_0_output_scale_frac_bits,
          std::enable_if_t<postVFuncId == 1, int> = 0>
INLINE void postVFunc(std::int8_t head,
                      NDArrayTypeV& NDArrayV,
                      OcmAllocatorType& ocmAllocator,
                      FixedPoint32<requantize_0_input_scale_frac_bits> requantize_0_input_scale,
                      std::int8_t requantize_0_input_zero_point,
                      FixedPoint32<requantize_0_output_scale_frac_bits> requantize_0_output_scale,
                      std::int8_t requantize_0_output_zero_point,
                      FixedPoint32<dequantize_0_input_scale_frac_bits> dequantize_0_input_scale,
                      std::int8_t dequantize_0_input_zero_point,
                      FixedPoint32<quantize_0_output_scale_frac_bits> quantize_0_output_scale,
                      std::int32_t quantize_0_output_zero_point) {
  static_assert(0, "User must fill in implementation of general_attention intermediate function");
}

The function above is the Value path—rightmost edge in the ONNX graph shown above—of the Query, Key and Value (QKV) paths for a self-attention mechanism. From the ONNX graph shown above, we can conclude that the correct path is: requantize, dequantize, quantize, then requantize. This can be implemented using the SDK functions in the following way:

  constexpr FracRepType scaleFB = 29;
  constexpr std::int32_t requantScale0Shift = scaleFB - (requantize_0_input_scale_frac_bits - requantize_0_output_scale_frac_bits);
  constexpr std::int32_t requantScale1Shift = scaleFB - (dequantize_0_input_scale_frac_bits - quantize_0_output_scale_frac_bits);

  qVar_t<FixedPoint32<scaleFB>> requantScale0 =
    FXDirectAssign(math::fxDiv<requantScale0Shift>(requantize_0_input_scale.value, requantize_0_output_scale.value));
  qVar_t<FixedPoint32<scaleFB>> requantScale1 =
    FXDirectAssign(math::fxDiv<requantScale1Shift>(dequantize_0_input_scale.value, quantize_0_output_scale.value));

  for(std::size_t i = 0; i < NDArrayTypeV::size(); ++i) {
    NDArrayV[i] = nn::quantizeLinear(NDArrayV[i], requantScale0);
    NDArrayV[i] = nn::quantizeLinear(NDArrayV[i], requantScale1);
  }

It is not necessary to reimplement all these functions after each compilation. The user can pass in a pre-filled file to the ChimeraJob object indicating to the compiler to use this instead of creating an empty one and continue without halting. For this demo, we have provided a pre-filled attention_stubs.hpp file called attention_stubs_solution.hpp.

The command below show how to invoke a compile-job. We have added a few extra parameters to the ChimeraJob object:

• attn_stub_src_path - Path for the the pre-filled attention_stub.hpp file.
• directives - A JSON file that specify a set of layout directives for specific nodes. For this demo, vit-quadric-pretrained_directives.json is provided.

NOTE: We have project-plans to remove these directive and compiler-flag usage by having the compiler recognize all these optimizations automatically.

from tvm.contrib.epu.chimera_job.hw_config import HWConfig, DEFAULT_32_ARRAY_SIZE

hw_cfg = DEFAULT_32_ARRAY_SIZE

cgc_job = ChimeraJob(
    hw_config=hw_cfg,
    model_p=custom_op_path,
    onnx_ort_override_p=int8_path,
    trange_file=trange_path,
    attn_stub_src_path=f"{os.getcwd()}/attention_stubs_solution.hpp",
    directives=directives_file,
)
cgc_job.compile(quiet=True)
print(cgc_job)

╒═════════════════════╤══════════════════════════════════════════════════════════════╕
│ Module Name         │ vit_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ ONNX File           │ ./vit-custom-op.onnx                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Product Target      │ QC-U                                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Number of Cores     │ 1                                                            │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ ISS Clock Frequency │ 1.700                                                        │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ L2M Size            │ 16MB                                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ LRM Size            │ 4kB                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ External Read BW    │ 128GBps                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ External Write BW   │ 128GBps                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ MACS per PE         │ 16                                                           │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max L2M             │ 10.990MB                                                     │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max LRM             │ 2.250kB                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max Temp Ext Bytes  │ 0.000MB                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Network GMACs       │ 17.564                                                       │
╘═════════════════════╧══════════════════════════════════════════════════════════════╛

╒════╤════════╤═══════════════════════════╤══════════════════╤══════════════════════════╤═══════╕
│    │ Type   │ Name                      │ shape            │ type                     │ mse   │
╞════╪════════╪═══════════════════════════╪══════════════════╪══════════════════════════╪═══════╡
│  0 │ Input  │ l_x_                      │ [1, 3, 224, 224] │ tensor[FixedPoint32<29>] │ n/a   │
├────┼────────┼───────────────────────────┼──────────────────┼──────────────────────────┼───────┤
│  1 │ Output │ dequantize_per_tensor_200 │ [1, 1000]        │ tensor[FixedPoint32<27>] │ n/a   │
╘════╧════════╧═══════════════════════════╧══════════════════╧══════════════════════════╧═══════╛

10. Demo

We first run inference, then show the statistics collected.

IMAGENET_DATA_DIR = Path("../../common/validation/imagenet-mini-quadric/val")
all_images_dict = {
    str(IMAGENET_DATA_DIR / "n03777568" / "ILSVRC2012_val_00049034.JPEG"): 661,
    str(IMAGENET_DATA_DIR / "n02877765" / "ILSVRC2012_val_00049979.JPEG"): 455,
    str(IMAGENET_DATA_DIR / "n02125311" / "ILSVRC2012_val_00015558.JPEG"): 286,
}
all_image_paths = list(all_images_dict.keys())
NUM_IMAGES = len(all_image_paths)

engines = {
    InferenceEngine.CHIMERA_ORT_INT8: cgc_job,
    InferenceEngine.CHIMERA_ISS_INT8: cgc_job,
}

all_images = []
for image_path in all_image_paths:
    all_images.append(np.expand_dims(test_transforms(Image.open(image_path)), axis=0))

outputs_per_inference_engine = {}
THREADS = min(NUM_IMAGES, 6)
for inference_engine, engine in engines.items():
    outputs_per_inference_engine[inference_engine] = batch_inference(
        inference_engine,
        engine,
        all_images,
        threads=THREADS,
    )

2026-07-18 15:09 - WARNING - epu - chimera_job - ORT is not threadsafe -- forcing single threaded batch execution
100%|██████████| 3/3 [00:01<00:00,  2.64it/s]
Processing: 100%|██████████| 3/3 [02:34<00:00, 51.46s/it]

from sdk_cli.utils.datasets.ImageNet import OPTIMIZED_IMAGENET_1K_LABELS
from sdk_cli.visualizers.layouter import ClassifierLayouter

%matplotlib inline

for i, image_path in enumerate(all_image_paths):
    title = f"{all_images_dict[image_path]}: {OPTIMIZED_IMAGENET_1K_LABELS.class_map[all_images_dict[image_path]]}"
    classifier_layouter = ClassifierLayouter(image_path, title)

    for inference_engine, all_outputs in outputs_per_inference_engine.items():
        classifier_layouter.add_data(all_outputs[i][0], f"ViT Top 5 (%): {str(inference_engine)}")
    classifier_layouter.display()

cgc_job.plot_run_statistics();

2026-07-18 15:11 - INFO - epu - chimera_job - Combined plots generated and saved to: 
/quadric/sdk-cli/examples/models/vit/ccl_build/vit_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/run/20260718_150915_eb0c2c/data/vit_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1.combined.png