TPU hangs when saving optimizer/scheduler
Author: LysandreJik
Pull request: yes

Similarly to when saving a model state dict, the optimizer and scheduler should be saved using `xm.save` and behind an `xm.rendezvous`.

Additional fix: `pl.ParallelLoader` is not a `torch.utils.data.DataLoader`, and, therefore, must be reinitialized at each epoch.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=h1) Report
> Merging [#4467](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/07dd7c2fd8996fec2979555437dfeff0d38cbf28&el=desc) will **decrease** coverage by `0.09%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4467/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #4467      +/-   ##
==========================================
- Coverage   78.07%   77.97%   -0.10%     
==========================================
  Files         123      123              
  Lines       20436    20439       +3     
==========================================
- Hits        15955    15937      -18     
- Misses       4481     4502      +21     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/4467/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `39.48% <0.00%> (-0.28%)` | :arrow_down: |
| [src/transformers/data/processors/squad.py](https://codecov.io/gh/huggingface/transformers/pull/4467/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL3Byb2Nlc3NvcnMvc3F1YWQucHk=) | `28.66% <0.00%> (-4.78%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4467/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `73.44% <0.00%> (-0.42%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4467/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `88.66% <0.00%> (-0.17%)` | :arrow_down: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4467/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `90.41% <0.00%> (-0.12%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=footer). Last update [07dd7c2...11186c1](https://codecov.io/gh/huggingface/transformers/pull/4467?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Solid work, as discussed 👍

Nice! I tried two days ago and just skipped the checkpoints but something with the eval seemed to be messing things up(ForMultipleChoice) as well. WandB created 8 runs. dont know if its useful, im still figuring out alot of the stuff.