Wrong outputs of converted mt5 models

Description

I'm trying to convert a mt5 model to TensorRT. I adapted the T5 demo notebook of the main branch, however, the outputs of the TensorRT model are not what they should be. (The original T5 notebook works as expected. The environment is a docker container build from the main branch of this repository.)

To be able do experiment, I basically copied the demo/HuggingFace/T5 folder to demo/HuggingFace/MT5. The changes I did in there are the usage of the MT5 classes of huggingface (e.g. MT5Config instead of T5Config) and I added the correct parameters for the mt5 models to the MT5ModelTRTConfig.

As you can see at the bottom of the adapted notebook, the outputs of the TensorRT model are not as expected, resulting in generation of nonsense-text: mt5-small-sum.ipynb. Further above, I also verified that the exported ONNX model does output the same as the pytorch model, so the problem seems to be with TensorRT.

As the model code of mT5 is basically the same as for T5, I don't understand why the conversion does not seem to work correctly for the mT5 models (besides the public T-Systems-onsite/mt5-small-sum-de-en-v2 model used in the commited notebook, I also tried it with a private fine-tuned mt5-small model, which didn't work either).

I would be very thankful for any ideas why this could be happening or any tips how to debug this further.
Cheers!

cc @vinhngx @parthchadha @rajeevsrao

Environment

TensorRT Version: 8.2.0.6
NVIDIA GPU: RTX 3090
NVIDIA Driver Version: 470.57.02
CUDA Version: 11.4, V11.4.120
CUDNN Version: 8.2.4
Operating System: ubuntu 20.04
Python Version (if applicable): 3.8.10
Tensorflow Version (if applicable): 2.5.1
PyTorch Version (if applicable): 1.9.1+cu111 (cudnn 8005)
Baremetal or Container (if so, version): Docker container build from the main branch of this repository

Relevant Files

• https://github.com/tobigue/TensorRT/commits/demo-mt5

Steps To Reproduce

• Clone demo-mt5 branch of https://github.com/tobigue/TensorRT
• Run the notebook demo/HuggingFace/notebooks/mt5-small-sum.ipynb

Start by running demo/HuggingFace/notebooks/mt5-small-sum.ipynb from the demo-mt5 branch and compare its PyTorch, ONNX, and TensorRT outputs. Read demo/HuggingFace/MT5/MT5ModelConfig.py alongside the adapted T5 demo and the linked mT5 model code. Done means identifying why TensorRT diverges and confirming that converted mT5 models produce the expected generated text.