When save a model on TPU, make a copy to be moved to CPU
Author: qihqi
Pull request: yes

# What does this PR do?

When we save a model on TPU, we first move it to CPU because TPU tensors have no 
storage. However, we should do it with a copy of the model, so that the original model 
still on TPU. Because otherwise `model.to('cpu')` would also modify the model in-place.
Then, it would raise the following error when that model is used in compute:

```
indices should be either on cpu or on the same device as the indexed tensor (XLA). When using XLA, the indexed tensor must be an XLA tensor.
```

Tested by running this command on TPU v4-8:
```
python3 examples/pytorch/text-classification/run_glue.py \
        --model_name_or_path=distilbert-base-uncased \
        --task_name=MNLI \
        --do_train=true \
        --num_train_epochs=1 \
        --max_seq_length=128 \
        --learning_rate=3e-5 \
        --overwrite_output_dir=true \
        --save_steps=3000 \
        --save_strategy=no --output_dir=/workspace/mnli
```

cc @muellerzr and @pacman100



Comments:
@qihqi this is done inside `_save_tpu` https://github.com/huggingface/transformers/pull/27799/files, or we missed something?

> @qihqi this is done inside `_save_tpu` https://github.com/huggingface/transformers/pull/27799/files, or we missed something?

I see, we need to make a copy. Should we do that inside `_save_tpu` or do it before here?

Thanks for the fix! Quick question: I used to see this error `indices should be either on cpu or on the same device as the indexed tensor (XLA). When using XLA, the indexed tensor must be an XLA tensor.` and it seems some index ops failed so I tried to fix the op instead of the model (and it didn't work out). How did you find out it's the model's problem?

> Thanks for the fix! Quick question: I used to see this error `indices should be either on cpu or on the same device as the indexed tensor (XLA). When using XLA, the indexed tensor must be an XLA tensor.` and it seems some index ops failed so I tried to fix the op instead of the model (and it didn't work out). How did you find out it's the model's problem?

Thanks! So I also started with assumption that the input is wrong, and traced where the input comes from, and arrived here: https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L2668 and here clearly the input is moving to the device. 

Also I ran the script the first few steps trains fine. So I start suspecting something happened after training is done. 

Hi @LysandreJik @ArthurZucker would you guys be able to help with a review? Thanks!

Pinging our TPU expert @muellerzr 

> Thanks for adding this!
> 
> I'm not very familiar with TPUs. My only concern with this is, if it behaves the same as the model being on GPU, then it requires enough space for 2 models on the device when `copy.deepcopy(self.model)` is called i.e. there is no offloading.

Good point! I have changed the code to first move it CPU then move it back. PTAL.

@qihqi Just to confirm that you were able to successfully run the script in the PR description with the most recent changes to the PR: 

```
python3 examples/pytorch/text-classification/run_glue.py \
        --model_name_or_path=distilbert-base-uncased \
        --task_name=MNLI \
        --do_train=true \
        --num_train_epochs=1 \
        --max_seq_length=128 \
        --learning_rate=3e-5 \
        --overwrite_output_dir=true \
        --save_steps=3000 \
        --save_strategy=no --output_dir=/workspace/mnli
```


> @qihqi Just to confirm that you were able to successfully run the script in the PR description with the most recent changes to the PR:
> 
> ```
> python3 examples/pytorch/text-classification/run_glue.py \
>         --model_name_or_path=distilbert-base-uncased \
>         --task_name=MNLI \
>         --do_train=true \
>         --num_train_epochs=1 \
>         --max_seq_length=128 \
>         --learning_rate=3e-5 \
>         --overwrite_output_dir=true \
>         --save_steps=3000 \
>         --save_strategy=no --output_dir=/workspace/mnli
> ```

Yes. (Well, to be exact it's this script: https://github.com/GoogleCloudPlatform/ml-testing-accelerators/blob/master/tests/pytorch/nightly/hf-glue.libsonnet#L41) 

Hi everyone,
One question, was this fix supposed to work with XLA FSDP? 

I'm having problems with the model.to ('cpu') part before unwrapping the model (memory access errors), and if I unwrap it before applying .to('cpu'), just one shard of the model is saved.