trt-torch-quickstart

Torch-TensorRT Quickstart

Convert a PyTorch nn.Module to a TensorRT engine using torch_tensorrt (the Torch-TensorRT frontend). Covers the AOT path (for production and C++ deploy) and the JIT path (for Python-only inference).

When to Use

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 — each release lists the exact torch / tensorrt / CUDA versions it was built against. Do not invent pins from memory.

For TensorRT 11.0:

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:

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.

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:

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:

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):

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:

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:

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:

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:

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:

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:

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.

   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:

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 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
• Related skills: trt-onnx-quickstart, trt-strong-typing-migration, trt-cpp-runtime-quickstart