Rework TPU checkpointing in Trainer
# What does this PR do?

This PR rewors a tiny bit the checkpointing mechanism in the `Trainer` and `PreTrainedModel` to get rid of the hack that stored something in the config (which was then forever present if the user decided to share their model on the hub).

The main problem is that the save on TPU has to be called on all processes because there is a synchronization inside (so if we only execute it in one process, the other do not reach the synchronization point and everything hangs). At the same time, the config should only be saved on one process to avoid race conditions. So this PR solves the problem by adding two new arguments to `save_pretrained`: `is_main` and `_save_function`.
Tested on TPUs and all worked well (training, checkpointing, reloading from checkpoint, evaluating a saved model), so I will merge this.