Repository: pytorch
Issue: Unable to launch CUDA Graph with DDP model 
Labels: oncall: distributed

### 🐛 Describe the bug

Hello! I am unable to get a CUDA Graph to replay when using a DDP model. I am unsure if this is because I am not doing the required steps to make both play nice with each other, or if this is a legitimate bug. I appreciate any help nonetheless! Here is a toy case which replicates my issue: 

```python
import os.path

import torch.multiprocessing as mp
import torch.distributed
import torch.nn as nn

import socket
from contextlib import closing


def train(rank: str,
          port: str,
          world_size: int,
          model: nn.Module):

    # Init for DDP
    MASTER_ADDR = '127.0.0.1'
    os.environ['MASTER_ADDR'] = MASTER_ADDR
    os.environ['MASTER_PORT'] = port
    torch.distributed.init_process_group('nccl', rank=rank, world_size=world_size)

    device = f'cuda:{rank}'

    model = model.to(device)

    static_input = torch.rand((1, 1, 300, 300, 20), device=device, dtype=torch.float)

    s = torch.cuda.Stream()
    with torch.cuda.stream(s):
        model = torch.nn.parallel.DistributedDataParallel(model)
        for _ in range(20):
            x = model(static_input)

    torch.cuda.current_stream().wait_stream(s)

    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)

    # Captures the graph
    # To allow capture, automatically sets a side stream as the current stream in the context
    g = torch.cuda.CUDAGraph()
    optimizer.zero_grad()
    with torch.cuda.graph(g):
        x = model(static_input)
        x.sum().backward()  # silly loss function...
        optimizer.step()

    # Fills the graph's input memory with new data to compute on
    static_input.copy_(torch.rand((1, 1, 300, 300, 20), device=device, dtype=torch.float))

    for _ in range(100):
        g.replay()

    print('we got through the stream...')

    # cleanup
    torch.distributed.destroy_process_group(rank)


if __name__ == '__main__':
    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

    def find_free_port():
        """ https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number """

        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
            s.bind(('', 0))
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return str(s.getsockname()[1])

    port = find_free_port()

    print(f'{torch.cuda.nccl.version() =}')
    print(f'{torch.__version__ =}')
    print(f'{torch.version.cuda=}')

    world_size = 2

    model = nn.Sequential(nn.Conv3d(1, 1, kernel_size=(3, 3, 3), padding=(1, 1, 1)),
                          nn.ReLU())

    mp.spawn(train, args=(port, world_size, model), nprocs=world_size, join=True)


```

This gives the error stack trace: 

```
torch.cuda.nccl.version() =(2, 14, 3)
torch.__version__ ='1.14.0.dev20221027'
torch.version.cuda='11.7'
Traceback (most recent call last):
  File "/home/chris/Desktop/CudaGraphDDP/main.py", line 83, in <module>
    mp.spawn(train, args=(port, world_size, model), nprocs=world_size, join=True)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 240, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 198, in start_processes
    while not context.join():
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/home/chris/Desktop/CudaGraphDDP/main.py", line 44, in train
    x = model(static_input)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1096, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1052, in _run_ddp_forward
    return module_to_run(*inputs, **kwargs)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/container.py", line 204, in forward
    input = module(input)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 613, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 608, in _conv_forward
    return F.conv3d(
RuntimeError: CUDA error: operation not permitted when stream is capturing

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/chris/Desktop/CudaGraphDDP/main.py", line 46, in train
    optimizer.step()
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/cuda/graphs.py", line 159, in __exit__
    self.cuda_graph.capture_end()
  File "/home/chris/anaconda3/envs/deeplearning/lib/python3.9/site-packages/torch/cuda/graphs.py", line 81, in capture_end
    super(CUDAGraph, self).capture_end()
RuntimeError: CUDA error: operation failed due to a previous error during capture

```

### Versions

Collecting environment information...
PyTorch version: 1.14.0.dev20221027
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.22.5
Libc version: glibc-2.31

Python version: 3.9.12 (main, Jun  1 2022, 11:38:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.14.0-1054-oem-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] adabelief-pytorch==0.2.1
[pip3] numpy==1.22.3
[pip3] torch==1.14.0.dev20221027+cu117
[pip3] torchaudio==0.14.0.dev20221027
[pip3] torchvision==0.15.0.dev20221027
[conda] adabelief-pytorch         0.2.1                    pypi_0    pypi
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.22.3           py39he7a7128_0  
[conda] numpy-base                1.22.3           py39hf524024_0  
[conda] pytorch                   1.14.0.dev20221027 py3.9_cuda11.7_cudnn8.5.0_0    pytorch-nightly
[conda] pytorch-cuda              11.7                 h67b0de4_0    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.14.0.dev20221027+cu117          pypi_0    pypi
[conda] torchaudio                0.14.0.dev20221027      py39_cu117    pytorch-nightly
[conda] torchvision               0.13.0                   pypi_0    pypi


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu