Repository: pytorch
Issue: BatchNorm with LSTM in DistributedDataParallel throws one of the variables needed for gradient computation has been modified by an inplace operation
Labels: oncall: distributed

### 🐛 Describe the bug

DistributedDataParallel instance of a neural network module that has BatchNorm computation  along with an LSTM cell throws the following error (when the NN module is run recursively for each timestep passing the hidden state of the LSTM cell):

```
[W python_anomaly_mode.cpp:104] Warning: Error detected in NativeBatchNormBackward.

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [64]] is at version 4; expected version 3 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
```

Please note:
1. Replacing BatchNorm with SyncBatchNorm does not help.
2. This error goes away when broadcast_buffers is set to False in the DDP call, but it is not the intended operation as the statistics need to be averaged across all the ranks. 

A minimal working example is pasted below.

```python
import os
import torch.nn as nn
import torch.distributed

os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "45678"
os.environ['TORCH_DISTRIBUTED_DEBUG'] = 'DETAIL'
world_size = 1
rank = 0

torch.distributed.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=world_size)
device = 'cuda'

class BatchNorm_LSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.bn = nn.BatchNorm1d(64)
        self.rnn = nn.LSTMCell(64, 64)  # input_size, hidden_size

    def forward(self, input, mem_h, mem_c):
        input = self.bn(input)
        mem_h, mem_c = self.rnn(input, (mem_h, mem_c))
        return (mem_h, mem_c)

net = BatchNorm_LSTM()
net.to(device=device)
net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = torch.nn.parallel.DistributedDataParallel(
    net,
    device_ids=[rank],
    output_device=rank,
)
net.train()

with torch.autograd.set_detect_anomaly(True):
    input = torch.rand(2, 64, 64, device=device) #time_steps x batch_size x input_size
    mem_h = torch.zeros(64, 64, device=device) # batch_size x hidden_size
    mem_c = torch.zeros(64, 64, device=device) # batch_size x hidden_size
    for i in range(input.size()[0]):
        mem_h, mem_c = net(input[i], mem_h, mem_c)
    out = mem_h.mean()
    out.backward() # throws error!
    print('Passed!') # never reached :(
```

### Versions

PyTorch version: 1.9.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.6.13 |Anaconda, Inc.| (default, Jun  4 2021, 14:25:59)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-163-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 2080 Ti
GPU 1: Quadro P400

Nvidia driver version: 470.86
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.10.0
[pip3] torch-ac==1.1.0
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] numpy                     1.19.5           py36hfc0c790_2    conda-forge
[conda] pytorch                   1.9.1           py3.6_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torch                     1.10.0                   pypi_0    pypi
[conda] torch-ac                  1.1.0                     dev_0    <develop>
[conda] torchvision               0.11.1                   pypi_0    pypi

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang