WIP Umt5
Author: agemagician
Pull request: yes

# What does this PR do?

It supports umt5 models, which need separate relative attention biases for each layer.
The current code will have backward compatibility with previous T5 and MT5 checkpoints.

Fixes # (issue)
https://github.com/huggingface/transformers/issues/22573

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


## Who can review?

Models:

- text models: @ArthurZucker and @younesbelkada 
- @stefan-it

Comments:
Many thanks @agemagician ! I could import the flax import for umT5 (and made some suggestions here) and will try to get the conversion script running (I think it only needs adjustments for this relative position embeddings thing).

cc @ArthurZucker 

So far the Pytorch version correctly uses "relative_attention_bias" across all layers. However, the flax version doesn't.
Any idea what could be the issue?

You can see this in line 224 on the flax file. I have even removed the if statement to force adding the "relative_attention_bias" layer but it doesn't 

@stefan-it  @sgugger @ArthurZucker

> So far the Pytorch version correctly uses "relative_attention_bias" across all layers. However, the flax version doesn't. Any idea what could be the issue?
> 
> You can see this in line 224 on the flax file. I have even removed the if statement to force adding the "relative_attention_bias" layer but it doesn't
> 
> @stefan-it @sgugger @ArthurZucker

To clarify, if I created a new flax model for UMT5, and printed the keys for the zero layer then:
```
test_flax_model.params["encoder"]["block"]["0"]["layer"]["0"]["SelfAttention"].keys()
dict_keys(['q', 'k', 'v', 'o', 'relative_attention_bias'])
```
but if I printed it for any follow layers:
```
test_flax_model.params["encoder"]["block"]["1"]["layer"]["0"]["SelfAttention"].keys()
dict_keys(['q', 'k', 'v', 'o'])
```

However, in pytorch it shows that all layers have a separate relative_attention_bias:
```
MT5ForConditionalGeneration(
  (shared): Embedding(256384, 512)
  (encoder): UMT5Stack(
    (embed_tokens): Embedding(256384, 512)
    (block): ModuleList(
      (0-7): 8 x UMT5Block(
        (layer): ModuleList(
          (0): UMT5LayerSelfAttention(
            (SelfAttention): UMT5Attention(
              (q): Linear(in_features=512, out_features=384, bias=False)
              (k): Linear(in_features=512, out_features=384, bias=False)
              (v): Linear(in_features=512, out_features=384, bias=False)
              (o): Linear(in_features=384, out_features=512, bias=False)
              (relative_attention_bias): Embedding(32, 6)
            )
            (layer_norm): UMT5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (1): UMT5LayerFF(
            (DenseReluDense): UMT5DenseGatedActDense(
              (wi_0): Linear(in_features=512, out_features=1024, bias=False)
              (wi_1): Linear(in_features=512, out_features=1024, bias=False)
              (wo): Linear(in_features=1024, out_features=512, bias=False)
              (dropout): Dropout(p=0.1, inplace=False)
              (act): NewGELUActivation()
            )
            (layer_norm): UMT5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
      )
    )
    (final_layer_norm): UMT5LayerNorm()
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (decoder): UMT5Stack(
    (embed_tokens): Embedding(256384, 512)
    (block): ModuleList(
      (0-7): 8 x UMT5Block(
        (layer): ModuleList(
          (0): UMT5LayerSelfAttention(
            (SelfAttention): UMT5Attention(
              (q): Linear(in_features=512, out_features=384, bias=False)
              (k): Linear(in_features=512, out_features=384, bias=False)
              (v): Linear(in_features=512, out_features=384, bias=False)
              (o): Linear(in_features=384, out_features=512, bias=False)
              (relative_attention_bias): Embedding(32, 6)
            )
            (layer_norm): UMT5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (1): UMT5LayerCrossAttention(
            (EncDecAttention): UMT5Attention(
              (q): Linear(in_features=512, out_features=384, bias=False)
              (k): Linear(in_features=512, out_features=384, bias=False)
              (v): Linear(in_features=512, out_features=384, bias=False)
              (o): Linear(in_features=384, out_features=512, bias=False)
              (relative_attention_bias): Embedding(32, 6)
            )
            (layer_norm): UMT5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (2): UMT5LayerFF(
            (DenseReluDense): UMT5DenseGatedActDense(
              (wi_0): Linear(in_features=512, out_features=1024, bias=False)
              (wi_1): Linear(in_features=512, out_features=1024, bias=False)
              (wo): Linear(in_features=1024, out_features=512, bias=False)
              (dropout): Dropout(p=0.1, inplace=False)
              (act): NewGELUActivation()
            )
            (layer_norm): UMT5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
      )
    )
    (final_layer_norm): UMT5LayerNorm()
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (lm_head): Linear(in_features=512, out_features=256384, bias=False)
)
```

I have added a script to convert the original t5x jax model directly to pytorch.
However, the results of the pytorch output are still garbage.

Pytoch automatically check for any missing keys + the shape. So, I think now the problem might be the merging of the KV into a single matrix.

I have also disabled relative bias for cross attention as it is not needed there.

@adarob could you please help us with the `joined_kv` "fusing" operation. 

T5 uses `joined_kv` for `num_heads` and `d_kv`. But scaled T5 has these two variables, so we need to "fuse" them manually. Can we just use a `.reshape(config.d_model, config.num_heads * config.d_kv)`, e.g. shape of (512, 6, 64) will be converted to (512, 384) :thinking: 

