From bfaeaf2bbaad1d6ca39c42baf919b838e6f814a5 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 29 May 2026 20:24:57 +0000 Subject: [PATCH 01/17] [6078291] Add ViT FP8/NVFP4 recipes + Torch-TRT example * modelopt_recipes/huggingface/vit/ptq/{fp8,nvfp4}.yaml -- self-contained ViT-tuned PTQ recipes targeting HuggingFace ViTForImageClassification. Encoder Linear weights/inputs quantized; attention Q/K/V BMMs, softmax, and per-block LayerNorm outputs at FP8; patch-embed nn.Conv2d, classifier, and the final vit.layernorm left FP16. NVFP4 variant runs encoder Linears in W4A4 NVFP4 (E2M1, block 16, FP8 scales) with AWQ-lite calibration. * examples/torch_trt/ -- end-to-end Torch-TensorRT deployment example (load HF model -> calibrate from tiny-imagenet -> mtq.quantize -> torch_tensorrt.compile(ir="dynamo") -> benchmark). Defaults to google/vit-large-patch16-224; --model_id + --recipe retarget any HF model + ModelOpt PTQ recipe. Attention softmax-P quantization uses the shared attention_qkv_fp8 recipe unit, which enables _QuantAttention.p_bmm_quantizer on HuggingFace attention. ImageNet-1k full-50k validation accuracy on google/vit-base-patch16-224 (batch=128, 49920/50000 samples): FP16 baseline: Top-1 81.769% Top-5 96.124% FP8 modelopt.onnx CLI: Top-1 81.707% Top-5 96.110% (-0.062 pp) FP8 torch path (this PR): Top-1 81.637% Top-5 96.140% (-0.132 pp) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.rst | 1 + examples/torch_trt/README.md | 102 ++++++ .../torch_trt/quantize_and_compile_vit.py | 317 ++++++++++++++++++ examples/torch_trt/requirements.txt | 3 + modelopt_recipes/huggingface/vit/ptq/fp8.yaml | 55 +++ .../huggingface/vit/ptq/nvfp4.yaml | 63 ++++ 6 files changed, 541 insertions(+) create mode 100644 examples/torch_trt/README.md create mode 100644 examples/torch_trt/quantize_and_compile_vit.py create mode 100644 examples/torch_trt/requirements.txt create mode 100644 modelopt_recipes/huggingface/vit/ptq/fp8.yaml create mode 100644 modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed03fec6ef7..6785a7e09b8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,6 +38,7 @@ Changelog - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. +- Add Torch-TensorRT FP8 / NVFP4 deployment example for HuggingFace ViT (``examples/torch_trt/``) covering ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``. Ships two ViT-tuned PTQ recipes under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``, ``nvfp4.yaml``) — encoder Linear weights+inputs quantized; attention Q/K/V BMMs, softmax, and per-block LayerNorm outputs at FP8; patch-embed ``nn.Conv2d``, ``classifier``, and the final ``vit.layernorm`` left FP16. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. **Bug Fixes** diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md new file mode 100644 index 00000000000..bed90d6c50f --- /dev/null +++ b/examples/torch_trt/README.md @@ -0,0 +1,102 @@ +# ModelOpt + Torch-TensorRT Deployment + +End-to-end examples that quantize a PyTorch model with NVIDIA ModelOpt and +then compile the quantized graph with +[Torch-TensorRT](https://docs.pytorch.org/TensorRT/) for deployment. + +The flow follows the +[Torch-TensorRT quantization guide](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html): +ModelOpt inserts Q/DQ nodes into the eager PyTorch graph, then +`torch_tensorrt.compile(ir="dynamo")` converts those Q/DQ nodes into native +TensorRT precision layers. + +## Setup + +```bash +# From the NVIDIA TensorRT docker image (recommended): +docker run --gpus all -it --rm -v $(pwd):/workspace -w /workspace nvcr.io/nvidia/tensorrt:26.02-py3 bash + +pip install -U "nvidia-modelopt[torch]" +pip install -r examples/torch_trt/requirements.txt +``` + +Torch-TensorRT itself follows the +[official install instructions](https://docs.pytorch.org/TensorRT/getting_started/installation.html) — +the version pulled by `pip` must match your installed PyTorch. + +## Usage + +```bash +# FP8 / NVFP4 default model is google/vit-large-patch16-224 +python examples/torch_trt/quantize_and_compile_vit.py \ + --precision fp8/nvfp4 \ + --calib_samples 128 \ + --batch_size 1 + +# Quantize but don't TRT-compile (handy on a non-TRT host) +python examples/torch_trt/quantize_and_compile_vit.py \ + --precision fp8/nvfp4 \ + --skip_trt + +# Custom model + custom recipe +python examples/torch_trt/quantize_and_compile_vit.py \ + --model_id \ + --recipe +``` + +## What the example does + +1. Loads a HuggingFace model (default: `google/vit-large-patch16-224`). +2. Builds a tiny calibration loader from `zh-plus/tiny-imagenet` (avoids the + gated `ILSVRC/imagenet-1k` repo so the example runs unauthenticated). +3. Runs `mtq.quantize` with one of the recipes shipped under + [`modelopt_recipes/`](../../modelopt_recipes/). The default recipes + target ViT; pass `--recipe ` to use a different one for a + different model. +4. Compiles the quantized model with `torch_tensorrt.compile` and prints a + median-latency benchmark against the BF16 eager baseline. + +## ViT-specific recipes shipped with the example + +These are the recipes the CLI selects by default when `--model_id` points +at a HF ViT classifier. They are **not** thin wrappers around the modelopt +defaults — they're tuned for the HF ViT module layout. + +| Flag | Recipe path | Key differences from the default | +|------|-------------|----------------------------------| +| `--precision fp8` | `huggingface/vit/ptq/fp8` | W8A8 FP8 **plus** MHA-aware FP8 on every per-block `nn.LayerNorm` output (shared Q/DQ feeds Q/K/V + MLP), FP8 attention Q/K/V BMM + softmax slots, patch-embedding `nn.Conv2d` left FP16, `classifier` head left in FP16, final `vit.layernorm` left FP16. | +| `--precision nvfp4` | `huggingface/vit/ptq/nvfp4` | Same skip list as the FP8 recipe; encoder Linear weights/inputs run NVFP4 W4A4 (E2M1, block 16, FP8 scales). Attention BMMs, softmax, and per-block LayerNorm outputs stay at FP8 — NVFP4 is too aggressive there. Uses `awq_lite` calibration. | + +Each recipe is self-contained (no `$import` of shared snippets) and uses +the "specific-enable" style: narrow `parent_class` + path scoping on the +enable rules means no `enable: false` carve-outs are needed. + +## Hardware requirements + +| Recipe | Minimum GPU | +|--------|-------------| +| `fp8` | Hopper (H100) / Ada (RTX 4090 / 6000 Ada) — compute capability 8.9+ | +| `nvfp4` | Blackwell (B100/B200) — TRT ≥ 10.8 | + +Older GPUs will still let `mtq.quantize` succeed (it emits fake-quant +nodes in PyTorch), but `torch_tensorrt.compile` will not find a real +low-precision kernel and the speedup column will be ~1×. + +### Resuming from a saved checkpoint + +Pass `--save_dir ` to persist the modelopt-quantized model +(`vit_modelopt_state.pt`). To reload without recalibrating, restore it +before the TRT compile step with: + +```python +import modelopt.torch.opt as mto +mto.restore(model, "vit_modelopt_state.pt") +``` + +## Custom recipes + +Use `--recipe ` to plug in a different recipe — either a path +relative to `modelopt_recipes/` (resolved against the built-in library) or +an absolute filesystem path to a YAML file. The recipe must declare +`metadata.recipe_type: ptq` and a `quantize:` section; see existing +`modelopt_recipes/huggingface/vit/ptq/*.yaml` for the patterns used here. diff --git a/examples/torch_trt/quantize_and_compile_vit.py b/examples/torch_trt/quantize_and_compile_vit.py new file mode 100644 index 00000000000..e986ec6ff2d --- /dev/null +++ b/examples/torch_trt/quantize_and_compile_vit.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Quantize a HuggingFace ViT model with ModelOpt and deploy it with Torch-TensorRT. + +Pipeline: + +1. Load ``google/vit-large-patch16-224`` (`ViTForImageClassification`) from HF. +2. Build a calibration loader from `zh-plus/tiny-imagenet` (same pattern as the + `torch_onnx` example) so the recipe runs end-to-end without ImageNet access. +3. Run ``mtq.quantize`` with one of the ViT-specific recipes + (`modelopt_recipes/huggingface/vit/ptq/`). Two non-default variants are + shipped: + + * ``fp8`` -> ``fp8_mha-classifier_skip``: W8A8 FP8 with an MHA-aware + LayerNorm output quantizer, FP8 attention BMM/softmax slots, and the + `classifier` head left in FP16. + * ``nvfp4`` -> ``nvfp4_linear-fp8_conv-classifier_skip``: NVFP4 W4A4 on + encoder Linear layers, FP8 override on the patch-embedding Conv2d (TRT + has no NVFP4 kernel for 4D Conv inputs), AWQ-lite calibration, and the + `classifier` head left in FP16. + +4. Compile the quantized model with ``torch_tensorrt.compile`` (Dynamo IR, + ``min_block_size=1``) and run an end-to-end sanity check + small benchmark + against the eager BF16 baseline. + +This script is intentionally CLI-driven and side-effect-free outside of the +optional ``--save_dir`` checkpoint. The quantized graph keeps Q/DQ nodes; the +TRT compile step is what turns them into TRT precision layers. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from datasets import load_dataset +from transformers import AutoImageProcessor, ViTForImageClassification + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.recipe import ModelOptPTQRecipe, load_recipe +from modelopt.torch.quantization.utils import export_torch_mode + +# Maps the user-facing precision flag to the ViT-specific recipe under +# `modelopt_recipes/huggingface/vit/ptq/`. The recipe loader resolves this +# relative path against the built-in recipe library. +PRECISION_TO_RECIPE: dict[str, str] = { + "fp8": "huggingface/vit/ptq/fp8", + "nvfp4": "huggingface/vit/ptq/nvfp4", +} + + +def load_model_and_processor(model_id: str, device: torch.device, dtype: torch.dtype): + """Pull the HF ViT classifier and its preprocessor.""" + print(f"Loading {model_id} (dtype={dtype})...") + processor = AutoImageProcessor.from_pretrained(model_id) + model = ViTForImageClassification.from_pretrained(model_id, torch_dtype=dtype) + model.eval().to(device) + return model, processor + + +def build_calibration_loader( + processor, + num_samples: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +): + """Build a calibration tensor stream from tiny-imagenet. + + tiny-imagenet avoids the gated `ILSVRC/imagenet-1k` repo so this example + runs unauthenticated. Images go through the HF processor (resize + center + crop + ImageNet normalization), which is exactly the eval-time transform + used by the released `vit-large-patch16-224` checkpoint. + """ + print(f"Loading calibration data ({num_samples} samples)...") + dataset = load_dataset("zh-plus/tiny-imagenet", split="train") + dataset = dataset.shuffle(seed=42).select(range(num_samples)) + + tensors: list[torch.Tensor] = [] + for sample in dataset: + image = sample["image"] + if image.mode != "RGB": + image = image.convert("RGB") + # HF image processors emit `pixel_values` of shape (1, 3, H, W). + pixel_values = processor(images=image, return_tensors="pt")["pixel_values"] + tensors.append(pixel_values.squeeze(0)) + + batched = torch.stack(tensors).to(device=device, dtype=dtype) + return torch.split(batched, batch_size) + + +def quantize_with_recipe(model, recipe_path: str, calib_batches): + """Resolve the YAML recipe and run `mtq.quantize`. + + Returns the quantized model. The graph still uses high-precision math at + this point — Q/DQ nodes have been inserted around weights and activations + and amax values populated, but no kernel substitution has happened yet. + """ + print(f"Loading recipe: {recipe_path}") + recipe = load_recipe(recipe_path) + if not isinstance(recipe, ModelOptPTQRecipe): + raise TypeError(f"Expected PTQ recipe, got {type(recipe).__name__}") + quant_cfg = recipe.quantize.model_dump() + + def forward_loop(model_): + with torch.no_grad(): + for batch in calib_batches: + model_(pixel_values=batch) + + print("Running mtq.quantize ...") + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + mtq.print_quant_summary(model) + return model + + +class ViTLogitsWrapper(torch.nn.Module): + """Returns raw logits as a single tensor. + + HF's `ViTForImageClassification.forward` returns an `ImageClassifierOutput` + dataclass. `torch_tensorrt.compile` (and `torch.export`) need a tensor-tree + return, so we unwrap it here. The wrapper holds the quantized model as a + submodule; Q/DQ nodes flow through unchanged. + """ + + def __init__(self, vit_model: torch.nn.Module): + super().__init__() + self.vit = vit_model + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + return self.vit(pixel_values=pixel_values).logits + + +def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Tensor): + """Compile the quantized model with Torch-TensorRT (Dynamo IR). + + `min_block_size=1` follows the Torch-TRT quantization guide — it makes the + partitioner accept single-node TRT subgraphs, which is what we want so the + Q/DQ + matmul pairs become TRT precision layers instead of falling back to + eager. The compile step expects fake-quant operators in the graph; we run + it under `export_torch_mode` so modelopt's Q/DQ are exported in the + TRT-friendly form. + """ + import torch_tensorrt + + print("Compiling with torch_tensorrt.compile (Dynamo IR)...") + with export_torch_mode(): + trt_model = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=[example_input], + min_block_size=1, + # The recipes export weights in BF16; TRT picks the FP8/NVFP4 + # kernel from the Q/DQ pattern, not from this list. + enabled_precisions={torch.bfloat16, torch.float16, torch.float32}, + truncate_double=True, + ) + return trt_model + + +def benchmark(model: torch.nn.Module, example_input: torch.Tensor, n_warmup: int, n_iters: int): + """Median-of-`n_iters` latency over `example_input`. CUDA-event timed.""" + torch.cuda.synchronize() + with torch.no_grad(): + for _ in range(n_warmup): + model(example_input) + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] + with torch.no_grad(): + for i in range(n_iters): + starts[i].record() + model(example_input) + ends[i].record() + torch.cuda.synchronize() + times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) + return times[len(times) // 2] + + +def _argmax_logits(out) -> torch.Tensor: + """Handle either an HF `ImageClassifierOutput` or a raw tensor.""" + logits = out.logits if hasattr(out, "logits") else out + return logits.argmax(dim=-1) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model_id", + default="google/vit-large-patch16-224", + help="HuggingFace model id of the ViT classifier to quantize.", + ) + parser.add_argument( + "--precision", + choices=sorted(PRECISION_TO_RECIPE), + default="fp8", + help="Which ViT recipe variant to apply.", + ) + parser.add_argument( + "--recipe", + default=None, + help="Override the recipe path (relative to modelopt_recipes/ or absolute). " + "If unset, the recipe is picked by --precision.", + ) + parser.add_argument( + "--calib_samples", + type=int, + default=128, + help="Number of tiny-imagenet samples to use for calibration.", + ) + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size for calibration / TRT compile / benchmarking.", + ) + parser.add_argument( + "--benchmark_iters", + type=int, + default=50, + help="Number of timed iterations (after warmup) per benchmark phase.", + ) + parser.add_argument( + "--save_dir", + type=str, + default=None, + help="If set, save the quantized modelopt state-dict here (BF16 weights " + "+ Q/DQ metadata) — re-usable across runs without recalibration.", + ) + parser.add_argument( + "--skip_trt", + action="store_true", + help="Quantize + run the BF16-fake-quant model only; skip torch_tensorrt.compile. " + "Useful for environments without torch_tensorrt installed.", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable GPU.") + device = torch.device("cuda") + # ViT-Large is a transformer in BF16 on the released checkpoint; the Q/DQ + # nodes operate on top of BF16 master weights either way. + dtype = torch.bfloat16 + + model, processor = load_model_and_processor(args.model_id, device, dtype) + image_size = model.config.image_size + num_channels = model.config.num_channels + example_input = torch.randn( + args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype + ) + + # Baseline forward + benchmark for a comparison number that survives + # quantization. argmax preserves the predicted-class check below. + print("\n=== Baseline (BF16) ===") + with torch.no_grad(): + baseline_pred = _argmax_logits(model(example_input)) + baseline_latency = benchmark( + lambda x: model(x), example_input, n_warmup=5, n_iters=args.benchmark_iters + ) + print(f"Baseline argmax class: {baseline_pred.tolist()}") + print(f"Baseline latency: {baseline_latency:.3f} ms (median over {args.benchmark_iters} iters)") + + calib_batches = build_calibration_loader( + processor, args.calib_samples, args.batch_size, device, dtype + ) + + recipe_path = args.recipe or PRECISION_TO_RECIPE[args.precision] + quantize_with_recipe(model, recipe_path, calib_batches) + + if args.save_dir: + save_path = Path(args.save_dir) + save_path.mkdir(parents=True, exist_ok=True) + ckpt = save_path / "vit_modelopt_state.pt" + mto.save(model, ckpt) + print(f"Saved quantized modelopt state to {ckpt}") + + print("\n=== Fake-quant (modelopt, BF16 math) ===") + with torch.no_grad(): + fq_pred = _argmax_logits(model(example_input)) + fq_match = (fq_pred == baseline_pred).all().item() + print(f"Quantized argmax class: {fq_pred.tolist()} (matches baseline: {fq_match})") + + if args.skip_trt: + print("\n--skip_trt set; not compiling with Torch-TensorRT.") + return + + wrapped = ViTLogitsWrapper(model).to(device).eval() + trt_model = compile_with_torch_tensorrt(wrapped, example_input) + + print("\n=== Torch-TensorRT compiled ===") + with torch.no_grad(): + trt_pred = trt_model(example_input).argmax(dim=-1) + trt_match = (trt_pred == baseline_pred).all().item() + trt_latency = benchmark(trt_model, example_input, n_warmup=5, n_iters=args.benchmark_iters) + print(f"TRT argmax class: {trt_pred.tolist()} (matches baseline: {trt_match})") + print(f"TRT latency: {trt_latency:.3f} ms (median over {args.benchmark_iters} iters)") + speedup = baseline_latency / trt_latency if trt_latency > 0 else float("inf") + print(f"\nSpeedup vs. BF16 baseline: {speedup:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/examples/torch_trt/requirements.txt b/examples/torch_trt/requirements.txt new file mode 100644 index 00000000000..8caf5b5fd93 --- /dev/null +++ b/examples/torch_trt/requirements.txt @@ -0,0 +1,3 @@ +datasets>=2.14.4 +torch-tensorrt>=2.4.0 +transformers>=4.40 diff --git a/modelopt_recipes/huggingface/vit/ptq/fp8.yaml b/modelopt_recipes/huggingface/vit/ptq/fp8.yaml new file mode 100644 index 00000000000..6158a99590a --- /dev/null +++ b/modelopt_recipes/huggingface/vit/ptq/fp8.yaml @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +metadata: + recipe_type: ptq + description: >- + HuggingFace ViT FP8 PTQ recipe matching modelopt.onnx FP8 layout: + per-tensor FP8 E4M3 on encoder Linear weights + inputs, attention QKV + BMMs and softmax outputs, MHA-aware per-block LayerNorm outputs; + patch-embed Conv2d, classifier head, and the final LayerNorm are + left FP16. Uses max calibration. +quantize: + algorithm: max + quant_cfg: + - quantizer_name: '*' + enable: false + + - parent_class: 'nn.Linear' + quantizer_name: 'vit.encoder.layer.*.weight_quantizer' + cfg: + num_bits: e4m3 + axis: + - parent_class: 'nn.Linear' + quantizer_name: 'vit.encoder.layer.*.input_quantizer' + cfg: + num_bits: e4m3 + axis: + + - quantizer_name: '*[qkv]_bmm_quantizer' + cfg: + num_bits: e4m3 + axis: + + - quantizer_name: '*softmax_quantizer' + cfg: + num_bits: e4m3 + axis: + + - parent_class: 'nn.LayerNorm' + quantizer_name: 'vit.encoder.layer.*.output_quantizer' + cfg: + num_bits: e4m3 + axis: diff --git a/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml b/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml new file mode 100644 index 00000000000..e30f7184dd6 --- /dev/null +++ b/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +metadata: + recipe_type: ptq + description: >- + HuggingFace ViT NVFP4 W4A4 PTQ recipe. Skip list mirrors modelopt.onnx + FP8 (patch-embed Conv2d, classifier, and final vit.layernorm stay + FP16). Encoder Linear weights/inputs run NVFP4; attention BMMs, + softmax, and per-block LayerNorm outputs stay at FP8. + Uses AWQ-lite calibration. +quantize: + algorithm: awq_lite + quant_cfg: + - quantizer_name: '*' + enable: false + + - parent_class: 'nn.Linear' + quantizer_name: 'vit.encoder.layer.*.weight_quantizer' + cfg: + num_bits: e2m1 + axis: + block_sizes: + -1: 16 + type: dynamic + scale_bits: e4m3 + - parent_class: 'nn.Linear' + quantizer_name: 'vit.encoder.layer.*.input_quantizer' + cfg: + num_bits: e2m1 + axis: + block_sizes: + -1: 16 + type: dynamic + scale_bits: e4m3 + + - quantizer_name: '*[qkv]_bmm_quantizer' + cfg: + num_bits: e4m3 + axis: + + - quantizer_name: '*softmax_quantizer' + cfg: + num_bits: e4m3 + axis: + + - parent_class: 'nn.LayerNorm' + quantizer_name: 'vit.encoder.layer.*.output_quantizer' + cfg: + num_bits: e4m3 + axis: From 0869610773c9ae3c62b261c26fe1821c9a8e087e Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 29 May 2026 20:27:58 +0000 Subject: [PATCH 02/17] Rename example script to torch_tensorrt_ptq.py + strip benchmarking comparison * examples/torch_trt/quantize_and_compile_vit.py -> torch_tensorrt_ptq.py * Drop the latency / speedup benchmarking comparison from the script and README; the script now only verifies that the compiled-model argmax matches the fake-quant argmax on a sample input. Accuracy comparison belongs in a separate harness, not in a "quantize + compile" example. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/README.md | 13 +- ...d_compile_vit.py => torch_tensorrt_ptq.py} | 112 ++++-------------- 2 files changed, 29 insertions(+), 96 deletions(-) rename examples/torch_trt/{quantize_and_compile_vit.py => torch_tensorrt_ptq.py} (62%) diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index bed90d6c50f..d826f7fc3d6 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -28,18 +28,18 @@ the version pulled by `pip` must match your installed PyTorch. ```bash # FP8 / NVFP4 default model is google/vit-large-patch16-224 -python examples/torch_trt/quantize_and_compile_vit.py \ +python examples/torch_trt/torch_tensorrt_ptq.py \ --precision fp8/nvfp4 \ --calib_samples 128 \ --batch_size 1 # Quantize but don't TRT-compile (handy on a non-TRT host) -python examples/torch_trt/quantize_and_compile_vit.py \ +python examples/torch_trt/torch_tensorrt_ptq.py \ --precision fp8/nvfp4 \ --skip_trt # Custom model + custom recipe -python examples/torch_trt/quantize_and_compile_vit.py \ +python examples/torch_trt/torch_tensorrt_ptq.py \ --model_id \ --recipe ``` @@ -53,8 +53,9 @@ python examples/torch_trt/quantize_and_compile_vit.py \ [`modelopt_recipes/`](../../modelopt_recipes/). The default recipes target ViT; pass `--recipe ` to use a different one for a different model. -4. Compiles the quantized model with `torch_tensorrt.compile` and prints a - median-latency benchmark against the BF16 eager baseline. +4. Compiles the quantized model with `torch_tensorrt.compile` and verifies + that the compiled-model argmax matches the fake-quant argmax on a sample + input. ## ViT-specific recipes shipped with the example @@ -80,7 +81,7 @@ enable rules means no `enable: false` carve-outs are needed. Older GPUs will still let `mtq.quantize` succeed (it emits fake-quant nodes in PyTorch), but `torch_tensorrt.compile` will not find a real -low-precision kernel and the speedup column will be ~1×. +low-precision kernel. ### Resuming from a saved checkpoint diff --git a/examples/torch_trt/quantize_and_compile_vit.py b/examples/torch_trt/torch_tensorrt_ptq.py similarity index 62% rename from examples/torch_trt/quantize_and_compile_vit.py rename to examples/torch_trt/torch_tensorrt_ptq.py index e986ec6ff2d..2e65e59a269 100644 --- a/examples/torch_trt/quantize_and_compile_vit.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -13,32 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Quantize a HuggingFace ViT model with ModelOpt and deploy it with Torch-TensorRT. +"""Quantize a HuggingFace ViT model with ModelOpt and compile with Torch-TensorRT. Pipeline: 1. Load ``google/vit-large-patch16-224`` (`ViTForImageClassification`) from HF. -2. Build a calibration loader from `zh-plus/tiny-imagenet` (same pattern as the - `torch_onnx` example) so the recipe runs end-to-end without ImageNet access. -3. Run ``mtq.quantize`` with one of the ViT-specific recipes - (`modelopt_recipes/huggingface/vit/ptq/`). Two non-default variants are - shipped: - - * ``fp8`` -> ``fp8_mha-classifier_skip``: W8A8 FP8 with an MHA-aware - LayerNorm output quantizer, FP8 attention BMM/softmax slots, and the - `classifier` head left in FP16. - * ``nvfp4`` -> ``nvfp4_linear-fp8_conv-classifier_skip``: NVFP4 W4A4 on - encoder Linear layers, FP8 override on the patch-embedding Conv2d (TRT - has no NVFP4 kernel for 4D Conv inputs), AWQ-lite calibration, and the - `classifier` head left in FP16. - -4. Compile the quantized model with ``torch_tensorrt.compile`` (Dynamo IR, - ``min_block_size=1``) and run an end-to-end sanity check + small benchmark - against the eager BF16 baseline. - -This script is intentionally CLI-driven and side-effect-free outside of the -optional ``--save_dir`` checkpoint. The quantized graph keeps Q/DQ nodes; the -TRT compile step is what turns them into TRT precision layers. +2. Build a calibration loader from `zh-plus/tiny-imagenet` so the recipe runs + end-to-end without ImageNet access. +3. Run ``mtq.quantize`` with one of the ViT-specific recipes under + `modelopt_recipes/huggingface/vit/ptq/` (FP8 or NVFP4). +4. Compile the quantized model with ``torch_tensorrt.compile(ir="dynamo", + min_block_size=1)`` and verify the compiled-model argmax matches the + fake-quant argmax on a sample input. + +The quantized graph keeps Q/DQ nodes; the TRT compile step is what turns +them into TRT precision layers. """ from __future__ import annotations @@ -80,13 +69,7 @@ def build_calibration_loader( device: torch.device, dtype: torch.dtype, ): - """Build a calibration tensor stream from tiny-imagenet. - - tiny-imagenet avoids the gated `ILSVRC/imagenet-1k` repo so this example - runs unauthenticated. Images go through the HF processor (resize + center - crop + ImageNet normalization), which is exactly the eval-time transform - used by the released `vit-large-patch16-224` checkpoint. - """ + """Build a calibration tensor stream from tiny-imagenet.""" print(f"Loading calibration data ({num_samples} samples)...") dataset = load_dataset("zh-plus/tiny-imagenet", split="train") dataset = dataset.shuffle(seed=42).select(range(num_samples)) @@ -96,7 +79,6 @@ def build_calibration_loader( image = sample["image"] if image.mode != "RGB": image = image.convert("RGB") - # HF image processors emit `pixel_values` of shape (1, 3, H, W). pixel_values = processor(images=image, return_tensors="pt")["pixel_values"] tensors.append(pixel_values.squeeze(0)) @@ -105,12 +87,7 @@ def build_calibration_loader( def quantize_with_recipe(model, recipe_path: str, calib_batches): - """Resolve the YAML recipe and run `mtq.quantize`. - - Returns the quantized model. The graph still uses high-precision math at - this point — Q/DQ nodes have been inserted around weights and activations - and amax values populated, but no kernel substitution has happened yet. - """ + """Resolve the YAML recipe and run `mtq.quantize`.""" print(f"Loading recipe: {recipe_path}") recipe = load_recipe(recipe_path) if not isinstance(recipe, ModelOptPTQRecipe): @@ -133,8 +110,7 @@ class ViTLogitsWrapper(torch.nn.Module): HF's `ViTForImageClassification.forward` returns an `ImageClassifierOutput` dataclass. `torch_tensorrt.compile` (and `torch.export`) need a tensor-tree - return, so we unwrap it here. The wrapper holds the quantized model as a - submodule; Q/DQ nodes flow through unchanged. + return, so we unwrap it here. """ def __init__(self, vit_model: torch.nn.Module): @@ -146,14 +122,11 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Tensor): - """Compile the quantized model with Torch-TensorRT (Dynamo IR). - - `min_block_size=1` follows the Torch-TRT quantization guide — it makes the - partitioner accept single-node TRT subgraphs, which is what we want so the - Q/DQ + matmul pairs become TRT precision layers instead of falling back to - eager. The compile step expects fake-quant operators in the graph; we run - it under `export_torch_mode` so modelopt's Q/DQ are exported in the - TRT-friendly form. + """Compile the quantized model with Torch-TensorRT (Dynamo IR, strongly-typed). + + `min_block_size=1` follows the Torch-TRT quantization guide so single-node + Q/DQ + matmul subgraphs become TRT precision layers. `export_torch_mode` + makes modelopt emit Q/DQ in the TRT-friendly form during `torch.export`. """ import torch_tensorrt @@ -164,34 +137,11 @@ def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Ten ir="dynamo", arg_inputs=[example_input], min_block_size=1, - # The recipes export weights in BF16; TRT picks the FP8/NVFP4 - # kernel from the Q/DQ pattern, not from this list. - enabled_precisions={torch.bfloat16, torch.float16, torch.float32}, truncate_double=True, ) return trt_model -def benchmark(model: torch.nn.Module, example_input: torch.Tensor, n_warmup: int, n_iters: int): - """Median-of-`n_iters` latency over `example_input`. CUDA-event timed.""" - torch.cuda.synchronize() - with torch.no_grad(): - for _ in range(n_warmup): - model(example_input) - torch.cuda.synchronize() - - starts = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] - with torch.no_grad(): - for i in range(n_iters): - starts[i].record() - model(example_input) - ends[i].record() - torch.cuda.synchronize() - times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) - return times[len(times) // 2] - - def _argmax_logits(out) -> torch.Tensor: """Handle either an HF `ImageClassifierOutput` or a raw tensor.""" logits = out.logits if hasattr(out, "logits") else out @@ -227,13 +177,7 @@ def main(): "--batch_size", type=int, default=1, - help="Batch size for calibration / TRT compile / benchmarking.", - ) - parser.add_argument( - "--benchmark_iters", - type=int, - default=50, - help="Number of timed iterations (after warmup) per benchmark phase.", + help="Batch size for calibration / TRT compile.", ) parser.add_argument( "--save_dir", @@ -245,7 +189,7 @@ def main(): parser.add_argument( "--skip_trt", action="store_true", - help="Quantize + run the BF16-fake-quant model only; skip torch_tensorrt.compile. " + help="Quantize + run the fake-quant model only; skip torch_tensorrt.compile. " "Useful for environments without torch_tensorrt installed.", ) args = parser.parse_args() @@ -253,8 +197,6 @@ def main(): if not torch.cuda.is_available(): raise SystemExit("This example requires a CUDA-capable GPU.") device = torch.device("cuda") - # ViT-Large is a transformer in BF16 on the released checkpoint; the Q/DQ - # nodes operate on top of BF16 master weights either way. dtype = torch.bfloat16 model, processor = load_model_and_processor(args.model_id, device, dtype) @@ -264,16 +206,10 @@ def main(): args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype ) - # Baseline forward + benchmark for a comparison number that survives - # quantization. argmax preserves the predicted-class check below. print("\n=== Baseline (BF16) ===") with torch.no_grad(): baseline_pred = _argmax_logits(model(example_input)) - baseline_latency = benchmark( - lambda x: model(x), example_input, n_warmup=5, n_iters=args.benchmark_iters - ) print(f"Baseline argmax class: {baseline_pred.tolist()}") - print(f"Baseline latency: {baseline_latency:.3f} ms (median over {args.benchmark_iters} iters)") calib_batches = build_calibration_loader( processor, args.calib_samples, args.batch_size, device, dtype @@ -289,7 +225,7 @@ def main(): mto.save(model, ckpt) print(f"Saved quantized modelopt state to {ckpt}") - print("\n=== Fake-quant (modelopt, BF16 math) ===") + print("\n=== Fake-quant (modelopt) ===") with torch.no_grad(): fq_pred = _argmax_logits(model(example_input)) fq_match = (fq_pred == baseline_pred).all().item() @@ -306,11 +242,7 @@ def main(): with torch.no_grad(): trt_pred = trt_model(example_input).argmax(dim=-1) trt_match = (trt_pred == baseline_pred).all().item() - trt_latency = benchmark(trt_model, example_input, n_warmup=5, n_iters=args.benchmark_iters) print(f"TRT argmax class: {trt_pred.tolist()} (matches baseline: {trt_match})") - print(f"TRT latency: {trt_latency:.3f} ms (median over {args.benchmark_iters} iters)") - speedup = baseline_latency / trt_latency if trt_latency > 0 else float("inf") - print(f"\nSpeedup vs. BF16 baseline: {speedup:.2f}x") if __name__ == "__main__": From b2b3f9047206a390474aaabbb1f6d3b081bf7874 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 29 May 2026 20:43:17 +0000 Subject: [PATCH 03/17] Add e2e integration test for torch_trt example * tests/examples/torch_trt/test_torch_tensorrt_ptq.py -- mirrors the tests/examples/torch_onnx/test_torch_quant_to_onnx.py pattern: invokes the example via run_example_command, parametrizes over (fp8, nvfp4), uses a 1-layer ViT config (--no_pretrained + --model_kwargs) so the test completes in ~30 s per parametrized case. Two variants: - test_torch_tensorrt_ptq[precision] -- full e2e through torch_tensorrt.compile (importorskip on torch_tensorrt). - test_torch_tensorrt_ptq_skip_trt[precision] -- quantize-only smoke test, useful on hosts without torch_tensorrt installed. * examples/torch_trt/torch_tensorrt_ptq.py: - Add --no_pretrained + --model_kwargs flags (mirroring torch_onnx) so the same script doubles as the test entry point. - Force aten.cat.default into PyTorch fallback inside compile_with_torch_tensorrt -- torch_tensorrt 2.10's cat converter chokes on the HF ViT cls-token + patch-embedding concat (BF16: "Got unsupported ScalarType BFloat16"; FP16: rank-(-1) TRT tensor that crashes the downstream `embeddings + position_embeddings` add). The cat is a tiny [1,1,H] + [1,N,H] op that runs once per forward, so PyTorch fallback costs essentially nothing. Verified locally: pytest tests/examples/torch_trt/test_torch_tensorrt_ptq.py -> 4 passed in 103 s on RTX 6000 Ada. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/torch_tensorrt_ptq.py | 61 +++++++++++++-- .../torch_trt/test_torch_tensorrt_ptq.py | 78 +++++++++++++++++++ 2 files changed, 133 insertions(+), 6 deletions(-) create mode 100644 tests/examples/torch_trt/test_torch_tensorrt_ptq.py diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index 2e65e59a269..da584928865 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -33,11 +33,12 @@ from __future__ import annotations import argparse +import json from pathlib import Path import torch from datasets import load_dataset -from transformers import AutoImageProcessor, ViTForImageClassification +from transformers import AutoImageProcessor, ViTConfig, ViTForImageClassification import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -53,11 +54,30 @@ } -def load_model_and_processor(model_id: str, device: torch.device, dtype: torch.dtype): - """Pull the HF ViT classifier and its preprocessor.""" - print(f"Loading {model_id} (dtype={dtype})...") +def load_model_and_processor( + model_id: str, + device: torch.device, + dtype: torch.dtype, + pretrained: bool = True, + config_overrides: dict | None = None, +): + """Pull the HF ViT classifier and its preprocessor. + + With ``pretrained=False`` the model is built from a config with random + weights (test path); ``config_overrides`` lets the caller shrink it + (e.g. ``{"num_hidden_layers": 1, "hidden_size": 64, ...}``). The + preprocessor is always loaded from ``model_id`` since it only carries + a small JSON config. + """ + print(f"Loading {model_id} (dtype={dtype}, pretrained={pretrained})...") processor = AutoImageProcessor.from_pretrained(model_id) - model = ViTForImageClassification.from_pretrained(model_id, torch_dtype=dtype) + if pretrained: + model = ViTForImageClassification.from_pretrained(model_id, torch_dtype=dtype) + else: + config = ViTConfig.from_pretrained(model_id) + for k, v in (config_overrides or {}).items(): + setattr(config, k, v) + model = ViTForImageClassification(config).to(dtype) model.eval().to(device) return model, processor @@ -131,6 +151,13 @@ def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Ten import torch_tensorrt print("Compiling with torch_tensorrt.compile (Dynamo IR)...") + # `aten.cat.default` is force-executed in PyTorch because torch_tensorrt + # 2.10's cat converter chokes on the cls-token + patch-embedding concat + # in HF ViT (BFloat16 path: `TypeError: Got unsupported ScalarType + # BFloat16`; FP16 path: rank-(-1) TRT tensor that trips the downstream + # `embeddings + position_embeddings` add). The cat is a tiny [1,1,H] + # + [1,N,H] concat that runs once per forward, so falling back to + # PyTorch costs essentially nothing. with export_torch_mode(): trt_model = torch_tensorrt.compile( model, @@ -138,6 +165,7 @@ def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Ten arg_inputs=[example_input], min_block_size=1, truncate_double=True, + torch_executed_ops={torch.ops.aten.cat.default}, ) return trt_model @@ -192,6 +220,20 @@ def main(): help="Quantize + run the fake-quant model only; skip torch_tensorrt.compile. " "Useful for environments without torch_tensorrt installed.", ) + parser.add_argument( + "--no_pretrained", + action="store_true", + help="Build the model from config with random weights instead of " + "downloading pretrained weights. Useful for fast e2e tests.", + ) + parser.add_argument( + "--model_kwargs", + type=str, + default=None, + help="JSON string of ViTConfig overrides applied when --no_pretrained " + 'is set (e.g. \'{"num_hidden_layers": 1, "hidden_size": 64, ' + '"intermediate_size": 128, "num_attention_heads": 2}\').', + ) args = parser.parse_args() if not torch.cuda.is_available(): @@ -199,7 +241,14 @@ def main(): device = torch.device("cuda") dtype = torch.bfloat16 - model, processor = load_model_and_processor(args.model_id, device, dtype) + config_overrides = json.loads(args.model_kwargs) if args.model_kwargs else None + model, processor = load_model_and_processor( + args.model_id, + device, + dtype, + pretrained=not args.no_pretrained, + config_overrides=config_overrides, + ) image_size = model.config.image_size num_channels = model.config.num_channels example_input = torch.randn( diff --git a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py new file mode 100644 index 00000000000..3b540cd00e8 --- /dev/null +++ b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +from _test_utils.examples.run_command import extend_cmd_parts, run_example_command + +# Recipe variants the example ships. Mirrors the parametrization style of +# ``tests/examples/torch_onnx/test_torch_quant_to_onnx.py``. +_PRECISIONS = ["fp8", "nvfp4"] + +# Tiny ViT config (~1 encoder block) so the test stays under a few seconds +# of GPU time while exercising every code path the recipe touches: encoder +# Linear weight/input quantizers, attention BMM + softmax quantizers, +# per-block LayerNorm output quantizer, and the patch-embed Conv / final +# vit.layernorm / classifier skip rules. +_TINY_VIT_KWARGS = { + "num_hidden_layers": 1, + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 2, +} + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_torch_tensorrt_ptq(precision): + """End-to-end: load tiny ViT -> mtq.quantize via recipe -> torch_tensorrt.compile. + + Runs against the smallest viable ``ViTForImageClassification`` config so + the test stays fast; ``--no_pretrained`` skips the multi-GB pretrained + download. The example's CLI exits non-zero if any step (calibration, + quantization, TRT compile) fails or if the compiled-model argmax doesn't + match the fake-quant argmax on the sample input. + """ + pytest.importorskip("torch_tensorrt") + + cmd_parts = extend_cmd_parts( + ["python", "torch_tensorrt_ptq.py"], + model_id="google/vit-base-patch16-224", + precision=precision, + calib_samples="4", + batch_size="1", + model_kwargs=json.dumps(_TINY_VIT_KWARGS), + ) + cmd_parts.append("--no_pretrained") + run_example_command(cmd_parts, "torch_trt") + + +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_torch_tensorrt_ptq_skip_trt(precision): + """Quantize-only smoke test (no torch_tensorrt.compile). + + Useful on hosts without ``torch_tensorrt`` installed and as a faster + sanity check that just exercises the recipe + ``mtq.quantize`` path. + """ + cmd_parts = extend_cmd_parts( + ["python", "torch_tensorrt_ptq.py"], + model_id="google/vit-base-patch16-224", + precision=precision, + calib_samples="4", + batch_size="1", + model_kwargs=json.dumps(_TINY_VIT_KWARGS), + ) + cmd_parts.extend(["--no_pretrained", "--skip_trt"]) + run_example_command(cmd_parts, "torch_trt") From 814c97068487092ae9ef7c7da8adaf034b414c92 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 29 May 2026 20:52:31 +0000 Subject: [PATCH 04/17] Remove quantize-only smoke variant from torch_trt e2e test Drops `test_torch_tensorrt_ptq_skip_trt` -- the full `test_torch_tensorrt_ptq` variant already exercises the same mtq.quantize path and goes further (torch_tensorrt.compile). The skip-variant added duplicate CI runtime without unique coverage. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .../torch_trt/test_torch_tensorrt_ptq.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py index 3b540cd00e8..f0350d525ad 100644 --- a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py +++ b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py @@ -57,22 +57,3 @@ def test_torch_tensorrt_ptq(precision): ) cmd_parts.append("--no_pretrained") run_example_command(cmd_parts, "torch_trt") - - -@pytest.mark.parametrize("precision", _PRECISIONS) -def test_torch_tensorrt_ptq_skip_trt(precision): - """Quantize-only smoke test (no torch_tensorrt.compile). - - Useful on hosts without ``torch_tensorrt`` installed and as a faster - sanity check that just exercises the recipe + ``mtq.quantize`` path. - """ - cmd_parts = extend_cmd_parts( - ["python", "torch_tensorrt_ptq.py"], - model_id="google/vit-base-patch16-224", - precision=precision, - calib_samples="4", - batch_size="1", - model_kwargs=json.dumps(_TINY_VIT_KWARGS), - ) - cmd_parts.extend(["--no_pretrained", "--skip_trt"]) - run_example_command(cmd_parts, "torch_trt") From 280a9ae4de2a1d27b7f11c3349932c8be89c37f6 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:02:11 +0000 Subject: [PATCH 05/17] [6078291] Simplify ViT FP8/NVFP4 recipes via shared config units Recompose the ViT PTQ recipes from the shared $import building blocks under modelopt_recipes/configs/ instead of inlined quant_cfg entries: - fp8: import w8a8_fp8_fp8 + attention_qkv_fp8 (resolves to the same config). - nvfp4: import numerics/nvfp4 for the nn.Linear weight/input cfg and attention_qkv_fp8 for the q/k/v BMM + softmax, and hold the patch-embed Conv2d and classifier head at per-tensor FP8. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt_recipes/huggingface/vit/ptq/fp8.yaml | 40 ++----------- .../huggingface/vit/ptq/nvfp4.yaml | 56 +++++++------------ 2 files changed, 25 insertions(+), 71 deletions(-) diff --git a/modelopt_recipes/huggingface/vit/ptq/fp8.yaml b/modelopt_recipes/huggingface/vit/ptq/fp8.yaml index 6158a99590a..6eb39351f62 100644 --- a/modelopt_recipes/huggingface/vit/ptq/fp8.yaml +++ b/modelopt_recipes/huggingface/vit/ptq/fp8.yaml @@ -15,41 +15,13 @@ metadata: recipe_type: ptq - description: >- - HuggingFace ViT FP8 PTQ recipe matching modelopt.onnx FP8 layout: - per-tensor FP8 E4M3 on encoder Linear weights + inputs, attention QKV - BMMs and softmax outputs, MHA-aware per-block LayerNorm outputs; - patch-embed Conv2d, classifier head, and the final LayerNorm are - left FP16. Uses max calibration. +imports: + w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8 + attention_qkv_fp8: configs/ptq/units/attention_qkv_fp8 quantize: algorithm: max quant_cfg: - - quantizer_name: '*' + - $import: w8a8_fp8_fp8 + - quantizer_name: '*output_quantizer' enable: false - - - parent_class: 'nn.Linear' - quantizer_name: 'vit.encoder.layer.*.weight_quantizer' - cfg: - num_bits: e4m3 - axis: - - parent_class: 'nn.Linear' - quantizer_name: 'vit.encoder.layer.*.input_quantizer' - cfg: - num_bits: e4m3 - axis: - - - quantizer_name: '*[qkv]_bmm_quantizer' - cfg: - num_bits: e4m3 - axis: - - - quantizer_name: '*softmax_quantizer' - cfg: - num_bits: e4m3 - axis: - - - parent_class: 'nn.LayerNorm' - quantizer_name: 'vit.encoder.layer.*.output_quantizer' - cfg: - num_bits: e4m3 - axis: + - $import: attention_qkv_fp8 diff --git a/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml b/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml index e30f7184dd6..3fb98da7366 100644 --- a/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml +++ b/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml @@ -15,49 +15,31 @@ metadata: recipe_type: ptq - description: >- - HuggingFace ViT NVFP4 W4A4 PTQ recipe. Skip list mirrors modelopt.onnx - FP8 (patch-embed Conv2d, classifier, and final vit.layernorm stay - FP16). Encoder Linear weights/inputs run NVFP4; attention BMMs, - softmax, and per-block LayerNorm outputs stay at FP8. - Uses AWQ-lite calibration. +imports: + nvfp4: configs/numerics/nvfp4 + fp8: configs/numerics/fp8 + attention_qkv_fp8: configs/ptq/units/attention_qkv_fp8 + w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 quantize: algorithm: awq_lite quant_cfg: - - quantizer_name: '*' + - $import: w4a4_nvfp4_nvfp4 + + - quantizer_name: '*output_quantizer' enable: false - - parent_class: 'nn.Linear' - quantizer_name: 'vit.encoder.layer.*.weight_quantizer' - cfg: - num_bits: e2m1 - axis: - block_sizes: - -1: 16 - type: dynamic - scale_bits: e4m3 - - parent_class: 'nn.Linear' - quantizer_name: 'vit.encoder.layer.*.input_quantizer' - cfg: - num_bits: e2m1 - axis: - block_sizes: - -1: 16 - type: dynamic - scale_bits: e4m3 + - $import: attention_qkv_fp8 - - quantizer_name: '*[qkv]_bmm_quantizer' + # Hold the patch-embed Conv2d and the classifier head at per-tensor FP8. + - quantizer_name: '*patch_embeddings.projection.weight_quantizer' cfg: - num_bits: e4m3 - axis: - - - quantizer_name: '*softmax_quantizer' + $import: fp8 + - quantizer_name: '*patch_embeddings.projection.input_quantizer' cfg: - num_bits: e4m3 - axis: - - - parent_class: 'nn.LayerNorm' - quantizer_name: 'vit.encoder.layer.*.output_quantizer' + $import: fp8 + - quantizer_name: '*classifier.weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*classifier.input_quantizer' cfg: - num_bits: e4m3 - axis: + $import: fp8 From a6c5707fc639aa491c1010fda7f70c2b3189ccf4 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:02:42 +0000 Subject: [PATCH 06/17] [6078291] Add Torch-TRT ViT ImageNet accuracy example examples/torch_trt/torch_tensorrt_accuracy.py quantizes and Torch-TensorRT- compiles a HuggingFace ViT (reusing torch_tensorrt_ptq.py) and reports ImageNet-1k top-1/top-5 via the onnx_ptq evaluate() harness. A thin adapter casts the float32 dataloader batches to the compiled model's compute dtype and unwraps logits; the Torch-TRT path is pinned to batch_size=1 (static engine). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/torch_tensorrt_accuracy.py | 252 ++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 examples/torch_trt/torch_tensorrt_accuracy.py diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py new file mode 100644 index 00000000000..85f9d0fa8c1 --- /dev/null +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Measure ImageNet top-1/top-5 accuracy of a Torch-TensorRT ViT. + +Pipeline: + +1. Quantize a HuggingFace ViT with a ModelOpt recipe and compile it with + ``torch_tensorrt.compile(ir="dynamo")`` — reusing the sibling example + ``torch_tensorrt_ptq.py``. +2. Score the compiled model on the ImageNet-1k validation split using the + ``onnx_ptq`` example's ``evaluate`` API (``examples/onnx_ptq/evaluation.py``). + +The compiled Torch-TRT module is a ``torch.nn.Module``, so ``evaluate`` runs it +exactly like an eager model. A thin :class:`_EvalAdapter` bridges the two +contracts: it casts the dataloader's float32 image batches to the model's +compute dtype and unwraps HF ``ImageClassifierOutput`` to a plain logits tensor. + +Example:: + + python torch_tensorrt_accuracy.py --precision fp8 --eval_data_size 5000 --baseline + +``--imagenet_path`` defaults to the gated ``ILSVRC/imagenet-1k`` HF dataset +(accept its license / set ``HF_TOKEN``), or point it at a local copy. Note the +``evaluate`` API shuffles the validation set, so a partial ``--eval_data_size`` +samples a different random subset each run; use the full set for a stable score. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import sys +from pathlib import Path + +import torch + +# Reuse the quantize -> torch_tensorrt.compile pipeline from the sibling example. +_THIS_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(_THIS_DIR)) +import torch_tensorrt_ptq as ttptq # noqa: E402 + +# Reuse the ImageNet accuracy harness from the onnx_ptq example (sibling dir). +_ONNX_PTQ_DIR = _THIS_DIR.parent / "onnx_ptq" +sys.path.insert(0, str(_ONNX_PTQ_DIR)) +from evaluation import evaluate # noqa: E402 + + +class _EvalAdapter(torch.nn.Module): + """Adapt a compiled/eager ViT to the ``onnx_ptq`` ``evaluate`` contract. + + ``evaluate_accuracy`` feeds float32 image batches, calls ``model(inputs)``, + and reads ``outputs.data``. This adapter casts inputs to the model's compute + dtype (the dataloader yields FP32) and unwraps an HF ``ImageClassifierOutput`` + to the bare logits tensor. + """ + + def __init__(self, model: torch.nn.Module, dtype: torch.dtype): + super().__init__() + self.model = model + self._dtype = dtype + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + out = self.model(pixel_values.to(self._dtype)) + return out.logits if hasattr(out, "logits") else out + + +def build_processor_transform(processor): + """Return a ``PIL.Image -> (C, H, W) float tensor`` transform from the HF processor. + + Using the model's own image processor keeps eval preprocessing (resize, + normalization mean/std) consistent with how the ViT was trained, which is + more faithful for a HuggingFace checkpoint than a generic timm transform. + The model and ``ILSVRC/imagenet-1k`` share the standard 1000-class ordering, + so predicted indices line up with the dataset labels. + """ + + def _transform(image): + return processor(images=image, return_tensors="pt")["pixel_values"][0] + + return _transform + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--model_id", + default="google/vit-large-patch16-224", + help="HuggingFace model id of the ViT classifier to quantize and score.", + ) + parser.add_argument( + "--precision", + choices=sorted(ttptq.PRECISION_TO_RECIPE), + default="fp8", + help="Which ViT recipe variant to apply.", + ) + parser.add_argument( + "--recipe", + default=None, + help="Override the recipe path (relative to modelopt_recipes/ or absolute). " + "If unset, the recipe is picked by --precision.", + ) + parser.add_argument( + "--calib_samples", + type=int, + default=128, + help="Number of tiny-imagenet samples to use for calibration.", + ) + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Calibration / compile / eval batch size. The Torch-TRT engine is " + "compiled for this single static batch shape and the onnx_ptq evaluate() " + "dataloader keeps the trailing partial batch, so the Torch-TRT path " + "requires --batch_size 1; larger batches are only allowed with --skip_trt.", + ) + parser.add_argument( + "--eval_data_size", + type=int, + default=None, + help="Number of ImageNet validation images to score (default: full 50k).", + ) + parser.add_argument( + "--imagenet_path", + default="ILSVRC/imagenet-1k", + help="HF dataset card or local path to the ImageNet validation set (gated).", + ) + parser.add_argument( + "--baseline", + action="store_true", + help="Also score the un-quantized BF16 model for a reference point.", + ) + parser.add_argument( + "--skip_trt", + action="store_true", + help="Score the fake-quant (modelopt) model; skip torch_tensorrt.compile. " + "Useful for environments without torch_tensorrt installed.", + ) + parser.add_argument( + "--no_pretrained", + action="store_true", + help="Build the model from config with random weights (fast smoke test).", + ) + parser.add_argument( + "--model_kwargs", + default=None, + help="JSON string of ViTConfig overrides applied when --no_pretrained is set.", + ) + parser.add_argument( + "--results_path", + default=None, + help="If set, write the accuracy results to this CSV path.", + ) + args = parser.parse_args() + + # The Torch-TRT engine is compiled for one static batch shape, and the reused + # onnx_ptq evaluate() dataloader does not drop the trailing partial batch, so a + # batch size that doesn't divide the validation set would crash the static + # engine mid-run. Fail fast. The fake-quant (--skip_trt) path is a plain eager + # module and tolerates any batch size. + if not args.skip_trt and args.batch_size != 1: + raise SystemExit( + "The Torch-TensorRT path requires --batch_size 1 (the engine is compiled " + "for a single static batch shape). Use --skip_trt to score the fake-quant " + "model at a larger batch size." + ) + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable GPU.") + device = torch.device("cuda") + dtype = torch.float16 + + config_overrides = json.loads(args.model_kwargs) if args.model_kwargs else None + model, processor = ttptq.load_model_and_processor( + args.model_id, + device, + dtype, + pretrained=not args.no_pretrained, + config_overrides=config_overrides, + ) + transform = build_processor_transform(processor) + + def run_eval(m: torch.nn.Module) -> tuple[float, float]: + top1, top5 = evaluate( + _EvalAdapter(m, dtype), + transform, + batch_size=args.batch_size, + num_examples=args.eval_data_size, + device="cuda", + dataset_path=args.imagenet_path, + ) + return top1, top5 + + results: list[list[str | float]] = [["Metric", "Top1 (%)", "Top5 (%)"]] + + # Baseline must run before in-place quantization mutates `model`. + if args.baseline: + prec = str(dtype).rsplit(".", 1)[-1] # e.g. "float16" + print(f"\n=== Baseline ({prec}) ===") + top1, top5 = run_eval(model) + print(f"baseline ({prec}) top1={top1:.2f}% top5={top5:.2f}%") + results.append([f"baseline_{prec}", top1, top5]) + + calib_batches = ttptq.build_calibration_loader( + processor, args.calib_samples, args.batch_size, device, dtype + ) + recipe_path = args.recipe or ttptq.PRECISION_TO_RECIPE[args.precision] + ttptq.quantize_with_recipe(model, recipe_path, calib_batches) + + wrapped = ttptq.ViTLogitsWrapper(model).to(device).eval() + if args.skip_trt: + print("\n--skip_trt set; scoring the fake-quant model (no Torch-TensorRT compile).") + eval_model: torch.nn.Module = wrapped + tag = f"{args.precision} (fake-quant)" + else: + image_size = model.config.image_size + num_channels = model.config.num_channels + example_input = torch.randn( + args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype + ) + eval_model = ttptq.compile_with_torch_tensorrt(wrapped, example_input) + tag = f"{args.precision} (torch-trt)" + + print(f"\n=== {tag} ===") + top1, top5 = run_eval(eval_model) + print(f"{tag} top1={top1:.2f}% top5={top5:.2f}%") + results.append([tag, top1, top5]) + + if args.results_path: + with open(args.results_path, "w", newline="") as f: + csv.writer(f).writerows(results) + print(f"\nWrote results to {args.results_path}") + + +if __name__ == "__main__": + main() From 712f13bf185a9ddf0b628ac90e4e2b9754ed2190 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:13:11 +0000 Subject: [PATCH 07/17] [6078291] Torch-TRT-compile the accuracy baseline + document the example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - torch_tensorrt_accuracy.py: the --baseline reference is now Torch-TensorRT- compiled like the quantized model (shared to_eval_model helper), so every reported number comes from the same TRT runtime. - README/CHANGELOG: document torch_tensorrt_accuracy.py and correct the recipe description — the recipes now compose from the shared modelopt_recipes/configs/ $import units, and the nvfp4 recipe holds the patch-embed Conv2d + classifier head at FP8. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/torch_trt/README.md | 45 ++++++++++++---- examples/torch_trt/torch_tensorrt_accuracy.py | 53 ++++++++++++------- 3 files changed, 68 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6785a7e09b8..85401ff738b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,7 +38,7 @@ Changelog - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. -- Add Torch-TensorRT FP8 / NVFP4 deployment example for HuggingFace ViT (``examples/torch_trt/``) covering ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``. Ships two ViT-tuned PTQ recipes under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``, ``nvfp4.yaml``) — encoder Linear weights+inputs quantized; attention Q/K/V BMMs, softmax, and per-block LayerNorm outputs at FP8; patch-embed ``nn.Conv2d``, ``classifier``, and the final ``vit.layernorm`` left FP16. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. +- Add Torch-TensorRT FP8 / NVFP4 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships two ViT-tuned PTQ recipes under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``, ``nvfp4.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: FP8 quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax; NVFP4 runs W4A4 on the encoder ``nn.Linear`` weights/inputs while holding the patch-embed ``nn.Conv2d``, ``classifier``, and attention BMMs/softmax at FP8. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. **Bug Fixes** diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index d826f7fc3d6..b599b20dfff 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -57,20 +57,43 @@ python examples/torch_trt/torch_tensorrt_ptq.py \ that the compiled-model argmax matches the fake-quant argmax on a sample input. -## ViT-specific recipes shipped with the example +## Measuring ImageNet accuracy + +`torch_tensorrt_accuracy.py` reuses the pipeline above and reports ImageNet-1k +top-1 / top-5 accuracy via the `onnx_ptq` example's `evaluate()` harness +([`examples/onnx_ptq/evaluation.py`](../onnx_ptq/evaluation.py)): + +```bash +python examples/torch_trt/torch_tensorrt_accuracy.py \ + --precision fp8 \ + --baseline \ + --eval_data_size 5000 +``` -These are the recipes the CLI selects by default when `--model_id` points -at a HF ViT classifier. They are **not** thin wrappers around the modelopt -defaults — they're tuned for the HF ViT module layout. +- `--baseline` also scores the unquantized model. It is Torch-TensorRT-compiled + the same way as the quantized model, so every reported number comes from the + same TRT runtime (pass `--skip_trt` to score the eager / fake-quant models). +- The Torch-TRT engine is compiled for one static batch shape, so the eval path + requires `--batch_size 1`. +- Validation uses the gated `ILSVRC/imagenet-1k` split (accept its license / set + `HF_TOKEN`), or point `--imagenet_path` at a local copy. `evaluate()` shuffles + the split, so a partial `--eval_data_size` draws a different random subset each + run — omit it (full set) for a stable, comparable score. +- `--results_path results.csv` writes the metrics table to CSV. -| Flag | Recipe path | Key differences from the default | -|------|-------------|----------------------------------| -| `--precision fp8` | `huggingface/vit/ptq/fp8` | W8A8 FP8 **plus** MHA-aware FP8 on every per-block `nn.LayerNorm` output (shared Q/DQ feeds Q/K/V + MLP), FP8 attention Q/K/V BMM + softmax slots, patch-embedding `nn.Conv2d` left FP16, `classifier` head left in FP16, final `vit.layernorm` left FP16. | -| `--precision nvfp4` | `huggingface/vit/ptq/nvfp4` | Same skip list as the FP8 recipe; encoder Linear weights/inputs run NVFP4 W4A4 (E2M1, block 16, FP8 scales). Attention BMMs, softmax, and per-block LayerNorm outputs stay at FP8 — NVFP4 is too aggressive there. Uses `awq_lite` calibration. | +## ViT-specific recipes shipped with the example -Each recipe is self-contained (no `$import` of shared snippets) and uses -the "specific-enable" style: narrow `parent_class` + path scoping on the -enable rules means no `enable: false` carve-outs are needed. +These are the recipes the CLI selects by default when `--model_id` points at a +HF ViT classifier. They are tuned for the HF ViT module layout and are composed +from the shared `$import` building blocks under +[`modelopt_recipes/configs/`](../../modelopt_recipes/configs/) +(`numerics/{fp8,nvfp4}`, `ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8}`) rather +than spelling out each `quant_cfg` entry. + +| Flag | Recipe path | What it quantizes | +|------|-------------|-------------------| +| `--precision fp8` | `huggingface/vit/ptq/fp8` | W8A8 FP8 (E4M3) on every weight + input quantizer — encoder Linears, the patch-embed `nn.Conv2d`, the `classifier` head, and per-block `nn.LayerNorm` inputs — plus FP8 on the attention Q/K/V BMMs and softmax. Output quantizers disabled. | +| `--precision nvfp4` | `huggingface/vit/ptq/nvfp4` | NVFP4 W4A4 (E2M1, block 16, FP8 scales) on the encoder `nn.Linear` weights/inputs, with the patch-embed `nn.Conv2d`, the `classifier` head, and the attention Q/K/V BMMs + softmax held at FP8. Uses `awq_lite` calibration. | ## Hardware requirements diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py index 85f9d0fa8c1..f896fcf0519 100644 --- a/examples/torch_trt/torch_tensorrt_accuracy.py +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -144,7 +144,9 @@ def main(): parser.add_argument( "--baseline", action="store_true", - help="Also score the un-quantized BF16 model for a reference point.", + help="Also score the unquantized model as a reference. It is " + "Torch-TensorRT-compiled like the quantized model (or run eager under " + "--skip_trt) so the comparison is apples-to-apples.", ) parser.add_argument( "--skip_trt", @@ -207,15 +209,38 @@ def run_eval(m: torch.nn.Module) -> tuple[float, float]: ) return top1, top5 + image_size = model.config.image_size + num_channels = model.config.num_channels + example_input = torch.randn( + args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype + ) + runtime = "fake-quant" if args.skip_trt else "torch-trt" + + def to_eval_model(m: torch.nn.Module, what: str) -> torch.nn.Module: + """Logits-wrap and (unless --skip_trt) Torch-TensorRT-compile ``m`` for eval. + + The baseline is compiled the same way as the quantized model so all + reported numbers come from the same Torch-TRT runtime. + """ + wrapped = ttptq.ViTLogitsWrapper(m).to(device).eval() + if args.skip_trt: + return wrapped + print(f"\nCompiling {what} with Torch-TensorRT ...") + return ttptq.compile_with_torch_tensorrt(wrapped, example_input) + results: list[list[str | float]] = [["Metric", "Top1 (%)", "Top5 (%)"]] - # Baseline must run before in-place quantization mutates `model`. + # Baseline must be built + scored before in-place quantization mutates `model`. if args.baseline: prec = str(dtype).rsplit(".", 1)[-1] # e.g. "float16" - print(f"\n=== Baseline ({prec}) ===") - top1, top5 = run_eval(model) - print(f"baseline ({prec}) top1={top1:.2f}% top5={top5:.2f}%") - results.append([f"baseline_{prec}", top1, top5]) + base_tag = f"baseline-{prec} ({runtime})" + base_eval = to_eval_model(model, "unquantized baseline") + print(f"\n=== {base_tag} ===") + top1, top5 = run_eval(base_eval) + print(f"{base_tag} top1={top1:.2f}% top5={top5:.2f}%") + results.append([base_tag, top1, top5]) + del base_eval + torch.cuda.empty_cache() calib_batches = ttptq.build_calibration_loader( processor, args.calib_samples, args.batch_size, device, dtype @@ -223,20 +248,8 @@ def run_eval(m: torch.nn.Module) -> tuple[float, float]: recipe_path = args.recipe or ttptq.PRECISION_TO_RECIPE[args.precision] ttptq.quantize_with_recipe(model, recipe_path, calib_batches) - wrapped = ttptq.ViTLogitsWrapper(model).to(device).eval() - if args.skip_trt: - print("\n--skip_trt set; scoring the fake-quant model (no Torch-TensorRT compile).") - eval_model: torch.nn.Module = wrapped - tag = f"{args.precision} (fake-quant)" - else: - image_size = model.config.image_size - num_channels = model.config.num_channels - example_input = torch.randn( - args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype - ) - eval_model = ttptq.compile_with_torch_tensorrt(wrapped, example_input) - tag = f"{args.precision} (torch-trt)" - + tag = f"{args.precision} ({runtime})" + eval_model = to_eval_model(model, f"{args.precision} model") print(f"\n=== {tag} ===") top1, top5 = run_eval(eval_model) print(f"{tag} top1={top1:.2f}% top5={top5:.2f}%") From 934e2c0af1357b60fa01e10ce1e94f2bec91285d Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:23:38 +0000 Subject: [PATCH 08/17] [6078291] Comment the local torch_tensorrt import in the example compile_with_torch_tensorrt imports torch_tensorrt inside the function (not at module scope) so the quantize-only --skip_trt path still runs on hosts without torch_tensorrt installed. Add a one-line comment making that intent explicit. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/torch_tensorrt_ptq.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index da584928865..ad6c7a926d0 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -148,6 +148,8 @@ def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Ten Q/DQ + matmul subgraphs become TRT precision layers. `export_torch_mode` makes modelopt emit Q/DQ in the TRT-friendly form during `torch.export`. """ + # Imported here (not at module scope) so the quantize-only `--skip_trt` path + # still runs on hosts without torch_tensorrt installed. import torch_tensorrt print("Compiling with torch_tensorrt.compile (Dynamo IR)...") From 5d47e12703794a081f80d0360a0bcdb41a08c5ee Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:39:41 +0000 Subject: [PATCH 09/17] [6078291] Address PR review: wire torch_trt into example CI + deps - example_tests.yml: add torch_trt to the ONNX/TensorRT example matrix. - .github/CODEOWNERS: add /examples/torch_trt. - README: install nvidia-modelopt (drop the [torch] extra). - requirements.txt: bump transformers to >=4.56. - torch_tensorrt_accuracy.py: default --calib_samples to 512. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .github/CODEOWNERS | 1 + .github/workflows/example_tests.yml | 2 +- examples/torch_trt/README.md | 4 ++-- examples/torch_trt/requirements.txt | 2 +- examples/torch_trt/torch_tensorrt_accuracy.py | 2 +- 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 7f40b270da1..fbcb9aa9915 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -59,6 +59,7 @@ modelopt_recipes @NVIDIA/modelopt-recipes-codeowners /examples/specdec_bench @NVIDIA/modelopt-examples-specdec_bench-codeowners /examples/speculative_decoding @NVIDIA/modelopt-torch-speculative-codeowners /examples/torch_onnx @NVIDIA/modelopt-onnx-codeowners +/examples/torch_trt @NVIDIA/modelopt-onnx-codeowners /examples/vllm_serve @NVIDIA/modelopt-examples-llm_ptq-codeowners /examples/windows @NVIDIA/modelopt-windows-codeowners diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 8564c772a49..d9e4fb71d9d 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -98,7 +98,7 @@ jobs: strategy: fail-fast: false matrix: - example: [diffusers, torch_onnx] + example: [diffusers, torch_onnx, torch_trt] uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index b599b20dfff..5d06804f0ae 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -16,7 +16,7 @@ TensorRT precision layers. # From the NVIDIA TensorRT docker image (recommended): docker run --gpus all -it --rm -v $(pwd):/workspace -w /workspace nvcr.io/nvidia/tensorrt:26.02-py3 bash -pip install -U "nvidia-modelopt[torch]" +pip install -U "nvidia-modelopt" pip install -r examples/torch_trt/requirements.txt ``` @@ -30,7 +30,7 @@ the version pulled by `pip` must match your installed PyTorch. # FP8 / NVFP4 default model is google/vit-large-patch16-224 python examples/torch_trt/torch_tensorrt_ptq.py \ --precision fp8/nvfp4 \ - --calib_samples 128 \ + --calib_samples 512 \ --batch_size 1 # Quantize but don't TRT-compile (handy on a non-TRT host) diff --git a/examples/torch_trt/requirements.txt b/examples/torch_trt/requirements.txt index 8caf5b5fd93..0da9517074e 100644 --- a/examples/torch_trt/requirements.txt +++ b/examples/torch_trt/requirements.txt @@ -1,3 +1,3 @@ datasets>=2.14.4 torch-tensorrt>=2.4.0 -transformers>=4.40 +transformers>=4.56 diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py index f896fcf0519..f9d9ac39ffa 100644 --- a/examples/torch_trt/torch_tensorrt_accuracy.py +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -118,7 +118,7 @@ def main(): parser.add_argument( "--calib_samples", type=int, - default=128, + default=512, help="Number of tiny-imagenet samples to use for calibration.", ) parser.add_argument( From cfa84448f59a7f420c05941b4e808e8af8596845 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 10 Jun 2026 16:13:49 +0000 Subject: [PATCH 10/17] [6078291] Document torch_trt vs torch_onnx/onnx_ptq in example README Address PR review: explain how the Torch-TensorRT PTQ example differs from the ONNX export examples with a 3-way comparison table covering starting point, where quantization happens, export step, intermediate artifact, and runtime. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/README.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index 5d06804f0ae..7d4c6bca321 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -10,6 +10,30 @@ ModelOpt inserts Q/DQ nodes into the eager PyTorch graph, then `torch_tensorrt.compile(ir="dynamo")` converts those Q/DQ nodes into native TensorRT precision layers. +## How this differs from the ONNX examples + +All three of these examples reach the same destination — a low-precision +TensorRT engine — but quantize at a different point in the pipeline and emit a +different artifact, so they suit different deployment stacks: + +| | Torch-TensorRT (this example) | [`torch_onnx`](../torch_onnx/) | [`onnx_ptq`](../onnx_ptq/) | +|---|---|---|---| +| Starting point | a PyTorch / HF model | a PyTorch / timm model | an already-exported ONNX model | +| Quantize on | the eager PyTorch graph (`mtq.quantize`) | the eager PyTorch graph (`mtq.quantize`) | the ONNX graph directly (ModelOpt ONNX PTQ) | +| Export step | none — the FX/Dynamo graph stays in-process | `torch.onnx.export` of the Q/DQ graph, postprocessed for TRT | none — Q/DQ inserted straight into the ONNX graph | +| Intermediate artifact | none | a Q/DQ ONNX file | a Q/DQ ONNX file | +| Compiler + runtime | `torch_tensorrt.compile(ir="dynamo")` → a `torch.nn.Module` you call from PyTorch | TensorRT builds a standalone engine from the ONNX | TensorRT builds a standalone engine from the ONNX | +| Best when | PyTorch-native serving; you want a drop-in compiled module | you quantize in PyTorch but deploy via a portable ONNX → TRT engine | you only have an ONNX model and never touch PyTorch | + +This example and [`torch_onnx`](../torch_onnx/) share the same PyTorch front end +(`mtq.quantize`), so the numerics are identical — they differ only in the back +end: this one keeps the graph in-process and hands it to Torch-TensorRT, while +`torch_onnx` exports a portable ONNX artifact for the standalone TensorRT +runtime. [`onnx_ptq`](../onnx_ptq/) instead quantizes the ONNX graph directly, +for when you start from an ONNX model rather than PyTorch. Pick this example +when your serving stack is PyTorch-native and you'd rather avoid an ONNX export +step. + ## Setup ```bash From c3eda36d4a0ba0c835dab12fb228357b5478db0c Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:10:38 +0000 Subject: [PATCH 11/17] [6078291] torch_trt example: default to dynamic FP8 engine, simplify CLI - Compile a dynamic-batch engine by default: the cls-token cat stays in TRT (one engine) so Q/DQ fuse into FP8 across the whole ViT (~121 vs ~24 GEMMs). Static keeps the bf16-cat PyTorch fallback (two engines). - hidden_act="gelu_fast"; defaults: --batch_size 128, --calib_samples 1024, --save_dir ./modelopt_quantized. - Replace --precision with --recipe (default ViT FP8 recipe); drop --no_pretrained / --model_kwargs. - Accuracy script: dynamic engine + batch 128, recipe-derived result label. - Update example test + README. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/README.md | 55 +++++-- examples/torch_trt/torch_tensorrt_accuracy.py | 67 ++------ examples/torch_trt/torch_tensorrt_ptq.py | 151 +++++++++++------- .../torch_trt/test_torch_tensorrt_ptq.py | 40 ++--- 4 files changed, 155 insertions(+), 158 deletions(-) diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index 7d4c6bca321..c64fa068615 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -51,16 +51,18 @@ the version pulled by `pip` must match your installed PyTorch. ## Usage ```bash -# FP8 / NVFP4 default model is google/vit-large-patch16-224 +# Default model is google/vit-large-patch16-224, default recipe is the ViT FP8 recipe python examples/torch_trt/torch_tensorrt_ptq.py \ - --precision fp8/nvfp4 \ - --calib_samples 512 \ - --batch_size 1 + --recipe huggingface/vit/ptq/fp8 \ + --calib_samples 1024 \ + --batch_size 128 -# Quantize but don't TRT-compile (handy on a non-TRT host) +# NVFP4 instead of FP8 python examples/torch_trt/torch_tensorrt_ptq.py \ - --precision fp8/nvfp4 \ - --skip_trt + --recipe huggingface/vit/ptq/nvfp4 + +# Quantize but don't TRT-compile (handy on a non-TRT host) +python examples/torch_trt/torch_tensorrt_ptq.py --skip_trt # Custom model + custom recipe python examples/torch_trt/torch_tensorrt_ptq.py \ @@ -81,6 +83,23 @@ python examples/torch_trt/torch_tensorrt_ptq.py \ that the compiled-model argmax matches the fake-quant argmax on a sample input. +## Dynamic vs static engine (and why FP8 cares) + +HF ViT concatenates a cls token onto the patch embeddings. In a **static** +compile, Torch-TensorRT's `cat` converter constant-folds that token through +numpy — which has no bfloat16 dtype — so the bf16 graph must run `aten.cat` in +PyTorch. That graph break splits the model into **two** TRT engines and stops +TRT from fusing the Q/DQ nodes into FP8 across the boundary, so only ~24 FP8 +GEMMs survive for ViT-large (one per layer). + +In a **dynamic** compile the `cat` operands are symbolic tensors (no numpy +materialization), so the model stays a **single** engine and TRT fuses Q/DQ into +~121 FP8 GEMMs — including bias+GELU epilogue-fused `e4m3` kernels — a ~5× jump +in FP8 coverage. The example therefore uses a **dynamic engine** for both fp8 +and nvfp4, built for `min=1, opt=--batch_size, max=1024` and serving any batch in +that range. (nvfp4's low-precision kernels require Blackwell — see *Hardware +requirements*.) + ## Measuring ImageNet accuracy `torch_tensorrt_accuracy.py` reuses the pipeline above and reports ImageNet-1k @@ -89,7 +108,8 @@ top-1 / top-5 accuracy via the `onnx_ptq` example's `evaluate()` harness ```bash python examples/torch_trt/torch_tensorrt_accuracy.py \ - --precision fp8 \ + --recipe huggingface/vit/ptq/fp8 \ + --batch_size 128 \ --baseline \ --eval_data_size 5000 ``` @@ -97,8 +117,9 @@ python examples/torch_trt/torch_tensorrt_accuracy.py \ - `--baseline` also scores the unquantized model. It is Torch-TensorRT-compiled the same way as the quantized model, so every reported number comes from the same TRT runtime (pass `--skip_trt` to score the eager / fake-quant models). -- The Torch-TRT engine is compiled for one static batch shape, so the eval path - requires `--batch_size 1`. +- The eval uses a **dynamic** engine (default `--batch_size 128`) for both + precisions, so it serves the trailing partial batch at any batch size — see + *Dynamic vs static engine*. - Validation uses the gated `ILSVRC/imagenet-1k` split (accept its license / set `HF_TOKEN`), or point `--imagenet_path` at a local copy. `evaluate()` shuffles the split, so a partial `--eval_data_size` draws a different random subset each @@ -114,10 +135,10 @@ from the shared `$import` building blocks under (`numerics/{fp8,nvfp4}`, `ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8}`) rather than spelling out each `quant_cfg` entry. -| Flag | Recipe path | What it quantizes | -|------|-------------|-------------------| -| `--precision fp8` | `huggingface/vit/ptq/fp8` | W8A8 FP8 (E4M3) on every weight + input quantizer — encoder Linears, the patch-embed `nn.Conv2d`, the `classifier` head, and per-block `nn.LayerNorm` inputs — plus FP8 on the attention Q/K/V BMMs and softmax. Output quantizers disabled. | -| `--precision nvfp4` | `huggingface/vit/ptq/nvfp4` | NVFP4 W4A4 (E2M1, block 16, FP8 scales) on the encoder `nn.Linear` weights/inputs, with the patch-embed `nn.Conv2d`, the `classifier` head, and the attention Q/K/V BMMs + softmax held at FP8. Uses `awq_lite` calibration. | +| `--recipe` value | What it quantizes | +|------------------|-------------------| +| `huggingface/vit/ptq/fp8` (default) | W8A8 FP8 (E4M3) on every weight + input quantizer — encoder Linears, the patch-embed `nn.Conv2d`, the `classifier` head, and per-block `nn.LayerNorm` inputs — plus FP8 on the attention Q/K/V BMMs and softmax. Output quantizers disabled. | +| `huggingface/vit/ptq/nvfp4` | NVFP4 W4A4 (E2M1, block 16, FP8 scales) on the encoder `nn.Linear` weights/inputs, with the patch-embed `nn.Conv2d`, the `classifier` head, and the attention Q/K/V BMMs + softmax held at FP8. Uses `awq_lite` calibration. | ## Hardware requirements @@ -132,9 +153,9 @@ low-precision kernel. ### Resuming from a saved checkpoint -Pass `--save_dir ` to persist the modelopt-quantized model -(`vit_modelopt_state.pt`). To reload without recalibrating, restore it -before the TRT compile step with: +The quantized modelopt model is saved to `--save_dir` (default +`./modelopt_quantized`) as `vit_modelopt_state.pt`. To reload without +recalibrating, restore it before the TRT compile step with: ```python import modelopt.torch.opt as mto diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py index f9d9ac39ffa..85c2f12f218 100644 --- a/examples/torch_trt/torch_tensorrt_accuracy.py +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -30,7 +30,7 @@ Example:: - python torch_tensorrt_accuracy.py --precision fp8 --eval_data_size 5000 --baseline + python torch_tensorrt_accuracy.py --batch_size 128 --eval_data_size 5000 --baseline ``--imagenet_path`` defaults to the gated ``ILSVRC/imagenet-1k`` HF dataset (accept its license / set ``HF_TOKEN``), or point it at a local copy. Note the @@ -42,7 +42,6 @@ import argparse import csv -import json import sys from pathlib import Path @@ -103,32 +102,25 @@ def main(): default="google/vit-large-patch16-224", help="HuggingFace model id of the ViT classifier to quantize and score.", ) - parser.add_argument( - "--precision", - choices=sorted(ttptq.PRECISION_TO_RECIPE), - default="fp8", - help="Which ViT recipe variant to apply.", - ) parser.add_argument( "--recipe", - default=None, - help="Override the recipe path (relative to modelopt_recipes/ or absolute). " - "If unset, the recipe is picked by --precision.", + default=ttptq.DEFAULT_RECIPE, + help="Recipe path (relative to modelopt_recipes/ or an absolute YAML). " + "Defaults to the ViT FP8 recipe; pass huggingface/vit/ptq/nvfp4 for NVFP4.", ) parser.add_argument( "--calib_samples", type=int, - default=512, + default=1024, help="Number of tiny-imagenet samples to use for calibration.", ) parser.add_argument( "--batch_size", type=int, - default=1, + default=128, help="Calibration / compile / eval batch size. The Torch-TRT engine is " - "compiled for this single static batch shape and the onnx_ptq evaluate() " - "dataloader keeps the trailing partial batch, so the Torch-TRT path " - "requires --batch_size 1; larger batches are only allowed with --skip_trt.", + "dynamic (min=1, opt=--batch_size, max=1024) and handles any batch incl. " + "the trailing partial batch, so any --batch_size (e.g. 128) works.", ) parser.add_argument( "--eval_data_size", @@ -154,16 +146,6 @@ def main(): help="Score the fake-quant (modelopt) model; skip torch_tensorrt.compile. " "Useful for environments without torch_tensorrt installed.", ) - parser.add_argument( - "--no_pretrained", - action="store_true", - help="Build the model from config with random weights (fast smoke test).", - ) - parser.add_argument( - "--model_kwargs", - default=None, - help="JSON string of ViTConfig overrides applied when --no_pretrained is set.", - ) parser.add_argument( "--results_path", default=None, @@ -171,31 +153,14 @@ def main(): ) args = parser.parse_args() - # The Torch-TRT engine is compiled for one static batch shape, and the reused - # onnx_ptq evaluate() dataloader does not drop the trailing partial batch, so a - # batch size that doesn't divide the validation set would crash the static - # engine mid-run. Fail fast. The fake-quant (--skip_trt) path is a plain eager - # module and tolerates any batch size. - if not args.skip_trt and args.batch_size != 1: - raise SystemExit( - "The Torch-TensorRT path requires --batch_size 1 (the engine is compiled " - "for a single static batch shape). Use --skip_trt to score the fake-quant " - "model at a larger batch size." - ) + dynamic = True if not torch.cuda.is_available(): raise SystemExit("This example requires a CUDA-capable GPU.") device = torch.device("cuda") dtype = torch.float16 - config_overrides = json.loads(args.model_kwargs) if args.model_kwargs else None - model, processor = ttptq.load_model_and_processor( - args.model_id, - device, - dtype, - pretrained=not args.no_pretrained, - config_overrides=config_overrides, - ) + model, processor = ttptq.load_model_and_processor(args.model_id, device, dtype) transform = build_processor_transform(processor) def run_eval(m: torch.nn.Module) -> tuple[float, float]: @@ -225,8 +190,8 @@ def to_eval_model(m: torch.nn.Module, what: str) -> torch.nn.Module: wrapped = ttptq.ViTLogitsWrapper(m).to(device).eval() if args.skip_trt: return wrapped - print(f"\nCompiling {what} with Torch-TensorRT ...") - return ttptq.compile_with_torch_tensorrt(wrapped, example_input) + print(f"\nCompiling {what} with Torch-TensorRT (dynamic={dynamic}) ...") + return ttptq.compile_with_torch_tensorrt(wrapped, example_input, dynamic=dynamic) results: list[list[str | float]] = [["Metric", "Top1 (%)", "Top5 (%)"]] @@ -245,11 +210,11 @@ def to_eval_model(m: torch.nn.Module, what: str) -> torch.nn.Module: calib_batches = ttptq.build_calibration_loader( processor, args.calib_samples, args.batch_size, device, dtype ) - recipe_path = args.recipe or ttptq.PRECISION_TO_RECIPE[args.precision] - ttptq.quantize_with_recipe(model, recipe_path, calib_batches) + ttptq.quantize_with_recipe(model, args.recipe, calib_batches) - tag = f"{args.precision} ({runtime})" - eval_model = to_eval_model(model, f"{args.precision} model") + label = Path(args.recipe).stem # e.g. "fp8" / "nvfp4" + tag = f"{label} ({runtime})" + eval_model = to_eval_model(model, f"{label} model") print(f"\n=== {tag} ===") top1, top5 = run_eval(eval_model) print(f"{tag} top1={top1:.2f}% top5={top5:.2f}%") diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index ad6c7a926d0..9d866bcc238 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -33,7 +33,6 @@ from __future__ import annotations import argparse -import json from pathlib import Path import torch @@ -45,13 +44,10 @@ from modelopt.recipe import ModelOptPTQRecipe, load_recipe from modelopt.torch.quantization.utils import export_torch_mode -# Maps the user-facing precision flag to the ViT-specific recipe under -# `modelopt_recipes/huggingface/vit/ptq/`. The recipe loader resolves this -# relative path against the built-in recipe library. -PRECISION_TO_RECIPE: dict[str, str] = { - "fp8": "huggingface/vit/ptq/fp8", - "nvfp4": "huggingface/vit/ptq/nvfp4", -} +# Default ViT PTQ recipe under `modelopt_recipes/huggingface/vit/ptq/`. The +# recipe loader resolves this relative path against the built-in recipe library; +# pass `--recipe` for a different one (e.g. `huggingface/vit/ptq/nvfp4`). +DEFAULT_RECIPE = "huggingface/vit/ptq/fp8" def load_model_and_processor( @@ -71,10 +67,14 @@ def load_model_and_processor( """ print(f"Loading {model_id} (dtype={dtype}, pretrained={pretrained})...") processor = AutoImageProcessor.from_pretrained(model_id) + # tanh-approximation GELU rather than the erf-based default. if pretrained: - model = ViTForImageClassification.from_pretrained(model_id, torch_dtype=dtype) + model = ViTForImageClassification.from_pretrained( + model_id, torch_dtype=dtype, hidden_act="gelu_fast" + ) else: config = ViTConfig.from_pretrained(model_id) + config.hidden_act = "gelu_fast" for k, v in (config_overrides or {}).items(): setattr(config, k, v) model = ViTForImageClassification(config).to(dtype) @@ -141,37 +141,84 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: return self.vit(pixel_values=pixel_values).logits -def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Tensor): +def compile_with_torch_tensorrt( + model: torch.nn.Module, example_input: torch.Tensor, dynamic: bool = True +): """Compile the quantized model with Torch-TensorRT (Dynamo IR, strongly-typed). `min_block_size=1` follows the Torch-TRT quantization guide so single-node Q/DQ + matmul subgraphs become TRT precision layers. `export_torch_mode` makes modelopt emit Q/DQ in the TRT-friendly form during `torch.export`. + + ``dynamic`` builds a dynamic-batch engine (min=1, opt=``example_input``'s + batch, max=1024); otherwise the engine is specialized to ``example_input``. + Dynamic keeps the model in one engine and fuses far more FP8 (see the `cat` + comment below), so it is the default. """ # Imported here (not at module scope) so the quantize-only `--skip_trt` path # still runs on hosts without torch_tensorrt installed. import torch_tensorrt - print("Compiling with torch_tensorrt.compile (Dynamo IR)...") - # `aten.cat.default` is force-executed in PyTorch because torch_tensorrt - # 2.10's cat converter chokes on the cls-token + patch-embedding concat - # in HF ViT (BFloat16 path: `TypeError: Got unsupported ScalarType - # BFloat16`; FP16 path: rank-(-1) TRT tensor that trips the downstream - # `embeddings + position_embeddings` add). The cat is a tiny [1,1,H] - # + [1,N,H] concat that runs once per forward, so falling back to - # PyTorch costs essentially nothing. - with export_torch_mode(): + print(f"Compiling with torch_tensorrt.compile (Dynamo IR, dynamic={dynamic})...") + if dynamic: + n, c, h, w = example_input.shape + # torch.export specializes a size-1 dynamic dim to a constant, so trace at + # opt batch >= 2; min=1 still serves batch 1 at runtime. + opt_n = max(int(n), 2) + compile_inputs = { + "inputs": [ + torch_tensorrt.Input( + min_shape=(1, c, h, w), + opt_shape=(opt_n, c, h, w), + max_shape=(1024, c, h, w), + dtype=example_input.dtype, + ) + ] + } + else: + compile_inputs = {"arg_inputs": [example_input]} + + # The cls-token `cat` decides how much FP8 TRT can fuse. Static: the bf16 cat + # converter materializes a constant via numpy (no bf16 dtype), so we run + # `aten.cat` in PyTorch -> the graph splits into two engines and FP8 fusion is + # capped (~24 GEMMs). Dynamic: the cat operands are symbolic, so it stays one + # engine (~121 FP8 GEMMs); forcing it to PyTorch there also breaks the engine + # at runtime. The Debugger keeps per-layer info for dump_trt_layer_info(). + cat_fallback = {} if dynamic else {"torch_executed_ops": {torch.ops.aten.cat.default}} + with export_torch_mode(), torch_tensorrt.dynamo.Debugger(log_level="error"): trt_model = torch_tensorrt.compile( model, ir="dynamo", - arg_inputs=[example_input], min_block_size=1, truncate_double=True, - torch_executed_ops={torch.ops.aten.cat.default}, + **cat_fallback, + **compile_inputs, ) return trt_model +def dump_trt_layer_info(trt_model: torch.nn.Module, path: Path) -> None: + """Write the per-layer engine info of every TRT submodule to ``path``. + + A Dynamo-compiled module can hold several ``TorchTensorRTModule`` subgraphs + (the parts that fell back to PyTorch sit between them), so we concatenate the + ``get_layer_info()`` JSON of each. + """ + import torch_tensorrt + + infos = [ + mod.get_layer_info() + for _, mod in trt_model.named_modules() + if isinstance(mod, torch_tensorrt.dynamo.runtime.TorchTensorRTModule) + ] + if not infos: + print("No TorchTensorRTModule found; nothing to dump (whole graph fell back to PyTorch?).") + return + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(infos)) + print(f"Wrote TRT layer info ({len(infos)} engine(s)) to {path}") + + def _argmax_logits(out) -> torch.Tensor: """Handle either an HF `ImageClassifierOutput` or a raw tensor.""" logits = out.logits if hasattr(out, "logits") else out @@ -185,35 +232,29 @@ def main(): default="google/vit-large-patch16-224", help="HuggingFace model id of the ViT classifier to quantize.", ) - parser.add_argument( - "--precision", - choices=sorted(PRECISION_TO_RECIPE), - default="fp8", - help="Which ViT recipe variant to apply.", - ) parser.add_argument( "--recipe", - default=None, - help="Override the recipe path (relative to modelopt_recipes/ or absolute). " - "If unset, the recipe is picked by --precision.", + default=DEFAULT_RECIPE, + help="Recipe path (relative to modelopt_recipes/ or an absolute YAML). " + "Defaults to the ViT FP8 recipe; pass huggingface/vit/ptq/nvfp4 for NVFP4.", ) parser.add_argument( "--calib_samples", type=int, - default=128, + default=1024, help="Number of tiny-imagenet samples to use for calibration.", ) parser.add_argument( "--batch_size", type=int, - default=1, + default=128, help="Batch size for calibration / TRT compile.", ) parser.add_argument( "--save_dir", type=str, - default=None, - help="If set, save the quantized modelopt state-dict here (BF16 weights " + default="./modelopt_quantized", + help="Directory to save the quantized modelopt state-dict (BF16 weights " "+ Q/DQ metadata) — re-usable across runs without recalibration.", ) parser.add_argument( @@ -223,18 +264,10 @@ def main(): "Useful for environments without torch_tensorrt installed.", ) parser.add_argument( - "--no_pretrained", - action="store_true", - help="Build the model from config with random weights instead of " - "downloading pretrained weights. Useful for fast e2e tests.", - ) - parser.add_argument( - "--model_kwargs", - type=str, + "--layer_info_path", default=None, - help="JSON string of ViTConfig overrides applied when --no_pretrained " - 'is set (e.g. \'{"num_hidden_layers": 1, "hidden_size": 64, ' - '"intermediate_size": 128, "num_attention_heads": 2}\').', + help="If set, write the compiled TRT engine's per-layer info " + "(get_layer_info()) to this file.", ) args = parser.parse_args() @@ -243,14 +276,7 @@ def main(): device = torch.device("cuda") dtype = torch.bfloat16 - config_overrides = json.loads(args.model_kwargs) if args.model_kwargs else None - model, processor = load_model_and_processor( - args.model_id, - device, - dtype, - pretrained=not args.no_pretrained, - config_overrides=config_overrides, - ) + model, processor = load_model_and_processor(args.model_id, device, dtype) image_size = model.config.image_size num_channels = model.config.num_channels example_input = torch.randn( @@ -266,15 +292,13 @@ def main(): processor, args.calib_samples, args.batch_size, device, dtype ) - recipe_path = args.recipe or PRECISION_TO_RECIPE[args.precision] - quantize_with_recipe(model, recipe_path, calib_batches) + quantize_with_recipe(model, args.recipe, calib_batches) - if args.save_dir: - save_path = Path(args.save_dir) - save_path.mkdir(parents=True, exist_ok=True) - ckpt = save_path / "vit_modelopt_state.pt" - mto.save(model, ckpt) - print(f"Saved quantized modelopt state to {ckpt}") + save_path = Path(args.save_dir) + save_path.mkdir(parents=True, exist_ok=True) + ckpt = save_path / "vit_modelopt_state.pt" + mto.save(model, ckpt) + print(f"Saved quantized modelopt state to {ckpt}") print("\n=== Fake-quant (modelopt) ===") with torch.no_grad(): @@ -287,7 +311,10 @@ def main(): return wrapped = ViTLogitsWrapper(model).to(device).eval() - trt_model = compile_with_torch_tensorrt(wrapped, example_input) + trt_model = compile_with_torch_tensorrt(wrapped, example_input, dynamic=True) + + if args.layer_info_path: + dump_trt_layer_info(trt_model, Path(args.layer_info_path)) print("\n=== Torch-TensorRT compiled ===") with torch.no_grad(): diff --git a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py index f0350d525ad..bc18a6f5e0c 100644 --- a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py +++ b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py @@ -13,47 +13,31 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json - import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -# Recipe variants the example ships. Mirrors the parametrization style of -# ``tests/examples/torch_onnx/test_torch_quant_to_onnx.py``. -_PRECISIONS = ["fp8", "nvfp4"] - -# Tiny ViT config (~1 encoder block) so the test stays under a few seconds -# of GPU time while exercising every code path the recipe touches: encoder -# Linear weight/input quantizers, attention BMM + softmax quantizers, -# per-block LayerNorm output quantizer, and the patch-embed Conv / final -# vit.layernorm / classifier skip rules. -_TINY_VIT_KWARGS = { - "num_hidden_layers": 1, - "hidden_size": 64, - "intermediate_size": 128, - "num_attention_heads": 2, -} +# Recipe variants the example ships. +_RECIPES = ["huggingface/vit/ptq/fp8", "huggingface/vit/ptq/nvfp4"] -@pytest.mark.parametrize("precision", _PRECISIONS) -def test_torch_tensorrt_ptq(precision): - """End-to-end: load tiny ViT -> mtq.quantize via recipe -> torch_tensorrt.compile. +@pytest.mark.parametrize("recipe", _RECIPES) +def test_torch_tensorrt_ptq(recipe, tmp_path): + """End-to-end: load ViT -> mtq.quantize via recipe -> torch_tensorrt.compile. - Runs against the smallest viable ``ViTForImageClassification`` config so - the test stays fast; ``--no_pretrained`` skips the multi-GB pretrained - download. The example's CLI exits non-zero if any step (calibration, - quantization, TRT compile) fails or if the compiled-model argmax doesn't - match the fake-quant argmax on the sample input. + Uses the pretrained ``google/vit-base-patch16-224`` (the example no longer + exposes a random-weight path). The CLI exits non-zero if any step + (calibration, quantization, TRT compile) fails; the printed argmax + comparison is informational only. NVFP4's low-precision kernels require a + Blackwell GPU, so the nvfp4 case only builds there. """ pytest.importorskip("torch_tensorrt") cmd_parts = extend_cmd_parts( ["python", "torch_tensorrt_ptq.py"], model_id="google/vit-base-patch16-224", - precision=precision, + recipe=recipe, calib_samples="4", batch_size="1", - model_kwargs=json.dumps(_TINY_VIT_KWARGS), + save_dir=str(tmp_path / "ckpt"), ) - cmd_parts.append("--no_pretrained") run_example_command(cmd_parts, "torch_trt") From 93bf5a4598a270d938781fc49b57efdf9071b256 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 11 Jun 2026 04:41:48 +0000 Subject: [PATCH 12/17] [6078291] torch_trt example: FP16, dynamic-only engine, simpler loader Run the ViT in float16 in both torch_tensorrt_ptq.py and torch_tensorrt_accuracy.py (the PTQ script previously used bfloat16), so both paths share one compute dtype. compile_with_torch_tensorrt now always builds a single dynamic-batch engine: remove the static-compile branch, its forced-cat fallback, and the dynamic= plumbing from both entry points. load_model_and_processor always loads pretrained weights; drop the pretrained/config_overrides args and the unused ViTConfig import (matches the test's CLI-only contract). README: restructure as a focused single-example guide; list the FP8/NVFP4 minimum compute capability (8.9+ / 10.0+) and architectures in Hardware Requirements; note both scripts run in float16; remove stale Dynamic Engine links. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/README.md | 332 +++++++++++------- examples/torch_trt/torch_tensorrt_accuracy.py | 6 +- examples/torch_trt/torch_tensorrt_ptq.py | 99 ++---- 3 files changed, 231 insertions(+), 206 deletions(-) diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index c64fa068615..0168ac9c344 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -1,171 +1,245 @@ -# ModelOpt + Torch-TensorRT Deployment +# Torch-TensorRT Quantization -End-to-end examples that quantize a PyTorch model with NVIDIA ModelOpt and -then compile the quantized graph with -[Torch-TensorRT](https://docs.pytorch.org/TensorRT/) for deployment. +[Torch-TensorRT](https://docs.pytorch.org/TensorRT/) compiles a PyTorch model into an optimized TensorRT engine with no separate export or runtime. This example quantizes a PyTorch / HuggingFace model with NVIDIA Model Optimizer and then compiles the quantized graph in-framework with Torch-TensorRT for deployment. -The flow follows the -[Torch-TensorRT quantization guide](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html): -ModelOpt inserts Q/DQ nodes into the eager PyTorch graph, then -`torch_tensorrt.compile(ir="dynamo")` converts those Q/DQ nodes into native -TensorRT precision layers. +Quantization is an effective model optimization technique that compresses your models. Model Optimizer inserts Q/DQ nodes into the eager PyTorch graph; `torch_tensorrt.compile(ir="dynamo")` then converts those Q/DQ nodes into native TensorRT precision layers — including FP8 and NVFP4 — following the [Torch-TensorRT quantization guide](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html). -## How this differs from the ONNX examples +This section focuses on the in-framework Torch-TensorRT path: a PyTorch front end (`mtq.quantize`) feeding a Dynamo-compiled TensorRT engine, demonstrated end-to-end on a HuggingFace ViT image classifier. If you instead want a portable ONNX → TensorRT artifact, or you start from an ONNX model, see the sibling [`torch_onnx`](../torch_onnx/) and [`onnx_ptq`](../onnx_ptq/) examples (compared in the [Support Matrix](#support-matrix)). -All three of these examples reach the same destination — a low-precision -TensorRT engine — but quantize at a different point in the pipeline and emit a -different artifact, so they suit different deployment stacks: +
+ +| **Section** | **Description** | **Link** | **Docs** | +| :------------: | :------------: | :------------: | :------------: | +| Pre-Requisites | Required packages and installation | \[[Link](#pre-requisites)\] | | +| Getting Started | Quantize and compile a ViT in a few lines | \[[Link](#getting-started)\] | \[[docs](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html)\] | +| Support Matrix | How this path compares to the ONNX examples | \[[Link](#support-matrix)\] | | +| ViT Recipes | The FP8 / NVFP4 recipes shipped with the example | \[[Link](#vit-recipes)\] | | +| Usage | CLI flags for the quantize and accuracy scripts | \[[Link](#usage)\] | | +| Evaluate Accuracy | Measure ImageNet top-1 / top-5 accuracy | \[[Link](#evaluate-accuracy)\] | | +| Custom Recipes | Plug in your own recipe / model | \[[Link](#custom-recipes)\] | | +| Resources | Roadmap, docs, benchmarks, and support | \[[Link](#resources)\] | | + +
+ +## Pre-Requisites + +### Docker + +Please use the TensorRT docker image (e.g., `nvcr.io/nvidia/tensorrt:26.02-py3`) or visit our [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more information. + +```bash +docker run --gpus all -it --rm -v $(pwd):/workspace -w /workspace nvcr.io/nvidia/tensorrt:26.02-py3 bash +``` + +Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies. + +### Local Installation + +```bash +pip install -U "nvidia-modelopt[hf]" +pip install -r requirements.txt +``` + +### Hardware Requirements + +The low-precision kernels Torch-TensorRT emits need a GPU that supports the target format: + +
+ +| Recipe | Minimum GPU | +| :---: | :---: | +| `fp8` | Ada / Hopper — compute capability 8.9+ | +| `nvfp4` | Blackwell — compute capability 10.0+ | + +
+ +> [!NOTE] +> Older GPUs still let `mtq.quantize` succeed — it emits fake-quant nodes in PyTorch — but `torch_tensorrt.compile` will not find a real low-precision kernel for an unsupported format. + +## Getting Started + +Quantize a HuggingFace ViT, then compile the Q/DQ graph with Torch-TensorRT into a single `torch.nn.Module` you call from PyTorch: + +```python +import torch +import torch_tensorrt + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.utils import export_torch_mode + +# 1. Quantize the eager PyTorch model with a Model Optimizer PTQ recipe. +recipe = load_recipe("huggingface/vit/ptq/fp8") # or huggingface/vit/ptq/nvfp4 +mtq.quantize(model, recipe.quantize.model_dump(), forward_loop=calibrate) + +# 2. Compile the quantized (Q/DQ) graph with Torch-TensorRT. +# export_torch_mode() makes Model Optimizer emit Q/DQ in the TRT-friendly form, +# and min_block_size=1 lets single-node Q/DQ + matmul subgraphs become TRT +# precision layers (per the Torch-TensorRT quantization guide). +with export_torch_mode(): + trt_model = torch_tensorrt.compile( + model, + ir="dynamo", + min_block_size=1, + truncate_double=True, + inputs=[torch_tensorrt.Input( + min_shape=(1, 3, 224, 224), + opt_shape=(128, 3, 224, 224), + max_shape=(1024, 3, 224, 224), + dtype=torch.float16, + )], + ) + +logits = trt_model(pixel_values) # call it like any nn.Module +``` + +The runnable script [`torch_tensorrt_ptq.py`](./torch_tensorrt_ptq.py) wraps this flow end-to-end. It: + +1. Loads a HuggingFace ViT classifier (default `google/vit-large-patch16-224`). +1. Builds a tiny calibration loader from `zh-plus/tiny-imagenet` (avoids the gated `ILSVRC/imagenet-1k` repo, so the example runs unauthenticated). +1. Runs `mtq.quantize` with one of the recipes under [`modelopt_recipes/`](../../modelopt_recipes/) (see [ViT Recipes](#vit-recipes)). +1. Saves the quantized Model Optimizer state (FP16 weights + Q/DQ metadata) to `/vit_modelopt_state.pt` for reuse without recalibration (see [Custom Recipes](#custom-recipes)). +1. Compiles the quantized model with `torch_tensorrt.compile` and verifies that the compiled-model argmax matches the fake-quant argmax on a sample input. + +```bash +# Default model is google/vit-large-patch16-224, default recipe is the ViT FP8 recipe. +python torch_tensorrt_ptq.py --calib_samples 1024 --batch_size 128 + +# NVFP4 instead of FP8. +python torch_tensorrt_ptq.py --recipe huggingface/vit/ptq/nvfp4 + +# Quantize but don't TRT-compile (handy on a non-TRT host). +python torch_tensorrt_ptq.py --skip_trt +``` + +> [!NOTE] +> Both `torch_tensorrt_ptq.py` and the accuracy script ([`torch_tensorrt_accuracy.py`](./torch_tensorrt_accuracy.py)) run the model in `float16`. + +## Support Matrix + +All three of these examples reach the same destination — a low-precision TensorRT engine — but quantize at a different point in the pipeline and emit a different artifact, so they suit different deployment stacks: + +
| | Torch-TensorRT (this example) | [`torch_onnx`](../torch_onnx/) | [`onnx_ptq`](../onnx_ptq/) | -|---|---|---|---| +| :---: | :---: | :---: | :---: | | Starting point | a PyTorch / HF model | a PyTorch / timm model | an already-exported ONNX model | -| Quantize on | the eager PyTorch graph (`mtq.quantize`) | the eager PyTorch graph (`mtq.quantize`) | the ONNX graph directly (ModelOpt ONNX PTQ) | +| Quantize on | the eager PyTorch graph (`mtq.quantize`) | the eager PyTorch graph (`mtq.quantize`) | the ONNX graph directly (ONNX PTQ) | | Export step | none — the FX/Dynamo graph stays in-process | `torch.onnx.export` of the Q/DQ graph, postprocessed for TRT | none — Q/DQ inserted straight into the ONNX graph | | Intermediate artifact | none | a Q/DQ ONNX file | a Q/DQ ONNX file | | Compiler + runtime | `torch_tensorrt.compile(ir="dynamo")` → a `torch.nn.Module` you call from PyTorch | TensorRT builds a standalone engine from the ONNX | TensorRT builds a standalone engine from the ONNX | | Best when | PyTorch-native serving; you want a drop-in compiled module | you quantize in PyTorch but deploy via a portable ONNX → TRT engine | you only have an ONNX model and never touch PyTorch | -This example and [`torch_onnx`](../torch_onnx/) share the same PyTorch front end -(`mtq.quantize`), so the numerics are identical — they differ only in the back -end: this one keeps the graph in-process and hands it to Torch-TensorRT, while -`torch_onnx` exports a portable ONNX artifact for the standalone TensorRT -runtime. [`onnx_ptq`](../onnx_ptq/) instead quantizes the ONNX graph directly, -for when you start from an ONNX model rather than PyTorch. Pick this example -when your serving stack is PyTorch-native and you'd rather avoid an ONNX export -step. +
-## Setup +This example and [`torch_onnx`](../torch_onnx/) share the same PyTorch front end (`mtq.quantize`), so the numerics are identical — they differ only in the back end: this one keeps the graph in-process and hands it to Torch-TensorRT, while `torch_onnx` exports a portable ONNX artifact for the standalone TensorRT runtime. [`onnx_ptq`](../onnx_ptq/) instead quantizes the ONNX graph directly, for when you start from an ONNX model rather than PyTorch. Pick this example when your serving stack is PyTorch-native and you'd rather avoid an ONNX export step. -```bash -# From the NVIDIA TensorRT docker image (recommended): -docker run --gpus all -it --rm -v $(pwd):/workspace -w /workspace nvcr.io/nvidia/tensorrt:26.02-py3 bash +## ViT Recipes -pip install -U "nvidia-modelopt" -pip install -r examples/torch_trt/requirements.txt -``` +These are the recipes the CLI selects by default when `--model_id` points at a HF ViT classifier. They are tuned for the HF ViT module layout and are composed from the shared `$import` building blocks under [`modelopt_recipes/configs/`](../../modelopt_recipes/configs/) (`numerics/{fp8,nvfp4}`, `ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8,w4a4_nvfp4_nvfp4}`) rather than spelling out each `quant_cfg` entry. + +
-Torch-TensorRT itself follows the -[official install instructions](https://docs.pytorch.org/TensorRT/getting_started/installation.html) — -the version pulled by `pip` must match your installed PyTorch. +| `--recipe` value | Calibration | What it quantizes | +| :---: | :---: | :--- | +| `huggingface/vit/ptq/fp8` (default) | `max` | Per-tensor FP8 (E4M3) on every weight + input quantizer matched by the `*weight_quantizer` / `*input_quantizer` globs — encoder Linears, the patch-embed `nn.Conv2d` projection, and the `classifier` head — plus FP8 on the attention Q/K/V BMMs and softmax. All output quantizers disabled. | +| `huggingface/vit/ptq/nvfp4` | `awq_lite` | Dynamic NVFP4 W4A4 (E2M1, block size 16, FP8 E4M3 scales) on all weight/input quantizers, with the patch-embed `nn.Conv2d` projection, the `classifier` head, and the attention Q/K/V BMMs + softmax held at per-tensor FP8. All output quantizers disabled. | + +
## Usage -```bash -# Default model is google/vit-large-patch16-224, default recipe is the ViT FP8 recipe -python examples/torch_trt/torch_tensorrt_ptq.py \ - --recipe huggingface/vit/ptq/fp8 \ - --calib_samples 1024 \ - --batch_size 128 +### `torch_tensorrt_ptq.py` -# NVFP4 instead of FP8 -python examples/torch_trt/torch_tensorrt_ptq.py \ - --recipe huggingface/vit/ptq/nvfp4 +[Script](./torch_tensorrt_ptq.py) — quantize and (optionally) Torch-TensorRT-compile a ViT. -# Quantize but don't TRT-compile (handy on a non-TRT host) -python examples/torch_trt/torch_tensorrt_ptq.py --skip_trt +
-# Custom model + custom recipe -python examples/torch_trt/torch_tensorrt_ptq.py \ +| Flag | Default | Description | +| :---: | :---: | :--- | +| `--model_id` | `google/vit-large-patch16-224` | HuggingFace model id of the ViT classifier to quantize. | +| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). Pass `huggingface/vit/ptq/nvfp4` for NVFP4. | +| `--calib_samples` | `1024` | Number of tiny-imagenet samples to use for calibration. | +| `--batch_size` | `128` | Batch size for calibration / TRT compile. | +| `--save_dir` | `./modelopt_quantized` | Directory the quantized Model Optimizer state-dict (FP16 weights + Q/DQ metadata) is always saved to, as `vit_modelopt_state.pt` — re-usable across runs without recalibration. | +| `--skip_trt` | off | Quantize + run the fake-quant model only; skip `torch_tensorrt.compile`. Useful for environments without Torch-TensorRT installed. | +| `--layer_info_path` | unset | If set, write the compiled TRT engine's per-layer info (`get_layer_info()`) to this file. | + +
+ +```bash +# Custom model + custom recipe, saving the quantized state elsewhere. +python torch_tensorrt_ptq.py \ --model_id \ - --recipe + --recipe \ + --save_dir ./my_quantized + +# Dump the compiled engine's per-layer info to inspect FP8/NVFP4 fusion. +python torch_tensorrt_ptq.py --layer_info_path ./vit_fp8_layers.txt ``` -## What the example does - -1. Loads a HuggingFace model (default: `google/vit-large-patch16-224`). -2. Builds a tiny calibration loader from `zh-plus/tiny-imagenet` (avoids the - gated `ILSVRC/imagenet-1k` repo so the example runs unauthenticated). -3. Runs `mtq.quantize` with one of the recipes shipped under - [`modelopt_recipes/`](../../modelopt_recipes/). The default recipes - target ViT; pass `--recipe ` to use a different one for a - different model. -4. Compiles the quantized model with `torch_tensorrt.compile` and verifies - that the compiled-model argmax matches the fake-quant argmax on a sample - input. - -## Dynamic vs static engine (and why FP8 cares) - -HF ViT concatenates a cls token onto the patch embeddings. In a **static** -compile, Torch-TensorRT's `cat` converter constant-folds that token through -numpy — which has no bfloat16 dtype — so the bf16 graph must run `aten.cat` in -PyTorch. That graph break splits the model into **two** TRT engines and stops -TRT from fusing the Q/DQ nodes into FP8 across the boundary, so only ~24 FP8 -GEMMs survive for ViT-large (one per layer). - -In a **dynamic** compile the `cat` operands are symbolic tensors (no numpy -materialization), so the model stays a **single** engine and TRT fuses Q/DQ into -~121 FP8 GEMMs — including bias+GELU epilogue-fused `e4m3` kernels — a ~5× jump -in FP8 coverage. The example therefore uses a **dynamic engine** for both fp8 -and nvfp4, built for `min=1, opt=--batch_size, max=1024` and serving any batch in -that range. (nvfp4's low-precision kernels require Blackwell — see *Hardware -requirements*.) - -## Measuring ImageNet accuracy - -`torch_tensorrt_accuracy.py` reuses the pipeline above and reports ImageNet-1k -top-1 / top-5 accuracy via the `onnx_ptq` example's `evaluate()` harness -([`examples/onnx_ptq/evaluation.py`](../onnx_ptq/evaluation.py)): +### `torch_tensorrt_accuracy.py` + +[Script](./torch_tensorrt_accuracy.py) — quantize, compile, and score on ImageNet (see [Evaluate Accuracy](#evaluate-accuracy)). + +
+ +| Flag | Default | Description | +| :---: | :---: | :--- | +| `--model_id` | `google/vit-large-patch16-224` | HuggingFace model id of the ViT classifier to quantize and score. | +| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). Pass `huggingface/vit/ptq/nvfp4` for NVFP4. | +| `--calib_samples` | `1024` | Number of tiny-imagenet samples to use for calibration. | +| `--batch_size` | `128` | Calibration / compile / eval batch size. The Torch-TRT engine is dynamic (`min=1`, `opt=max(--batch_size, 2)`, `max=1024`) and handles any batch including the trailing partial batch. | +| `--eval_data_size` | full 50k | Number of ImageNet validation images to score. | +| `--imagenet_path` | `ILSVRC/imagenet-1k` | HF dataset card or local path to the ImageNet validation set (gated). | +| `--baseline` | off | Also score the unquantized model as a reference. It is Torch-TensorRT-compiled like the quantized model (or run eager under `--skip_trt`) so the comparison is apples-to-apples. | +| `--skip_trt` | off | Score the fake-quant (Model Optimizer) model; skip `torch_tensorrt.compile`. Useful for environments without Torch-TensorRT installed. | +| `--results_path` | unset | If set, write the accuracy results to this CSV path. | + +
+ +## Evaluate Accuracy + +[`torch_tensorrt_accuracy.py`](./torch_tensorrt_accuracy.py) reuses the quantize → compile pipeline above and reports ImageNet-1k top-1 / top-5 accuracy via the `onnx_ptq` example's `evaluate()` harness ([`examples/onnx_ptq/evaluation.py`](../onnx_ptq/evaluation.py)): ```bash -python examples/torch_trt/torch_tensorrt_accuracy.py \ +python torch_tensorrt_accuracy.py \ --recipe huggingface/vit/ptq/fp8 \ --batch_size 128 \ --baseline \ - --eval_data_size 5000 + --eval_data_size 5000 \ + --results_path results.csv ``` -- `--baseline` also scores the unquantized model. It is Torch-TensorRT-compiled - the same way as the quantized model, so every reported number comes from the - same TRT runtime (pass `--skip_trt` to score the eager / fake-quant models). -- The eval uses a **dynamic** engine (default `--batch_size 128`) for both - precisions, so it serves the trailing partial batch at any batch size — see - *Dynamic vs static engine*. -- Validation uses the gated `ILSVRC/imagenet-1k` split (accept its license / set - `HF_TOKEN`), or point `--imagenet_path` at a local copy. `evaluate()` shuffles - the split, so a partial `--eval_data_size` draws a different random subset each - run — omit it (full set) for a stable, comparable score. -- `--results_path results.csv` writes the metrics table to CSV. - -## ViT-specific recipes shipped with the example - -These are the recipes the CLI selects by default when `--model_id` points at a -HF ViT classifier. They are tuned for the HF ViT module layout and are composed -from the shared `$import` building blocks under -[`modelopt_recipes/configs/`](../../modelopt_recipes/configs/) -(`numerics/{fp8,nvfp4}`, `ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8}`) rather -than spelling out each `quant_cfg` entry. - -| `--recipe` value | What it quantizes | -|------------------|-------------------| -| `huggingface/vit/ptq/fp8` (default) | W8A8 FP8 (E4M3) on every weight + input quantizer — encoder Linears, the patch-embed `nn.Conv2d`, the `classifier` head, and per-block `nn.LayerNorm` inputs — plus FP8 on the attention Q/K/V BMMs and softmax. Output quantizers disabled. | -| `huggingface/vit/ptq/nvfp4` | NVFP4 W4A4 (E2M1, block 16, FP8 scales) on the encoder `nn.Linear` weights/inputs, with the patch-embed `nn.Conv2d`, the `classifier` head, and the attention Q/K/V BMMs + softmax held at FP8. Uses `awq_lite` calibration. | - -## Hardware requirements +- `--baseline` also scores the unquantized model. It is Torch-TensorRT-compiled the same way as the quantized model, so every reported number comes from the same TRT runtime (pass `--skip_trt` to score the eager / fake-quant models instead). +- The eval uses a **dynamic** engine (default `--batch_size 128`) for both precisions, so it serves the trailing partial batch at any batch size. +- `--results_path results.csv` writes the metrics table (`Metric`, `Top1 (%)`, `Top5 (%)`) to CSV. -| Recipe | Minimum GPU | -|--------|-------------| -| `fp8` | Hopper (H100) / Ada (RTX 4090 / 6000 Ada) — compute capability 8.9+ | -| `nvfp4` | Blackwell (B100/B200) — TRT ≥ 10.8 | +> [!NOTE] +> Validation uses the gated `ILSVRC/imagenet-1k` split: accept its license / set `HF_TOKEN`, or point `--imagenet_path` at a local copy. `evaluate()` shuffles the split, so a partial `--eval_data_size` draws a different random subset each run — omit it (full 50k set) for a stable, comparable score. -Older GPUs will still let `mtq.quantize` succeed (it emits fake-quant -nodes in PyTorch), but `torch_tensorrt.compile` will not find a real -low-precision kernel. +## Custom Recipes -### Resuming from a saved checkpoint +Use `--recipe ` to plug in a different recipe — either a path relative to `modelopt_recipes/` (resolved against the built-in recipe library) or an absolute filesystem path to a YAML file. The recipe is loaded via `modelopt.recipe.load_recipe`, must declare `metadata.recipe_type: ptq` and a `quantize:` section, and its `quantize` config is passed straight to `mtq.quantize`. See the existing [`modelopt_recipes/huggingface/vit/ptq/*.yaml`](../../modelopt_recipes/huggingface/vit/ptq/) for the patterns used here. -The quantized modelopt model is saved to `--save_dir` (default -`./modelopt_quantized`) as `vit_modelopt_state.pt`. To reload without -recalibrating, restore it before the TRT compile step with: +### Resuming From a Saved Checkpoint + +`torch_tensorrt_ptq.py` always saves the quantized Model Optimizer state to `/vit_modelopt_state.pt` (default `--save_dir ./modelopt_quantized`) via `mto.save`. To reload it without recalibrating, restore it onto a freshly-loaded model before the TRT compile step: ```python import modelopt.torch.opt as mto -mto.restore(model, "vit_modelopt_state.pt") + +mto.restore(model, "./modelopt_quantized/vit_modelopt_state.pt") ``` -## Custom recipes +> [!NOTE] +> See the [save / restore guide](https://nvidia.github.io/Model-Optimizer/guides/2_save_load.html) for the full `mto.save` / `mto.restore` workflow. + +## Resources -Use `--recipe ` to plug in a different recipe — either a path -relative to `modelopt_recipes/` (resolved against the built-in library) or -an absolute filesystem path to a YAML file. The recipe must declare -`metadata.recipe_type: ptq` and a `quantize:` section; see existing -`modelopt_recipes/huggingface/vit/ptq/*.yaml` for the patterns used here. +- 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/146) +- 📖 [Documentation](https://nvidia.github.io/Model-Optimizer) +- 🎯 [Benchmarks](../benchmark.md) +- 💡 [Release Notes](https://nvidia.github.io/Model-Optimizer/reference/0_changelog.html) +- 🐛 [File a bug](https://github.com/NVIDIA/Model-Optimizer/issues/new?template=1_bug_report.md) +- ✨ [File a Feature Request](https://github.com/NVIDIA/Model-Optimizer/issues/new?template=2_feature_request.md) diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py index 85c2f12f218..f973842f103 100644 --- a/examples/torch_trt/torch_tensorrt_accuracy.py +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -153,8 +153,6 @@ def main(): ) args = parser.parse_args() - dynamic = True - if not torch.cuda.is_available(): raise SystemExit("This example requires a CUDA-capable GPU.") device = torch.device("cuda") @@ -190,8 +188,8 @@ def to_eval_model(m: torch.nn.Module, what: str) -> torch.nn.Module: wrapped = ttptq.ViTLogitsWrapper(m).to(device).eval() if args.skip_trt: return wrapped - print(f"\nCompiling {what} with Torch-TensorRT (dynamic={dynamic}) ...") - return ttptq.compile_with_torch_tensorrt(wrapped, example_input, dynamic=dynamic) + print(f"\nCompiling {what} with Torch-TensorRT ...") + return ttptq.compile_with_torch_tensorrt(wrapped, example_input) results: list[list[str | float]] = [["Metric", "Top1 (%)", "Top5 (%)"]] diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index 9d866bcc238..d77aa6c1706 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -37,7 +37,7 @@ import torch from datasets import load_dataset -from transformers import AutoImageProcessor, ViTConfig, ViTForImageClassification +from transformers import AutoImageProcessor, ViTForImageClassification import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -50,34 +50,14 @@ DEFAULT_RECIPE = "huggingface/vit/ptq/fp8" -def load_model_and_processor( - model_id: str, - device: torch.device, - dtype: torch.dtype, - pretrained: bool = True, - config_overrides: dict | None = None, -): - """Pull the HF ViT classifier and its preprocessor. - - With ``pretrained=False`` the model is built from a config with random - weights (test path); ``config_overrides`` lets the caller shrink it - (e.g. ``{"num_hidden_layers": 1, "hidden_size": 64, ...}``). The - preprocessor is always loaded from ``model_id`` since it only carries - a small JSON config. - """ - print(f"Loading {model_id} (dtype={dtype}, pretrained={pretrained})...") +def load_model_and_processor(model_id: str, device: torch.device, dtype: torch.dtype): + """Pull the HF ViT classifier and its preprocessor.""" + print(f"Loading {model_id} (dtype={dtype})...") processor = AutoImageProcessor.from_pretrained(model_id) # tanh-approximation GELU rather than the erf-based default. - if pretrained: - model = ViTForImageClassification.from_pretrained( - model_id, torch_dtype=dtype, hidden_act="gelu_fast" - ) - else: - config = ViTConfig.from_pretrained(model_id) - config.hidden_act = "gelu_fast" - for k, v in (config_overrides or {}).items(): - setattr(config, k, v) - model = ViTForImageClassification(config).to(dtype) + model = ViTForImageClassification.from_pretrained( + model_id, torch_dtype=dtype, hidden_act="gelu_fast" + ) model.eval().to(device) return model, processor @@ -141,58 +121,31 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: return self.vit(pixel_values=pixel_values).logits -def compile_with_torch_tensorrt( - model: torch.nn.Module, example_input: torch.Tensor, dynamic: bool = True -): - """Compile the quantized model with Torch-TensorRT (Dynamo IR, strongly-typed). - - `min_block_size=1` follows the Torch-TRT quantization guide so single-node - Q/DQ + matmul subgraphs become TRT precision layers. `export_torch_mode` - makes modelopt emit Q/DQ in the TRT-friendly form during `torch.export`. - - ``dynamic`` builds a dynamic-batch engine (min=1, opt=``example_input``'s - batch, max=1024); otherwise the engine is specialized to ``example_input``. - Dynamic keeps the model in one engine and fuses far more FP8 (see the `cat` - comment below), so it is the default. - """ +def compile_with_torch_tensorrt(model: torch.nn.Module, example_input: torch.Tensor): + """Compile the quantized model with Torch-TensorRT (Dynamo IR, strongly-typed).""" # Imported here (not at module scope) so the quantize-only `--skip_trt` path # still runs on hosts without torch_tensorrt installed. import torch_tensorrt - print(f"Compiling with torch_tensorrt.compile (Dynamo IR, dynamic={dynamic})...") - if dynamic: - n, c, h, w = example_input.shape - # torch.export specializes a size-1 dynamic dim to a constant, so trace at - # opt batch >= 2; min=1 still serves batch 1 at runtime. - opt_n = max(int(n), 2) - compile_inputs = { - "inputs": [ - torch_tensorrt.Input( - min_shape=(1, c, h, w), - opt_shape=(opt_n, c, h, w), - max_shape=(1024, c, h, w), - dtype=example_input.dtype, - ) - ] - } - else: - compile_inputs = {"arg_inputs": [example_input]} - - # The cls-token `cat` decides how much FP8 TRT can fuse. Static: the bf16 cat - # converter materializes a constant via numpy (no bf16 dtype), so we run - # `aten.cat` in PyTorch -> the graph splits into two engines and FP8 fusion is - # capped (~24 GEMMs). Dynamic: the cat operands are symbolic, so it stays one - # engine (~121 FP8 GEMMs); forcing it to PyTorch there also breaks the engine - # at runtime. The Debugger keeps per-layer info for dump_trt_layer_info(). - cat_fallback = {} if dynamic else {"torch_executed_ops": {torch.ops.aten.cat.default}} + print("Compiling with torch_tensorrt.compile (Dynamo IR, dynamic batch)...") + n, c, h, w = example_input.shape + # torch.export specializes a size-1 dynamic dim to a constant, so trace at + # opt batch >= 2; min=1 still serves batch 1 at runtime. + opt_n = max(int(n), 2) with export_torch_mode(), torch_tensorrt.dynamo.Debugger(log_level="error"): trt_model = torch_tensorrt.compile( model, ir="dynamo", min_block_size=1, truncate_double=True, - **cat_fallback, - **compile_inputs, + inputs=[ + torch_tensorrt.Input( + min_shape=(1, c, h, w), + opt_shape=(opt_n, c, h, w), + max_shape=(1024, c, h, w), + dtype=example_input.dtype, + ) + ], ) return trt_model @@ -254,7 +207,7 @@ def main(): "--save_dir", type=str, default="./modelopt_quantized", - help="Directory to save the quantized modelopt state-dict (BF16 weights " + help="Directory to save the quantized modelopt state-dict (FP16 weights " "+ Q/DQ metadata) — re-usable across runs without recalibration.", ) parser.add_argument( @@ -274,7 +227,7 @@ def main(): if not torch.cuda.is_available(): raise SystemExit("This example requires a CUDA-capable GPU.") device = torch.device("cuda") - dtype = torch.bfloat16 + dtype = torch.float16 model, processor = load_model_and_processor(args.model_id, device, dtype) image_size = model.config.image_size @@ -283,7 +236,7 @@ def main(): args.batch_size, num_channels, image_size, image_size, device=device, dtype=dtype ) - print("\n=== Baseline (BF16) ===") + print("\n=== Baseline (FP16) ===") with torch.no_grad(): baseline_pred = _argmax_logits(model(example_input)) print(f"Baseline argmax class: {baseline_pred.tolist()}") @@ -311,7 +264,7 @@ def main(): return wrapped = ViTLogitsWrapper(model).to(device).eval() - trt_model = compile_with_torch_tensorrt(wrapped, example_input, dynamic=True) + trt_model = compile_with_torch_tensorrt(wrapped, example_input) if args.layer_info_path: dump_trt_layer_info(trt_model, Path(args.layer_info_path)) From 80134ebc4ebe64f4fad6964e3ca4e825200ac7b6 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:02:15 +0000 Subject: [PATCH 13/17] [6078291] torch_trt example: load ViT with eager attention for softmax PTQ Load the ViT classifier with attn_implementation="eager" so softmax runs through F.softmax instead of the fused SDPA kernel. The recipe's *softmax_quantizer is then exercised during calibration (otherwise it stays uncalibrated, amax=dynamic) and emits Q/DQ around the softmax output on export. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_trt/torch_tensorrt_ptq.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index d77aa6c1706..8edab142ec0 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -54,9 +54,15 @@ def load_model_and_processor(model_id: str, device: torch.device, dtype: torch.d """Pull the HF ViT classifier and its preprocessor.""" print(f"Loading {model_id} (dtype={dtype})...") processor = AutoImageProcessor.from_pretrained(model_id) - # tanh-approximation GELU rather than the erf-based default. + # `gelu_fast` selects the tanh-approximation GELU rather than the erf-based + # default. Eager attention runs softmax through `F.softmax` instead of the + # fused SDPA kernel, so the recipe's `*softmax_quantizer` is exercised during + # calibration and emits Q/DQ around the softmax output on export. model = ViTForImageClassification.from_pretrained( - model_id, torch_dtype=dtype, hidden_act="gelu_fast" + model_id, + torch_dtype=dtype, + hidden_act="gelu_fast", + attn_implementation="eager", ) model.eval().to(device) return model, processor From 406ca1a2b6ca93d2b416fa8fa1ecb3f642eb458c Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:18:19 +0000 Subject: [PATCH 14/17] [6078291] torch_trt test: skip nvfp4 case on non-Blackwell GPUs NVFP4's TRT kernels are only generated on Blackwell (SM >= 100); on older GPUs torch_tensorrt.compile fails at NVRTC/Myelin codegen. Guard the nvfp4 parametrization with skipif(not fp4_compatible()) so the suite skips rather than fails on non-Blackwell hardware. The fp8 case is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .../examples/torch_trt/test_torch_tensorrt_ptq.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py index bc18a6f5e0c..abd28415b3d 100644 --- a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py +++ b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py @@ -16,8 +16,19 @@ import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -# Recipe variants the example ships. -_RECIPES = ["huggingface/vit/ptq/fp8", "huggingface/vit/ptq/nvfp4"] +from modelopt.torch.quantization.backends.utils import fp4_compatible + +# Recipe variants the example ships. NVFP4's TRT kernels are only generated on +# Blackwell (compute capability >= 10.0), so that case is skipped elsewhere. +_RECIPES = [ + "huggingface/vit/ptq/fp8", + pytest.param( + "huggingface/vit/ptq/nvfp4", + marks=pytest.mark.skipif( + not fp4_compatible(), reason="NVFP4 requires a Blackwell GPU (SM >= 100)" + ), + ), +] @pytest.mark.parametrize("recipe", _RECIPES) From 59fa56858827e261e0a5cb7d752cff11b029c6d1 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:49:29 +0000 Subject: [PATCH 15/17] [6078291] torch_trt test: use tiny offline ViT, drop NVFP4 case Build the e2e test's ViT from a tiny randomly-initialized config saved with its image processor (new create_tiny_vit_dir helper in tests/_test_utils/torch/transformers_models.py) instead of downloading google/vit-base-patch16-224, so the test runs offline and fast while exercising the same module structure (3-channel patch conv, attention q/k/v, classifier). Drop the NVFP4/Blackwell-gated parametrization (and the now-unused fp4_compatible import) so the test covers the FP8 recipe only. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .../_test_utils/torch/transformers_models.py | 34 +++++++++++++++++++ .../torch_trt/test_torch_tensorrt_ptq.py | 27 ++++++--------- 2 files changed, 45 insertions(+), 16 deletions(-) diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index a946ebf2a74..90a5f266546 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -24,6 +24,7 @@ transformers = pytest.importorskip("transformers") from transformers import ( AutoModelForCausalLM, + AutoModelForImageClassification, AutoModelForImageTextToText, AutoModelForQuestionAnswering, AutoProcessor, @@ -39,6 +40,8 @@ Qwen3MoeConfig, T5Config, T5ForConditionalGeneration, + ViTConfig, + ViTImageProcessor, ) import modelopt.torch.opt as mto @@ -644,6 +647,37 @@ def create_tiny_bert_dir(tmp_path: Path | str, **config_kwargs) -> Path: return _create_tiny_llm_dir(Path(tmp_path) / "tiny_bert", get_tiny_bert, **config_kwargs) +##### ViT (vision) ##### +def get_tiny_vit(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + + # Keep num_channels=3 and a 16x16 patch so the patch-embedding stem conv matches a + # real ViT; image_size=32 gives a 2x2 patch grid, keeping the model tiny. + kwargs = { + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "image_size": 32, + "patch_size": 16, + "num_channels": 3, + "num_labels": 2, + } + kwargs.update(**config_kwargs) + return AutoModelForImageClassification.from_config(ViTConfig(**kwargs)) + + +def create_tiny_vit_dir(tmp_path: Path | str, **config_kwargs) -> Path: + vit_dir = Path(tmp_path) / "tiny_vit" + tiny_vit = get_tiny_vit(**config_kwargs) + tiny_vit.save_pretrained(vit_dir) + # A vision model also needs a saved image processor so the example's + # AutoImageProcessor.from_pretrained(dir) resolves; size it to the tiny image. + image_size = tiny_vit.config.image_size + ViTImageProcessor(size={"height": image_size, "width": image_size}).save_pretrained(vit_dir) + return vit_dir + + ##### TESTERS ##### def tf_output_tester(model_ref, model_test): inputs = model_ref.dummy_inputs diff --git a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py index abd28415b3d..5bc3962a911 100644 --- a/tests/examples/torch_trt/test_torch_tensorrt_ptq.py +++ b/tests/examples/torch_trt/test_torch_tensorrt_ptq.py @@ -15,19 +15,11 @@ import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command +from _test_utils.torch.transformers_models import create_tiny_vit_dir -from modelopt.torch.quantization.backends.utils import fp4_compatible - -# Recipe variants the example ships. NVFP4's TRT kernels are only generated on -# Blackwell (compute capability >= 10.0), so that case is skipped elsewhere. +# Recipe variants the example ships. _RECIPES = [ "huggingface/vit/ptq/fp8", - pytest.param( - "huggingface/vit/ptq/nvfp4", - marks=pytest.mark.skipif( - not fp4_compatible(), reason="NVFP4 requires a Blackwell GPU (SM >= 100)" - ), - ), ] @@ -35,17 +27,20 @@ def test_torch_tensorrt_ptq(recipe, tmp_path): """End-to-end: load ViT -> mtq.quantize via recipe -> torch_tensorrt.compile. - Uses the pretrained ``google/vit-base-patch16-224`` (the example no longer - exposes a random-weight path). The CLI exits non-zero if any step - (calibration, quantization, TRT compile) fails; the printed argmax - comparison is informational only. NVFP4's low-precision kernels require a - Blackwell GPU, so the nvfp4 case only builds there. + Uses a tiny randomly-initialized ViT (saved locally with its image processor) + rather than downloading a full pretrained checkpoint, so the test stays offline + and fast while exercising the same module structure (3-channel patch conv, + attention q/k/v, classifier). The CLI exits non-zero if any step (calibration, + quantization, TRT compile) fails; the printed argmax comparison is informational + only. """ pytest.importorskip("torch_tensorrt") + model_dir = create_tiny_vit_dir(tmp_path) + cmd_parts = extend_cmd_parts( ["python", "torch_tensorrt_ptq.py"], - model_id="google/vit-base-patch16-224", + model_id=str(model_dir), recipe=recipe, calib_samples="4", batch_size="1", From 3e4d4ec532738fce393f59cd66b1f73b74a2a879 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:49:43 +0000 Subject: [PATCH 16/17] [6078291] torch_trt: remove the ViT NVFP4 recipe and its references Delete modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml and scrub every remaining reference to it: the README (recipe/hardware table rows, usage examples, building-block list), the --recipe help text in torch_tensorrt_ptq.py / torch_tensorrt_accuracy.py, and the CHANGELOG entry. The ViT Torch-TRT example is now FP8-only. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/torch_trt/README.md | 19 +++----- examples/torch_trt/torch_tensorrt_accuracy.py | 4 +- examples/torch_trt/torch_tensorrt_ptq.py | 8 ++-- .../huggingface/vit/ptq/nvfp4.yaml | 45 ------------------- 5 files changed, 14 insertions(+), 64 deletions(-) delete mode 100644 modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 85401ff738b..7960dac388b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,7 +38,7 @@ Changelog - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. -- Add Torch-TensorRT FP8 / NVFP4 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships two ViT-tuned PTQ recipes under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``, ``nvfp4.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: FP8 quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax; NVFP4 runs W4A4 on the encoder ``nn.Linear`` weights/inputs while holding the patch-embed ``nn.Conv2d``, ``classifier``, and attention BMMs/softmax at FP8. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. +- Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. **Bug Fixes** diff --git a/examples/torch_trt/README.md b/examples/torch_trt/README.md index 0168ac9c344..fcd233e6713 100644 --- a/examples/torch_trt/README.md +++ b/examples/torch_trt/README.md @@ -2,7 +2,7 @@ [Torch-TensorRT](https://docs.pytorch.org/TensorRT/) compiles a PyTorch model into an optimized TensorRT engine with no separate export or runtime. This example quantizes a PyTorch / HuggingFace model with NVIDIA Model Optimizer and then compiles the quantized graph in-framework with Torch-TensorRT for deployment. -Quantization is an effective model optimization technique that compresses your models. Model Optimizer inserts Q/DQ nodes into the eager PyTorch graph; `torch_tensorrt.compile(ir="dynamo")` then converts those Q/DQ nodes into native TensorRT precision layers — including FP8 and NVFP4 — following the [Torch-TensorRT quantization guide](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html). +Quantization is an effective model optimization technique that compresses your models. Model Optimizer inserts Q/DQ nodes into the eager PyTorch graph; `torch_tensorrt.compile(ir="dynamo")` then converts those Q/DQ nodes into native TensorRT FP8 precision layers, following the [Torch-TensorRT quantization guide](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html). This section focuses on the in-framework Torch-TensorRT path: a PyTorch front end (`mtq.quantize`) feeding a Dynamo-compiled TensorRT engine, demonstrated end-to-end on a HuggingFace ViT image classifier. If you instead want a portable ONNX → TensorRT artifact, or you start from an ONNX model, see the sibling [`torch_onnx`](../torch_onnx/) and [`onnx_ptq`](../onnx_ptq/) examples (compared in the [Support Matrix](#support-matrix)). @@ -13,7 +13,7 @@ This section focuses on the in-framework Torch-TensorRT path: a PyTorch front en | Pre-Requisites | Required packages and installation | \[[Link](#pre-requisites)\] | | | Getting Started | Quantize and compile a ViT in a few lines | \[[Link](#getting-started)\] | \[[docs](https://docs.pytorch.org/TensorRT/user_guide/shapes_precision/quantization.html)\] | | Support Matrix | How this path compares to the ONNX examples | \[[Link](#support-matrix)\] | | -| ViT Recipes | The FP8 / NVFP4 recipes shipped with the example | \[[Link](#vit-recipes)\] | | +| ViT Recipes | The FP8 recipe shipped with the example | \[[Link](#vit-recipes)\] | | | Usage | CLI flags for the quantize and accuracy scripts | \[[Link](#usage)\] | | | Evaluate Accuracy | Measure ImageNet top-1 / top-5 accuracy | \[[Link](#evaluate-accuracy)\] | | | Custom Recipes | Plug in your own recipe / model | \[[Link](#custom-recipes)\] | | @@ -49,7 +49,6 @@ The low-precision kernels Torch-TensorRT emits need a GPU that supports the targ | Recipe | Minimum GPU | | :---: | :---: | | `fp8` | Ada / Hopper — compute capability 8.9+ | -| `nvfp4` | Blackwell — compute capability 10.0+ | @@ -69,7 +68,7 @@ from modelopt.recipe import load_recipe from modelopt.torch.quantization.utils import export_torch_mode # 1. Quantize the eager PyTorch model with a Model Optimizer PTQ recipe. -recipe = load_recipe("huggingface/vit/ptq/fp8") # or huggingface/vit/ptq/nvfp4 +recipe = load_recipe("huggingface/vit/ptq/fp8") mtq.quantize(model, recipe.quantize.model_dump(), forward_loop=calibrate) # 2. Compile the quantized (Q/DQ) graph with Torch-TensorRT. @@ -105,9 +104,6 @@ The runnable script [`torch_tensorrt_ptq.py`](./torch_tensorrt_ptq.py) wraps thi # Default model is google/vit-large-patch16-224, default recipe is the ViT FP8 recipe. python torch_tensorrt_ptq.py --calib_samples 1024 --batch_size 128 -# NVFP4 instead of FP8. -python torch_tensorrt_ptq.py --recipe huggingface/vit/ptq/nvfp4 - # Quantize but don't TRT-compile (handy on a non-TRT host). python torch_tensorrt_ptq.py --skip_trt ``` @@ -136,14 +132,13 @@ This example and [`torch_onnx`](../torch_onnx/) share the same PyTorch front end ## ViT Recipes -These are the recipes the CLI selects by default when `--model_id` points at a HF ViT classifier. They are tuned for the HF ViT module layout and are composed from the shared `$import` building blocks under [`modelopt_recipes/configs/`](../../modelopt_recipes/configs/) (`numerics/{fp8,nvfp4}`, `ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8,w4a4_nvfp4_nvfp4}`) rather than spelling out each `quant_cfg` entry. +This is the recipe the CLI selects by default when `--model_id` points at a HF ViT classifier. It is tuned for the HF ViT module layout and is composed from the shared `$import` building blocks under [`modelopt_recipes/configs/`](../../modelopt_recipes/configs/) (`ptq/units/{w8a8_fp8_fp8,attention_qkv_fp8}`) rather than spelling out each `quant_cfg` entry.
| `--recipe` value | Calibration | What it quantizes | | :---: | :---: | :--- | | `huggingface/vit/ptq/fp8` (default) | `max` | Per-tensor FP8 (E4M3) on every weight + input quantizer matched by the `*weight_quantizer` / `*input_quantizer` globs — encoder Linears, the patch-embed `nn.Conv2d` projection, and the `classifier` head — plus FP8 on the attention Q/K/V BMMs and softmax. All output quantizers disabled. | -| `huggingface/vit/ptq/nvfp4` | `awq_lite` | Dynamic NVFP4 W4A4 (E2M1, block size 16, FP8 E4M3 scales) on all weight/input quantizers, with the patch-embed `nn.Conv2d` projection, the `classifier` head, and the attention Q/K/V BMMs + softmax held at per-tensor FP8. All output quantizers disabled. |
@@ -158,7 +153,7 @@ These are the recipes the CLI selects by default when `--model_id` points at a H | Flag | Default | Description | | :---: | :---: | :--- | | `--model_id` | `google/vit-large-patch16-224` | HuggingFace model id of the ViT classifier to quantize. | -| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). Pass `huggingface/vit/ptq/nvfp4` for NVFP4. | +| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). | | `--calib_samples` | `1024` | Number of tiny-imagenet samples to use for calibration. | | `--batch_size` | `128` | Batch size for calibration / TRT compile. | | `--save_dir` | `./modelopt_quantized` | Directory the quantized Model Optimizer state-dict (FP16 weights + Q/DQ metadata) is always saved to, as `vit_modelopt_state.pt` — re-usable across runs without recalibration. | @@ -174,7 +169,7 @@ python torch_tensorrt_ptq.py \ --recipe \ --save_dir ./my_quantized -# Dump the compiled engine's per-layer info to inspect FP8/NVFP4 fusion. +# Dump the compiled engine's per-layer info to inspect FP8 fusion. python torch_tensorrt_ptq.py --layer_info_path ./vit_fp8_layers.txt ``` @@ -187,7 +182,7 @@ python torch_tensorrt_ptq.py --layer_info_path ./vit_fp8_layers.txt | Flag | Default | Description | | :---: | :---: | :--- | | `--model_id` | `google/vit-large-patch16-224` | HuggingFace model id of the ViT classifier to quantize and score. | -| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). Pass `huggingface/vit/ptq/nvfp4` for NVFP4. | +| `--recipe` | `huggingface/vit/ptq/fp8` | Recipe path (relative to `modelopt_recipes/` or an absolute YAML). | | `--calib_samples` | `1024` | Number of tiny-imagenet samples to use for calibration. | | `--batch_size` | `128` | Calibration / compile / eval batch size. The Torch-TRT engine is dynamic (`min=1`, `opt=max(--batch_size, 2)`, `max=1024`) and handles any batch including the trailing partial batch. | | `--eval_data_size` | full 50k | Number of ImageNet validation images to score. | diff --git a/examples/torch_trt/torch_tensorrt_accuracy.py b/examples/torch_trt/torch_tensorrt_accuracy.py index f973842f103..c450b458993 100644 --- a/examples/torch_trt/torch_tensorrt_accuracy.py +++ b/examples/torch_trt/torch_tensorrt_accuracy.py @@ -106,7 +106,7 @@ def main(): "--recipe", default=ttptq.DEFAULT_RECIPE, help="Recipe path (relative to modelopt_recipes/ or an absolute YAML). " - "Defaults to the ViT FP8 recipe; pass huggingface/vit/ptq/nvfp4 for NVFP4.", + "Defaults to the ViT FP8 recipe.", ) parser.add_argument( "--calib_samples", @@ -210,7 +210,7 @@ def to_eval_model(m: torch.nn.Module, what: str) -> torch.nn.Module: ) ttptq.quantize_with_recipe(model, args.recipe, calib_batches) - label = Path(args.recipe).stem # e.g. "fp8" / "nvfp4" + label = Path(args.recipe).stem # e.g. "fp8" tag = f"{label} ({runtime})" eval_model = to_eval_model(model, f"{label} model") print(f"\n=== {tag} ===") diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index 8edab142ec0..cfa2ff9bdf5 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -20,8 +20,8 @@ 1. Load ``google/vit-large-patch16-224`` (`ViTForImageClassification`) from HF. 2. Build a calibration loader from `zh-plus/tiny-imagenet` so the recipe runs end-to-end without ImageNet access. -3. Run ``mtq.quantize`` with one of the ViT-specific recipes under - `modelopt_recipes/huggingface/vit/ptq/` (FP8 or NVFP4). +3. Run ``mtq.quantize`` with the ViT-specific FP8 recipe under + `modelopt_recipes/huggingface/vit/ptq/`. 4. Compile the quantized model with ``torch_tensorrt.compile(ir="dynamo", min_block_size=1)`` and verify the compiled-model argmax matches the fake-quant argmax on a sample input. @@ -46,7 +46,7 @@ # Default ViT PTQ recipe under `modelopt_recipes/huggingface/vit/ptq/`. The # recipe loader resolves this relative path against the built-in recipe library; -# pass `--recipe` for a different one (e.g. `huggingface/vit/ptq/nvfp4`). +# pass `--recipe` for a different one. DEFAULT_RECIPE = "huggingface/vit/ptq/fp8" @@ -195,7 +195,7 @@ def main(): "--recipe", default=DEFAULT_RECIPE, help="Recipe path (relative to modelopt_recipes/ or an absolute YAML). " - "Defaults to the ViT FP8 recipe; pass huggingface/vit/ptq/nvfp4 for NVFP4.", + "Defaults to the ViT FP8 recipe.", ) parser.add_argument( "--calib_samples", diff --git a/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml b/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml deleted file mode 100644 index 3fb98da7366..00000000000 --- a/modelopt_recipes/huggingface/vit/ptq/nvfp4.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -metadata: - recipe_type: ptq -imports: - nvfp4: configs/numerics/nvfp4 - fp8: configs/numerics/fp8 - attention_qkv_fp8: configs/ptq/units/attention_qkv_fp8 - w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 -quantize: - algorithm: awq_lite - quant_cfg: - - $import: w4a4_nvfp4_nvfp4 - - - quantizer_name: '*output_quantizer' - enable: false - - - $import: attention_qkv_fp8 - - # Hold the patch-embed Conv2d and the classifier head at per-tensor FP8. - - quantizer_name: '*patch_embeddings.projection.weight_quantizer' - cfg: - $import: fp8 - - quantizer_name: '*patch_embeddings.projection.input_quantizer' - cfg: - $import: fp8 - - quantizer_name: '*classifier.weight_quantizer' - cfg: - $import: fp8 - - quantizer_name: '*classifier.input_quantizer' - cfg: - $import: fp8 From 3b0a30214d907a8644f648d8c90bbaec5ddbed46 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 26 Jun 2026 16:16:00 +0000 Subject: [PATCH 17/17] [6078291] torch_trt: quantize non-causal (ViT) attention softmax-P via eager p_bmm_quantizer wrapper ModelOpt's built-in softmax-P quant path (``_QuantAttention.p_bmm_quantizer``) runs through the Triton flash-attention kernel, which only supports causal attention. ViT attention is non-causal, so the FP8 ViT recipe's p_bmm_quantizer tripped ``NotImplementedError: ... does not support non-causal attention`` during ``mtq.quantize`` / ``torch_tensorrt.compile``. Route non-causal attention (``is_causal=False``) to an eager softmax wrapper that applies p_bmm_quantizer by temporarily swapping ``torch.nn.functional.softmax`` for a Q/DQ-emitting version. This keeps the softmax-P quantizer in the traced graph so ONNX / Torch-TRT export sees Q/DQ around the softmax output, while attention is still computed by the original (eager) interface. Requires an eager attention implementation; SDPA-fused softmax is unaffected. The causal LLM path still uses the Triton kernel, and the other out-of-envelope cases (sliding window, sinks, softcapping, dropout, KV cache, padded masks) still raise. Verified on google/vit-base-patch16-224: the torch_trt e2e example test (quantize -> torch_tensorrt.compile) passes, the causal Triton path is unchanged, and a new GPU test covers the non-causal eager fallback. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/torch_trt/torch_tensorrt_ptq.py | 5 +- .../torch/quantization/plugins/huggingface.py | 30 ++++++++++++ .../plugins/test_attention_quant.py | 47 ++++++++++++++++++- 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/examples/torch_trt/torch_tensorrt_ptq.py b/examples/torch_trt/torch_tensorrt_ptq.py index cfa2ff9bdf5..dcd60534de7 100644 --- a/examples/torch_trt/torch_tensorrt_ptq.py +++ b/examples/torch_trt/torch_tensorrt_ptq.py @@ -56,8 +56,9 @@ def load_model_and_processor(model_id: str, device: torch.device, dtype: torch.d processor = AutoImageProcessor.from_pretrained(model_id) # `gelu_fast` selects the tanh-approximation GELU rather than the erf-based # default. Eager attention runs softmax through `F.softmax` instead of the - # fused SDPA kernel, so the recipe's `*softmax_quantizer` is exercised during - # calibration and emits Q/DQ around the softmax output on export. + # fused SDPA kernel, so the recipe's attention softmax-P quantizer + # (`p_bmm_quantizer` on HF attention) is exercised during calibration and + # emits Q/DQ around the softmax output on export. model = ViTForImageClassification.from_pretrained( model_id, torch_dtype=dtype, diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 9f94e1a67c5..6779c3f9ade 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -194,6 +194,29 @@ def _triton_qdq_attention(self, p_qdq, query_states, key_states, value_states, * p_qdq_amax=p_qdq_amax, ) + def _eager_p_qdq_attention( + self, original_attention_interface, query_states, key_states, value_states, **kwargs + ): + """Apply ``p_bmm_quantizer`` to the softmax output via an eager wrapper. + + For attention outside the causal-only Triton kernel's envelope (e.g. ViT's + non-causal attention). Swapping ``F.softmax`` for a quantized version keeps + the quantizer in the traced graph so ONNX / Torch-TRT export emits Q/DQ + around the softmax probabilities. Requires an eager attention + implementation; SDPA-fused softmax (computed inside the C++ kernel) is + unaffected. + """ + _pq = self.p_bmm_quantizer + _orig_softmax = torch.nn.functional.softmax + + def _quantized_softmax(*s_args, **s_kwargs): + return _pq(_orig_softmax(*s_args, **s_kwargs)) + + with replace_function(torch.nn.functional, "softmax", _quantized_softmax): + return original_attention_interface( + self, query_states, key_states, value_states, **kwargs + ) + @staticmethod def _quantized_attention( original_attention_interface, @@ -216,6 +239,13 @@ def _quantized_attention( # positional argument after q/k/v; everything else is a kwarg. if args: kwargs["attention_mask"] = args[0] + # The built-in Triton P kernel is causal-only. Non-causal attention + # (e.g. ViT) applies p_bmm_quantizer through an eager softmax wrapper + # that stays export-traceable for ONNX / Torch-TRT instead. + if kwargs.get("is_causal") is False or getattr(self, "is_causal", True) is False: + return self._eager_p_qdq_attention( + original_attention_interface, query_states, key_states, value_states, **kwargs + ) return self._triton_qdq_attention( p_qdq, query_states, key_states, value_states, **kwargs ) diff --git a/tests/gpu/torch/quantization/plugins/test_attention_quant.py b/tests/gpu/torch/quantization/plugins/test_attention_quant.py index d18274bed75..79d541147bd 100644 --- a/tests/gpu/torch/quantization/plugins/test_attention_quant.py +++ b/tests/gpu/torch/quantization/plugins/test_attention_quant.py @@ -134,8 +134,6 @@ def run(seq_q=seqlen, seq_k=seqlen, **kwargs): run(s_aux=torch.zeros(num_q_heads, device="cuda")) with pytest.raises(NotImplementedError, match="softcapping"): run(softcap=50.0) - with pytest.raises(NotImplementedError, match="non-causal"): - run(is_causal=False) with pytest.raises(NotImplementedError, match="dropout"): run(dropout=0.1) with pytest.raises(NotImplementedError, match="KV cache"): @@ -167,6 +165,51 @@ def run(seq_q=seqlen, seq_k=seqlen, **kwargs): assert torch.isfinite(output).all() +@pytest.mark.skipif(not TRITON_FA_AVAILABLE, reason="Triton attention kernel unavailable") +def test_p_qdq_non_causal_falls_back_to_eager(): + """Non-causal attention (e.g. ViT) is outside the causal-only Triton kernel's + envelope, so p_bmm_quantizer is applied through the eager softmax wrapper + instead of raising -- keeping the softmax-P quant in an export-traceable graph.""" + batch_size, num_q_heads, num_kv_heads, seqlen, head_dim = 2, 4, 2, 32, 64 + + quant_attention = _make_quant_attention(num_q_heads=num_q_heads, num_kv_heads=num_kv_heads) + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + getattr(quant_attention, name).disable() + + torch.manual_seed(29) + q = torch.randn(batch_size, num_q_heads, seqlen, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.randn(batch_size, num_kv_heads, seqlen, head_dim, dtype=torch.bfloat16, device="cuda") + v = torch.randn(batch_size, num_kv_heads, seqlen, head_dim, dtype=torch.bfloat16, device="cuda") + + # Eager interface so the wrapper's F.softmax swap actually fires (SDPA fuses softmax). + module = inspect.getmodule(quant_attention.get_attn_type(quant_attention)) + eager_fn = module.eager_attention_forward + + def run(): + return quant_attention._quantized_attention( + eager_fn, + quant_attention, + q, + k, + v, + attention_mask=None, + scaling=head_dim**-0.5, + is_causal=False, + )[0] + + quant_attention.p_bmm_quantizer.disable() + expected = run() + + quant_attention.p_bmm_quantizer.enable() + quant_attention.p_bmm_quantizer.num_bits = (4, 3) # FP8 + quant_attention.p_bmm_quantizer.amax = torch.tensor(1.0, device="cuda") # softmax P in [0, 1] + output = run() + + assert output.shape == expected.shape + assert not torch.equal(output, expected), "softmax qdq should perturb the output" + torch.testing.assert_close(output, expected, atol=0.1, rtol=0.1) + + @pytest.mark.skipif(kitchen is None, reason="kitchen is not installed.") def test_kitchen_fa(): batch_size = 2