TPU training is stuck using T5 with PyTorch Lightning
## Environment info

- `transformers` version: 4.6.1
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.8.1+cu101 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Using GPU in script?: False
- Using distributed or parallel set-up in script?: False

### Who can help

@patrickvonplaten, @patil-suraj

## Information

I'm fine-tuning the `t5-small` model using PyTorch Lightning on TPU v3 (Google Colab) on the `imdb` dataset. The training is stuck using an 8-cores setup and works well using a 1-core setup. It seems super weird since the `roberta-base` model works just fine using all 8-cores. 

I've filed a similar issue https://github.com/PyTorchLightning/pytorch-lightning/issues/7984, but would be great to receive some feedback if `t5` models are proved to work on TPUs using Lightning trainer.

## To reproduce

Please, use this Google Colab Notebook:
https://colab.research.google.com/drive/1FbWfkho3Otfl19y5ybkrK5Jw_tWNqV-M?usp=sharing

## Expected behavior

`t5` models should work fine using TPU 8-cores training setup.

Hello! We don't have CI running with pytorch lightning so we would recommend opening an issue on their repository.

Did you try a TPU training with the `Trainer` that comes with `transformers`? It should work fine on 8 cores.<|||||>Hello @LysandreJik, thanks for replying! 

Yes, `🤗/Trainer` works perfectly fine with 8-cores, `t5-small` gets fine-tuned in under 10 minutes on TPU-v3. Nevertheless, it's still unclear whether this issue happens due to some lightning trainer internals or T5 model – `roberta-base` works fine with lightning using 8 TPU cores. 

The latter makes me think that it might be some T5-specific issue, mightn't be?

P.S. Lightning Trainer reports, that `lm_head.weight` parameter isn't tied (it seems missing prior to moving the model to the XLA device). Just in case.<|||||>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.