Pretraing T5 model with run_t5_mlm_flax.py script does not support distributed training with deepspeed
### System Info

- `transformers` version: 4.24.0.dev0
- Platform: Linux-3.10.0-1127.18.2.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.8.2+cu111
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.6.1
- Jax version: 0.3.25
- JaxLib version: 0.3.25
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@patil-suraj 
@patrickvonplaten 
@stas

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

sh file to pretrain T5 model

 deepspeed --hostfile hostfile \
    --master_port <fill in> \
     run_t5_mlm_flax.py \
    --deepspeed deepspeed_configs.json \
    --train_file <fill in> \
    --output_dir <fill in> \
    --model_name_or_path=t5-small \
    --do_train \
    --max_seq_length="512" \
    --num_train_epochs=1 \
    --save_steps=100 \
    --per_device_train_batch_size=4 \
    --warmup_steps=100 \
    --logging_steps=100 \
    --overwrite_output_dir
    
```
ValueError: Some specified arguments are not used by the HfArgumentParser: ['--local_rank=0']
```

### Expected behavior

It seems run_t5_mlm_flax.py uses its own TrainingArguments file which does not define "local_rank" and "deepspeed"attributes (unlike transformers.TrainingArguments which defined these variables).

run_t5_mlm_flax.py should be configured for these attributes in order to train in a distributed manner. 
@ArthurZucker could you take a look? :-) <|||||>DeepSpeed only supports PyTorch and the script you mention is for Flax. I don't think there is anything that can be done ;-)<|||||>Thanks for letting me know. Is there anything for distributed training on Flax models 