❓ [Question] Examples not working in nvcr.io/nvidia/pytorch:23.09-py3.

[sayakpaul]

❓ Question

I am within the nvcr.io/nvidia/pytorch:23.09-py3 container. Trying out some snippets from:
https://youtu.be/eGDMJ3MY4zk?si=MhkbgwAPVQSFZEha.

Both JIT and AoT examples failed. For JIT, it complained that "tensorrt" backend isn't available, for AoT, it complained that "The user code is using a feature we don't support. Please try torchdynamo.explain() to get possible the reasons".

I am on an A100. What's going on?

[gs-olive]

Hello - for the backend="tensorrt" issue, please ensure the line import torch_tensorrt is in the script (even though it may not be used directly). This is important because the import itself registers the backend, as here:

TensorRT/py/torch_tensorrt/dynamo/backend/backends.py

Line 27 in 0112bb4

@td.register_backend(name="tensorrt") # type: ignore[misc]

If this does not address the issue, could you please also share the version of Torch-TensorRT being used?

Regarding the AoT approach, @peri044 may be able to help further, but based on the error it sounds like the model may not be natively trace-able. Which specific example from the presentation are you referencing for the AoT sample?

[sayakpaul]

This one:

[peri044]

Hello @sayakpaul , I would suggest you to try our latest main branch. 23.09 is old now. Let us know if you see any issues with the main branch.
Alternatively, you can try our nightly container if you don't want to build from source : https://github.com/pytorch/TensorRT/pkgs/container/tensorrt%2Ftorch_tensorrt

[sayakpaul]

With the latest container (https://github.com/pytorch/TensorRT/pkgs/container/tensorrt%2Ftorch_tensorrt), the JIT workflow works. Have a few questions. So, let me use this thread to ask them. Maybe the community might find them to be useful.

[sayakpaul]

I used the following code snippet I am using to benchmark the JIT workflow:

import torch
import torch_tensorrt
import torch.utils.benchmark as benchmark

class MyModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = torch.nn.Conv2d(64, 128, 3)
    def forward(self, a):
        if a.sum() > 3:
            b = torch.nn.functional.elu(a)
        else:
            b = a + 1
        return self.conv(b)


# Taken from
# https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html
def benchmark_fn(f, *args, **kwargs):
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
    )
    return t0.blocked_autorange().mean * 1e6


model = MyModel().eval().cuda()
inputs = torch.nn.randn(32, 64, 64, 64, device="cuda")

print("Using regular model:")
print(benchmark_fn(model, inputs))


optimized_model = torch.compile(
    model,
    backend="tensorrt",
    dynamic=False,
    options={
        "debug": True,
        "enabled_precisions": {torch.half},
        "min_block_size": 1,
    }
)
_ = optimized_model(inputs)
print("Using optimized model:")
print(benchmark_fn(optimized_model, inputs))

I am getting:

620.6526302793078 microseconds (for the non-compiled model)
2095.269737765193 (for the compiled model)

I am on an 80 GB A100. I can understand the compute load might not be enough for me to see evident speedups here, but is this expected? Wanted to know your opinion.

[sayakpaul]

The AOT workflow still fails:

import torch
import torch_tensorrt
import torch.utils.benchmark as benchmark

class MyModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = torch.nn.Conv2d(64, 128, 3)
    def forward(self, a):
        if a.sum() > 3:
            b = torch.nn.functional.elu(a)
        else:
            b = a + 1
        return self.conv(b)


# Taken from
# https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html
def benchmark_fn(f, *args, **kwargs):
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
    )
    return t0.blocked_autorange().mean * 1e6


inputs = torch.randn((256, 64, 32, 32)).cuda()
model = MyModel().eval().cuda()

print("Using regular model:")
print(benchmark_fn(model, inputs))

exp_program = torch_tensorrt.dynamo.trace(model, [inputs])
trt_model = torch_tensorrt.dynamo.compile(exp_program, inputs=[inputs])

_ = trt_model(input)
print("Using optimized model:")
print(benchmark_fn(optimized_model, inputs))

print("Serializing...")
trt_ser_model = torch_tensorrt.dynamo.serialize(trt_model, *inputs)
torch.save(trt_ser_model, "trt_model.pt")

Trace:

