[`Umt5`]  Add google's umt5 to `transformers`
Author: ArthurZucker
Pull request: yes

# What does this PR do?
Superseeds #22626 which has been stale for quite some time

A kaggle notebook for reproducing and running the originial model:
https://www.kaggle.com/arthurzucker/umt5-inference 

- Tokenizer is a BertGenerationTokenizer. Here is how to convert it: 

```python 
!wget https://storage.googleapis.com/t5-data/vocabs/umt5.256000/sentencepiece.model
from transformers import T5Tokenizer

umt5 = T5Tokenizer("/Users/arthurzucker/Work/transformers/sentencepiece.model")
to_add = []
for i in range(300):
    to_add.append(f"<extra_id_{i}>")
tokenizer.add_tokens(list(reversed(to_add)), True)  
```
84 tokens are free to use apparently.

- The modeling code is just MT5, addapted to have more relative bias that is not the same for all layers. 
For a first conversion I'll be using this:
```bash 
python src/transformers/models/t5/convert_t5x_checkpoint_to_pytorch.py --t5x_checkpoint_path "/Users/arthurzucker/Work/transformers/checkpoint_1000000" --config_file "/Users/arthurzucker/Work/transformers/checkpoint_1000000/config.json" --pytorch_dump_path ./ArthurZ --scalable_attention
```


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Hi @ArthurZucker , thanks for updating this! As far as we can tell, it is not just mT5, because of joined/separate key-value in attention. Was this problem solved in latest conversion script of this PR :thinking: 

/cc @agemagician 

The conversion went well, the outputs are still a bit gibberish but didn’t have problem of un matching shape. 
They mentioned that the model is closer to MT5, which is why if we we can have minimal changes, it will look like this + adapted conversion 

> The conversion went well, the outputs are still a bit gibberish but didn’t have problem of un matching shape. They mentioned that the model is closer to MT5, which is why if we we can have minimal changes, it will look like this + adapted conversion

So far, I can see you made similar changes as we did before, which led to gibberish output.
The one addition change you did is allowing fall back to byte for the tokenizer.

I belive the issue still exist because of the way we reshape and convert the q,k and v for the attention as @stefan-it mentioned.

There is also a different logic for `postion_bias` which seems to be missing. 
Joint 3D matrix vs what we have now can be linked to the new sharding scheme, It probably the last thing to check. 



Regarding the split / merge, I don't really see a problem with the code. The checkpoints are split, and the actual code is similar to mt5 with the difference being `scanning` I believe. However feel free to check and I hope we can get coherent outputs! 

Update, the outputs match 🔥 The issue was : the tokenizer

@ArthurZucker Awesome news! I will check downstream performance as well soon :hugs: 

Wait wait 😅 I have to update and push the tokenizers 😉 

> Update, the outputs match 🔥 The issue was : the tokenizer

"The outputs match" Do you mean you have tested both the original t5x inference pipeline against the converted transformer Pytorch version ?

Yes @agemagician. If you read the PR description there’s a link to the reproducing script for generating with the t5x repo

> Yes @agemagician. If you read the PR description there’s a link to the reproducing script for generating with the t5x repo

Awesome work :)

Currently setting up an instance ton convert an upload the `xxl` model, other models are available [here](https://huggingface.co/models?search=umt5)