Convert T5x "Scalable T5" models to PyTorch
### Feature request

Google T5X library has a new model architecture called "scalable_t5" based on T5, the main difference is that it supports giant model training.

It supports Jax Scan and Rematerialization / Checkpointing, allowing it to load and train giant models on TPU or GPU.

Links:
https://github.com/google-research/t5x/tree/main/t5x/examples/scalable_t5

 It has the same architecture as T5, but the checkpoints are stored differently, rather than storing each layer in a separate folder. It stores all layers for the encoder or the decoder in a single folder.
 This makes the current T5 convertors doesn't work.

### Motivation

Training large models like PaLM on TPU pods requires this specific architecture. Unfortunately, trained models using this architecture are not yet convertible to hugging face. This means the community can't use such models at Hugging Face.

Pinging:
@patrickvonplaten
@stefan-it
@bastings 
@ArthurZucker


### Your contribution

None
Hi @agemagician ,

thanks for pinging! Could you confirm that the new umT5 checkpoints have this new scalable format:

https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/umt5_small/checkpoint_1000000?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&prefix=&forceOnObjectsSortingFiltering=false

(Yeah, they put the wrong link in the [overview table](https://github.com/google-research/t5x/blob/main/docs/models.md#umt5-checkpoints) for umT5, I will also prepare a PR fixing that...)<|||||>Update:

I wrote the umT5X conversion script - and it conversion seems to work.

Here's the initial draft:

https://gist.github.com/stefan-it/5d6a4ec89e7ad97181983881434cb4eb

@agemagician Could you please check if it's working with your checkpoints?

I placed that file in `/home/stefan/Repositories/transformers/src/transformers/models/t5`.

And installed latest `t5x` version (Git main branch) and latest `jaxlib`.

I tried it with umT5 Small checkpoints:

```bash
gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r gs://t5-data/pretrained_models/t5x/umt5_small/checkpoint_1000000 .
```

Then you can call the script with:

```bash
python3 convert_umt5x_checkpoint_to_flax.py --config_name google/mt5-small --t5x_checkpoint_path ~/Dokumente/umt5/checkpoint_1000000 --flax_dump_folder_path ./exported
```

-> It's important that the `config_name` matches architecture size of the checkpoint.

Caveats: I will of course do some downstream tasks experiments to see if conversion works. If @agemagician has a working evaluation pipeline it would be great to hear some feedback of the performance! I will work on the conversion script later - need some sleep now.<|||||>Hi @stefan-it ,

Thanks a lot for your quick reply.

I have created a small random model based on the new scalable architecture like umT5X to check the conversion script, which was converted successfully.
However, I debugged the code to make sure it was converted correctly, and I think it was converted incorrectly.

I created a small model based on this configuration:
```
{
  "_name_or_path": "./",
  "architectures": [
    "T5ForConditionalGeneration"
  ],
  "d_ff": 16,
  "d_kv": 6,
  "d_model": 8,
  "decoder_start_token_id": 0,
  "dense_act_fn": "silu",
  "dropout_rate": 0.0,
  "eos_token_id": 1,
  "feed_forward_proj": "gated-silu",
  "initializer_factor": 1.0,
  "is_encoder_decoder": true,
  "is_gated_act": true,
  "layer_norm_epsilon": 1e-06,
  "model_type": "t5",
  "num_decoder_layers": 3,
  "num_heads": 4,
  "num_layers": 3,
  "output_past": true,
  "pad_token_id": 0,
  "relative_attention_max_distance": 128,
  "relative_attention_num_buckets": 64,
  "tie_word_embeddings": false,
  "torch_dtype": "float32",
  "transformers_version": "4.26.0",
  "use_cache": true,
  "vocab_size": 256
}
```

I also checked the gin file to make sure it is similar :
```
network.T5Config.dropout_rate = %DROPOUT_RATE
network.T5Config.dtype = 'bfloat16'
network.T5Config.emb_dim = 8
network.T5Config.head_dim = 6
network.T5Config.logits_via_embedding = False
network.T5Config.mlp_activations = ('silu', 'linear')
network.T5Config.mlp_dim = 16
network.T5Config.num_decoder_layers = 3
network.T5Config.num_encoder_layers = 3
network.T5Config.num_heads = 4
network.T5Config.remat_policy = 'minimal'
network.T5Config.scan_layers = True
network.T5Config.vocab_size = 256
```

Then I added the following print statement to check the dimensions :
```
   config = T5Config.from_pretrained(config_name)
    flax_model = FlaxT5ForConditionalGeneration(config=config)
    t5x_model = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)

    print(config.num_layers)
    print(len(t5x_model["target"]["encoder"]["encoder"]["attention"]["key"]["kernel"]))
    print(t5x_model["target"]["encoder"]["encoder"]["attention"]["key"]["kernel"][0].shape)
    print(flax_model.params["encoder"]["block"][str(0)]["layer"]["0"]["SelfAttention"]["k"][
            "kernel"
        ].shape)
```

The output is:
```
3
8
(3, 4, 6)
(8, 24)
```
As you can see, the current conversion tries to copy three layers while the checkpoint shows 8.
Also, you can see the dimension of the selected layer doesn't match (3, 4, 6) vs (8, 24).

It seems the new checkpoints are stored based on the following order:
1. emb_dim - > num_layers -> num_heads -> head_dim

So, I think we either need to copy a slice during every iteration or rearrange the parameters.
I think we should also have an assert check to make sure both destination and original parameters have the same size.
<|||||>To make it easy for you to debug it, I have created a repo contains this small model which should accelerate testing and debugging :
[agemagician/scalable_t5x_tiny_test](https://huggingface.co/agemagician/scalable_t5x_tiny_test/tree/main)<|||||>Ah, yeah length/shape checks would be the next thing that I would have tried! Many thanks for your feedback, I will work on it today! Also thanks for uploading your checkpoints!!<|||||>Great, I will wait for your updated version.

Maybe at the end, we could meet and celebrate since we both live in Munich 😉 <|||||>Hey @agemagician ,

yeah really good idea :hugs: 

I read through the code and compared the "normal" t5x `layer.py` vs. scaled t5x `layer.py`.

As you already noticed in the (4, 6) vs (24) notation: old t5x used a `joined_kv`, whereas scaled t5 uses `heads` and `kv` in separate variables. This joining stuff is also "explained" in the [readme](https://github.com/google-research/t5x/blob/main/docs/usage/partitioning.md#canonical-logical-axis-names) - markdown in that table is broken, here's a better overview:

```bash
For "heads * kv" fused dimension of attention  matrices,
when the kernel is reshaped such that "heads" and "kv"
are packed in the same dimension. 
```

So I will try to reshape it to get a `joined_kv`!

**Update**: yes! Transposing and reshaping yields the correct shape now!!<|||||>Hi @agemagician ,

I updated the gist: https://gist.github.com/stefan-it/5d6a4ec89e7ad97181983881434cb4eb

Conversion script has now a shape check (it compares the shape of the init. FLAX model with the shape of read T5X checkpoint model).

I will do some more tests after I got some sleep -> hopefully on downstream tasks to test the performance.

Please also test the new version of the script :hugs: <|||||>Amazing work @stefan-it 👍 

I went through the code and tested it, and I believe it should lead to a correct conversion.

It was a smart idea to use transpose to correct the order and then reshape. The only drawback is that we have to store 0.5x additional memory of the model, either encoder or decoder, during the weights copying process. So, this might be a bit problematic with very large models. However, given this is readable code, I think we should stick with it :)

Yes, I agree that the next step should be a downstream task test, before merge this script to HF.<|||||>I think we should definitely support this is the t5x -> pytorch conversion script! Maybe @ArthurZucker here as well <|||||>Great, thanks a lot @patrickvonplaten for joining forces 😄 

@stefan-it , I have created a small Colab example to test the model output vs mt5, which should give somehow a similar output:
https://colab.research.google.com/drive/1QrqxNdIK7ugQ3FC8tqxUqZZwP0zdvYE4?usp=sharing

However, the output from the umt5 model is garbage compared to mt5 for the following input:
`"Wikipedia is a <extra_id_0>"`
umt5:
`<pad>xictoarelor nhaulated辙ktör betroffen syntet Undesagrado硼颤oplasm betroffen nhau痍剖 المختلفةrieks`
mt5:
`<pad> <extra_id_0> political encyclopedia</s>`

I have checked the paper in case there is something different, and indeed, there is a difference in the architecture:
```
C ADDITIONAL TRAINING DETAILS
The model architectures used in this study are the same as mT5 models, except that relative position embeddings are not shared across layers. In all of our models, the vocabulary size is 256,000
subwords, and byte-level fallback is enabled, so unknown tokens are broken down into UTF-8 bytes.
We use the T5X library (Roberts et al., 2022) to train the models using Google Cloud TPUs. For
pretraining, we use Adafactor optimizer (Shazeer & Stern, 2018) with a constant learning rate of
0.01 in the first 10,000 steps and inverse square root decay afterwards. For finetuning, we use
Adafactor with a constant learning rate of 5e−5. Unlike mT5, we do not use loss normalization
factor. Instead we use the number of real target tokens as the effective loss normalization.
Finally, we do not factorize the second moment of the Adafactor states and we also use momentum,
neither of which are used in T5 and mT5 studies.
```

So we can't simply use the current HF mt5 model architecture as it is.

@patrickvonplaten, any thoughts on how not to share relative position embeddings across layers on mt5 model script ?<|||||>Many thanks for that Notebook! Make things a bit easier - I've also converted the model to PyTorch incl. vocab and uploaded it on the hub:

https://huggingface.co/stefan-it/umt5-small/tree/main

I noticed one difference - I think it was in `t5x_model["target"]["decoder"]["decoder"]["relpos_bias"]["rel_embedding"]` and yeah... it corresponds to the relative position embeddings, oh no!<|||||>Looking more into "network.py" for both t5 and scalable_t5, I found it is true what is mentioned in the paper.

On t5, they define the relative embedding once, then they call it on each encoder layer:
https://github.com/google-research/t5x/blob/main/t5x/examples/t5/network.py#L56

On scalable_t5, they define the relative embedding on each encoder layer separately:
https://github.com/google-research/t5x/blob/main/t5x/examples/scalable_t5/network.py#L64

The same goes for the decoder.

So the current implementation of mt5 at huggingface can't work directly with the new umt5 because at mt5 we only have a single shared relative bias, while on umt5 we have a separate relative bias for each layer.<|||||>Yeah, this architecture breaking change is really annoying! It means a lot of copying of code from T5 I guess...<|||||>But this issue is a good pointer where to perform some modifications (in a new umT5 model implementation):

https://github.com/huggingface/transformers/issues/13397<|||||>> But this issue is a good pointer where to perform some modifications (in a new umT5 model implementation):
> 
> #13397

yes, I am already working on a solution for that :)
I will make a PR today that allows umt5 to work with separate relative bias using mt5 code base without the need of a new model.<|||||>I have made the pull request :
https://github.com/huggingface/transformers/pull/22613

All we need is to set the following parameter in the config :
share_relative_attention_bias = False<|||||>Hi @agemagician ,

do you see the `relative_attention_bias` in all layers? I'm using the PR and it shows:

```
T5Config {
  "_name_or_path": "./",
  "architectures": [
    "T5ForConditionalGeneration"
  ],
  "d_ff": 1024,
  "d_kv": 64,
  "d_model": 512,
  "decoder_start_token_id": 0,
  "dense_act_fn": "gelu_new",
  "dropout_rate": 0.1,
  "eos_token_id": 1,
  "feed_forward_proj": "gated-gelu",
  "initializer_factor": 1.0,
  "is_encoder_decoder": true,
  "is_gated_act": true,
  "layer_norm_epsilon": 1e-06,
  "model_type": "t5",
  "num_decoder_layers": 8,
  "num_heads": 6,
  "num_layers": 8,
  "pad_token_id": 0,
  "relative_attention_max_distance": 128,
  "relative_attention_num_buckets": 32,
  "share_relative_attention_bias": false,
  "tie_word_embeddings": false,
  "tokenizer_class": "T5Tokenizer",
  "torch_dtype": "float32",
  "transformers_version": "4.28.0.dev0",
  "use_cache": true,
  "vocab_size": 256384
}

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Layer: 0
dict_keys(['k', 'o', 'q', 'relative_attention_bias', 'v'])
Layer: 1
dict_keys(['k', 'o', 'q', 'v'])
Layer: 2
dict_keys(['k', 'o', 'q', 'v'])
Layer: 3
dict_keys(['k', 'o', 'q', 'v'])
Layer: 4
dict_keys(['k', 'o', 'q', 'v'])
Layer: 5
dict_keys(['k', 'o', 'q', 'v'])
Layer: 6
dict_keys(['k', 'o', 'q', 'v'])
Layer: 7
dict_keys(['k', 'o', 'q', 'v'])
```<|||||>But it is there when I'm using PyTorch! I can see a difference between `config.share_relative_attention_bias` = `True` or `False`, but not with Flax implementation at the moment!<|||||>> share_relative_attention_bias

hmmm, checking ..<|||||>---- Replied Message ----
| From | Ahmed ***@***.***> |
| Date | 04/06/2023 18:18 |
| To | ***@***.***> |
| Cc | ***@***.***> |
| Subject | Re: [huggingface/transformers] Convert T5x "Scalable T5" models to PyTorch (Issue #22573) |

Great, thanks a lot @patrickvonplaten for joining forces 😄

@stefan-it , I have created a small Colab example to test the model output vs mt5, which should give somehow a similar output:
https://colab.research.google.com/drive/1QrqxNdIK7ugQ3FC8tqxUqZZwP0zdvYE4?usp=sharing

However, the output from the umt5 model is garbage compared to mt5 for the following input:
"Wikipedia is a <extra_id_0>"
umt5:
<pad>xictoarelor nhaulated辙ktör betroffen syntet Undesagrado硼颤oplasm betroffen nhau痍剖 المختلفةrieks
mt5:
<pad> <extra_id_0> political encyclopedia</s>

I have checked the paper in case there is something different and indeed there is a difference in the architecture:

C ADDITIONAL TRAINING DETAILS
The model architectures used in this study are the same as mT5 models, except that relative position embeddings are not shared across layers. In all of our models, the vocabulary size is 256,000
subwords, and byte-level fallback is enabled, so unknown tokens are broken down into UTF-8 bytes.
We use the T5X library (Roberts et al., 2022) to train the models using Google Cloud TPUs. For
pretraining, we use Adafactor optimizer (Shazeer & Stern, 2018) with a constant learning rate of
0.01 in the first 10,000 steps and inverse square root decay afterwards. For finetuning, we use
Adafactor with a constant learning rate of 5e−5. Unlike mT5, we do not use loss normalization
factor. Instead we use the number of real target tokens as the effective loss normalization.
Finally, we do not factorize the second moment of the Adafactor states and we also use momentum,
neither of which are used in T5 and mT5 studies.


So we can't simply use the current HF mt5 model architecture as it is.

—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you are subscribed to this thread.Message ID: ***@***.***><|||||>Hmm, so @sgugger asked to create a separate model for that.
I will do it and share a new PR with u.<|||||>@agemagician please let me know if you need some help with that :hugs: <|||||>Closed by #24477 