Skip to content

Torch-TensorRT v2.15.0

Latest
Latest

Choose a tag to compare

View all tags
@lanluo-nvidia lanluo-nvidia released this 13 Nov 00:34
· 129 commits to main since this release
v2.15.0
955b9fa
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

Torch-TensorRT 2.15.0 Linux x86-64 and Windows targets

PyTorch 2.15, CUDA 13.2/13.4, TensorRT 11.3, Python 3.11~3.14

Torch-TensorRT Wheels are available:

x86-64 Linux and Windows:
CUDA 13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

CUDA 13.2/13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

aarch64 SBSA Linux and Jetson Thor and Orin:
CUDA 13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

CUDA 13.2/13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

Torch-TensorRT-RTX 2.15.0 Linux x86-64 and Windows targets

PyTorch 2.15, CUDA 13.2/13.4, TensorRT-RTX 1.6, Python 3.11~3.14

Torch-TensorRT-RTX Wheels are available:

x86-64 Linux and Windows:
CUDA 13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT-RTX 1.6

CUDA 13.2/13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT-RTX 1.6

Torch-TensorRT-Executorch-Runtime 0.2.1 Linux x86-64 and Windows targets

PyTorch 2.15, CUDA 13.2/13.4, TensorRT 11.3, Python 3.11~3.14

Torch-TensorRT-Executorch-Runtime Wheels are available:

x86-64 Linux and Windows:
CUDA 13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

CUDA 13.2/13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

aarch64 SBSA Linux and Jetson Thor and Orin:
CUDA 13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

CUDA 13.2/13.4 + Python 3.11-3.14 + Torch 2.15 + TensorRT 11.4

New Features:

Highlights

  • Efficient Deployment of LeRobot Policies: Compile and run policies from LeRobot without PyTorch at TensorRT speeds using lerobot-export and Executorch-TensorRT
  • ExecuTorch on Windows: compile with Torch-TensorRT, save a .pte, and run it with ExecuTorch's runtime on Windows x86-64.
  • Compile and save large models with far less memory: A number of improvements have been made to the host memory / unified memory efficiency of Torch-TensorRT compilation. Some of these features are exposed through a new API torch_tensorrt.dynamo.low_memory, an experimental API that manages weight and engine streaming to/from disk during compilation to reduce max memory utilization. Now a LeRobot pi0.5 ExecuTorch export peaks at 5GB max utilization.
  • TorchAO Quantization Support: compile FP8, static FP8, INT4, NVFP4 and MXFP4 models quantized with TorchAO
  • Bring your own kernels (experimental): triton_op and cutile_op turn a Triton or cuTile kernel into a TensorRT plugin in one call, and custom operators that update their inputs in place can be compiled into the engine, bypassing the overhead of running the kernels outside of the engine .
  • More control over compilation: module-level PyTorch fallback, Global Performance Tuning (sweeps TensorRT builder options for a faster engine), and improvements to torch_tensorrt.autocast.

LeRobot on ExecuTorch + TensorRT: from checkpoint to robot, without PyTorch

Take a trained LeRobot policy from a checkpoint to a compiled, self-checking program on your robot's NVIDIA GPU. With Torch-TensorRT 2.15, ExecuTorch 1.6 and LeRobot 0.7.1, lerobot-export --backend.type=executorch_tensorrt compiles a policy once, ahead of time, into a bundle: a plain folder of ExecuTorch programs with TensorRT engines inside. lerobot-rollout --bundle.path=<bundle> then drives the robot from it, in a process that never imports PyTorch.
.

  • Light weight deployment The run-time process never loads PyTorch, so that memory stays free for the model.
  • Fit large policies on small boards. pi0, pi0.5, GR00T N1.7, SmolVLA, X-VLA, ACT, Diffusion Policy and more export and run on an 8 GB Jetson Orin Nano
  • Currently over 20 Policies Supported SmolVLA, pi0, pi0.5, pi0-FAST, GR00T N1.7, X-VLA, ACT, Diffusion Policy, Multi-Task DiT, VQ-BeT, TD-MPC, FastWAM, EO-1, VLA-JEPA, Wall-OSS and more have export recipes.
  • Verification tooling The export replays a test frame through the bundle and through PyTorch and refuses a bundle that disagrees. lerobot-rollout replays it again before the robot connects, lerobot-check-bundle verifies a bundle on the robot's own computer, and a NaN or infinite action stops the run before it reaches the motors.
  • Serving models lerobot-policy-server can serve a executorch-tensorrt bundle from a GPU workstation to a small robot computer to be integrated into greater pipelines.
