Torch.compile dependancy on system architecture / environment

torch.compile
1

TLDR: Does the PyTorch-TensorRT compiler depend on the (a) underlying system architecture and/or (b) available system resources?

More details: I want to evaluate the performance of a compiled PyTorch model, for the TensorRT framework on the Jetson Orin Nano in Jetpack 6.2.

A minimal working example is avaliable in the README.md for pytorch/TensorRT:

import torch
import torch_tensorrt

model = MyModel().eval().cuda() # define your model here
x = torch.randn((1, 3, 224, 224)).cuda() # define what the inputs to the model will look like

optimized_model = torch.compile(model, backend="tensorrt")
optimized_model(x) # compiled on first run

optimized_model(x) # this will be fast!

I think that the underlying TensorRT engine runs an optimisation to search the compilation space for the ~best compilation when we run optimized_model(x) for the first time.

optimization_level (0–5, default 3) controls how long TRT spends searching for
faster kernel implementations. Higher values produce faster engines at the cost of
longer compile time.

Source: TensorRT/docs/_sources/user_guide/performance_tuning.rst.txt

In practise, this means for effective performative use of torch.compile() we should:

  • Be on the same device (or as close as possible) that we wish to deploy the compiled model
  • Ensure the working environment is the same (or as close as possible) to the deployed environment. As an example, the jetson_clocks and nvpmodel power state should be the same.

If the two above items are not respected, then the TensorRT compiler is NOT optimising for our target deployment scenario.

Is my understanding correct? And where can I find more information about this? I have read the (very helpful :sparkles:) pytorch/TensorRT repo and the NVIDIA TensorRT documentation but perhaps I have missed something?

2

This matches my understanding, yes.

CC @narendasan as the TorchTRT expert to chime in

3

Yes this is correct. You should use the deployment device in the power mode you will deploy when you compile to build engines tuned for that scenario. More information on controlling the determinism of the engine construction can be found here: Optimizing TensorRT Performance — NVIDIA TensorRT

4

Perfect, I understand now. Thanks!:sparkles:

As a gentle suggestion, it would be helpful to people like me if you could add your sentence to the Deterministic Tactic Selection section. Like this:

TensorRT runs through all the possible tactics in the engine-building phase and selects the fastest ones. Since the selection is based on the tactics’ latency measurements, TensorRT can select different tactics across different runs if some have similar latencies. You should use the deployment device in the power mode you will deploy when you compile to build engines tuned for that scenario.…