t5: add conversion script for T5X to FLAX
Author: stefan-it
Pull request: yes

Hi,

this PR adds the (long awaited) conversion script from T5X to HF FLAX, previously available in this [GIST](https://gist.github.com/stefan-it/30e4998ef159f33696e377a46f699d9f).

This conversion scripts allows to convert models that were trained with [T5X](https://github.com/google-research/t5x) to a FLAX model, so it can be used with Transformers.

Script was road-tested and performance was compared against official T5 (v1.1) checkpoint (because T5 checkpoints can be converted into T5X checkpoints). More information can be found in [this issue](https://github.com/google-research/t5x/issues/198) in T5X upstream repo.

Comments:
/cc @patrickvonplaten @patil-suraj :hugs: 

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

I have just tested it with v1.0 and there's a problem with the lm head:

```
1.0: dict_keys(['decoder_norm', 'layers_0', 'layers_1', 'layers_2', 'layers_3', 'layers_4', 'layers_5', 'relpos_bias'])
1.1: dict_keys(['decoder_norm', 'layers_0', 'layers_1', 'layers_2', 'layers_3', 'layers_4', 'layers_5', 'layers_6', 'layers_7', 'logits_dense', 'relpos_bias'])
```

So `logits_dense` is missing in the 1.0 checkpoints and the conversion script can't handle it. I will try to find a solution here and post a short conversion-pipeline guidline soon.

v1.0 checkpoints can also be converted now. Here are some checks (with final evaluation on downstream task):

Requirements:

```bash
pip3 install git+https://github.com/google-research/t5x.git
pip3 install --upgrade tensorstore==0.1.13
```

Pinned `tensorstore` version fixes a strange `zarr` error when loading the checkpoints.

Then clone 1.0 and 1.1 T5X checkpoints:

```bash
gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r -n gs://t5-data/pretrained_models/t5x/t5_small .
gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp -r -n gs://t5-data/pretrained_models/t5x/t5_1_1_small .
```

Transformer configs can be downloaded from model hub:

```bash
curl --silent https://huggingface.co/t5-small/resolve/main/config.json > config_1_0.json
curl --silent https://huggingface.co/google/t5-v1_1-small/resolve/main/config.json > config_1_1.json
```

Models can be converted via:

```bash
python3 convert_t5x_checkpoint_to_flax.py --t5x_checkpoint_path ./t5_small --config_name ./config_1_0.json --flax_dump_folder_path ./t5x_1_0_exported

python3 convert_t5x_checkpoint_to_flax.py --t5x_checkpoint_path ./t5_1_1_small --config_name ./config_1_1.json --flax_dump_folder_path ./t5x_1_1_exported
```

Then I ran downstream evaluation with the original T5 models (from model hub) and the converted ones on the summarization task, e.g.:

```bash
python examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path t5-small \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir ./t5_1_0_original \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --num_train_epochs 1
```

and:

```bash
python examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path /mnt/transformers/src/transformers/models/t5/t5x_1_0_exported \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir ./t5_1_0_converted \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate \
    --num_train_epochs 1
```

And compared the training losses. Training losses are identical (original vs. converted model).

Hi @stancld , thanks for adding the longt5 variant :hugs: Could you also add the patch for v1.0 checkpoints from this commit:

https://github.com/huggingface/transformers/pull/16853/commits/4f36d429fa03fd85011a0b6bd5f5b04c9077b852

Would be awesome :)

Does this script support the transformation of XL or XXL models?

> Does this script support the transformation of XL or XXL models?
I tried to generate the following files in /content/flan_t5x_xl_exported, and then I used this below code to load and happen error. How do I solve it?
```python
model = T5ForConditionalGeneration.from_pretrained("/content/flan_t5x_xl_exported", from_flax=True)
# Error no file named pytorch_model.bin, tf_model.h5, model.ckpt.index or flax_model.msgpack found in 
# directory /content/flan_t5x_xl_exported.
```
/content/flan_t5x_xl_exported:
"
*model-00001-of-00002.msgpack
*model-00002-of-00002.msgpack
*model.msgpack.index.json
config.json
"

@joytianya Currently, you cannot use cross-platform loading when the large model is split into multiple files. But this feature is planned soon -- please see #19965 

@stefan-it 
@stancld 
Does the script support T5X converted into pytorch?
if not, Is there any other solution?