Not sure it follow, in transformers, T5 has 
```python        
self.key_value_proj_dim = config.d_kv
self.n_heads = config.num_heads
self.inner_dim = self.n_heads * self.key_value_proj_dim
```
are you looking for something like this? 

> @adarob could you please help us with the `joined_kv` "fusing" operation.
> 
> T5 uses `joined_kv` for `num_heads` and `d_kv`. But scaled T5 has these two variables, so we need to "fuse" them manually. Can we just use a `.reshape(config.d_model, config.num_heads * config.d_kv)`, e.g. shape of (512, 6, 64) will be converted to (512, 384) 🤔

@@adarob, we will highly appreciate your feedback as this is the current bottleneck to finalize the model integration.

> Not sure it follow, in transformers, T5 has
> 
> ```python
> self.key_value_proj_dim = config.d_kv
> self.n_heads = config.num_heads
> self.inner_dim = self.n_heads * self.key_value_proj_dim
> ```
> 
> are you looking for something like this?

The current issue is related to the weights, in T5  the attention k,v,q,v are fused into 2d matrix, while the new umt5x version has a 3d matrix.
Please check the difference in "t5x_attention_lookup" function in both T5 and UMT5. We need to use reshape to fix it, but it seems it doesn't provide the correct results.

Also pinging the main author of umT5 paper @hwchung27 for help :hugs: 

Maybe we have more luck with pinging @cpgaffney1 as t5x contributor here :hugs: 

adarob and hwchung27 no longer have this project as their main focus. I am also not a T5X owner. Tagging @gauravmishra  

When are you planning to merge this MR? I would like to test this new model

> When are you planning to merge this MR? I would like to test this new model

Unfortunately, we are stuck in the conversion process and we didn't get help yet from either the authors or the hugging face team.

@agemagician Could you please point out where you were asking a Hugging Face team member for help? I'm sorry we missed that message, but looking at the conversation I only see pings directed at persons out of the team.

cc @younesbelkada since Arthur is on vacation.

Thanks for the PR !
@agemagician , could you point me on the exact issue you are facing, and on which conversion process?
Is there something that is reproducible so that I can try to have a look locally?

> @agemagician Could you please point out where you were asking a Hugging Face team member for help? I'm sorry we missed that message, but looking at the conversation I only see pings directed at persons out of the team.
> 
> cc @younesbelkada since Arthur is on vacation.

Hi @sgugger No worries, It was mainly when we created a new issue for integrating the model and then we created this PR ourselves.
https://github.com/huggingface/transformers/issues/22573

It would be great if HF team could help us to finalize the integration of this model, since it is almost finished.

> Thanks for the PR ! @agemagician , could you point me on the exact issue you are facing, and on which conversion process? Is there something that is reproducible so that I can try to have a look locally?

Hi @younesbelkada ,

Thanks for offering your help.

Here is a summary of the current state:
UMT5 model is almost similar to T5 model except the following:
1. The original T5X checkpoint does not merge kv for k, o, q, and v values.
2. They use separate relative attention for each layer.
3. They use byte fallback in case of OOV for the tokenizer.

What we have done:
1. We created the script for converting the original t5x to pytorch and jax.
2. We replicated T5 model and separated the relative attention for each layer.
3. We converted the tokenizer.

What does not work:
1. The results of the output model are rubbish.

Where do we think the problem:
1. It is how we join the the separated kv values.
https://github.com/agemagician/transformers/blob/UMT5/src/transformers/models/umt5/convert_umt5x_checkpoint_to_pytorch.py#L51

How you can replicate the process:
1. You can use the pytorch conversion script here:
https://github.com/agemagician/transformers/blob/UMT5/src/transformers/models/umt5/convert_umt5x_checkpoint_to_pytorch.py 
2. Then you can use the pytorch model from here:
https://huggingface.co/agemagician/umt5-small

@stefan-it please, let me know if I missed something here.

@younesbelkada  Please, let me know if you need any additional information :)

Thanks for the detailed pointers! 
Can you point me to the t5x checkpoint for umt5-small so that I can try to convert the weights myself? 

> Thanks for the detailed pointers! Can you point me to the t5x checkpoint for umt5-small so that I can try to convert the weights myself?

Sure, here is the link:
https://github.com/google-research/t5x/blob/main/docs/models.md#umt5-checkpoints

I appreciate everyone's hard work in this thread. I really hope that umT5 gets merged to HF, since the last time we got an LM that supports ~100 languages was in 2020 with mT5.

it seems they moved out the checkpoints to another place, this should be the correct path:
```bash
gcloud storage cp -r gs://t5-data/pretrained_models/t5x/umt5_small/checkpoint_1000000 ./
```


Thank you for all the work in this thread. Are you aware of [mLongT5](https://arxiv.org/pdf/2305.11129.pdf)? 
Do you know when this is going to be merged?

Any update on the status of integrating umT5 and mLong?


Hi, I am ramping up on taking over the PR, will update on the progress whenever possible. thanks

I am currently working on the tokenizer, we'll be linking a new PR for this model addition in the coming days! 

Hi @ArthurZucker ,

The tokenizer in my repo should be working:
https://huggingface.co/agemagician/umt5-small/tree/main



Mmm I meant the fast tokenizer, I am mostly working on adding byte fallback support for fast tokenizers, which doesn’t exist yet 👌