# Install on the machine that will run the policy (Linux x86-64 or aarch64, CUDA 13)
pip install 'lerobot[export]' 'torch-tensorrt[executorch]>=2.15,<2.16' torch-tensorrt-executorch-runtime
pip install 'cuda-toolkit[nvcc]==13.2.*'   # the export compiles CUDA code; JetPack already ships nvcc
export CUDA_HOME=$(python -c "import os, nvidia; print(os.path.join(list(nvidia.__path__)[0], 'cu13'))")
export PATH=$CUDA_HOME/bin:$PATH

# Compile a trained policy into a bundle. The export checks it against PyTorch
# ("Bundle test case: largest difference ... (tolerance ...)") and stops if it disagrees.
lerobot-export \
    --policy.path=outputs/train/act_so101/checkpoints/last/pretrained_model \
    --backend.type=executorch_tensorrt

# Drive the robot from the bundle: PyTorch is never imported
lerobot-rollout \
    --bundle.path=outputs/export/act_executorch_tensorrt \
    --robot.type=so101_follower \
    --robot.port=/dev/ttyACM0 \
    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
    --duration=30

Docs: LeRobot export guide (policies and backends, a Jetson Orin Nano walkthrough, remote serving) · ExecuTorch Deployment guide

ExecuTorch on Windows

