[Umt5] Add google's umt5 to transformers

[ArthurZucker]

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:

!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:

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

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[stefan-it]

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 🤔

/cc @agemagician

[ArthurZucker]

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

[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

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.

[ArthurZucker]

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.

[ArthurZucker]

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!

[ArthurZucker]

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

[sgugger]

Thanks for adding the modeling file. Have a couple more nits.

[ArthurZucker]

Currently setting up an instance ton convert an upload the xxl model, other models are available here