run_t5_mlm_flax.py Multi GPU Training
Author: ToluClassics
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version:  4.15.0
- Platform: Ubuntu
- Python version: Python 3.9.9
- PyTorch version (GPU?): 1.10.1

### Who can help
@patil-suraj 
@patrickvonplaten
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- T5, BART, Marian, Pegasus, EncoderDecoder: @patrickvonplaten
- Blenderbot, MBART: @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet, BigBird: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder, UniSpeech, UniSpeechSAT, SEW, SEW-D, Speech2Text: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten @narsil
- Tokenizers: @SaulLu
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (T5):

I am currently trying to pretrain T5 from scratch on a custom dataset using the [`run_t5_mlm_flax.py`](https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_t5_mlm_flax.py) on multiple gpus but when I log the GPU count using `jax.local_device_count()` and `jax.local_device_count()` it returns 1.

Does this mean the script is training with just one GPU and if so how can I make It distributed?
Thank you!


The tasks I am working on is:
* [ ] Pretraining T5 from scratch on a custom dataset


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->


Comments:
Hey @ToluClassics,

Could you please open an issue on https://github.com/google/jax as your problem does not seem to be related to Transformers

@ToluClassics @patrickvonplaten @patil-suraj @sgugger 
I also have the same issue with the exact same script.

How is it resolved finally?
How can I make run_t5_mlm_flax.py use GPU instead of CPU?


This is how I run it:


```
python run_t5_mlm_flax.py \
    --output_dir="models/my_mini_trial-t5-base" \
    --model_type="t5" \
    --config_name="models/my_mini_trial-t5-base" \
    --tokenizer_name="models/my_mini_trial-t5-base" \
    --dataset_name="oscar" \
    --dataset_config_name="unshuffled_deduplicated_no" \
    --max_seq_length="512" \
    --per_device_train_batch_size="32" \
    --per_device_eval_batch_size="32" \
    --adafactor \
    --learning_rate="0.005" \
    --weight_decay="0.001" \
    --warmup_steps="2000" \
    --overwrite_output_dir \
    --logging_steps="500" \
    --save_steps="10000" \
    --eval_steps="2500"
  ```
  
  And I confirm that JAX recognize the GPU:
  
  ```
  In [1]: import jax
   ...:
   ...: print("Number of available GPUs:", jax.device_count())
   ...: print("Default GPU:", jax.default_backend())
Number of available GPUs: 1
Default GPU: gpu
```

G.V.