torch_tensorrt.executorch#

ExecuTorch compilation and export integration.

Runtime loading is provided by the optional torch-tensorrt-executorch-runtime distribution and dispatched through torch_tensorrt.load(..., format="executorch").

Note

This module requires ExecuTorch and is only supported on Linux.

Overview#

The executorch module lowers exported programs into an ExecuTorch Edge program with TensorRT engines delegated to the TensorRT backend. Use it when you want to inspect or customize the Edge program before serialization, or when a single .pte has to carry more than one method.

import torch_tensorrt.executorch

edge = torch_tensorrt.executorch.export({"forward": exported_program})
program = edge.to_executorch()

with open("model.pte", "wb") as f:
    program.write_to_file(f)

For the simpler case where no inspection is needed, torch_tensorrt.save(..., output_format="executorch") writes the file in one call.

Functions#

torch_tensorrt.executorch.export(source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], *, inputs: Sequence[Any] | None = None, arg_inputs: Sequence[Any] | None = None, kwarg_inputs: Mapping[str, Any] | None = None, dynamic_shapes: Any | None = None, retrace: bool | None = None, transform_passes: Any | None = None, partitioners: Sequence['Partitioner'] | Mapping[str, Sequence['Partitioner'] | None] | None = None, compile_specs: Sequence['CompileSpec'] | Mapping[str, Sequence['CompileSpec'] | None] | None = None, compile_config: EdgeCompileConfig | None = None, constant_methods: Mapping[str, Any] | None = None, generate_etrecord: bool = False, weight_streaming_budget_per_engine: int | None = None) → EdgeProgramManager[source]#

Prepare TensorRT-compiled programs for composable ExecuTorch lowering.

TensorRT claims engine nodes first. Additional partitioners run afterward in caller-provided order. The returned EdgeProgramManager is the standard ExecuTorch inspection and customization boundary; call to_executorch() on it when ready to perform final memory planning and serialization.

Export stages independent graph and metadata containers while sharing tensor payload storage with the source programs, so transform passes must treat a shared payload value as immutable. Engines are never deep copied, since that would serialize and deserialize them, but each one is decoded into a byte buffer the returned program owns, so the bytes of a multi-gigabyte engine are resident twice while both programs are alive. Method mappings preserve independent entry points but do not imply shared mutable state between them.

When exporting more than one method, give each method its own partitioner instances via partitioners={"method": [...]}. A partitioner may carry method-specific state. Give each instance the compile spec for the method it serves, since a backend that reads its method name from the specs cannot find it otherwise. Sharing one instance across methods is rejected when its specs name a method, because every method sharing it would be tagged with the same name. Sharing an instance whose specs name no method is not rejected here, but a backend that reads its own method name from its specs, such as the CUDA backend, then raises during lowering.

generate_etrecord=True is outside the payload sharing described above. It makes ExecuTorch deep copy the whole program, so peak memory grows by roughly the size of the program including engines.

Each engine is serialized once, into base64 text. Every method is serialized before the first one is rewritten, and a method’s text is released once that method is rewritten. On top of the engine bytes the exported program itself carries, peak memory adds at most roughly 1.33x the engine bytes of every method together, which is the base64 text, plus about 2.3x the bytes of the largest engine, which is what decoding one method’s text holds at once. For methods whose engines are the same size that is about 3.7x for one method, 2.5x for two, and 1.9x for four.

constant_methods keys are restricted to valid Python identifiers here, which is narrower than ExecuTorch itself accepts.

Parameters:

source (Union(torch.export.ExportedProgram, torch.fx.GraphModule, Dict[str, torch.export.ExportedProgram])) – A TensorRT-compiled source. A GraphModule is exported first, and a mapping becomes one method per key. An ExportedProgram and a mapping both need their engines already compiled, so inputs, arg_inputs, kwarg_inputs, dynamic_shapes and retrace=True are rejected for them.

Keyword Arguments:

inputs (Sequence[Union(Input, torch.Tensor)]) – Example positional inputs used to export a GraphModule source. Mutually exclusive with arg_inputs, which is the same argument under the newer name.
arg_inputs (Sequence[Union(Input, torch.Tensor)]) – See inputs.
kwarg_inputs (Dict[str, Any]) – Example keyword inputs used to export a GraphModule source.
dynamic_shapes (Any) – Dynamic shape specification passed through to torch.export. Inferred from any torch_tensorrt.Input in the example inputs when omitted, which requires every input to be an Input.
retrace (bool) – Re-trace a GraphModule source instead of wrapping its existing graph. Unset behaves as False, and True requires example inputs.
transform_passes (Union(Sequence[Any], Dict[str, Sequence[Any]], PassManager)) – ExecuTorch transform passes, either for every method or per method.
partitioners (Union(Sequence[Partitioner], Dict[str, Sequence[Partitioner]])) – Extra partitioners to run after the TensorRT one, either for every method or per method. Give each method its own instances.
compile_specs (Union(Sequence[CompileSpec], Dict[str, Sequence[CompileSpec]])) – Compile specs for the TensorRT partitioner, either for every method or per method.
compile_config (executorch.exir.EdgeCompileConfig) – Edge compile config. Defaults to get_edge_compile_config().
constant_methods (Dict[str, Any]) – Methods returning a constant, such as a vocab size. Keys must be valid Python identifiers and must not name a method of source.
generate_etrecord (bool) – Ask ExecuTorch for an ETRecord for later debugging. This copies the whole program, engines included.
weight_streaming_budget_per_engine (Optional[python:int]) – Bytes of engine weights that may stay resident in GPU memory, with the rest streamed from host memory. It applies to each TensorRT engine separately, not as a total for the program, so a program with N engines can hold up to N times this value resident: every delegate is initialized when its method loads and they stay resident together. Requires the engine to have been compiled with enable_weight_streaming=True; it is ignored, with a log message, on an engine that cannot stream. Leave it unset, the default, so each engine gets TensorRT’s own automatic budget, sized against free memory on the device it actually loads on. This is a different unit from torch_tensorrt.runtime.weight_streaming(...).device_budget, which is a program total split proportionally across engines.

Returns:

The Edge program, ready for inspection, further transformation, or to_executorch().

Return type:

executorch.exir.EdgeProgramManager

torch_tensorrt.executorch.get_edge_compile_config() → EdgeCompileConfig[source]#: Return the EdgeCompileConfig used for Torch-TensorRT ExecuTorch export.

Classes#

class torch_tensorrt.executorch.TensorRTPartitioner(compile_specs: List[CompileSpec] | None = None)[source]#

Partitions the graph for TensorRT delegation.

Only nodes that are torch.ops.tensorrt.execute_engine are supported; each such node becomes its own partition so the backend can serialize the engine to the same format as the TRT runtime.

If compile_specs does not already contain a target_device entry, one defaulting to cuda:0 is auto-appended (mirroring CudaPartitioner). Callers targeting a non-default GPU should pre-populate compile_specs with the desired CompileSpec("target_device", b"cuda:<index>") to override the default.

Note: target_device is AOT metadata only – it drives ExecuTorch’s PropagateDevicePass tagging at export time. At runtime the C++ backend selects the GPU from the device baked into the serialized engine blob, not from this value.

final class torch_tensorrt.executorch.TensorRTBackend[source]#

Backend that serializes TensorRT engines for the native ExecuTorch runtime.

The partition contains a single execute_engine node; we extract the engine and metadata and encode them as a standalone blob. The C++ runtime backend parses that blob directly without the legacy Torch-TensorRT C++ runtime.