seq2seq_trainer optimization issue on TPU 
Hi
I am running seq2seq_trainer.py model on TPU v3-8 instance with pytorch xla 1.7, using adafactor,  here are the logs, could you please assist? thanks

```
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
  0%|                                                                                                                                                                         | 1/19380 [00:02<13:07:23,  2.44s/it]/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
/root/anaconda3/envs/pytorch/lib/python3.6/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
{'loss': 10383.4560546875, 'learning_rate': 6e-07, 'epoch': 0.0010319917440660474}    

```

on GPU also I cannot use adafactor with this error:

```
/opt/conda/envs/internship/lib/python3.7/site-packages/transformers/optimization.py:506: UserWarning: This overload of add_ is deprecated:
        add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
        add_(Tensor other, *, Number alpha) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
Traceback (most recent call last):
  File "finetune_t5_trainer.py", line 223, in <module>
    main()
  File "finetune_t5_trainer.py", line 159, in main
    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
  File "/opt/conda/envs/internship/lib/python3.7/site-packages/transformers/trainer.py", line 797, in train
    self.optimizer.step()
  File "/opt/conda/envs/internship/lib/python3.7/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/conda/envs/internship/lib/python3.7/site-packages/transformers/optimization.py", line 510, in step
    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
  File "/opt/conda/envs/internship/lib/python3.7/site-packages/transformers/optimization.py", line 441, in _approx_sq_grad
    return torch.mm(r_factor.unsqueeze(-1), c_factor.unsqueeze(0))
RuntimeError: tensors must be 2-D
```


@patrickvonplaten

I'm  not sure your first stacktrace shows an actual error? Just deprecation warnings.

Also, without your code we have no way of understanding what might have happened here.<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.<|||||>Hello, I recently got a very similar problem when trying to implement (manually) a self-attention module for images on a model which is trained using adafactor. I'm also using PyTorch lightning, but I don't think that makes a difference, since I tried with the "default" optimizers coming with torch.optim (RMSProp, Adam), and they work. This means that the problem might possibly be caused by the adafactor implementation. I'm running the latest, stable version of 'transformers' as of now (4.9.1). Here are the detailed logs:
```
File "train.py", line 125, in <module>
    run(args)
  File "train.py", line 90, in run
    trainer.fit(model, train_dataloader=train_loader, val_dataloaders=test_loader)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 473, in fit
    results = self.accelerator_backend.train()
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py", line 66, in train
    results = self.train_or_test()
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 69, in train_or_test
    results = self.trainer.train()
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 524, in train
    self.train_loop.run_training_epoch()
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 572, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 730, in run_training_batch
    self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 505, in optimizer_step
    model_ref.optimizer_step(
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1261, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 286, in step
    self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 144, in __optimizer_step
    optimizer.step(closure=closure, *args, **kwargs)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/transformers/optimization.py", line 576, in step
    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
  File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/transformers/optimization.py", line 507, in _approx_sq_grad
    return torch.mm(r_factor.unsqueeze(-1), c_factor.unsqueeze(0))
RuntimeError: tensors must be 2-D
```<|||||>> 
> 
> Hello, I recently got a very similar problem when trying to implement (manually) a self-attention module for images on a model which is trained using adafactor. I'm also using PyTorch lightning, but I don't think that makes a difference, since I tried with the "default" optimizers coming with torch.optim (RMSProp, Adam), and they work. This means that the problem might possibly be caused by the adafactor implementation. I'm running the latest, stable version of 'transformers' as of now (4.9.1). Here are the detailed logs:
> 
> ```
> File "train.py", line 125, in <module>
>     run(args)
>   File "train.py", line 90, in run
>     trainer.fit(model, train_dataloader=train_loader, val_dataloaders=test_loader)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 473, in fit
>     results = self.accelerator_backend.train()
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu_accelerator.py", line 66, in train
>     results = self.train_or_test()
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 69, in train_or_test
>     results = self.trainer.train()
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 524, in train
>     self.train_loop.run_training_epoch()
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 572, in run_training_epoch
>     batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 730, in run_training_batch
>     self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 505, in optimizer_step
>     model_ref.optimizer_step(
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1261, in optimizer_step
>     optimizer.step(closure=optimizer_closure)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 286, in step
>     self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 144, in __optimizer_step
>     optimizer.step(closure=closure, *args, **kwargs)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
>     return wrapped(*args, **kwargs)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/transformers/optimization.py", line 576, in step
>     update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
>   File "/home/gskenderi/anaconda3/envs/geri/lib/python3.8/site-packages/transformers/optimization.py", line 507, in _approx_sq_grad
>     return torch.mm(r_factor.unsqueeze(-1), c_factor.unsqueeze(0))
> RuntimeError: tensors must be 2-D
> ```

Upon further inspection in these last few minutes, it seems that the Adafactor optimizer has difficulties optimizing PyTorch's nn.Conv* layers. If I try to use some Conv1d or to fine-tune a Resnet model, I get the error indicated above but if not my model works fine. In both these cases, the only thing that has changed is the task of updating the weights of some Convolutional layers. I urge you to take a look at this, as it is quite bizarre. Since I did not mentioned this above, this is the optimizer I'm currently trying to use:
```
optimizer = Adafactor(self.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
lr_scheduler = AdafactorSchedule(optimizer)
```
