Slow gpt2 training on TPU with run_clm_flax.py
## Environment info

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cpu (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
- gpt2: @patrickvonplaten, @LysandreJik 

 -->

## Information

Model I am using (Bert, XLNet ...): gpt2

The problem arises when using:
* [X] the official example scripts: (give details below)
https://huggingface.co/flax-community/papuGaPT2/blob/main/run_clm_flax.py
with only small modification added print(jax.device_count()) in main() to see if TPU is being used.

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)
Oscar, polish dataset 47GB:
--dataset_name="oscar" \
--dataset_config_name="unshuffled_deduplicated_pl" \

## To reproduce

Steps to reproduce the behavior:

1. login to remote VM:
```./google-cloud-sdk/bin/gcloud alpha compute tpus tpu-vm ssh dishcloth --zone us-central1-a --project hf-flax```
2. activate venv:
 source ~/papugapt2/bin/activate
3. run pretraining:
```
cd papuGaPT2
bash pretrain_model.sh
```

## Expected behavior

Expected to get training speed around 1s/it on this dataset (as this was speed achieved before updating to latest scripts from master)

But instead the speed is around 10-20s/it

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.