FP8 ViT on NVIDIA Jetson AGX Thor

Thor’s Blackwell GPU runs FP8 GEMMs, so a Vision Transformer whose attention and MLP matmuls are quantized to FP8 beats the FP16 model on latency without touching the rest of the network. This tutorial takes the HuggingFace google/vit-base-patch16-224 classifier through the whole path:

1. transform() fuses it with the TensorRT patterns under the tensorrt_fp8 backend.

2. quantize() inserts FP8 Q/DQ stubs and calibrates them on ImageNette.

3. convert_float_to_float16() turns everything that is not FP8 into float16 in PyTorch, so torch.onnx.export writes the final graph.

4. trtexec --stronglyTyped builds an engine that keeps every tensor at the precision the ONNX says, and the engine is profiled and scored on the device.

The FP16 baseline goes through exactly the same steps minus quantize(), so the two engines differ only in the Q/DQ nodes. There is no ONNX-level rewriting anywhere: what torch.onnx.export produces is what TensorRT builds.

The script runs end to end wherever trtexec is found, so the simplest way to use it is to run the whole file on the Thor: the PyTorch half (steps 1–3) takes a few minutes there and writes two ONNX files under artifacts/, and step 4 builds, profiles and scores them in place. On a workstation without TensorRT the script stops after the export.

Measured on a Jetson AGX Thor with TensorRT 10.13.3 (batch 1, GPU compute time with CUDA graphs, ImageNette validation, 3925 images):

engine                          median ms    top-1
fp16, weakly typed --fp16           1.399    86.0 %
fp16, strongly typed                1.402    86.0 %
fp8 + fp16, strongly typed          1.065    86.2 %

FP8 is 1.32x faster than the FP16 engine at the same accuracy.

Install dependencies:

pip install "embedl-deploy[tensorrt]" torchvision transformers onnx onnxscript

On the Thor itself, the PyPI torch wheel runs on the GPU as is (torch 2.14, cu130), and the tensorrt Python module comes with JetPack – create the virtualenv with --system-site-packages to see it.

Setup: the model and ImageNette

The HuggingFace classifier takes pixel_values and returns a logits-carrying output; a thin wrapper gives it the plain tensor -> tensor signature every step below expects. Attention is set to the eager implementation so the fuser sees the individual matmuls and softmax rather than a fused SDPA call.

import copy
import itertools
import json
import re
import shutil
import subprocess
import tarfile
import urllib.request
from pathlib import Path

import torch
import torchvision
from torch import fx, nn
from torchvision import transforms
from transformers import (  # type: ignore[import-not-found]
    AutoModelForImageClassification,
)

MODEL_ID = "google/vit-base-patch16-224"
IMAGENETTE_URL = (
    "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz"
)
DATA_DIR = Path("artifacts/data")
IMAGENETTE_DIR = DATA_DIR / "imagenette2-320"
OUT_DIR = Path("artifacts/thor_fp8_vit")
OUT_DIR.mkdir(parents=True, exist_ok=True)

#: ImageNette's ten classes, as ImageNet-1k indices.
IMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701]
CALIBRATION_IMAGES = 128
BATCH_SIZE = 1

torch.manual_seed(0)


