{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# FP8 ViT on NVIDIA Jetson AGX Thor\n\nThor's Blackwell GPU runs FP8 GEMMs, so a Vision Transformer whose attention\nand MLP matmuls are quantized to FP8 beats the FP16 model on latency without\ntouching the rest of the network. This tutorial takes the HuggingFace\n``google/vit-base-patch16-224`` classifier through the whole path:\n\n1. ``transform()`` fuses it with the TensorRT patterns under the\n   ``tensorrt_fp8`` backend.\n2. ``quantize()`` inserts FP8 Q/DQ stubs and calibrates them on ImageNette.\n3. ``convert_float_to_float16()`` turns everything that is not FP8 into\n   float16 *in PyTorch*, so ``torch.onnx.export`` writes the final graph.\n4. ``trtexec --stronglyTyped`` builds an engine that keeps every tensor at the\n   precision the ONNX says, and the engine is profiled and scored on the\n   device.\n\nThe FP16 baseline goes through exactly the same steps minus ``quantize()``, so\nthe two engines differ only in the Q/DQ nodes. There is no ONNX-level rewriting\nanywhere: what ``torch.onnx.export`` produces is what TensorRT builds.\n\nThe script runs end to end wherever ``trtexec`` is found, so the simplest way\nto use it is to run the whole file on the Thor: the PyTorch half (steps 1--3)\ntakes a few minutes there and writes two ONNX files under ``artifacts/``, and\nstep 4 builds, profiles and scores them in place. On a workstation without\nTensorRT the script stops after the export.\n\nMeasured on a Jetson AGX Thor with TensorRT 10.13.3 (batch 1, GPU compute time\nwith CUDA graphs, ImageNette validation, 3925 images):\n\n```text\nengine                          median ms    top-1\nfp16, weakly typed --fp16           1.399    86.0 %\nfp16, strongly typed                1.402    86.0 %\nfp8 + fp16, strongly typed          1.065    86.2 %\n```\nFP8 is 1.32x faster than the FP16 engine at the same accuracy.\n\nInstall dependencies:\n\n```bash\npip install \"embedl-deploy[tensorrt]\" torchvision transformers onnx onnxscript\n```\nOn the Thor itself, the PyPI ``torch`` wheel runs on the GPU as is (``torch\n2.14, cu130``), and the ``tensorrt`` Python module comes with JetPack -- create\nthe virtualenv with ``--system-site-packages`` to see it.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup: the model and ImageNette\n\nThe HuggingFace classifier takes ``pixel_values`` and returns a\n``logits``-carrying output; a thin wrapper gives it the plain\n``tensor -> tensor`` signature every step below expects. Attention is set\nto the ``eager`` implementation so the fuser sees the individual matmuls\nand softmax rather than a fused SDPA call.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import copy\nimport itertools\nimport json\nimport re\nimport shutil\nimport subprocess\nimport tarfile\nimport urllib.request\nfrom pathlib import Path\n\nimport torch\nimport torchvision\nfrom torch import fx, nn\nfrom torchvision import transforms\nfrom transformers import (  # type: ignore[import-not-found]\n    AutoModelForImageClassification,\n)\n\nMODEL_ID = \"google/vit-base-patch16-224\"\nIMAGENETTE_URL = (\n    \"https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz\"\n)\nDATA_DIR = Path(\"artifacts/data\")\nIMAGENETTE_DIR = DATA_DIR / \"imagenette2-320\"\nOUT_DIR = Path(\"artifacts/thor_fp8_vit\")\nOUT_DIR.mkdir(parents=True, exist_ok=True)\n\n#: ImageNette's ten classes, as ImageNet-1k indices.\nIMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701]\nCALIBRATION_IMAGES = 128\nBATCH_SIZE = 1\n\ntorch.manual_seed(0)\n\n\nclass ImageClassifier(nn.Module):\n    \"\"\"``pixel_values -> logits`` view of a HuggingFace image classifier.\"\"\"\n\n    def __init__(self, model_id: str) -> None:\n        \"\"\"Load the pretrained classifier `model_id` with eager attention.\"\"\"\n        super().__init__()\n        self.model = AutoModelForImageClassification.from_pretrained(\n            model_id, attn_implementation=\"eager\"\n        )\n\n    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:\n        \"\"\"Return the ImageNet-1k logits for `pixel_values`.\"\"\"\n        logits: torch.Tensor = self.model(pixel_values=pixel_values).logits\n        return logits\n\n\ndef download_imagenette() -> None:\n    \"\"\"Download and extract ImageNette if not already present.\"\"\"\n    if IMAGENETTE_DIR.exists():\n        return\n    DATA_DIR.mkdir(parents=True, exist_ok=True)\n    tgz_path = DATA_DIR / \"imagenette2-320.tgz\"\n    print(f\"Downloading ImageNette to {tgz_path} ...\")\n    urllib.request.urlretrieve(IMAGENETTE_URL, str(tgz_path))\n    with tarfile.open(tgz_path) as tar:\n        tar.extractall(DATA_DIR, filter=\"data\")\n    tgz_path.unlink()\n\n\ndownload_imagenette()\n\n# ViT-B/16 was trained with 0.5 mean/std and a straight resize to 224.\npreprocess = transforms.Compose(\n    [\n        transforms.Resize(224),\n        transforms.CenterCrop(224),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]\n)\n\n\ndef remap(label: int) -> int:\n    \"\"\"Map an ImageNette class index to its ImageNet class index.\"\"\"\n    return IMAGENETTE_TO_IMAGENET[label]\n\n\ntrain_set = torchvision.datasets.ImageFolder(\n    str(IMAGENETTE_DIR / \"train\"), transform=preprocess, target_transform=remap\n)\nval_set = torchvision.datasets.ImageFolder(\n    str(IMAGENETTE_DIR / \"val\"), transform=preprocess, target_transform=remap\n)\nval_loader = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE)\n\ncalibration_loader = torch.utils.data.DataLoader(\n    train_set, batch_size=BATCH_SIZE, shuffle=True\n)\ncalibration_batches = [\n    images\n    for images, _ in itertools.islice(calibration_loader, CALIBRATION_IMAGES)\n]\n\nexample_input = torch.randn(1, 3, 224, 224)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 1: fuse for TensorRT with the FP8 policy\n\n``tensorrt_fp8`` is the TensorRT backend with the FP8 precision policy\nmeasured on Thor: the same conversions and fusions as ``tensorrt``, plus\nthe knowledge of which kernels take FP8 profitably and which stay in\nfloating point. ``transform()`` rewrites the model into fused modules\naround those kernels -- the attention projections, the MLP linears -- and\nleaves everything else as it is.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from embedl_deploy import transform\nfrom embedl_deploy.backend import set_backend\n\nset_backend(\"tensorrt_fp8\")\n\nmodel = ImageClassifier(MODEL_ID).eval()\nfused = transform(model, (example_input,)).model\n\nwith torch.no_grad():\n    drift = (model(example_input) - fused(example_input)).abs().max().item()\nprint(f\"max |logit drift| after fusion: {drift:.2e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 2: quantize to FP8\n\nActivations and weights both get E4M3 FP8 Q/DQ. Calibration runs the 128\ntraining images through the model to pick each stub's scale. Layers the\nFP8 policy rules out -- the three-channel patch-embedding stem, for one --\nare left in floating point on its own. ``freeze_weights=True`` bakes the\nweight scales in, which the export in step 3 requires.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from embedl_deploy.quantize import (\n    Precision,\n    QuantConfig,\n    TensorQuantConfig,\n    quantize,\n)\n\n\ndef forward_loop(m: nn.Module) -> None:\n    \"\"\"Feed the calibration images through `m`.\"\"\"\n    with torch.no_grad():\n        for images in calibration_batches:\n            m(images)\n\n\nfp8 = TensorQuantConfig(Precision.FP8)\n# ``quantize()`` works in place, so the baseline keeps its own copy.\nquantized = quantize(\n    copy.deepcopy(fused),\n    (example_input,),\n    config=QuantConfig(activation=fp8, weight=fp8),\n    forward_loop=forward_loop,\n    freeze_weights=True,\n)\n\n\ndef accuracy(m: nn.Module, limit: int | None = None) -> float:\n    \"\"\"Top-1 accuracy of `m` on the validation split (or its first `limit`).\"\"\"\n    correct = total = 0\n    with torch.no_grad():\n        for images, labels in val_loader:\n            correct += int((m(images).argmax(dim=1) == labels).sum())\n            total += labels.numel()\n            if limit is not None and total >= limit:\n                break\n    return correct / total\n\n\n# A quick sanity check in PyTorch: the FP8 fake-quantized model on the first\n# 256 validation images. The full-split numbers come from the engines.\nprint(f\"FP32 top-1 (256 images):     {accuracy(fused, 256):.1%}\")\nprint(f\"FP8 QDQ top-1 (256 images):  {accuracy(quantized, 256):.1%}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 3: float16 in PyTorch, then export\n\n``convert_float_to_float16()`` returns a copy whose floating parameters and\nbuffers are float16, keeps the quantizer scales usable, and pins the few\noperations that need float32 (interpolation, cumulative sums, top-k) with\ncasts at their boundaries. The Q/DQ nodes are emitted during export with\nfloat16 scales, so the exported graph is float16 end to end with FP8 in\nthe quantized matmuls. The same half model is also saved as a ``.pt2``\nthrough ``torch.export`` for inspection or a later re-export. The baseline\nis the same conversion applied to the fused, unquantized model.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from embedl_deploy.quantize import convert_float_to_float16\n\n\ndef export_onnx(m: fx.GraphModule, path: Path) -> Path:\n    \"\"\"Export the float16 copy of `m` to ONNX, and to ``.pt2`` alongside.\"\"\"\n    half = convert_float_to_float16(m)\n    with torch.no_grad():\n        exported = torch.export.export(half, (example_input.half(),))\n        torch.export.save(exported, str(path.with_suffix(\".pt2\")))\n        torch.onnx.export(\n            half,\n            (example_input.half(),),\n            str(path),\n            dynamo=True,\n            opset_version=20,\n            input_names=[\"input\"],\n            output_names=[\"logits\"],\n        )\n    # The exporter keeps the weights in a ``<name>.data`` file next to the ONNX.\n    size = sum(p.stat().st_size for p in path.parent.glob(f\"{path.name}*\"))\n    print(f\"wrote {path} ({size / 1e6:.0f} MB)\")\n    return path\n\n\nfp16_onnx = export_onnx(fused, OUT_DIR / \"vit_fp16.onnx\")\nfp8_onnx = export_onnx(quantized, OUT_DIR / \"vit_fp8.onnx\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 4: build, profile and score on the Thor\n\n``--stronglyTyped`` tells TensorRT to honour the dtypes in the ONNX exactly:\nfloat16 tensors stay float16, the ``QuantizeLinear`` outputs stay FP8, and\nnothing is re-decided by the autotuner. That is what makes the comparison\nclean -- both engines run precisely the graph the export wrote. The FP16\nengine is also built once weakly typed (``--fp16``), the way most people\nrun a float model, for reference.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "TRTEXEC = shutil.which(\"trtexec\") or \"/usr/src/tensorrt/bin/trtexec\"\nif not Path(TRTEXEC).exists():\n    print(\n        f\"trtexec not found -- copy {OUT_DIR} to the Thor and run this \"\n        \"script there to build, profile and score the engines.\"\n    )\n    raise SystemExit(0)\n\n_COMPUTE_TIME = re.compile(\n    r\"GPU Compute Time: min = ([\\d.]+) ms, max = ([\\d.]+) ms, \"\n    r\"mean = ([\\d.]+) ms, median = ([\\d.]+) ms\"\n)\n\n\ndef build_engine(onnx_path: Path, label: str, *flags: str) -> dict[str, float]:\n    \"\"\"Build `onnx_path` with `flags`, time it, return the latency summary.\"\"\"\n    engine = OUT_DIR / f\"{label}.engine\"\n    log = OUT_DIR / f\"trtexec_{label}.log\"\n    cmd = [\n        TRTEXEC,\n        f\"--onnx={onnx_path}\",\n        *flags,\n        f\"--saveEngine={engine}\",\n        f\"--timingCacheFile={OUT_DIR / 'timing.cache'}\",\n        \"--useCudaGraph\",\n        \"--noDataTransfers\",\n        \"--useSpinWait\",\n        \"--warmUp=1000\",\n        \"--iterations=500\",\n        \"--avgRuns=100\",\n    ]\n    print(f\"[{label}] {' '.join(cmd)}\")\n    with log.open(\"w\") as fh:\n        subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, check=True)\n    match = _COMPUTE_TIME.search(log.read_text())\n    if match is None:\n        raise RuntimeError(f\"no GPU compute time summary in {log}\")\n    minimum, maximum, mean, median = (float(v) for v in match.groups())\n    print(f\"[{label}] median {median:.3f} ms (mean {mean:.3f} ms)\")\n    return {\n        \"median_ms\": median,\n        \"mean_ms\": mean,\n        \"min_ms\": minimum,\n        \"max_ms\": maximum,\n    }\n\n\nlatency = {\n    \"fp16 (weakly typed --fp16)\": build_engine(\n        fp16_onnx, \"fp16_weak\", \"--fp16\"\n    ),\n    \"fp16 (strongly typed)\": build_engine(\n        fp16_onnx, \"fp16_strong\", \"--stronglyTyped\"\n    ),\n    \"fp8 + fp16 (strongly typed)\": build_engine(\n        fp8_onnx, \"fp8_strong\", \"--stronglyTyped\"\n    ),\n}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Accuracy is measured by running the engines over the whole ImageNette\nvalidation split (3925 images) through the TensorRT Python API. The inputs\nare float16 because the strongly typed engines' input tensor is.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tensorrt as trt\n\n\nclass Engine:\n    \"\"\"Minimal single-input, single-output TensorRT engine runner.\"\"\"\n\n    def __init__(self, path: Path) -> None:\n        \"\"\"Deserialize the engine at `path` and allocate its output buffer.\"\"\"\n        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))\n        self.engine = runtime.deserialize_cuda_engine(path.read_bytes())\n        self.context = self.engine.create_execution_context()\n        self.input_name = self.engine.get_tensor_name(0)\n        self.output_name = self.engine.get_tensor_name(1)\n        self.stream = torch.cuda.Stream()\n        self.output = torch.empty(\n            tuple(self.engine.get_tensor_shape(self.output_name)),\n            dtype=_torch_dtype(self.engine.get_tensor_dtype(self.output_name)),\n            device=\"cuda\",\n        )\n        self.input_dtype = _torch_dtype(\n            self.engine.get_tensor_dtype(self.input_name)\n        )\n\n    def __call__(self, images: torch.Tensor) -> torch.Tensor:\n        \"\"\"Run one batch through the engine and return float32 logits.\"\"\"\n        images = images.to(\"cuda\", self.input_dtype).contiguous()\n        self.context.set_tensor_address(self.input_name, images.data_ptr())\n        self.context.set_tensor_address(\n            self.output_name, self.output.data_ptr()\n        )\n        self.context.execute_async_v3(self.stream.cuda_stream)\n        self.stream.synchronize()\n        return self.output.float().cpu()\n\n\ndef _torch_dtype(dtype: \"trt.DataType\") -> torch.dtype:\n    \"\"\"Map a TensorRT tensor dtype to the torch dtype of its buffer.\"\"\"\n    return {trt.float32: torch.float32, trt.float16: torch.float16}[dtype]\n\n\ndef engine_accuracy(label: str) -> float:\n    \"\"\"Top-1 accuracy of the engine saved under `label`.\"\"\"\n    engine = Engine(OUT_DIR / f\"{label}.engine\")\n    correct = total = 0\n    for images, labels in val_loader:\n        correct += int((engine(images).argmax(dim=1) == labels).sum())\n        total += labels.numel()\n    return correct / total\n\n\ntop1 = {\n    \"fp16 (weakly typed --fp16)\": engine_accuracy(\"fp16_weak\"),\n    \"fp16 (strongly typed)\": engine_accuracy(\"fp16_strong\"),\n    \"fp8 + fp16 (strongly typed)\": engine_accuracy(\"fp8_strong\"),\n}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Results\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(f\"\\n{'engine':<32}{'median ms':>12}{'top-1':>10}\")\nfor label, stats in latency.items():\n    print(f\"{label:<32}{stats['median_ms']:>12.3f}{top1[label]:>10.1%}\")\nspeedup = (\n    latency[\"fp16 (strongly typed)\"][\"median_ms\"]\n    / latency[\"fp8 + fp16 (strongly typed)\"][\"median_ms\"]\n)\nprint(f\"\\nFP8 speed-up over the FP16 engine: {speedup:.2f}x\")\n\nresults = {\n    \"model\": MODEL_ID,\n    \"tensorrt\": trt.__version__,\n    \"latency\": latency,\n    \"top1\": top1,\n    \"latency_is\": \"GPU compute time, batch 1, CUDA graphs, H2D/D2H excluded\",\n    \"accuracy_is\": \"ImageNette validation, 3925 images, ImageNet-1k labels\",\n}\n(OUT_DIR / \"results.json\").write_text(json.dumps(results, indent=2))\nprint(f\"wrote {OUT_DIR / 'results.json'}\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}