[Hot-Fix][XLA] Re-enable broken _tpu_save for XLATensors
Author: yeounoh
Pull request: yes

# What does this PR do?

This re-enables checkpointing model on `xla` device, since that wasn't allowed in some cases where XLATensors have invalid storage.data_ptr(). This explicitly moves the model to `cpu` before checkpointing to ensure the storage is valid. 

This is a *hot-fix* to unblock all failing HF model checkpointing on XLA device.

This fixes the existing `_save_tpu` in `trainer.py` implementation, which is currently broken:
```
100%|███████████████████████████████████████████| 10/10 [00:44<00:00,  4.49s/it]
[INFO|trainer.py:2821] 2023-12-01 23:04:53,098 >> Saving model checkpoint to /workspace/MNLI
[INFO|configuration_utils.py:458] 2023-12-01 23:05:09,794 >> Configuration saved in /workspace/MNLI/config.json
/usr/local/lib/python3.8/site-packages/safetensors-0.4.1rc1-py3.8-linux-x86_64.egg/safetensors/torch.py:17: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  return tensor.storage().data_ptr()
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/safetensors-0.4.1rc1-py3.8-linux-x86_64.egg/safetensors/torch.py", line 13, in storage_ptr
    return tensor.untyped_storage().data_ptr()
RuntimeError: Attempted to access the data pointer on an invalid python storage.
```

With this change, it works,
```
100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  2.28it/s][INFO|trainer.py:1975] 2023-12-01 23:44:53,563 >> 

Training completed. Do not forget to share your model on huggingface.co/models =)


{'train_runtime': 7.1392, 'train_samples_per_second': 1.121, 'train_steps_per_second': 0.14, 'train_loss': 1.0944234132766724, 'epoch': 1.0}
100%|█████████████████████████████████████████████| 1/1 [00:07<00:00,  7.14s/it]
[INFO|trainer.py:2821] 2023-12-01 23:45:00,259 >> Saving model checkpoint to /workspace/MNLI
[INFO|configuration_utils.py:458] 2023-12-01 23:45:18,980 >> Configuration saved in /workspace/MNLI/config.json
[INFO|modeling_utils.py:1851] 2023-12-01 23:45:20,612 >> Model weights saved in /workspace/MNLI/pytorch_model.bin
[INFO|tokenization_utils_base.py:2210] 2023-12-01 23:45:20,613 >> tokenizer config file saved in /workspace/MNLI/tokenizer_config.json
[INFO|tokenization_utils_base.py:2217] 2023-12-01 23:45:20,613 >> Special tokens file saved in /workspace/MNLI/special_tokens_map.json
```

Comments:
I see this in the CI test run, `tests_torch`
```
 if cls.token is None:
>           raise ValueError("Cannot run tests as secret isn't setup.")
E           ValueError: Cannot run tests as secret isn't setup.
```

Also, `run_tests` is blocked, with "This workflow is awaiting approval from a maintainer in https://github.com/huggingface/transformers/pull/27799"

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_27799). All of your documentation changes will be reflected on that endpoint.

No worries, the test is skipped on main feel free to rebase! 

There was an issue with this recently #27578 that should be fixed by this I think 

@yeounoh @ArthurZucker  please help check if https://github.com/huggingface/transformers/pull/29703 would work for TPU?