PTQ support for ViT models

Description

I am trying to figure out if TensoRT and the pytorch_quantization module support post-training quantization for vision transformers.

The following piece of code follows the pytorch_quantization docs almost verbatim (with small changes for compatibility):

import torch
import timm
import torchvision
import pytorch_quantization.quant_modules

pytorch_quantization.quant_modules.initialize()

model = timm.create_model("vit_tiny_patch16_224", pretrained=False, num_classes=0)
# or
# model = torchvision.models.vit_b_16()
# model.heads = torch.nn.Identity()
model = model.eval()
data = torch.randn(1, 3, 224, 224)

for name, module in model.named_modules():
    if name.endswith("_quantizer"):
        module.enable_calib()
        module.disable_quant()

model(data)

for name, module in model.named_modules():
    if name.endswith("_quantizer"):
        module.load_calib_amax()
        module.disable_calib()
        module.enable_quant()

with pytorch_quantization.enable_onnx_export():
    torch.onnx.export(
        model,
        data,
        "timm_vit.onnx",
        opset_version=14,
        # opset_version=10, # scaled_dot_product_attention is not supported in opset 10
        # enable_onnx_checker=False, # unexpected keyword argument
    ) 

After that, I visualize the resulting engine graph with trex:

trex process timm_vit.onnx results

The conversion succeeds, however, the graph barely uses any INT8 operations. I would have expected almost the whole graph to consist of Int8 operators, but instead most edges in the graph are labeled as Float with only a few Int8s.

Is this expected? My understanding was that most operators in transformers were supposed to be quantizable (with the notable exception of LayerNorm and Softmax, which would require special custom layers for quantization).

Relevant Files

Environment

TensorRT Version: 10.0.0.6

NVIDIA GPU: NVIDIA RTX A6000

NVIDIA Driver Version: 535.171.04

CUDA Version: 12.2

CUDNN Version: 8

Operating System: Ubuntu 22.04

Python Version (if applicable): 3.10.12

PyTorch Version (if applicable): 2.3.0

Baremetal or Container (if so, version): nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 docker container

Start by reproducing the provided timm and torchvision ViT examples with TensorRT 10.0.0.6, then inspect the generated timm_vit.onnx graph using trex. Compare the observed Float and Int8 operations with the expected post-training quantization behavior, and document whether ViT PTQ is supported or what limitations apply.