aitune.torch.backend.torch_tensorrt_aot_backend

TorchTensorRT backend with AOT compilation and intermediate model save.

Module Contents

Classes

Name	Description
`TorchTensorRTAotBackend`	Backend that compiles model using TensorRT.
`TorchTensorRTAotBackendConfig`	Configuration for TorchTensorRTAotBackend.
`TorchTensorRTAotBuildStep`	Identifiers for discrete sub-steps of a TorchTensorRTAot backend build.
`TorchTensorRTConfig`	Fallback settings matching upstream defaults when torch_tensorrt is unavailable.

Functions

Name	Description
`_is_primitive`	Check if a value is a primitive type (int, float, str, bool).
`assert_torch_tensorrt`	Check if torch_tensorrt is installed.

Data

logger

API

class aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend(
    config: aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig | None = None
)

Bases: Backend

Backend that compiles model using TensorRT.

EXPORTED_MODEL_PATH_KEY

= 'exported_model_path'

STATE_DEVICE

= 'device'

STATE_TYPE

= 'type'

_config

= config or TorchTensorRTAotBackendConfig()

_devices

list[str] = ['cuda']

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._activate()

Load compiled module.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Build the model with TensorRT.

Exports the module via torch.export.export with shared Dim instances across inputs, then hands the resulting ExportedProgram to torch_tensorrt.dynamo.compile. This replaces the previous flow that called torch_tensorrt.compile and let it run its own export, which built one independent Dim per Input and raised ConstraintViolationError whenever two inputs had to share a runtime axis (e.g. HF encoders’ shared batch on input_ids and attention_mask).

Workaround pending a dynamic_shapes passthrough on torch_tensorrt.compile; revert to torch_tensorrt.compile(...) once that lands upstream.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._create_exported_model_path(
    cache_dir: pathlib.Path
) -> pathlib.Path

Create path to exported model.

In order to avoid name clashes we create a unique path for each model and graph spec.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._deactivate()

Deactivate backend and cleanup.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._deploy()

Deploys the backend.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run inference with TensorRT engine.

Parameters:

*args

AnyDefaults to ()

Inference arguments

**kwargs

AnyDefaults to {}

Inference keyword arguments

Returns: Any

Model outputs

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
) -> aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackend.to_dict() -> dict

Returns the state_dict of the backend.

class aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig(
    compile_config: aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTConfig = TorchTensorRTConfig(),
    pickle_protocol: int = 5
)

Dataclass

Bases: BackendConfig

Configuration for TorchTensorRTAotBackend.

See torch_tensorrt/dynamo/_settings.py CompilationSettings for compile_config(TorchTensorRTConfig)

compile_config

TorchTensorRTConfig = field(default_factory=TorchTensorRTConfig)

pickle_protocol

int = 5

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig.describe() -> str

Describe the backend configuration. Display only changed fields.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig.from_dict(
    data: dict
) -> aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig

classmethod

Initialise config from a plain dict (e.g. parsed from YAML).

compile_config may be passed as a nested dict and will be reconstructed into a TorchTensorRTConfig instance automatically.

aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBackendConfig.to_dict()

Saves the backend configuration to a file.

class aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTAotBuildStep()

Bases: BackendBuildStep

Identifiers for discrete sub-steps of a TorchTensorRTAot backend build.

COMPILED_MODEL_SAVE

= 'Compiled model save'

TORCHTRT_COMPILE

= 'Torch-TensorRT compile'

TORCH_EXPORT

= 'Torch export'

class aitune.torch.backend.torch_tensorrt_aot_backend.TorchTensorRTConfig(
    enabled_precisions: set[torch.dtype] | None = None,
    use_explicit_typing: bool = True,
    workspace_size: int = 0
)

Dataclass

Fallback settings matching upstream defaults when torch_tensorrt is unavailable.

enabled_precisions

set[dtype] | None = None

use_explicit_typing

bool = True

workspace_size

int = 0

aitune.torch.backend.torch_tensorrt_aot_backend._is_primitive(
    value: typing.Any
) -> bool

Check if a value is a primitive type (int, float, str, bool).

aitune.torch.backend.torch_tensorrt_aot_backend.assert_torch_tensorrt()

Check if torch_tensorrt is installed.

aitune.torch.backend.torch_tensorrt_aot_backend.logger = getLogger(__name__)