Traceback (most recent call last):
  File "/opt/torch_tensorrt/trt_aot.py", line 33, in <module>
    exp_program = torch_tensorrt.dynamo.trace(model, [inputs])
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch_tensorrt/dynamo/aten_tracer.py", line 51, in trace
    if input.shape_mode == Input._ShapeMode.DYNAMIC:
AttributeError: 'Tensor' object has no attribute 'shape_mode'

[sayakpaul]

When I swapped the inputs with the following:

inputs = [torch_tensorrt.Input(
    min_shape=(1, 512, 1, 1),
    opt_shape=(4, 512, 1, 1),
    max_shape=(8, 512, 1, 1)
)]

it led to the following:

Traceback (most recent call last):
  File "/opt/torch_tensorrt/trt_aot.py", line 37, in <module>
    exp_program = torch_tensorrt.dynamo.trace(model, inputs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch_tensorrt/dynamo/aten_tracer.py", line 86, in trace
    exp_program = export(model, tuple(trace_inputs), constraints=constraints)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_export/__init__.py", line 556, in export
    return _export(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_export/__init__.py", line 596, in _export
    gm_torch_level = _export_to_torch_ir(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_export/__init__.py", line 517, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1226, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 410, in _fn
    return fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 558, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 148, in _fn
    return fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 402, in _convert_frame_assert
    return _compile(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 610, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 221, in time_wrapper
    r = func(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 527, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 497, in transform
    tracer.run()
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2127, in run
    super().run()
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 751, in run
    and self.step()
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 714, in step
    getattr(self, inst.opname)(inst)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 382, in inner
    raise exc.UserError(
torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

from user code:
   File "/opt/torch_tensorrt/trt_aot.py", line 10, in forward
    if a.sum() > 3:

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

I changed the code accordingly as well:

exp_program = torch_tensorrt.dynamo.trace(model, inputs)
trt_model = torch_tensorrt.dynamo.compile(exp_program, inputs=inputs)

[gs-olive]

Hi @sayakpaul - thank you for the comments. Regarding the JIT workflow, this specific example was intended to illustrate how graph breaks are handled in torch.compile with the Torch-TRT backend. As you suggested, the model itself does not have sufficient operations to make the overhead of using a TRT engine worthwhile here. Additionally, this graph has an intentionally-added control-flow break, which means that a minimum of two TRT engines will be generated, further increasing the overhead for this small model. Since this model only has one computational operator per TRT block, the default options would not have converted these operators to TRT at all ("min_block_size": 1 overrides this default). As a result, I believe the metrics are roughly as expected here.

The AOT workflow is not expected to work for this specific model (the one with the conditional). It would only work for the model without the conditional, as here. This is because the Torch ATen tracer cannot handle this sort of Python control flow. In order to use the AOT workflow on a model like this, the code would need to be modified to either remove the conditional or use tracer-allowed conditionals, like torch.cond

[sayakpaul]

Will try and get back. Thanks for your inputs!

Could you maybe provide an example for JIT where the improvements are evident? FWIW, we did try something with SD but the benefits were not there: huggingface/diffusers#5564.

Probably we are missing something out?

[gs-olive]

Thanks for the follow-up! I will look into that sample you provided a bit more - it looks like the TRT conversion errored out for some reason, which would cause compilation to fall back to Torch eager. One other to try might be the Stable Diffusion version highlighted here.

[gs-olive]

Also, regarding the error cited in huggingface/diffusers#5564, I retried your sample on the latest nightly version of the repository, with the following small code modifications, and it is working:

    elif run_compile and with_tensorrt:
        print("Run torch compile with TensorRT backend")
        pipe.unet = torch.compile(
            pipe.unet, fullgraph=True, backend="tensorrt", dynamic=False,
            options={"min_block_size": 1,
                     "truncate_long_and_double": True,
                     "enabled_precisions": {torch.half}}
        )

Please let me know if this resolves the issue for you as well.

[sayakpaul]

I retried your sample on the latest nightly version of the repository

Don't want to build from source. Will https://github.com/pytorch/TensorRT/pkgs/container/tensorrt%2Ftorch_tensorrt work here? (I will prune the previous container before mounting)

[gs-olive]

Yes, I believe the Docker container should also work. You could alternatively also use the following for installing the latest nightly distribution of Torch-TRT:

pip install --pre torch torch-tensorrt --extra-index-url https://download.pytorch.org/whl/nightly/cu121