Torch-TensorRT 2.15 adds ExecuTorch support on Windows x86-64 (#4835): compile with Torch-TensorRT, save a .pte with output_format="executorch", and run it with ExecuTorch's runtime. All that is required to run compiled .pte files is ExecuTorch (the CUDA build), the TensorRT delegate (the separate torch-tensorrt-executorch-runtime wheel) and TensorRT.

# Linux
pip install "torch_tensorrt[executorch]"

# Windows: PyPI only has CPU builds of PyTorch and ExecuTorch there, so add the PyTorch CUDA index
pip install "torch_tensorrt[executorch]" --extra-index-url https://download.pytorch.org/whl/cu132

Docs: ExecuTorch Deployment guide · torch_tensorrt.executorch API reference

TorchAO support for FP8 weight-only, static FP8, symmetric INT4 weight-only, NVFP4 weight-only, and MXFP4 dynamic-activation/weight quantization configurations

Torch-TensorRT 2.15 compiles models quantized with TorchAO through the Dynamo backend. TorchAO's quantize/dequantize ops are preserved through torch.export and converted to TensorRT IQuantizeLayer / IDequantizeLayer, so the engine stores weights as low-precision constants instead of upcasting them at build time. Install with pip install "torch_tensorrt[quantization]" (or pip install torchao).

• FP8 weight-only (Float8WeightOnlyConfig): FP8 weight constants with BF16/FP16 activations; no calibration needed.
• Static FP8: weights and activations are quantized using observer-based calibration, and TensorRT fuses the Q/DQ pairs into FP8 GEMMs. Requires Hopper or newer.
• Symmetric INT4 weight-only: group-wise INT4 with a zero zero-point (TensorRT rejects asymmetric INT4 DQ). Engines keep Int4 weight constants; compile with immutable_weights=True.
• NVFP4 weight-only: packed FP4 (E2M1) weights with FP8 block scales, lowered to a two-level IDequantizeLayer. Weight dims must be divisible by 16; compile with immutable_weights=True.
• MXFP4 (MXDynamicActivationMXWeightConfig): packed FP4 weights with E8M0 block scales (block size 32), with activations kept in BF16. TorchAO has no MXFP4 weight-only config.

INT4, NVFP4, and MXFP4 are currently storage + dequantize paths: they reduce weight memory and engine size, but GEMMs run at the activation precision rather than on native low-bit MMA kernels. These paths require TensorRT ≥ 10.8.

End-to-end examples in examples/dynamo/torchao/ cover a Linear model for every config, FLUX.1 for FP8 weight-only, INT4, NVFP4, and MXFP4, and Qwen3 for INT4. See the Quantization user guide for the workflow and hardware requirements.

In Place Custom Operators

KV-cache writes, running statistics and recurrent state are often updated in place. In the case that these operators are backed by a user supplied kernel, Torch-TensorRT can now compile such a custom operator into a TensorRT engine as a Quick Deployable Plugin and writes the updates in place.

Declare the op with torch.library.custom_op(..., mutates_args=...) and register_fake, then register it with torch_tensorrt.dynamo.conversion.plugins.custom_op.

import torch
import torch_tensorrt
from torch_tensorrt.dynamo.runtime import TorchTensorRTModule


@torch.library.custom_op("demo::scale_add_", mutates_args=("state",))
def scale_add_(state: torch.Tensor, update: torch.Tensor) -> None:
    state.mul_(0.9).add_(update)  # in place: state = 0.9 * state + update


scale_add_.register_fake(lambda state, update: None)  # nothing is returned
torch_tensorrt.dynamo.conversion.plugins.custom_op(
    "demo::scale_add_", supports_dynamic_shapes=True
)


class Model(torch.nn.Module):
    def forward(self, state, update):
        scale_add_(state, update)
        return state * 2


state, update = torch.zeros(4, 8, device="cuda"), torch.randn(4, 8, device="cuda")
# min_block_size=1: this graph is smaller than the default block size of 5
trt_model = torch_tensorrt.compile(
    Model(), ir="dynamo", inputs=[state, update], min_block_size=1
)
engine = next(m for m in trt_model.modules() if isinstance(m, TorchTensorRTModule))
print("mutated engine inputs:", [i for i, _ in engine.aliased_io.values()])

ref_state = state.clone()
for step in range(3):  # the caller's state accumulates across calls
    update = torch.randn(4, 8, device="cuda")
    out = trt_model(state, update)  # TensorRT updates `state` in place
    ref_out = Model()(ref_state, update)  # eager reference
    print(step, torch.allclose(out, ref_out), torch.allclose(state, ref_state))

Docs: Plugins tutorial: Mutating Custom Operators

Limits and requirements

  • Mutating a module buffer (register_buffer) is not supported

Custom cuTile and Triton Kernel Support (Experimental)

If your model calls a Triton kernel through a PyTorch custom op, Torch-TensorRT has no converter for it: the op runs in PyTorch and splits the engine in two. torch_tensorrt.kernels.triton_op (#4435) and torch_tensorrt.kernels.cutile_op (#4463) turn a kernel you write into a TensorRT plugin in one call, so the kernel runs inside the engine. One call registers the PyTorch custom op, the Torch-TensorRT converter and a TensorRT plugin (a Quick Deployable Plugin, QDP) that carries the compiled kernel (PTX) inside the engine.

import tensorrt.plugin as trtp
import torch
import triton
import triton.language as tl

import torch_tensorrt
import torch_tensorrt.kernels as ttk

BLOCK_SIZE = 256


@triton.jit
def relu_sq_kernel(x_ptr, n, y_ptr, BLOCK_SIZE: tl.constexpr):
    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    x = tl.maximum(tl.load(x_ptr + offs, mask=offs < n), 0.0)
    tl.store(y_ptr + offs, x * x, mask=offs < n)


def relu_sq_meta(x: torch.Tensor) -> torch.Tensor:
    return torch.empty_like(x)


# One call compiles the kernel to PTX and registers torch.ops.demo.relu_sq, the
# TensorRT plugin and the Torch-TensorRT converter. The grid and n are symbolic
# expressions of the input shape, so one engine serves every batch size from 2 to 32.
ttk.triton_op(
    "demo::relu_sq",
    kernel=relu_sq_kernel,
    signature={"x_ptr": "*fp32", "n": "i32", "y_ptr": "*fp32"},
    constexprs={"BLOCK_SIZE": BLOCK_SIZE},
    grid=lambda inputs, outputs: (trtp.cdiv(inputs[0].shape_expr.numel(), BLOCK_SIZE),),
    meta_fn=relu_sq_meta,
    extra_args_fn=lambda inputs, outputs: [trtp.SymInt32(inputs[0].shape_expr.numel())],
)


class Net(torch.nn.Module):
    def forward(self, x):
        return torch.ops.demo.relu_sq(x - 0.5)


spec = torch_tensorrt.Input(
    min_shape=(2, 64), opt_shape=(8, 64), max_shape=(32, 64), dtype=torch.float32
)
# The default min_block_size (5) would leave this two-op graph in PyTorch.
trt_model = torch_tensorrt.compile(Net().cuda().eval(), inputs=[spec], min_block_size=1)
for batch in (2, 5, 32):
    x = torch.randn(batch, 64, device="cuda")
    torch.testing.assert_close(trt_model(x), torch.relu(x - 0.5) ** 2)
print("TensorRT engine with the Triton kernel matches PyTorch for batch 2, 5, 32")
torch_tensorrt.save(trt_model, "relu_sq.ep", arg_inputs=[spec])

Deploying the saved engine needs neither Triton nor the kernel definition:

import sys

import torch

import torch_tensorrt

# No kernel definition and no triton_op(...) call in this process.
loaded = torch_tensorrt.load("relu_sq.ep").module()
for n in (2, 5, 32):
    x = torch.randn(n, 64, device="cuda")
    torch.testing.assert_close(loaded(x), torch.relu(x - 0.5) ** 2)
print("ran without Triton; op registered here:", hasattr(torch.ops.demo, "relu_sq"))

cuTile works the same way: register the kernel, then use the op in a model and compile it as above. It needs cuda-tile>=1.3,<2, cuda-python, a tileiras compiler from CUDA Toolkit 13.2 or newer and a driver that supports CUDA 13.2 or newer.

import cuda.tile as ct
import tensorrt.plugin as trtp
import torch

import torch_tensorrt.kernels as ttk

TILE = 128


@ct.kernel
def reglu_kernel(gate, up, out, tile_size: ct.Constant[int]):
    pid = ct.bid(0)
    g = ct.load(gate, index=(pid,), shape=(tile_size,))
    u = ct.load(up, index=(pid,), shape=(tile_size,))
    ct.store(out, index=(pid,), tile=ct.maximum(g, 0.0) * u)


def reglu_meta(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    return torch.empty_like(gate)


# One call compiles the kernel with cuda.tile and registers torch.ops.demo.reglu, the
# TensorRT plugin and the converter. signature lists the arrays (inputs, then outputs);
# constants holds the ct.Constant values; grid is a tile count built from symbolic shapes.
ttk.cutile_op(
    "demo::reglu",
    kernel=reglu_kernel,
    signature={"gate": "fp32", "up": "fp32", "out": "fp32"},
    constants={"tile_size": TILE},
    grid=lambda inputs, outputs: (trtp.cdiv(inputs[0].shape_expr.numel(), TILE),),
    meta_fn=reglu_meta,
)
# From here, use torch.ops.demo.reglu(gate, up) in a model and compile it exactly as
# in the Triton example above.

Docs: API reference for triton_op and cutile_op · Plugins tutorial · example: triton_op.py

Limits and requirements

  • You supply the grid: give it as a symbolic expression of the input shapes. Both frontends derive the block size (triton_op from num_warps, cutile_op from the compiled kernel's .reqntid); only cutile_op also derives the extents and strides launch arguments, while for triton_op you pass any i32 scalar arguments through extra_args_fn. meta_fn must be fully type-annotated.
  • Validation: the kernel signature, dtypes and compiled launch features are validated at registration. Grid and scalar values are range-checked only when they are constants, so keep symbolic ones within Grid limits for every shape in the profile.
  • Known limitations: a Triton kernel that needs more than 48 KiB of dynamic shared memory per block builds but fails when the engine runs (#4877, not yet merged, will make registration reject such kernels with an explicit error); keep block sizes and num_stages small. AOT-compiled Triton kernels get no pointer-alignment specialization, so loads and stores are not vectorized (fix pending in #4875).
  • Portability: the embedded PTX targets one GPU architecture and one PTX ISA version. Therefore similar to TensorRT, it is important to build on the target platform.

Global Performance Tuning

TensorRT's builder has over 200 internal knobs that can change how fast an engine runs. Global Performance Tuning lets you sweep them, and Torch-TensorRT 2.15 brings it to the TensorRT partitions of Thorch-TensorRT compilations, so you can search for a faster engine without changing the model. get_all_build_routes() (in torch_tensorrt.dynamo) lists the knobs of the installed TensorRT, the build_route setting applies one route, and the tune_build_routes setting sweeps routes (tuning_search: fast, the default, builds the defaults and then varies one knob at a time; full builds every combination; mixed runs fast and then combines only the knobs that helped), timing each on the GPU and, if you set accuracy_threshold, rejecting routes whose output drifts from the eager model. The compile call returns the module built with the fastest route that passes the accuracy check; the winning route is logged at INFO and every trial is in the cache file, so you can pin it later with build_route. Each trial runs in its own spawned Python process, so a TensorRT or driver crash costs one trial. Gains depend on the model, GPU and TensorRT version; re-tune after changing any of them.

import glob, json, os, tempfile

import torch
import torch_tensorrt
from torch import nn
from torch_tensorrt.dynamo import get_all_build_routes, is_global_perf_tuning_available


def main():
    assert is_global_perf_tuning_available(), "needs TensorRT >= 11.1 (not TensorRT-RTX, not Windows)"
    # allowed_values is already in sweep syntax, e.g. "-kgen:codegen:cuda_tile=[0|1|2|3]"
    expr = get_all_build_routes("-kgen:codegen:cuda_tile")["tuner_options"][0]["allowed_values"]
    torch.manual_seed(0)
    model = nn.Sequential(
        nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
        nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, 10),
    ).eval().cuda().half()
    x = torch.randn(8, 3, 64, 64, device="cuda", dtype=torch.half)
    cache = os.path.join(tempfile.mkdtemp(), "tune.jsonl")

    torch_tensorrt.compile(
        model, ir="dynamo", arg_inputs=[x], min_block_size=1,
        tune_build_routes=expr, tuning_search="fast",  # one build per allowed value
        accuracy_threshold=0.01, accuracy_algorithm="cos",  # reject routes with 1 - cos_sim > 0.01
        tuning_timeout_s=120, tuning_cache_file=cache,  # 120 s budget per TRT partition
    )

    path = glob.glob(cache.replace(".jsonl", ".*.jsonl"))[0]  # one cache file per TRT partition
    rows = [json.loads(line) for line in open(path)][1:]  # line 0 is the sweep header
    for r in rows:
        print(r["build_route"], "->", f"{r['gpu_time']:.4f} ms" if r["gpu_time"] else r["error_message"])
    best = min((r for r in rows if r["gpu_time"]), key=lambda r: r["gpu_time"])["build_route"]

    tuned = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=[x], min_block_size=1, build_route=best)
    print("best:", best, "| max abs err vs eager:", (tuned(x) - model(x)).abs().max().item())


if __name__ == "__main__":  # required: each trial runs in a spawned child that re-imports this file
    main()

Docs: Performance Tuning guide: Global Performance Tuning · NVIDIA's Global Performance Tuning

Module-level PyTorch fallback for Dynamo Frontend

Keep a particular submodule in PyTorch while the rest of the model compiles to TensorRT. Dynamo compilation now honors torch_executed_modules, with the same convention as the TorchScript frontend, the fully qualified class name (torchvision.models.resnet.BasicBlock; a class defined in your own script is __main__.Gate to Dynamo, where TorchScript uses Gate) (#4639; #4958). Every operation traced from an instance of a listed class, including its nested submodules, stays out of the TensorRT partitions.

import torch
import torch.nn as nn
import torch_tensorrt


class Gate(nn.Module):  # a custom layer that should stay in PyTorch
    def forward(self, x):
        return x * torch.sigmoid(x)


torch.manual_seed(0)
model = nn.Sequential(nn.Conv2d(3, 8, 3, padding=1), Gate(), nn.Conv2d(8, 8, 3, padding=1))
model = model.eval().cuda()
x = torch.randn(1, 3, 16, 16, device="cuda")

trt_model = torch_tensorrt.compile(
    model, ir="dynamo", inputs=[x],
    min_block_size=1, disable_tf32=True,  # toy model: allow one-op engines, compare at full FP32
    torch_executed_modules=[f"{Gate.__module__}.{Gate.__qualname__}"],
)

for name, child in trt_model.named_children():
    print(name, type(child).__name__)
    if name.startswith("_run_on_gpu"):
        print("  runs in PyTorch:", [str(n.target) for n in child.graph.nodes if n.op == "call_function"])
print("max abs diff vs eager:", (trt_model(x) - model(x)).abs().max().item())

Output:

_run_on_acc_0 TorchTensorRTModule
_run_on_gpu_1 GraphModule
  runs in PyTorch: ['aten.sigmoid.default', 'aten.mul.Tensor']
_run_on_acc_2 TorchTensorRTModule
max abs diff vs eager: 1.7881393432617188e-07

Limits and requirements

  • Names are strings, in a list, set or tuple: f"{cls.__module__}.{cls.__qualname__}" for a class you define. A class object or a short name matches nothing; the partitioner logs a "not found in the graph" warning, and with require_full_compilation=True, where a matching module raises an AssertionError, a misspelled name gives no message at all, so check the spelling.

Low host memory compilation (Experimental)

Compiling a large model used to hold several copies of its weights in memory at once. That limited compiling models on host memory constrained devices, in particular Jetson and DGX Spark where memory is shared between the CPU and GPU Torch-TensorRT 2.15 lowers the peak memory consumption, some features are automatically applied, Automatic savings and new opt-in options go further when memory is particularly constrained.

New options
  • Engines kept on disk until the save: torch_tensorrt.compile_and_save() (#4863). Compiles an ExportedProgram and saves the result as an ExportedProgram (.pt2), an AOTInductor package or an ExecuTorch .pte in one call. Each engine is written to disk as soon as it is built, so finished engines stay out of memory while the rest build, and a .pte save streams them from disk without loading them.
  • Weights released once they are built: release_weights_after_build (#4853). A new compile setting that frees each weight as soon as the last engine that reads it is built, and frees parameters that no partition reads before the first build starts.
  • Compiling from memory-mapped weights (#4854; experimental, the API will change in 2.16). save_weights_for_mapping(model, path, source=checkpoint_source(checkpoint)) writes the file that map_weights reads from a safetensors checkpoint, one chunk at a time, and torch_tensorrt.dynamo.tools.map_weights() points a model's weights at that file instead of loading them, so each weight is read from disk only while the engine that uses it builds.
  • The above features can be used together using the experimental torch_tensorrt.dynamo.low_memory() (#4854, #4942; experimental, the API may change in 2.16). This API builds the model on as if it was on the virtual meta device, maps its weights, and supplies the compile settings that lower memory, passed as **low_mem.compile_settings.
  • A fixed engine size for the resource partitioner: max_partition_size (#4855). Sets the weight bytes per TensorRT engine, so a model shards are deterministic. With cpu_memory_budget, which moves with the process's memory, the same pi0.5 policy split into 27 engines in one run and 13 in another; a fixed 0.5 GiB size gave 16 engines in both.
  • Shared activation memory in the ExecuTorch runtime: use_shared_activation_scratch (#4739). A process option of the TensorRT delegate that makes all TensorRT execution contexts on a device share one scratch buffer, so a model split into many engines no longer pays for each engine's scratch separately.
import os
import tempfile

import torch
import torch_tensorrt
from torch_tensorrt.dynamo.tools import save_weights_for_mapping


def mlp():  # 8 x (Linear + ReLU), 32 MiB of float32 weights
    layers = []
    for _ in range(8):
        layers += [torch.nn.Linear(1024, 1024), torch.nn.ReLU()]
    return torch.nn.Sequential(*layers).eval()


folder = tempfile.mkdtemp()
weights, program = os.path.join(folder, "mlp.safetensors"), os.path.join(folder, "mlp.pt2")

reference = mlp()  # stands in for a trained model
save_weights_for_mapping(reference, weights)  # once: write its weights as a safetensors file

x = torch.randn(4, 1024)
# Builds mlp() on the meta device, so it holds no weights, and maps the file into memory.
with torch_tensorrt.dynamo.low_memory(mlp, weights) as low_mem:
    # CPU inputs: torch.export reads only shapes, so the weights stay on disk.
    exported = torch.export.export(low_mem.model, (x,))
    torch_tensorrt.compile_and_save(  # each engine goes to disk as soon as it is built
        exported,
        program,
        arg_inputs=[x.cuda()],
        max_partition_size=10 << 20,  # engines of at most 10 MiB of weights
        **low_mem.compile_settings,
    )
    print(f"{low_mem.report.mapped_bytes >> 20} MiB of weights mapped from disk")

loaded = torch_tensorrt.load(program).module()
print("max |reloaded - eager| =", (loaded(x.cuda()).cpu() - reference(x)).abs().max().item())

Limits and requirements

  • compile_and_save: not supported on Windows. Only output_format="executorch" avoids loading the engines; for the other formats every engine is loaded once the last one is built, so only the compile phase gets the lower peak. The temporary files (<TMPDIR>/torch_tensorrt_<user>/offloaded_engines) need disk space for potentially all engines concurrently.
  • Memory-mapped weights: export the model with CPU inputs. checkpoint_source needs the safetensors package, which Torch-TensorRT does not install (save_weights_for_mapping without source= and map_weights do not). On Linux the pages are released after each engine builds; on Windows the mapping works but the pages stay resident until the file is unmapped.
  • max_partition_size: a smaller size lowers the peak of each build at the cost of more engines (a weight read by several engines is stored in each). It replaces cpu_memory_budget (setting both raises ValueError; cpu_memory_budget itself is unchanged) and turns resource partitioning on. The size counts weight bytes, not build memory, which has been measured at 3 to 10 times the weights; give it in bytes.
  • use_shared_activation_scratch: off by default, and it affects run-time memory only: set it before loading the programs that should share the buffer, with Runtime.get().backend_registry.set_option("TensorRTBackend", {"use_shared_activation_scratch": True}). Pooled calls never run concurrently on one GPU.
Automatic savings
  • Compiling (#4792–#4796). With the C++ runtime, engines are built, set up and refit from TensorRT's serialized plan without copying it, and modules no longer keep a host copy of the plan once the engine is set up (#4794, #4795, #4796). On Linux, the host memory the TensorRT builder freed after each engine build goes back to the OS (#4792), and offload_module_to_cpu=True no longer holds the weights on the GPU and the CPU at the same time (#4793). Together, on a 3 GiB FP16 model (24 × nn.Linear(8192, 8192), RTX 5090, TensorRT 11.3), the peak host + GPU growth of compiling fell by about 40%, from 4.7× to 2.8× the weight size (2.2× with offload_module_to_cpu=True), measured over what the process held before compiling, which includes the PyTorch model; once the PyTorch model is deleted, the process holds about 0.2× the weight size more than before compiling, down from 1.7×.
  • Saving a .pte (#4821, #4774, #4775, #4864, #4865). A .pte stores each engine once, as named data that identical engines in different methods share, instead of holding two host copies until the file is written (pi0.5 with 13 engines: peak host memory of the export 17.8 → 11.6 GiB, for #4821). torch_tensorrt.save shares the compiled module's tensors instead of deep-copying them (#4864) and frees each program and its engines when it returns (#4865), so saving several programs in one process no longer stacks their memory (pi0.5's three-program .pte save: 23.99 → 12.07 GiB host + GPU, then 7.62 → 6.59 GiB).
  • Resource partitioner (#4963, #4961). With resource partitioning on, long graphs no longer hit Python's recursion limit, a multi-output node stays in the same engine as its outputs, and engines stay within the size budget (#4963). Values larger than 64 KiB that depend only on the model's inputs, such as attention masks and rotary cos/sin, are recomputed inside each engine instead of passed between engines, which avoids Gemma attention engines failing to build on Jetson AGX Thor (#4961).

Limits and requirements

  • Saving: torch_tensorrt.save shares tensors when it inlines the module without retracing (retrace=False, the default for output_format="executorch"), not when you pass your own ExecuTorch partitioners.

Docs: the Resource Management guide covers offload_module_to_cpu, builder trimming and compiling from memory-mapped weights, where it also describes the resource partitioner and release_weights_after_build as the settings low_memory supplies. API references: compile_and_save, low_memory and, on the same page, map_weights and save_weights_for_mapping. max_partition_size is documented in the CompilationSettings reference and in Partitioning: Fixing the Engine Size, and use_shared_activation_scratch in the ExecuTorch guide. Worked examples: low-CPU-memory compilation and compiling from mapped weights.

Fixes and Refinements

Rule-based mixed-precision autocast: reliability fixes and new options

Rule-based autocast (available since 2.10) runs a model in FP16 or BF16 and keeps in FP32 the operations you exclude and the ones an optional calibration run flags next to an out-of-range value, so you do not cast layers by hand. In 2.15 it compiles models that use list-returning operations such as chunk, split and unbind, or integer indices such as nn.Embedding, which failed under autocast before; saved programs no longer carry the calibration data; and it has two new opt-in settings.

Set enable_autocast=True and autocast_low_precision_type (torch.float16 or torch.bfloat16), exclude operations with autocast_excluded_ops, autocast_excluded_nodes or the new autocast_excluded_node_names, and let a calibration run over autocast_calibration_dataloader flag overflow-prone nodes (autocast_max_output_threshold, plus the opt-in reduction-depth check autocast_max_depth_of_reduction).

In FP16 the -1e9 mask fill in the model below becomes -inf, and multiplying it by a zero mask value gives NaN (observed on TensorRT 11.3). With calibration data, every node of this small model touches the -1e9 value, so all of them stay in FP32 and the result matches the FP32 model (the point is that the NaN is avoided); without it, compilation logs a warning and the FP16 engine returns NaN.

import torch
import torch.nn as nn
import torch_tensorrt
from torch.utils.data import DataLoader, TensorDataset


class Scores(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(16, 16)

    def forward(self, x, mask):
        # the -1e9 fill does not fit in FP16
        return torch.softmax(self.linear(x) + (1.0 - mask) * -1e9, dim=-1)


torch.manual_seed(0)
model = Scores().cuda().eval()
x = torch.randn(2, 8, 16, device="cuda")
mask = (torch.rand(2, 8, 16, device="cuda") > 0.3).float()
data = TensorDataset(torch.randn(8, 8, 16, device="cuda"), mask.repeat(4, 1, 1))

for loader in (DataLoader(data, batch_size=2), None):  # batches match the export
    compiled = torch_tensorrt.dynamo.compile(
        torch.export.export(model, (x, mask)),
        arg_inputs=[x, mask],
        enable_autocast=True,
        autocast_low_precision_type=torch.float16,
        autocast_calibration_dataloader=loader,
    )
    error = (compiled(x, mask).float() - model(x, mask)).abs().max().item()
    print(f"calibrated: {loader is not None}, max error vs FP32 eager: {error:.1e}")

Output:

calibrated: True, max error vs FP32 eager: 4.5e-08
calibrated: False, max error vs FP32 eager: nan

What changed in 2.15:

  • Calibration data stays out of saved programs (#4980, #4978). Saved programs and engine-cache entries no longer embed the calibration dataloader or its inputs (in one measured case a saved program was 2.96 MB with the calibration set embedded and 160 KB without), an unpicklable dataloader no longer fails compilation, and refit after loading reuses the recorded precision decisions.
  • Operator exclusions accept ATen op objects (#4980). {torch.ops.aten.softmax.int} used to fail with a pickling error; strings already worked.
  • List-returning operations such as chunk, split and unbind work (#4980, #4978). They used to fail with 'list' object has no attribute 'to'.
  • Integer indices, bool masks and factory operations are no longer cast (#4978). Models with nn.Embedding used to fail under autocast, and an integer comparison could give a silently wrong mask.
  • Missing calibration is reported (#4980). The output-range and reduction-depth checks need autocast_calibration_dataloader; without it they are skipped and compilation logs a warning.
  • New opt-in settings (#4995). autocast_excluded_node_names keeps nodes in FP32 by graph node name, and autocast_model_regions="reclassify" lets the rules choose precisions inside the model's own torch.autocast blocks; this overrides choices of the model's author, so check accuracy.
  • Calibration uses far less GPU memory (#4995). It keeps one maximum per node instead of every output, with no setting needed (peak extra GPU memory during calibration of four Linear + ReLU layers over 128 batches of 64 × 1024: 448.2 → 1.8 MiB).

Docs: Compile Mixed Precision models with Torch-TensorRT · CompilationSettings reference

Limits and requirements

  • Matching: autocast_excluded_nodes is a regex against module paths such as block.norm, and autocast_excluded_node_names a regex against graph node names; both use re.match, so they must match from the start of the name (norm does not match block.norm, .*norm$ does)
  • The checks: the output-range check flags a node when one of its inputs, or the output of a node that consumes it, exceeded autocast_max_output_threshold (default 512), so a node that itself produces a value beyond the FP16 range is only protected when one of its inputs or consumers is also out of range: check the output for inf and NaN. The reduction-depth check only handles matmul, mm, bmm, dot, mv, attention and convolution nodes and cannot see model inputs or weights, so it never flags an nn.Linear (aten.linear) or an nn.Conv layer.
  • Calibration batches must be tuples or lists with one tensor per positional input, on the model's device, with exactly the exported shapes (use drop_last=True).
  • Autocast applies on the exported-program path (torch_tensorrt.dynamo.compile, torch_tensorrt.compile(ir="dynamo"), refit); torch.compile(backend="torch_tensorrt") ignores it, and models with torch.cond are rejected. MutableTorchTensorRTModule.save still stores the dataloader.

Tensor-parallel Llama and Qwen export with static KV cache

This release adds an end-to-end multi-GPU workflow for exporting, saving, and loading tensor-parallel Llama and Qwen models as one serialized TensorRT engine per rank. The example supports FP16 and FP32 execution, KV caches, dynamic sequence profiles, and preserves symbolic bounds for scalar cache indices across save/load. Reproducible prefill/decode benchmark support is provided.

What's Changed

New Contributors

Full Changelog: v2.14.0...v2.15.0

Contributors

shoumikhin, huydhn, and 18 other contributors
Assets 29
Loading