torch_tensorrt.save for aot_inductor format compatibility with PyTorch 2.10

[leimao]

Hello, I have a question. Is torch_tensorrt.save for aot_inductor format compatible with PyTorch 2.10 (say, the latest nightly build)? There is no unit test testing torch_tensorrt.save based on my observations. Thanks.

[narendasan]

Not sure what you mean by compatibility. Yes it should work as in you can create a AOTI PT2 package using torch_tensorrt.save with PyTorch 2.10.

Tests here:

TensorRT/tests/py/dynamo/models/test_export_serde.py

Line 661 in a80572d

def test_save_load_aoti(ir, tmp_path):

[leimao]

@narendasan Thank you. I am building and using Torch-TensorRT 2.9.0 on top of the latest PyTorch (2.10 nightly) via some custom method. So far everything seems to be working fine, including torch_tensorrt.compile. The resulted engine works fine when I executed it and compared with the reference outputs. However, when I tried to save it as AOTInductor engine (using the exact unit test you pointed to), I encountered the following error:

self = <torch._export.non_strict_utils._NonStrictTorchFunctionHandler object at 0x7f9c12566360>
func = <method 'data_ptr' of 'torch._C.TensorBase' objects>
types = (<class 'torch._subclasses.fake_tensor.FakeTensor'>,)
args = (FakeTensor(..., device='cuda:0', size=(1, 3, 224, 224)),), kwargs = {}

    def __torch_function__(self, func, types, args=(), kwargs=None):
        kwargs = kwargs or {}
        if torch.compiler.is_dynamo_compiling():
            return func(*args, **kwargs)
    
        if log.isEnabledFor(logging.DEBUG) and config.extended_debug_current_loc:
            frame = _find_user_code_frame()
            if frame is not None:
                log.debug(
                    "%s called at %s:%s in %s",
                    func.__qualname__,
                    frame.f_code.co_filename,
                    frame.f_lineno,
                    frame.f_code.co_name,
                )
    
        func, args, kwargs = self._override(func, args, kwargs)
        try:
>           return func(*args, **kwargs)
E           RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor). If you're using torch.compile/export/fx, it is likely that we are erroneously tracing into a custom kernel. To fix this, please wrap the custom kernel into an opaque custom op. Please see the following for details: https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html
E           
E           While executing %_run_on_acc_0 : [num_users=1] = call_module[target=_run_on_acc_0](args = (%x,), kwargs = {})
E           Original traceback:
E           None
E           Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)

I know it's not your responsibility to support this custom build. But any suggestion for fixing this is appreciated. Thank you very much.

[leimao]

The error seems to happen in this line:
https://github.com/pytorch/TensorRT/blob/v2.9.0/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py#L402
where a data pointer is being get from a (fake) tensor.

[narendasan]

You need to use the C++ based runtime from Torch-TensorRT (TorchTensorRTModule). The Python one does not fully support serialization

[leimao]

Thanks @narendasan. I still have a few noob questions.
There are two files, _PythonTorchTensorRTModule.py and _TorchTensorRTModule.py. The former one contains the PythonTorchTensorRTModule and the latter one contains the TorchTensorRTModule you referred to. Both runtime modules were included in the init.py. How do I select which runtime to use?

[leimao]

So it seems that Torch-TensorRT does runtime check if Torch-TensorRT runtime is available. Because I am doing my custom build setup using a different build toolchain, these things were broken?

FeatureSet = namedtuple(
    "FeatureSet",
    [
        "torchscript_frontend",
        "torch_tensorrt_runtime",
        "dynamo_frontend",
        "fx_frontend",
        "refit",
        "qdp_plugin",
        "windows_cross_compile",
        "tensorrt_rtx",
    ],
)

trtorch_dir = os.path.dirname(__file__)
linked_file = os.path.join(
    "lib", "torchtrt.dll" if sys.platform.startswith("win") else "libtorchtrt.so"
)
linked_file_runtime = os.path.join(
    "lib",
    (
        "torchtrt_runtime.dll"
        if sys.platform.startswith("win")
        else "libtorchtrt_runtime.so"
    ),
)
linked_file_full_path = os.path.join(trtorch_dir, linked_file)
linked_file_runtime_full_path = os.path.join(trtorch_dir, linked_file_runtime)

_TENSORRT_RTX = tensorrt._package_name == "tensorrt_rtx"
_TS_FE_AVAIL = os.path.isfile(linked_file_full_path)
_TORCHTRT_RT_AVAIL = _TS_FE_AVAIL or os.path.isfile(linked_file_runtime_full_path)
_DYNAMO_FE_AVAIL = version.parse(sanitized_torch_version()) >= version.parse("2.1.dev")
_FX_FE_AVAIL = False if _TENSORRT_RTX else True
_REFIT_AVAIL = True
_WINDOWS_CROSS_COMPILE = check_cross_compile_trt_win_lib()

[narendasan]

Theres a compile option call something like use_python_runtime. You must have a build with the C++ runtime. If you are using python only builds this will not be there

[leimao]

@narendasan Thank you very much.

I do have built a C++ runtime using the code from the core directory. This runtime library is built as a C++ extension dependency for the Python torch_tensorrt.

I did some hack by setting _TORCHTRT_RT_AVAIL = True and changed some torch extension registration logics (although I really don't want to change the source code), the unit test worked.