Enable AMP for xla:gpu device in trainer class
Author: ymwangg
Pull request: yes

# What does this PR do?
This PR enables AMP in trainer class for xla:gpu device.

# Discussion
It looks like the torch_xla support in trainer class is primarily for xla:tpu device. 
I found the following features may be useful but not essential and I can include them in this PR if necessary:
1. Rename `tpu` to `xla` in the codebase.
2. Currently xla device is always turned on when torch_xla is installed. It may be useful to allow users to optionally turn it off without uninstalling torch_xla.
3. Currently users need to set `GPU_NUM_DEVICES` manually when using xla:gpu device. It may be useful to set a default value for it when torch_xla and cuda devices are available.

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Comments:
Oh, interesting! Thanks for your contribution, pinging @sgugger on the issue.

> I'm not entirely sure about this PR in the sense that PyTorch XLA support is mainly for TPU, and I don't know if traditional mixed precision training with the gradient scaler will work on TPUs.
> 
> So we should probably split the test to detect if we have GPUs available or TPUs. Some of the logic will stay common between the two, but the mixed precision part might only work for XLA GPUs?

Right, XLA:TPU does not support AMP and only XLA:GPU support it.

> Right, XLA:TPU does not support AMP and only XLA:GPU support it.

So as I said in my previous comment, could you add a new test `is_gpu_xla_available` and use this one for the part where you add grad scalers? Otherwise the changes will make the Trainer stop working on TPU.

@sgugger Maybe I'm missing something. Could you elaborate why the changes will make the Trainer stop working on TPU? The code inside `self.do_grad_scaling` is unreachable if running on TPU since either `--fp16` or `--bf16` option will raise error on TPU. I tested the following training script on TPU:
```bash
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_device_train_batch_size 10 \
    --do_eval \
    --do_train
```
With master branch:
```bash
WARNING:root:TPU has started up successfully with version pytorch-1.9
WARNING:__main__:Process rank: -1, device: xla:1, n_gpu: 0distributed training: False, 16-bits training: False
...
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     1.7568
  train_runtime            = 0:12:23.47
  train_samples            =       4627
  train_samples_per_second =      18.67
  train_steps_per_second   =      1.868
```
With this PR:
```bash
WARNING:root:TPU has started up successfully with version pytorch-1.9
WARNING:__main__:Process rank: -1, device: xla:1, n_gpu: 0distributed training: False, 16-bits training: False
...
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     1.7577
  train_runtime            = 0:10:19.70
  train_samples            =       4627
  train_samples_per_second =     22.399
  train_steps_per_second   =      2.241
```

Ah yes, you're right. Thanks for testing!