trt-torch-quickstart

# Torch-TensorRT Quickstart Convert a PyTorch `nn.Module` to a TensorRT engine using `torch_tensorrt` (the [Torch-TensorRT](https://github.com/pytorch/TensorRT) frontend). Covers the AOT path (for production and C++ deploy) and the JIT path (for Python-only inference). ## When to Use | Scenario | Use this skill? | |----------|-----------------| | PyTorch `nn.Module`; want a TensorRT engine without writing an ONNX intermediate | Yes | | Model uses ops that don't export cleanly to ONNX (custom autograd, dynamic control flow) | Yes | | Want Torch-TensorRT's automatic fallback of unsupported subgraphs to PyTorch | Yes | | Need a serialized engine for C++ deploy from a PyTorch source | Yes — use AOT path here, then `trt-cpp-runtime-quickstart` for the C++ load | | Have a working ONNX file already | No — use `trt-onnx-quickstart` | | LLM token generation (Llama, Mistral, Qwen text gen) | **No** — route to [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) | | Have a `.plan` file already and want to run it from C++ | No — use `trt-cpp-runtime-quickstart` | | Migrating an existing weakly-typed TRT network to strongly-typed | No — use `trt-strong-typing-migration` | ## Prerequisites Torch-TensorRT versions are tightly coupled to a (torch, tensorrt-cu13, CUDA) triple; mismatches produce ABI errors at `import torch_tensorrt`. **Pin matrix source of truth:** the [pytorch/TensorRT releases page](https://github.com/pytorch/TensorRT/releases) — each release lists the exact `torch` / `tensorrt` / CUDA versions it was built against. Do not invent pins from memory. For TensorRT 11.0: | Component | Required | Notes | |-----------|----------|-------| | TensorRT | 11.x (`tensorrt-cu13` wheel) | Mixing `-cu12` and `-cu13` wheels breaks imports. | | CUDA toolkit | 13.x | Driver R590+. | | Python | ≥ 3.10 | TRT 11 dropped 3.9 and earlier. | | `torch` | Match the `torch_tensorrt` release notes for the chosen `torch_tensorrt` version. | | | `torch_tensorrt` | The release matching the TRT 11.0 RC — see [pytorch/TensorRT releases](https://github.com/pytorch/TensorRT/releases). | | **Recommended environment**: NGC TensorRT container with PyTorch (`nvcr.io/nvidia/pytorch:<tag>` or `nvcr.io/nvidia/tensorrt:<tag>` with `torch_tensorrt` pip-installed on top). Run with `--gpus all`. Verification: ```bash python3 -c "import torch, torch_tensorrt, tensorrt; print(torch.__version__, torch_tensorrt.__version__, tensorrt.__version__)" ``` ## Step 1 — Load and prepare the model Use a `torch.nn.Module` in eval mode. Torch-TensorRT will trace through it. ```python import torch import torchvision.models as models model = models.resnet50(weights=None).eval().cuda() example = torch.randn(1, 3, 224, 224, device="cuda") ``` Notes: - `weights=None` skips the download — fine for shape/perf testing. For accuracy work, load real weights. - Model must be on CUDA before compile. Torch-TensorRT does not move it for you. - `.eval()` matters: BatchNorm and Dropout behave differently in train mode and can produce different engines. ## Step 2 — Compile to a TensorRT engine (AOT) Two compile paths — pick by deployment target: | Goal | Path | API | |------|------|-----| | Serialized engine (production, C++ deploy, reuse) | **AOT** — this section | `torch_tensorrt.dynamo.compile` → `torch_tensorrt.save` | | In-process Python callable only (no serialized engine) | **JIT** — end of this section | `torch.compile(backend="torch_tensorrt")` | The AOT path uses `torch_tensorrt.dynamo.compile` and returns a serializable `ExportedProgram`. **Important — strong typing is the default in torch_tensorrt ≥ 2.12.** With `use_explicit_typing=True` (the default), engine precision is **inferred from the exported model's dtype** — passing `enabled_precisions={torch.float16}` raises `AssertionError`. To compile in FP16, cast the model and example tensor to FP16 before exporting: ```python import torch import torch_tensorrt model = MyModel().eval().cuda().half() # cast to FP16 example = torch.randn(1, 3, 224, 224, device="cuda", dtype=torch.float16) trt_gm = torch_tensorrt.dynamo.compile( torch.export.export(model, (example,)), inputs=[example], # No enabled_precisions — precision comes from the model dtype. truncate_double=True, min_block_size=1, ) # Serialize for later / C++ loading — use torch_tensorrt.save, NOT raw bytes. torch_tensorrt.save(trt_gm, "resnet50_trt.ep", inputs=[example]) ``` **Do not extract raw engine bytes via `submod.engine` and write them with `open(...).write(...)`.** That bypasses Torch-TensorRT's metadata wrapper: the blob will not round-trip through `torch_tensorrt.load(...)` and silently drops the dispatch graph that handles partial-fallback subgraphs. `torch_tensorrt.save(trt_gm, path, inputs=...)` is the only supported serialization path. For mixed precision or to override the model dtype, set `use_explicit_typing=False` and `enabled_precisions` applies as before (the weakly-typed path, deprecated in TRT 11): ```python trt_gm = torch_tensorrt.dynamo.compile( torch.export.export(model, (example,)), inputs=[example], use_explicit_typing=False, # weakly-typed (deprecated) enabled_precisions={torch.float16, torch.float32}, truncate_double=True, ) ``` Key arguments: - `truncate_double=True`: silently downcasts FP64 constants. Without it, any FP64 op forces a partition boundary. - `min_block_size`: smallest subgraph (in node count) worth handing to TRT. `1` is aggressive; default `5` avoids tiny TRT subgraphs that don't pay back their launch overhead. - `workspace_size`: bytes of scratch memory the TRT builder can use. Leave unset to let TRT decide. **JIT alternative (Python-only inference):** if you just want a callable and not a serialized engine, use `torch.compile` with the Torch-TensorRT backend. Under the strong-typing default, precision comes from the model dtype — `.half()` the model for FP16 rather than passing `enabled_precisions`: ```python model = model.half() # FP16 inferred from model dtype (strong typing default) example = example.half() trt_model = torch.compile(model, backend="torch_tensorrt") out = trt_model(example) # compiles lazily on the first call, then runs through TRT in-process ``` This skips `torch.export` and gives a `torch.compile`-wrapped callable. It cannot be loaded from C++ — use the AOT path above if you need that. ## Step 3 — Dynamic shapes Wrap inputs in `torch_tensorrt.Input` with `min_shape`/`opt_shape`/`max_shape`: ```python dynamic_input = torch_tensorrt.Input( min_shape=(1, 3, 224, 224), opt_shape=(8, 3, 224, 224), max_shape=(32, 3, 224, 224), dtype=torch.float16, ) trt_gm = torch_tensorrt.dynamo.compile( torch.export.export(model, (example,), dynamic_shapes={"x": {0: torch.export.Dim("batch", min=1, max=32)}}), inputs=[dynamic_input], # Strong typing: FP16 comes from the input/model dtype, not enabled_precisions. ) ``` Pitfalls: - The `dynamic_shapes=` argument to `torch.export.export` and the `Input` ranges must agree. A mismatch produces a builder error 2–3 minutes into compilation, not at export time. - `opt_shape` is what TRT tunes kernels for. Set it to the most common runtime shape, not the midpoint. - Only mark dimensions dynamic that actually vary at runtime; unnecessary dynamism degrades performance. ## Step 4 — Run inference and verify Raw-tensor comparison (`assert_close`) is fine for FP32 but **misleading for FP16** — a few outlier positions exceed the tolerance even when downstream behavior is identical. Use a **semantic check** appropriate to the model type: ```python with torch.inference_mode(): torch_out = model(example) trt_out = trt_gm(example) ``` **Image classification** — top-1 (or top-5) class match, plus mean-softmax-distance: ```python assert torch.equal(torch_out.argmax(-1), trt_out.argmax(-1)), "top-1 mismatch" mean_prob_diff = (torch.softmax(torch_out.float(), -1) - torch.softmax(trt_out.float(), -1)).abs().mean() assert mean_prob_diff < 1e-3 ``` **Sentence / token embeddings** — cosine similarity per token, mean across the batch: ```python cos = torch.nn.functional.cosine_similarity(torch_out.flatten(0, -2), trt_out.flatten(0, -2), dim=-1) assert cos.mean() > 0.999 ``` **Detection / regression heads** — bounded raw-tensor closeness on the meaningful sub-tensor (e.g., bbox coordinates), softmax check on the class logits. **FP32 sanity** — for any model, an FP32 (no `.half()`) run should pass `torch.testing.assert_close(trt_out, torch_out, rtol=1e-4, atol=1e-4)`. If that fails, the issue is in the compilation, not numeric precision. ## Loading the serialized engine **From Python**: ```python import torch_tensorrt loaded = torch_tensorrt.load("resnet50_trt.ep").module() out = loaded(example) ``` **From C++**: this skill does NOT cover C++ loading. The `.ep` file `torch_tensorrt.save` produces by default is a `torch.export` archive for the **Python** `torch_tensorrt.load(...)` path — it is *not* directly consumable by the plain TensorRT C++ runtime (`IRuntime` / `deserializeCudaEngine`), which expects a serialized TensorRT engine (`.plan`). Two supported C++ routes: - Save with `output_format="torchscript"` and deploy the TorchScript module with libtorch + the Torch-TensorRT C++ runtime. - If you only need the raw TensorRT engine in C++, build it through the ONNX path (`trt-onnx-quickstart`) and load the resulting `.plan` with `trt-cpp-runtime-quickstart` (modern `IRuntime` + `enqueueV3` + `setTensorAddress`). Do not inline a C++ snippet here — it will diverge from the canonical pattern. ## Common issues **`RuntimeError: Trying to create tensor with negative dimension`** during export — the model has a shape-dependent control flow path that `torch.export` can't trace. Either rewrite with `torch.cond` / `torch.where`, or fall back to `ir="dynamo"` + `min_block_size=1` so the unsupported region runs in PyTorch. **`Unsupported operator` warnings** — Torch-TensorRT partitions around them and runs those nodes in PyTorch. Set `require_full_compilation=True` to turn this into an error; use during development to find what's actually falling back. ## Diagnosing excessive PyTorch fallback If the compile completes but most of the model fell back to PyTorch (only a few subgraphs run on TRT), three knobs in order: 1. **Identify what fell back.** Run with `TORCH_LOGS="graph_breaks"` and `require_full_compilation=True` to escalate fallbacks to errors that name the offending op. Without this you don't know what the autotuner rejected. ```bash TORCH_LOGS="graph_breaks" python3 your_compile.py ``` 2. **Reconsider `min_block_size`.** This is a real trade-off — not just a knob. - Default `min_block_size=5`: Torch-TensorRT only hands a subgraph to TRT if it has ≥5 nodes. Tiny TRT subgraphs cost more in kernel-launch overhead than they save. - `min_block_size=1`: aggressive — every supported op goes to TRT, even single-node subgraphs. Useful to see what's *theoretically* supported, but commonly slower at inference due to the per-subgraph launch tax. - Recommended: leave at default `5` for production; drop to `1` only during diagnosis to see maximum TRT coverage. 3. **Look for shape-dependent control flow.** `torch.export` traces only one branch of `if x.shape[0] > 0:` style code; rewrite with `torch.where`, `torch.cond`, or lift the condition out of the model. The trace error names the offending op. If fallback persists after all three, you have a genuine unsupported op — write a converter (see Torch-TensorRT upstream docs) or switch to the ONNX path (`trt-onnx-quickstart`). **Engine builds but output is garbage** — almost always a dtype issue. Check: 1. Model/input dtype matches your pipeline — under strong typing precision follows the model dtype (`.half()` for FP16); `enabled_precisions` applies only under weak typing (`use_explicit_typing=False`). 2. `truncate_double=True` if the model has any FP64 constants (common in positional encodings). 3. The model was `.eval()` before export. **Compile takes >10 minutes** — turn on the builder log to see what's going on: ```python import torch_tensorrt.logging as ttlog ttlog.set_reportable_log_level(ttlog.Level.Info) ``` Usually it's the autotuner exploring kernel variants for a heavy GEMM/conv. To cap exploration, set `workspace_size` smaller (less scratch → fewer candidates). **`ImportError: cannot import name 'XYZ' from 'torch_tensorrt'`** — version mismatch. Confirm the pin matrix; the `torch_tensorrt` Python API moves between minor versions. ## Numerical debugging If `assert_close` fails by more than ~5× tolerance: 1. Re-run in FP32 — drop the model's `.half()` (or under weak typing, `enabled_precisions={torch.float32}`). If that matches PyTorch, the issue is FP16 accumulator drift in a specific op — usually a softmax or LayerNorm. The pragmatic fix is to wrap the offending submodule in a module excluded from the compile (`min_block_size` won't help; explicit exclusion is needed). See the [Torch-TensorRT lowering guide](https://pytorch.org/TensorRT/) for the version-specific exclusion API. 2. Sanity-check at FP32 first, then re-introduce FP16 and bisect — almost always reveals one offending op. 3. For FP8: check that the calibration distribution actually covers the runtime distribution. FP8 has no headroom for outliers. See `trt-strong-typing-migration` for the broader weak-vs-strong typing discussion. By default Torch-TensorRT produces strongly-typed engines in TRT 11.0 — types inferred from the exported program — but you can override. ## What this skill is not - Not a guide to writing custom Torch-TensorRT converters. Use the upstream Torch-TensorRT docs for that. - Not for TensorRT-LLM / LLM inference — use the TRT-LLM `examples/` flow instead. - Not for QAT (quantization-aware training) — Torch-TensorRT consumes the QAT model but the training loop is upstream. ## References - Upstream Torch-TensorRT: https://pytorch.org/TensorRT/ - Version pin matrix: [pytorch/TensorRT releases](https://github.com/pytorch/TensorRT/releases) - Related skills: `trt-onnx-quickstart`, `trt-strong-typing-migration`, `trt-cpp-runtime-quickstart`