🐛 [Bug] Compilation Error on GPT-2

[gs-olive]

Bug Description

When converting the GPT-2 network (https://huggingface.co/gpt2) from TorchScript to Torch-TRT, the following error is encountered:

compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
RuntimeError: [Error thrown at core/partitioning/shape_analysis.cpp:167] Unsupported input data type unsigned char

To Reproduce

Steps to reproduce the behavior:

1. Run torch_tensorrt.compile with GPT-2 model as input, using fp32 precision.
2. Choose fixed input size of [1, 128] and enable truncate_long_and_double with 12 GB workspace.
3. Pass in model keyword args to disable attention and hidden state outputs

Expected behavior

Model should successfully compile to Torch-TRT. Specifically, internal (non-user-provided) type-casting issues should not cause errors.

Environment

• Torch-TensorRT Version: 1.3.0a0+e3b99294
• PyTorch Version: 1.13.0.dev20220921+cu116
• CPU Architecture: Intel Xeon CPU
• OS: Ubuntu 20.04
• How you installed PyTorch: pip
• Build command you used: python setup.py develop
• Are you using local sources or building from archives: local
• Python version: 3.8.13
• CUDA version: 11.6

Additional context

The problematic data in GPT-2 seems to be this bias term, instantiated in the attention module, which has type uint8. In both the TorchScript IR and the model code (example 1, example 2), it seems that this bias term is generally cast to a bool. The error is thrown in this code segment:

TensorRT/core/partitioning/shape_analysis.cpp

Lines 258 to 260 in e3b9929

	c10::optional<nvinfer1::DataType> dtype = util::optTypeMetaToTRTDataType(cur_ivalue.toTensor().dtype());
	if (dtype == c10::nullopt) {
	TORCHTRT_THROW_ERROR("Unsupported input data type " << cur_ivalue.toTensor().dtype());

The conversion of a uint8 type to a TRT Data Type fails, however simply patching this conversion also does not fix the issue, as an out-of-bounds error later follows.

Temporary Solution

A temporary fix to this problem is to add the following to the compilation arguments in torch_tensorrt.compile:

torch_tensorrt.compile( ..., torch_executed_ops=["aten::where"], ...)

This solution works as it happens to exclude the code which uses and processes the uint8 tensor, however it is only a temporary fix and does not resolve the underlying issue.

Steps to a Solution

• Fix mismatched dimension issue in aten::where
• Make at::kByte a valid input type

[eric8607242]

Hi, I face the same issues when I try to compile the GPT2 into the tensorrt format.

The corresponding environments:

ubuntu 22.04

Driver Version: 515.65.01
CUDA Version: 11.7
GPU: rtx 3090

torch == 1.13.1+cu117
torch_tensorrt == 1.3.0
transformer == 4.25.1

The sample code to produce the error:

tokens=tokenizer('The cat is on the table.', return_tensors='pt')
for k in tokens.keys():
    tokens[k] = tokens[k].cuda()

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
gpt2_model = GPT2LMHeadModel.from_pretrained('gpt2', return_dict=False, torchscript=True)
gpt2_model.eval()
gpt2_model.cuda()

traced_gpt2_model = torch.jit.trace(gpt2_model, tokens["input_ids"])

trt_gpt2_model = torch_tensorrt.compile(traced_gpt2_model, 
    inputs= [torch_tensorrt.Input(shape=[batch_size, 7], dtype=torch.int32)], # Input ids
    enabled_precisions= {torch.float32}, # Run with 32-bit precision
    workspace_size=2000000000,
    truncate_long_and_double=True
)

I try to compile with the temporary solution as:

trt_gpt2_model = torch_tensorrt.compile(traced_gpt2_model, 
    inputs= [torch_tensorrt.Input(shape=[batch_size, 7], dtype=torch.int32)], # Input ids
    enabled_precisions= {torch.float32}, # Run with 32-bit precision
    workspace_size=2000000000,
    truncate_long_and_double=True,
    torch_executed_ops=["aten::where"]
)

However, the other issue is raised as:

RuntimeError: forward() Expected a value of type 'Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]' for argument '1' but instead found type 'Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]'.
Position: 2
Declaration: forward.forward(__torch__.module self_1, Tensor 0, ((Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor)) 1) -> ((Tensor, ((Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor), (Tensor, Tensor))) 2)
 [Torch-TensorRT] - Truncating graph input type from at::kLong to at::kInt
WARNING: [Torch-TensorRT] - Truncating graph input type from at::kLong to at::kInt

Is there any error in my sample code?
And when this commit will be released?

Big thanks!

[gs-olive]

Hello - I am able to reproduce the issue you are describing. In the initial test, I had also added the argument: use_cache=False in the GPT2LMHeadModel.from_pretrained(...) call. With this addition, the compilation succeeds on my machine (using the main branch of Torch-TensorRT) - please let me know if this also resolves the issue on your end.