Adafactor gives RuntimeError: tensors must be 2-D
Author: TJKlein
Pull request: no

## Environment info
- `transformers` version: 4.2.2 (also tried with the latest version v.4.5.1)
- Platform: Linux-4.4.0-1127-aws-x86_64-with-debian-stretch-sid
- Python version: 3.6.13
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: False

### Who can help
@sgugger @patrickvonplaten

## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce



Steps to reproduce the behavior:

In my code, I replaced AdamW (which is working just fine) with **Adafactor** and then I get an error (see below). The code is using also gradient checkpointing. Using **Adafactor from FairSeq** works **well**

```
# Replacing AdamW
# optimizer = AdamW([{'params': model.parameters()}], lr=args.lr, eps=args.epsilon)
# with Adafactor

optimizer = Adafactor(
        [{'params': model.parameters()}], lr=None,
        eps=(1e-30, 1e-3),
        clip_threshold=1.0,
        decay_rate=-0.8,
        beta1=None,
        weight_decay=0.0,
        relative_step=True,
        scale_parameter=True,
        warmup_init=True
    )
```

Output:
```
home/ubuntu/transformers/src/transformers/optimization.py:557: 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  /opt/conda/conda-bld/pytorch_1607370116979/work/torch/csrc/utils/python_arg_parser.cpp:882.)
  exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
  0%|▎                                                                                                                   | 19/6858 [00:37<3:42:15,  1.95s/it]
Traceback (most recent call last):
  File "main.py", line 519, in <module>
    main()
  File "main.py", line 510, in main
    train(allincl_model, epoch, optimizer, scheduler, criterion)
  File "main.py", line 384, in train
    optimizer.step()
  File "/home/ubuntu/transformers/src/transformers/optimization.py", line 561, in step
    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
  File "/home/ubuntu/transformers/src/transformers/optimization.py", line 492, in _approx_sq_grad
    return torch.mm(r_factor.unsqueeze(-1), c_factor.unsqueeze(0))
RuntimeError: tensors must be 2-D
```




Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

Got the same problem. Have you solved it yet?

Finally I got to solve this problem. This error is caused by 3-D parameters. When the optimizer gets a `[dim1, dim2, dim3]` parameter, [transformers/optimization.py Line 544](https://github.com/huggingface/transformers/blob/master/src/transformers/optimization.py#L544) sets `state["exp_avg_sq_row"]` as `[dim1, dim2]` and `state["exp_avg_sq_col"]` as `[dim1, dim3]`. Then the two parameters in [line 508](https://github.com/huggingface/transformers/blob/master/src/transformers/optimization.py#L508) become `[dim1, dim2, 1]` and `[1, dim1, dim3]`, and the error occurs.

To solve this issue, I create my own adafactor optimizer and change line 506-508 to 
```
r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
return torch.mul(r_factor, c_factor)
```
according to [fairseq's implementation](https://github.com/pytorch/fairseq/blob/main/fairseq/optim/adafactor.py#L159).

Actually having the same problem

@ybch14 - do you think this could also be fixed in `transformers` Adafactor implementation?

> @ybch14 - do you think this could also be fixed in `transformers` Adafactor implementation?

Definitely, just change line 506-508 of [transformers/optimization.py](https://github.com/huggingface/transformers/blob/master/src/transformers/optimization.py#506) as I mentioned above then all done! I'm creating my custom optimizer just because I'm not familiar with pull request process and in a hurry with my development needs. I would really appreciate it if you can help initiate a pull request.

I will attach my local test code here to help your local test:
```
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.optimization import Adafactor

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.w = nn.Parameter(torch.randn(2, 3, 4), requires_grad=True)

    def forward(self):
        return self.w.mean().sigmoid()

device = torch.device("cuda")
target = torch.tensor(1.).to(device)
model = Model().to(device)
y = model()
loss = F.binary_cross_entropy(y, target)
loss.backward()
optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
optimizer.step()
```


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

Thanks a lot for your help here @ybch14 ! I've opened a PR to fix it just like you suggested and it seems to work just fine :-)

BTW, we have some guidelines here on how you can open pull requests: 
https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md

@patrickvonplaten Thank you for your PR and hope pytorch gets better :)