Repository: pytorch
Issue: Inplace error of BatchNorm layer in DistributedDataParallel module
Labels: oncall: distributed

## 🐛 Bug

If you have a DistributedDataParallel module which contains a BatchNorm layer, the extra call of model forward after the one w.r.t. derivative variable will incur the inplace error. 
If the `grad` takes right after the forward about the derivative variable , there will be no error.
This is quite related to the issue about the buffer
#22095

## To Reproduce
1. run the minimal example with `python -m torch.distributed.run`
2. The first `grad` function run without errors
3. The second  `grad`: Observe `one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [2]] is at version 6; expected version 5 instead.`
```
import os
import torch 
from models.wideresnet_update import *

torch.autograd.set_detect_anomaly(True)

class mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.bn = torch.nn.BatchNorm2d(2,affine=True)

    def forward(self, z):
        return (self.bn(z)).sum()

def run(distributed=False, local_rank=None):
    model = mod().to(local_rank)

    if distributed:
        model = torch.nn.parallel.DistributedDataParallel(model,
                device_ids=[local_rank], output_device=local_rank,broadcast_buffers=True)

    model.eval()
    x1 = torch.zeros(1,2,3,4).to(local_rank).requires_grad_()
    x2 = torch.zeros(1,2,3,4).to(local_rank).requires_grad_()
    print(model.module.bn.running_mean)
    #############################
    # No Error
    with torch.enable_grad():
        out1 = model(x1)  
        out2 = model(x2) 
        loss = (out2-out1).mean()
    grad = torch.autograd.grad(loss, x2)[0] # grad on x2, grad right after the forward model(x2) 
    print("1st Stepped.")   

    #############################  
    # Error
    with torch.enable_grad():
        out1 = model(x1) 
        out2 = model(x2) 
        loss = (out2-out1).mean()
    grad = torch.autograd.grad(loss, x1)[0]   # grad on x1, other forward after model(x1) 
    print("2nd Stepped.")

if __name__ == "__main__":
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ.get("WORLD_SIZE", 0))

    distributed = world_size >= 1
    if distributed:
        print(f"[Main] Distributed across {world_size} GPUs")
        torch.cuda.set_device(local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")

    run(distributed=distributed, local_rank=local_rank)
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

Message:
1st Stepped.
1st Stepped.
[W python_anomaly_mode.cpp:104] Warning: Error detected in CudnnBatchNormBackward. Traceback of forward call that caused the error:
  File "minimal_examples.py", line 58, in <module>
    run(distributed=distributed, local_rank=local_rank)
  File "minimal_examples.py", line 41, in run
    out1 = model(x1)
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 799, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "minimal_examples.py", line 13, in forward
    return (self.bn(z)).sum()
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 744, in forward
    self.eps,
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/functional.py", line 2282, in batch_norm
    input, weight, bias, running_mean, running_var, training, momentum, eps, torch.backends.cudnn.enabled
 (function _print_stack)
Traceback (most recent call last):
  File "minimal_examples.py", line 58, in <module>
    run(distributed=distributed, local_rank=local_rank)
  File "minimal_examples.py", line 44, in run
    grad = torch.autograd.grad(loss, x1)[0]  #extra forward after the one w.r.t. derivative variable 
  File "/data/anonym1/anaconda3/envs/torch/lib/python3.7/site-packages/torch/autograd/__init__.py", line 228, in grad
    inputs, allow_unused, accumulate_grad=False)
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [2]] is at version 6; expected version 5 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!

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): 1.9.1
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:  11.1
 - GPU models and configuration:
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang