Fix bug in using TPU
Author: pphuc25
Pull request: no

### System Info

transformers==4.28

### Who can help?
@sgugger and @sanchit-gandhi

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```
if (args.fp16 or args.bf16) and self.sharded_ddp is not None:
    if args.half_precision_backend == "auto":
        if args.device == torch.device("cpu"):
            if args.fp16:
                raise ValueError("Tried to use `fp16` but it is not supported on cpu")
            elif _is_native_cpu_amp_available:
                args.half_precision_backend = "cpu_amp"
            else:
                raise ValueError("Tried to use cpu amp but native cpu amp is not available")
        else:
            args.half_precision_backend = "cuda_amp"

    logger.info(f"Using {args.half_precision_backend} half precision backend")
```
in file trainer.py code from 615 to 627, I see that when using TPU with this logic so half_precision_backend will auto set to "cuda_amp", which TPU can not works seen the in TPU can not call some function like torch.cuda (I am traning Whisper using TPU)
 
 So in my suggest, we should have the condition if is_torch_tpu_available() then args.half_precision_backend = "cpu_amp"
 
 The new code:
 ```
 if (args.fp16 or args.bf16) and self.sharded_ddp is not None:
    if args.half_precision_backend == "auto":
        if args.device == torch.device("cpu"):
            if args.fp16:
                raise ValueError("Tried to use `fp16` but it is not supported on cpu")
            elif _is_native_cpu_amp_available:
                args.half_precision_backend = "cpu_amp"
            else:
                raise ValueError("Tried to use cpu amp but native cpu amp is not available")
        elif is_torch_tpu_available() :
            args.half_precision_backend = "cpu_amp"
        else:
            args.half_precision_backend = "cuda_amp"

    logger.info(f"Using {args.half_precision_backend} half precision backend")
```
However, I do not sure that my new changing is suitable so I post to this to have a better solution, can you review for me is this a good solution so that I can make a pull requests to contribute

I would like to advice @sgugger and @sanchit-gandhi to review it

### Expected behavior

There's no bug in training using TPU

Comments:
Could you please give us a reproducer? Also cc @muellerzr and @pacman100 

As an idea workflow, this representation may be limited due to the inability to showcase the complete setting. Consequently, there could potentially be bugs in the way it is set up. The following points outline the reproducers:


1. Download file [run_speech_recognition_seq2seq_streaming.py](https://github.com/huggingface/community-events/blob/main/whisper-fine-tuning-event/run_speech_recognition_seq2seq_streaming.py) and [xla_spawn.py](https://github.com/huggingface/transformers/blob/main/examples/legacy/seq2seq/xla_spawn.py)
2. Adding script
```
echo 'python xla_spawn.py \
        --num_cores 8 \
        run_speech_recognition_seq2seq_streaming.py \
        --tpu_num_cores 8 \
	--model_name_or_path="openai/whisper-tiny" \
	--dataset_name="mozilla-foundation/common_voice_11_0" \
	--dataset_config_name="vi" \
	--language="vi" \
	--train_split_name="train+validation" \
	--eval_split_name="test" \
	--model_index_name="Whisper Small Spanish" \
	--max_steps="5000" \
	--output_dir="./" \
	--per_device_train_batch_size="64" \
	--per_device_eval_batch_size="32" \
	--logging_steps="25" \
	--learning_rate="1e-5" \
	--warmup_steps="500" \
	--evaluation_strategy="steps" \
	--eval_steps="1000" \
	--save_strategy="steps" \
	--save_steps="1000" \
	--generation_max_length="225" \
	--length_column_name="input_length" \
	--max_duration_in_seconds="30" \
	--text_column_name="sentence" \
	--freeze_feature_encoder="False" \
	--report_to="tensorboard" \
	--metric_for_best_model="loss" \
	--greater_is_better="False" \
	--load_best_model_at_end False \
	--gradient_checkpointing \
	--bf16 \
	--overwrite_output_dir \
	--do_train \
	--do_eval \
	--predict_with_generate False \
	--do_normalize_eval \
	--streaming \
	--use_auth_token \
	--push_to_hub \
        --bf16_full_eval True' >> run.sh
```
3. run command line `bash run.sh`

--- 
After the execution reaches a certain point, a bug occurs when calling `torch.cuda` on the TPU. To prevent this issue, one possible solution is to include the argument `--half_precision_backend "cpu_amp"` in the script. 

To implement this fix, I suggest modifying the trainer.py file by adding a condition that checks if the code is running on a TPU (as I mentioned before). If it is, the half_precision_backend should be set to "cpu_amp".

I would like to cc @muellerzr and @pacman100 

The flags `--bf16` and `--bf16_full_eval` are not supported on TPU. I'm not sure using the CPU autocast is a good idea since it will trigger copy of the data to the CPU.

During my experiment, I observed that the removal of bf16 led to satisfactory results. However, the training speed was slower compared to using bf16 in transformers==4.27.0. Could you please suggest some methods to enhance the training speed using TPU?

There is some information in the accelerate docs: https://huggingface.co/docs/accelerate/concept_guides/training_tpu

Probably using JAX would be faster here on TPU. The script at #21764 works for non-streaming mode fine-tuning (the PR just needs tests + docs before merge), so you can use this already if you want

thank you so much for your information <3