Repository: pytorch
Issue: nn.Embedding backwards slow under high row contention
Labels: module: performance,module: nn,module: cuda,triaged

## 🐛 Bug

nn.Embedding backward is very slow when there is high contention for a small number of rows in the embedding matrix. This can be ameliorated with different approaches to backward in this case.

## To Reproduce

Steps to reproduce the behavior:

```
import torch
import torch.nn as nn
from torch.nn import Parameter

class ParameterEmb(nn.Module):
    def __init__(self, n, dim):
        super().__init__()
        self.emb = Parameter(torch.zeros(n, dim))

    def forward(self, input):
        return self.emb.index_select(0, input)

class OneHotEmb(nn.Module):
    def __init__(self, n, dim):
        super().__init__()
        self.n = n
        self.emb = nn.Linear(n, dim, bias=False)

    def forward(self, input):
        onehot = torch.zeros(input.size(0), self.n, device=input.device)
        onehot.scatter_(1, input.unsqueeze(1), 1)
        return self.emb(onehot)

def forward_backward(M, x):
    return M(x).sum().backward()


N, D = 100, 100
E = nn.Embedding(N, D)
PE = ParameterEmb(N, D)
OHE = OneHotEmb(N, D)
```

Lets first time everything with a uniform distribution over input rows:
```
  x = (torch.arange(10000, dtype=torch.long) % N)

  %timeit forward_backward(E, x)
  %timeit forward_backward(PE, x)
  %timeit forward_backward(OHE, x)

  E.cuda()
  PE.cuda()
  OHE.cuda()
  x = x.cuda()

  %timeit forward_backward(E, x)
  %timeit forward_backward(PE, x)
  %timeit forward_backward(OHE, x)
```
Everything is fine here
```
  CPU Emb     2.87 ms ± 53.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  CPU Param   12.8 ms ± 496 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
  CPU OneHOt  2.9 ms ± 288 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  GPU Emb     343 µs ± 660 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  GPU Param   150 µs ± 238 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
  GPU OneHOt  231 µs ± 983 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```
But if I have high contention on a single row, Embedding becomes 20x slower while other approaches do not.

```
  x = torch.zeros(10000, dtype=torch.long)


  CPU Emb     31.5 ms ± 248 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
  CPU Param   12.8 ms ± 66.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
  CPU OneHOt  2.52 ms ± 95.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  GPU Emb     6.34 ms ± 348 ns per loop (mean ± std. dev. of 7 runs, 100 loops each)
  GPU Param   513 µs ± 198 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
  GPU OneHOt  205 µs ± 500 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

## Environment
```
PyTorch version: 1.0.1.post2
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 7.3.0-27ubuntu1~18.04) 7.3.0
CMake version: version 3.12.1

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Quadro GP100
...

Nvidia driver version: 410.79
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.15.2
[pip3] torch==1.0.1.post2
[pip3] torchbiggraph==1.dev0
[pip3] torchvision==0.2.1
[conda] blas                      1.0                         mkl
[conda] mkl                       2019.0                      118
[conda] mkl_fft                   1.0.6            py37h7dd41cf_0
[conda] mkl_random                1.0.1            py37h4414c95_1
[conda] torch                     1.0.1.post2              pypi_0    pypi
[conda] torchvision               0.2.1                    py37_1    pytorch
 ```