Unable to use kwarg_inputs in torch_tensorrt.compile

[derMart]

I am trying to compile a CLIP text encoder model using:

pipe.text_encoder = torch_tensorrt.compile(
    pipe.text_encoder,
    ir="dynamo",
    inputs=[torch.zeros(1, 77, dtype=torch.int32, device="cuda")],
)

This compiles fine, but the model forward function later fails:

Traceback (most recent call last):
    image = pipe(
  File "\lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "\lib\site-packages\diffusers\pipelines\stable_diffusion_xl\pipeline_stable_diffusion_xl.py", line 1086, in __call__
    ) = self.encode_prompt(
  File "\lib\site-packages\diffusers\pipelines\stable_diffusion_xl\pipeline_stable_diffusion_xl.py", line 406, in encode_prompt
    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
  File "\lib\site-packages\torch\fx\graph_module.py", line 822, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "\lib\site-packages\torch\fx\graph_module.py", line 400, in __call__
    raise e
  File "\lib\site-packages\torch\fx\graph_module.py", line 387, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "\lib\site-packages\torch\nn\modules\module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "\lib\site-packages\torch\nn\modules\module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() got an unexpected keyword argument 'output_hidden_states'

So the text_encoder forward function needs to be compiled with the output_hidden_states as kwarg.
The torch_tensorrt.compile function seems to have a parameter for this (see also https://pytorch.org/TensorRT/py_api/torch_tensorrt.html?highlight=compile):

kwarg_inputs (dict[Any, ...]) – Optional, kwarg inputs to the module forward function.

Which is of type dict[Any,Any]. However if I supply this to compile:

pipe.text_encoder = torch_tensorrt.compile(
    pipe.text_encoder,
    ir="dynamo",
    inputs=[torch.zeros(1, 77, dtype=torch.int32, device="cuda")],
    kwarg_inputs=dict(output_hidden_states=True)
)

I get the following error at compile time:

Traceback (most recent call last):
    pipe.text_encoder = torch_tensorrt.compile(
  File "lib\site-packages\torch_tensorrt\_compile.py", line 284, in compile
    torchtrt_kwarg_inputs = prepare_inputs(kwarg_inputs)
  File "lib\site-packages\torch_tensorrt\dynamo\utils.py", line 272, in prepare_inputs
    torchtrt_input = prepare_inputs(
  File "lib\site-packages\torch_tensorrt\dynamo\utils.py", line 280, in prepare_inputs
    raise ValueError(
ValueError: Invalid input type <class 'bool'> encountered in the dynamo_compile input parsing. Allowed input types: {torch_tensorrt.Input, torch.Tensor, list, tuple, dict}

So obviously kwarg_inputs is not of type dict[Any,Any] but more restricted which does not make any sense to me.

How to use kwarg_inputs? Or is this a bug?

best

[narendasan]

Can you open a bug for this?

[derMart]

Done here #3377 :)

[cehongwang]

Hi, the correct usage of kwarg should be something similar to this:

kwarg_inputs={
        "output_hidden_states": torch.zeros(**someshape**, dtype=torch.int32, device="cuda"), 
        "other_argument_of_forward_function": **other_kwarg_sample_inputs**,
        ...
}

It might be a bit tricky to get all inputs in the forward function in this scenario. You can also try torch.compile, which is JIT compile and does not need sample input.
https://github.com/pytorch/TensorRT/blob/main/examples/dynamo/torch_compile_stable_diffusion.py

[derMart]

I am confused by what you mean with "something similar to this". I would think

kwarg_inputs={
        "output_hidden_states": True, 
        "other_argument_of_forward_function": True,
        ...
}

is "similar" to the example you showed, but it will bail. As it seems to be a bug, that providing only tensor values at the leaf nodes of kwarg_inputs will compile fine, I do not know how to interpret your answer.
For the record: torch.compile is not a solution sadly (I tried it), as the result is not serializable to disk. It may be worth for testing purposes, but for anything deployment or production related it is sadly of no use therefore.

[cehongwang]

What I mean is that you need to include everything the model takes in the forward function in the kwarg_inputs. If a kwarg input is a tensor, you give it a tensor of the same shape. If the kwarg has a bool, you can try torch.tensor(True)

You can also try MutableTorchTensorRTModule. That handles the sample input for you and supports extra functionalities such as weight updates.

TensorRT/examples/dynamo/mutable_torchtrt_module_example.py

Line 76 in 54e36db

from diffusers import DiffusionPipeline

This includes an example of compiling the unet of stable diffusion pipeline.