[Optimization] AdaFactor not working on TPU but works on GPU.
Author: prikmm
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.2 / 4.10.0.dev0
- Platform: Kaggle / Colab
- Python version: 3.7.10
- PyTorch version (GPU?):  1.9.0a0+git1a7c23c (Kaggle) / 1.9.0+cu102 (Colab)
- Tensorflow version (GPU?): 2.4.1
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes, TPU

## Information

Model I am using (Bert, XLNet ...): T5

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (XSum)
* [ ] my own task or dataset: (give details below)

I am trying to finetune t5-small on XSum using `AdaFactor` and `get_linear_schedule_with_warmup`. I am able to do this when I use GPU but when using TPU, the model doesn't converge. The train loss varies but doesn't decrease and validation loss stays constant. Linear Schedule works properly, I saw my `comet_ml` graph, and `lr` was changing the way it should. It's like the loss is not modifying the weights at all.

Code for initializing optimizer and lr_scheduler:
```python
optimizer = Adafactor(model.parameters(), scale_parameter=False, 
                      relative_step=False, warmup_init=False,
                      lr=1e-3)
lr_scheduler = get_linear_schedule_with_warmup(optimizer,
                                               num_training_steps=Config.total_train_steps,
                                               num_warmup_steps =Config.warmup_steps )
```

## To reproduce

The below given colabs are similar but I have provided GPU and TPU notebooks so that running and waiting for results is not needed:
TPU: [Colab Link](https://colab.research.google.com/drive/1MAID8RhaLSevIyhhotUmxZAgjZj9IXR_?usp=sharing)
GPU: [Colab Link](https://colab.research.google.com/drive/111i6_P7PTtpuQLMU26NBCx9qam-q7WSW?usp=sharing)

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
Training Loss should decrease and eventually converge (like it does on GPU)


#### Edit:

Provided links of two notebooks, for  GPU and TPU each.

Comments:
Gently pinging @sgugger here

Hey @patrickvonplaten, can I get an update on this issue? 

Thanks!

Adding Adafactor in the Transformers library was a mistake, Transformers is a library for models, not optimizers.
I don't think this will be addressed @prikmm so you should look for another implementation of this optimizer to use.

@sgugger It worked. I was initializing the `optimizer` and `lr_scheduler` in global scope.
```python
model = AutoModelForSeq2SeqLM.from_pretrained(....)
WRAPPED_MODEL = xmp.MpModelWrapper(model)
optimizer = Adafactor(model.parameters(), scale_parameter=False, 
                      relative_step=False, warmup_init=False,
                      lr=1e-3)
lr_scheduler = get_linear_schedule_with_warmup(optimizer,
                                               num_training_steps=Config.total_train_steps,
                                               num_warmup_steps =Config.warmup_steps )

def _mp_fn():
      .....
      trainer = Trainer(....., optimizers=(optimizer, lr_scheduler))
      trainer.train()
      .....
xmp.spawn(_mp_fn,  start_methods="fork")
```
When I initialized them inside `_mp_fn()`, everything worked fine.
```python
def _mp_fn():
      ......
      optimizer, lr_scheduler = get_optim_lr(model)
      trainer = Trainer(....., optimizers=(optimizer, lr_scheduler))
      trainer.train()
      ......

xmp.spawn(_mp_fn, start_method="fork")
```
I think in method-1, the optimizer gets linked to model weights present in host memory. And when the optimizer gets copied to each TPU device. It will still be linked to model weights present in host memory (or to nothing), and the loss will update the model weights in host memory (or it won't, I have not been able to check  that), and not the model present in each TPU device. 

Whereas , in method-2, since, the optimizer is defined in TPU device scope, and uses the model present there. It is able to update model weights using the loss of that device. I tried `AdamW` using both the methods, and found `AdamW` too doesn't work in method-1 but works in method-2.

For GPU, I perform:
```python
model = model.to("cuda")
```
before initializing the optimizer. So, here the optimizer is linked to right model weights (present in GPU). Hence, everything worked while training on single GPU.

Generally, when using GPU, if two or more variables in use are not on the same device, it will throw an error. This is not the case with TPU, it throws no error, because of which it took such a long time to solve.

This is what I have theorised. If I am wrong, please let me know?

I use a single GPU majority of the time, so pardon me for my lack of TPU knowledge (It's increasing everyday) :)

Ah yes, you should always define your optimizer after transferring your model on the TPU when working on TPUs, because moving the model to the TPU actually creates new tensors for each parameter. So in your case 1, the optimizer was completely disconnected from the model.