class ImageClassifier(nn.Module):
    """``pixel_values -> logits`` view of a HuggingFace image classifier."""

    def __init__(self, model_id: str) -> None:
        """Load the pretrained classifier `model_id` with eager attention."""
        super().__init__()
        self.model = AutoModelForImageClassification.from_pretrained(
            model_id, attn_implementation="eager"
        )

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Return the ImageNet-1k logits for `pixel_values`."""
        logits: torch.Tensor = self.model(pixel_values=pixel_values).logits
        return logits


def download_imagenette() -> None:
    """Download and extract ImageNette if not already present."""
    if IMAGENETTE_DIR.exists():
        return
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    tgz_path = DATA_DIR / "imagenette2-320.tgz"
    print(f"Downloading ImageNette to {tgz_path} ...")
    urllib.request.urlretrieve(IMAGENETTE_URL, str(tgz_path))
    with tarfile.open(tgz_path) as tar:
        tar.extractall(DATA_DIR, filter="data")
    tgz_path.unlink()


download_imagenette()

# ViT-B/16 was trained with 0.5 mean/std and a straight resize to 224.
preprocess = transforms.Compose(
    [
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


def remap(label: int) -> int:
    """Map an ImageNette class index to its ImageNet class index."""
    return IMAGENETTE_TO_IMAGENET[label]


train_set = torchvision.datasets.ImageFolder(
    str(IMAGENETTE_DIR / "train"), transform=preprocess, target_transform=remap
)
val_set = torchvision.datasets.ImageFolder(
    str(IMAGENETTE_DIR / "val"), transform=preprocess, target_transform=remap
)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE)

calibration_loader = torch.utils.data.DataLoader(
    train_set, batch_size=BATCH_SIZE, shuffle=True
)
calibration_batches = [
    images
    for images, _ in itertools.islice(calibration_loader, CALIBRATION_IMAGES)
]

example_input = torch.randn(1, 3, 224, 224)

Step 1: fuse for TensorRT with the FP8 policy

tensorrt_fp8 is the TensorRT backend with the FP8 precision policy measured on Thor: the same conversions and fusions as tensorrt, plus the knowledge of which kernels take FP8 profitably and which stay in floating point. transform() rewrites the model into fused modules around those kernels – the attention projections, the MLP linears – and leaves everything else as it is.

from embedl_deploy import transform
from embedl_deploy.backend import set_backend

set_backend("tensorrt_fp8")

model = ImageClassifier(MODEL_ID).eval()
fused = transform(model, (example_input,)).model

with torch.no_grad():
    drift = (model(example_input) - fused(example_input)).abs().max().item()
print(f"max |logit drift| after fusion: {drift:.2e}")

Step 2: quantize to FP8

Activations and weights both get E4M3 FP8 Q/DQ. Calibration runs the 128 training images through the model to pick each stub’s scale. Layers the FP8 policy rules out – the three-channel patch-embedding stem, for one – are left in floating point on its own. freeze_weights=True bakes the weight scales in, which the export in step 3 requires.

from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    TensorQuantConfig,
    quantize,
)


def forward_loop(m: nn.Module) -> None:
    """Feed the calibration images through `m`."""
    with torch.no_grad():
        for images in calibration_batches:
            m(images)


fp8 = TensorQuantConfig(Precision.FP8)
# ``quantize()`` works in place, so the baseline keeps its own copy.
quantized = quantize(
    copy.deepcopy(fused),
    (example_input,),
    config=QuantConfig(activation=fp8, weight=fp8),
    forward_loop=forward_loop,
    freeze_weights=True,
)


def accuracy(m: nn.Module, limit: int | None = None) -> float:
    """Top-1 accuracy of `m` on the validation split (or its first `limit`)."""
    correct = total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            correct += int((m(images).argmax(dim=1) == labels).sum())
            total += labels.numel()
            if limit is not None and total >= limit:
                break
    return correct / total


# A quick sanity check in PyTorch: the FP8 fake-quantized model on the first
# 256 validation images. The full-split numbers come from the engines.
print(f"FP32 top-1 (256 images):     {accuracy(fused, 256):.1%}")
print(f"FP8 QDQ top-1 (256 images):  {accuracy(quantized, 256):.1%}")

Step 3: float16 in PyTorch, then export

convert_float_to_float16() returns a copy whose floating parameters and buffers are float16, keeps the quantizer scales usable, and pins the few operations that need float32 (interpolation, cumulative sums, top-k) with casts at their boundaries. The Q/DQ nodes are emitted during export with float16 scales, so the exported graph is float16 end to end with FP8 in the quantized matmuls. The same half model is also saved as a .pt2 through torch.export for inspection or a later re-export. The baseline is the same conversion applied to the fused, unquantized model.

from embedl_deploy.quantize import convert_float_to_float16


def export_onnx(m: fx.GraphModule, path: Path) -> Path:
    """Export the float16 copy of `m` to ONNX, and to ``.pt2`` alongside."""
    half = convert_float_to_float16(m)
    with torch.no_grad():
        exported = torch.export.export(half, (example_input.half(),))
        torch.export.save(exported, str(path.with_suffix(".pt2")))
        torch.onnx.export(
            half,
            (example_input.half(),),
            str(path),
            dynamo=True,
            opset_version=20,
            input_names=["input"],
            output_names=["logits"],
        )
    # The exporter keeps the weights in a ``<name>.data`` file next to the ONNX.
    size = sum(p.stat().st_size for p in path.parent.glob(f"{path.name}*"))
    print(f"wrote {path} ({size / 1e6:.0f} MB)")
    return path


fp16_onnx = export_onnx(fused, OUT_DIR / "vit_fp16.onnx")
fp8_onnx = export_onnx(quantized, OUT_DIR / "vit_fp8.onnx")

Step 4: build, profile and score on the Thor

--stronglyTyped tells TensorRT to honour the dtypes in the ONNX exactly: float16 tensors stay float16, the QuantizeLinear outputs stay FP8, and nothing is re-decided by the autotuner. That is what makes the comparison clean – both engines run precisely the graph the export wrote. The FP16 engine is also built once weakly typed (--fp16), the way most people run a float model, for reference.

TRTEXEC = shutil.which("trtexec") or "/usr/src/tensorrt/bin/trtexec"
if not Path(TRTEXEC).exists():
    print(
        f"trtexec not found -- copy {OUT_DIR} to the Thor and run this "
        "script there to build, profile and score the engines."
    )
    raise SystemExit(0)

_COMPUTE_TIME = re.compile(
    r"GPU Compute Time: min = ([\d.]+) ms, max = ([\d.]+) ms, "
    r"mean = ([\d.]+) ms, median = ([\d.]+) ms"
)


def build_engine(onnx_path: Path, label: str, *flags: str) -> dict[str, float]:
    """Build `onnx_path` with `flags`, time it, return the latency summary."""
    engine = OUT_DIR / f"{label}.engine"
    log = OUT_DIR / f"trtexec_{label}.log"
    cmd = [
        TRTEXEC,
        f"--onnx={onnx_path}",
        *flags,
        f"--saveEngine={engine}",
        f"--timingCacheFile={OUT_DIR / 'timing.cache'}",
        "--useCudaGraph",
        "--noDataTransfers",
        "--useSpinWait",
        "--warmUp=1000",
        "--iterations=500",
        "--avgRuns=100",
    ]
    print(f"[{label}] {' '.join(cmd)}")
    with log.open("w") as fh:
        subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, check=True)
    match = _COMPUTE_TIME.search(log.read_text())
    if match is None:
        raise RuntimeError(f"no GPU compute time summary in {log}")
    minimum, maximum, mean, median = (float(v) for v in match.groups())
    print(f"[{label}] median {median:.3f} ms (mean {mean:.3f} ms)")
    return {
        "median_ms": median,
        "mean_ms": mean,
        "min_ms": minimum,
        "max_ms": maximum,
    }


latency = {
    "fp16 (weakly typed --fp16)": build_engine(
        fp16_onnx, "fp16_weak", "--fp16"
    ),
    "fp16 (strongly typed)": build_engine(
        fp16_onnx, "fp16_strong", "--stronglyTyped"
    ),
    "fp8 + fp16 (strongly typed)": build_engine(
        fp8_onnx, "fp8_strong", "--stronglyTyped"
    ),
}

Accuracy is measured by running the engines over the whole ImageNette validation split (3925 images) through the TensorRT Python API. The inputs are float16 because the strongly typed engines’ input tensor is.

import tensorrt as trt


class Engine:
    """Minimal single-input, single-output TensorRT engine runner."""

    def __init__(self, path: Path) -> None:
        """Deserialize the engine at `path` and allocate its output buffer."""
        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
        self.engine = runtime.deserialize_cuda_engine(path.read_bytes())
        self.context = self.engine.create_execution_context()
        self.input_name = self.engine.get_tensor_name(0)
        self.output_name = self.engine.get_tensor_name(1)
        self.stream = torch.cuda.Stream()
        self.output = torch.empty(
            tuple(self.engine.get_tensor_shape(self.output_name)),
            dtype=_torch_dtype(self.engine.get_tensor_dtype(self.output_name)),
            device="cuda",
        )
        self.input_dtype = _torch_dtype(
            self.engine.get_tensor_dtype(self.input_name)
        )

    def __call__(self, images: torch.Tensor) -> torch.Tensor:
        """Run one batch through the engine and return float32 logits."""
        images = images.to("cuda", self.input_dtype).contiguous()
        self.context.set_tensor_address(self.input_name, images.data_ptr())
        self.context.set_tensor_address(
            self.output_name, self.output.data_ptr()
        )
        self.context.execute_async_v3(self.stream.cuda_stream)
        self.stream.synchronize()
        return self.output.float().cpu()


def _torch_dtype(dtype: "trt.DataType") -> torch.dtype:
    """Map a TensorRT tensor dtype to the torch dtype of its buffer."""
    return {trt.float32: torch.float32, trt.float16: torch.float16}[dtype]


def engine_accuracy(label: str) -> float:
    """Top-1 accuracy of the engine saved under `label`."""
    engine = Engine(OUT_DIR / f"{label}.engine")
    correct = total = 0
    for images, labels in val_loader:
        correct += int((engine(images).argmax(dim=1) == labels).sum())
        total += labels.numel()
    return correct / total


top1 = {
    "fp16 (weakly typed --fp16)": engine_accuracy("fp16_weak"),
    "fp16 (strongly typed)": engine_accuracy("fp16_strong"),
    "fp8 + fp16 (strongly typed)": engine_accuracy("fp8_strong"),
}

Results

print(f"\n{'engine':<32}{'median ms':>12}{'top-1':>10}")
for label, stats in latency.items():
    print(f"{label:<32}{stats['median_ms']:>12.3f}{top1[label]:>10.1%}")
speedup = (
    latency["fp16 (strongly typed)"]["median_ms"]
    / latency["fp8 + fp16 (strongly typed)"]["median_ms"]
)
print(f"\nFP8 speed-up over the FP16 engine: {speedup:.2f}x")

results = {
    "model": MODEL_ID,
    "tensorrt": trt.__version__,
    "latency": latency,
    "top1": top1,
    "latency_is": "GPU compute time, batch 1, CUDA graphs, H2D/D2H excluded",
    "accuracy_is": "ImageNette validation, 3925 images, ImageNet-1k labels",
}
(OUT_DIR / "results.json").write_text(json.dumps(results, indent=2))
print(f"wrote {OUT_DIR / 'results.json'}")