Torch-TensorRT AOT Backend Guide

The Torch-TensorRT AOT (Ahead-Of-Time) backend exports models with torch.export.export(), compiles the exported program with torch_tensorrt.dynamo.compile(), and saves the compiled model for later use. This approach is ideal for production deployments where compilation happens once during tuning.

Overview

• AOT Compilation: Compiles during tuning, not at runtime
• Model Persistence: Compiled model is saved and loaded
• Fast Startup: No compilation overhead at inference time
• Production Ready: Deterministic performance
• Dynamo Export Path: Uses torch.export and the Torch-TensorRT Dynamo frontend

Quick Start

1
from aitune.torch.backend import TorchTensorRTAotBackend, TorchTensorRTAotBackendConfig
2
from torch_tensorrt.dynamo import CompilationSettings
3
import aitune.torch as ait
4
5
# Configure backend
6
config = TorchTensorRTAotBackendConfig(
7
    compile_config=CompilationSettings(),
8
)
9
backend = TorchTensorRTAotBackend(config)
10
11
# Use in tuning
12
from aitune.torch.tune_strategy import OneBackendStrategy
13
strategy = OneBackendStrategy(backend=backend)
14
15
model = ait.Module(model, "my-model", strategy=strategy)
16
ait.tune(model, input_data)
17
18
# Save compiled model
19
ait.save(model, "model.ait")
20
21
# Later: Load and use
22
ait.load(model, "model.ait")

Configuration Options

TorchTensorRTAotBackendConfig

1
@dataclass
2
class TorchTensorRTAotBackendConfig(BackendConfig):
3
    compile_config: TorchTensorRTConfig
4
    pickle_protocol: int = 5

compile_config

Compilation settings:

1
from torch_tensorrt.dynamo import CompilationSettings
2
3
config = TorchTensorRTAotBackendConfig(
4
    compile_config=CompilationSettings(
5
        workspace_size=1 << 30,
6
    )
7
)

By default, the engine matches the model’s loaded dtype. To request FP16 kernels via the legacy weak-typing path (deprecated in TensorRT 10.12), pair enabled_precisions with use_explicit_typing=False:

1
config = TorchTensorRTAotBackendConfig(
2
    compile_config=CompilationSettings(
3
        enabled_precisions={torch.float16},
4
        use_explicit_typing=False,
5
    )
6
)

pickle_protocol

Protocol for saving compiled model:

1
config = TorchTensorRTAotBackendConfig(
2
    pickle_protocol=5,  # Default
3
)

AOT vs JIT Comparison

For a detailed explanation of JIT vs AOT backends, see the JIT vs AOT Torch-TensorRT section.

Best Practices

1. Use Representative Data: Tune with inputs that match production shapes and dtypes
2. Verify After Load: Test loaded model before deployment
3. Version Control: Track both source code and .ait files
4. GPU Compatibility: Compile on the same or a compatible GPU as deployment

Troubleshooting

Issue: Load fails on different GPU

Cause: Engine compiled for different GPU architecture.

Solution: Recompile on target GPU or use hardware compatibility level in TensorRT backend.

Next Steps

• Learn about TorchTensorRT JIT Backend
• Explore Deployment Guide
• Review TensorRT Backend for more options