Repository: pytorch
Issue: Large performance difference of loss.backward() between torch-1.9.0 and torch-1.8.0
Labels: module: autograd,triaged,module: regression

### 🐛 Describe the bug

When I reimplemented a CNN model called [PatchNet ](https://github.com/soarsmu/PatchNet/blob/master/model.py) which is a classification model used for Linux stable patch identification (text classification actually), I found different PyTorch versions have totally different performance. Specifically, when I used torch=1.9.0+cu111 or higher versions, loss.backward() can take even 3 seconds or longer which is abnormal. I'm using Python-3.8, Ubuntu-20.04 and RTX3090 GPU. Detailed backward time is shown as following measured by torch.autograd.profiler.profile() as context manager:

```
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
             aten::embedding_dense_backward        85.31%        2.921s        85.34%        2.922s        1.461s             2
     aten::cudnn_convolution_backward_input        11.32%     387.625ms        11.32%     387.681ms      96.920ms             4
    aten::cudnn_convolution_backward_weight         1.90%      65.218ms         1.96%      67.092ms      16.773ms             4
                                 aten::div_         0.33%      11.390ms         0.33%      11.390ms      11.390ms             1
           aten::cudnn_convolution_backward         0.26%       8.921ms        13.54%     463.695ms     115.924ms             4
                                aten::fill_         0.25%       8.672ms         0.25%       8.672ms       1.239ms             7
                                  aten::sum         0.15%       4.981ms         0.15%       5.008ms     834.603us             6
                     aten::sigmoid_backward         0.13%       4.286ms         0.13%       4.303ms       4.303ms             1
                                   aten::mm         0.11%       3.663ms         0.11%       3.687ms     921.713us             4
                                aten::empty         0.06%       2.048ms         0.06%       2.048ms      45.503us            45
                   aten::threshold_backward         0.04%       1.253ms         0.04%       1.253ms     250.683us             5
        aten::binary_cross_entropy_backward         0.01%     484.898us         0.35%      11.923ms      11.923ms             1
                  aten::_local_scalar_dense         0.01%     481.430us         0.01%     481.430us     120.358us             4
                        aten::_masked_scale         0.01%     478.522us         0.01%     500.683us     250.342us             2
                                 aten::item         0.01%     475.891us         0.03%     957.321us     239.330us             4
                 BinaryCrossEntropyBackward         0.01%     414.821us         0.36%      12.338ms      12.338ms             1
         aten::adaptive_max_pool2d_backward         0.01%     365.964us         0.02%     570.998us     142.750us             4
                            aten::unsqueeze         0.01%     345.295us         0.01%     349.640us     349.640us             1
                                aten::slice         0.01%     344.157us         0.01%     346.699us     173.350us             2
                   CudnnConvolutionBackward         0.01%     240.269us        13.55%     463.935ms     115.984ms             4
                                  aten::add         0.01%     187.079us         0.01%     187.079us      37.416us             5
                                aten::copy_         0.00%     127.850us         0.00%     127.850us      31.962us             4
                                 aten::view         0.00%     110.744us         0.00%     110.744us       6.921us            16
                  AdaptiveMaxPool2DBackward         0.00%      93.294us         0.02%     664.292us     166.073us             4
                                    aten::t         0.00%      69.087us         0.00%     102.355us      12.794us             8
                              AddmmBackward         0.00%      68.350us         0.11%       3.842ms       1.921ms             2
            torch::autograd::AccumulateGrad         0.00%      56.038us         0.00%     108.763us      12.085us             9
                              aten::reshape         0.00%      45.049us         0.00%     129.478us      10.790us            12
                           aten::empty_like         0.00%      44.461us         0.00%     134.177us      11.181us            12
                        aten::empty_strided         0.00%      40.160us         0.00%      40.160us      20.080us             2
                            SigmoidBackward         0.00%      38.578us         0.13%       4.342ms       4.342ms             1
                              ReluBackward0         0.00%      38.034us         0.04%       1.291ms     258.290us             5
                               ViewBackward         0.00%      37.363us         0.00%     166.841us      13.903us            12
                           aten::as_strided         0.00%      37.272us         0.00%      37.272us       2.329us            16
                                     detach         0.00%      33.490us         0.00%      33.490us       3.721us             9
                              aten::resize_         0.00%      30.846us         0.00%      30.846us       1.928us            16
                          EmbeddingBackward         0.00%      28.907us        85.34%        2.922s        1.461s             2
                                aten::zero_         0.00%      27.715us         0.00%     127.761us      21.294us             6
                       FusedDropoutBackward         0.00%      24.200us         0.02%     524.883us     262.441us             2
                            aten::ones_like         0.00%      24.024us         0.25%       8.624ms       8.624ms             1
                               aten::detach         0.00%      19.235us         0.00%      52.725us       5.858us             9
                               aten::narrow         0.00%      18.714us         0.01%     365.413us     182.707us             2
                            aten::transpose         0.00%      18.336us         0.00%      33.268us       4.159us             8
                                aten::zeros         0.00%      17.007us         0.00%      58.904us      29.452us             2
                                CatBackward         0.00%      16.036us         0.01%     381.449us     381.449us             1
                           SqueezeBackward0         0.00%      14.738us         0.01%     364.378us     364.378us             1
                   aten::embedding_backward         0.00%      13.686us        85.34%        2.922s        1.461s             2
                           aten::contiguous         0.00%      12.796us         0.00%     105.813us      52.906us             2
                               aten::select         0.00%       9.268us         0.00%      13.854us       3.464us             4
                                  TBackward         0.00%       9.061us         0.00%      27.529us      13.764us             2
                               aten::expand         0.00%       7.026us         0.00%      17.893us      17.893us             1
                               AddBackward0         0.00%       5.713us         0.00%       5.713us       1.428us             4
                            aten::expand_as         0.00%       3.943us         0.00%      21.836us      21.836us             1
                                 aten::conj         0.00%       2.699us         0.00%       2.699us       0.675us             4
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 3.424s
```
However, it will not happen using torch==1.8.0+cu111 or lower versions of torch. Such a large performance difference between torch-1.8.0 and torch-1.9.0 is weird. It can be seen most of the time is taken on backward of embedding, but I have no idea about why torch-1.9.0 makes this happen.

```
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                       Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
             aten::embedding_dense_backward        50.80%      45.219ms        50.99%      45.393ms      22.696ms             2
     aten::cudnn_convolution_backward_input        23.77%      21.159ms        23.79%      21.176ms       5.294ms             4
    aten::cudnn_convolution_backward_weight        12.82%      11.412ms        16.26%      14.472ms       3.618ms             4
                                   aten::mm         6.64%       5.909ms         6.65%       5.919ms       1.480ms             4
                                aten::empty         3.92%       3.486ms         3.92%       3.486ms      72.615us            48
        aten::binary_cross_entropy_backward         0.33%     294.571us         0.42%     376.645us     376.645us             1
                  AdaptiveMaxPool2DBackward         0.21%     183.868us         0.91%     806.103us     201.526us             4
                   CudnnConvolutionBackward         0.18%     161.157us        40.27%      35.852ms       8.963ms             4
                                aten::fill_         0.13%     116.972us         0.13%     116.972us      10.634us            11
         aten::adaptive_max_pool2d_backward         0.11%      93.573us         0.70%     622.235us     155.559us             4
                                 aten::div_         0.07%      61.209us         0.07%      61.209us      61.209us             1
                                aten::copy_         0.07%      61.074us         0.07%      61.074us      15.268us             4
                   aten::threshold_backward         0.06%      54.029us         0.07%      65.142us      13.028us             5
                              AddmmBackward         0.05%      47.408us         6.76%       6.020ms       3.010ms             2
                 BinaryCrossEntropyBackward         0.05%      46.689us         0.48%     423.334us     423.334us             1
                                 aten::view         0.05%      45.229us         0.05%      45.229us       3.231us            14
           aten::cudnn_convolution_backward         0.05%      42.895us        40.09%      35.691ms       8.923ms             4
                     aten::sigmoid_backward         0.05%      41.255us         0.05%      46.324us      46.324us             1
                  aten::_local_scalar_dense         0.04%      37.912us         0.04%      37.912us       9.478us             4
                        aten::_masked_scale         0.04%      36.593us         0.05%      43.002us      21.501us             2
                                    aten::t         0.04%      33.992us         0.06%      52.985us       6.623us             8
                            SigmoidBackward         0.04%      33.941us         0.09%      80.265us      80.265us             1
                           aten::empty_like         0.03%      28.548us         0.46%     412.877us      25.805us            16
            torch::autograd::AccumulateGrad         0.03%      28.405us         0.05%      45.688us       5.076us             9
                                aten::zero_         0.03%      27.028us         0.12%     110.864us      11.086us            10
                                 aten::item         0.03%      26.328us         0.07%      64.240us      16.060us             4
                              ReluBackward0         0.03%      22.520us         0.10%      87.662us      17.532us             5
                           aten::as_strided         0.02%      20.204us         0.02%      20.204us       1.263us            16
                              aten::reshape         0.02%      20.096us         0.07%      62.875us       5.240us            12
                          EmbeddingBackward         0.02%      19.888us        51.03%      45.429ms      22.715ms             2
                                CatBackward         0.02%      18.975us         0.04%      36.059us      36.059us             1
                           aten::zeros_like         0.02%      18.486us         0.49%     439.077us     109.769us             4
                   aten::embedding_backward         0.02%      16.817us        51.01%      45.410ms      22.705ms             2
                               ViewBackward         0.02%      16.787us         0.09%      79.662us       6.639us            12
                           SqueezeBackward0         0.02%      16.350us         0.03%      28.827us      28.827us             1
                            aten::ones_like         0.02%      15.894us         0.07%      58.942us      58.942us             1
                        aten::empty_strided         0.02%      14.348us         0.02%      14.348us       7.174us             2
                       FusedDropoutBackward         0.01%      13.260us         0.06%      56.262us      28.131us             2
                              aten::resize_         0.01%      12.042us         0.01%      12.042us       0.602us            20
                            aten::unsqueeze         0.01%      11.046us         0.01%      12.477us      12.477us             1
                                     detach         0.01%      10.257us         0.01%      10.257us       1.140us             9
                                 aten::conj         0.01%      10.131us         0.01%      10.131us       2.533us             4
                               aten::narrow         0.01%       9.928us         0.02%      17.084us       8.542us             2
                           aten::contiguous         0.01%       9.541us         0.06%      50.057us      25.028us             2
                           aten::resize_as_         0.01%       8.658us         0.01%      11.016us       2.754us             4
                            aten::transpose         0.01%       8.057us         0.02%      18.993us       2.374us             8
                               aten::detach         0.01%       7.026us         0.02%      17.283us       1.920us             9
                               aten::select         0.01%       6.387us         0.01%      11.719us       2.930us             4
                               AddBackward0         0.01%       6.185us         0.01%       6.185us       1.546us             4
                                aten::slice         0.01%       6.029us         0.01%       7.156us       3.578us             2
                                aten::zeros         0.01%       5.534us         0.03%      29.306us      14.653us             2
                            aten::expand_as         0.01%       5.161us         0.01%       9.502us       9.502us             1
                               aten::expand         0.00%       2.963us         0.00%       4.341us       4.341us             1
                                  TBackward         0.00%       2.064us         0.01%      11.314us       5.657us             2
-------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 89.022ms
```

This model is a simple CNN architecture composed of word embedding, Conv2D and Conv3D modules and input is a 5D word index matrix with shape: `[batch, file, hunk, line, seq]`.
Input indexes will first be embedded by embedding to produce tensors and reshaped to `[batch*file*hunk*line, 1, seq, emb_dim]`, then it will be input to a Conv2D module with input_channel=1 and a max pooling operation will be applied to feature maps to produce features with shape: `[batch*file*hunk,*line, feature]`.
Then extracted features will be input to one more Conv2D module with shape `[batch, 1, file*hunk*line, feature]` and similar max pooling operation within each feature map will also applied to compress the feature vectors to a fixed size, namely number of feature maps (out_channel). Two input word matrices will be equally processed by this CNN encoder and their features will concatenated to produce the final feature vector.
At last, a MLP classifier with sigmoid activation will give predicted probabilities based on the feature vector and binary cross entropy loss will be used as loss function. 

Here is the demo code that can help replicate this:
``` python
import torch

class ReshapeLayer(torch.nn.Module):
    def __init__(self, shape):
        super(ReshapeLayer, self).__init__()
        self.shape = shape

    def forward(self, x):
        return x.reshape(*self.shape)

cuda_device = 0
torch.cuda.set_device(cuda_device)

inp_1 = torch.randint(low=0, high=43583, size=(32, 2, 5, 8, 32), dtype=torch.long).cuda(cuda_device)
inp_2 = torch.randint(low=0, high=43583, size=(32, 2, 5, 8, 32), dtype=torch.long).cuda(cuda_device)
labels = torch.ones((32,)).cuda(cuda_device)

# shape: [batch, 1, file, hunk, line, seq, dim] -> [batch*file*hunk*line, 1, seq, dim]
encoder = torch.nn.Sequential(
    ReshapeLayer((32*2*5*8, 1, 32, -1)),
    torch.nn.Conv2d(1, 64, (2, 128)),
    torch.nn.ReLU(),
    torch.nn.AdaptiveMaxPool2d(1),
    ReshapeLayer((32, 1, 2*5*8, 64)),
    torch.nn.Conv2d(1, 64, (2, 64)),
    torch.nn.ReLU(),
    torch.nn.AdaptiveMaxPool2d(1),
    torch.nn.Flatten(1, -1)
).cuda(cuda_device)

embedding = torch.nn.Embedding(43584, 128, 0).cuda(cuda_device)
classifier = torch.nn.Sequential(
    torch.nn.Dropout(0.5),
    torch.nn.Linear(128, 128),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.5),
    torch.nn.Linear(128, 1),
    torch.nn.Sigmoid()
).cuda(cuda_device)

inp_1 = embedding(inp_1)
inp_2 = embedding(inp_2)

inp_1 = inp_1.view(32, 1, 2, 5, 8, 32, 128)
inp_2 = inp_2.view(32, 1, 2, 5, 8, 32, 128)
encoded_1 = encoder(inp_1)
encoded_2 = encoder(inp_2)

encoded = torch.cat((encoded_1, encoded_2), dim=-1)

out = classifier(encoded)
loss_func = torch.nn.BCELoss().cuda(cuda_device)
loss = loss_func(out.squeeze(), labels)

with torch.autograd.profiler.profile() as prof:
    loss.backward()

print(prof.key_averages().table(sort_by="self_cpu_time_total"))
```


### Versions

```
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-42-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration: 
GPU 0: GeForce RTX 3090
GPU 1: GeForce RTX 3090
GPU 2: GeForce RTX 3090
GPU 3: GeForce RTX 3090

Nvidia driver version: 460.67
cuDNN version: Probably one of the following:
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.4
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.2
[pip3] torch==1.9.0+cu111
[pip3] torchaudio==0.9.0
[pip3] torchvision==0.10.0+cu111
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.22.2                   pypi_0    pypi
[conda] torch                     1.9.0+cu111              pypi_0    pypi
[conda] torchaudio                0.9.0                    pypi_0    pypi
[conda] torchvision               0.10.0+cu111             pypi_0    pypi
```

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7