Restore fp16 support on xla gpu device
https://github.com/huggingface/transformers/pull/20684 accidentally disabled fp16 support on xla gpu device, which leads to significant performance regression. This PR restores this feature.

cc @jeffhataws @sgugger @Lokiiiiii

Tested with 
```
GPU_NUM_DEVICES=1 python run_mlm.py \
    --model_name_or_path bert-base-uncased \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --overwrite_output_dir true \
    --output_dir /tmp/test-mlm \
    --per_gpu_train_batch_size 24 \
    --do_eval \
    --fp16 true \
    --do_train \
    --num_train_epochs 3 \
    --optim adamw_torch_xla
```

```
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     1.7725
  train_runtime            = 0:04:58.00
  train_samples            =       4627
  train_samples_per_second =      46.58
  train_steps_per_second   =      1.943
INFO:__main__:*** Evaluate ***
[INFO|trainer.py:739] 2023-03-21 19:05:53,483 >> The following columns in the evaluation set don't have a corresponding argument in `BertForMaskedLM.forward` and have been ignored: special_tokens_mask. If special_tokens_mask are not expected by `BertForMaskedLM.forward`,  you can safely ignore this message.
[INFO|trainer.py:3072] 2023-03-21 19:05:53,487 >> ***** Running Evaluation *****
[INFO|trainer.py:3074] 2023-03-21 19:05:53,487 >>   Num examples = 479
[INFO|trainer.py:3077] 2023-03-21 19:05:53,487 >>   Batch size = 8
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:07<00:00,  8.38it/s]
***** eval metrics *****
  epoch                   =        3.0
  eval_loss               =     1.5811
  eval_runtime            = 0:00:29.83
  eval_samples            =        479
  eval_samples_per_second =     16.055
  eval_steps_per_second   =      2.011
  perplexity              =     4.8601
```
_The documentation is not available anymore as the PR was closed or merged._<|||||>Failure is unrelated and due to the branch being old, it's fixed on main, so merging.