CUDA error: operation not permitted when stream is capturing (2 GPUs)
Labels: module: multi-gpu, triaged

### 🐛 Describe the bug

The error occurs when capturing a CUDA graph with a small network on a machine with more than 1 GPU.
The MWE indicates that some temporary requests or allocations are silently made on a different (default) device than where the network and the data reside and may not necessarily be an issue with capturing itself.

The minimal setup to reproduce is as follows. At least 2 GPUs, The MWE is:
```
import torch.nn as nn
import torch
import torch.optim as optim
import torch.nn.functional as F

dev = torch.device('cuda:1') # we intend to create data network and capture the graph on this GPU only

D = 32
input_shape = torch.Size([4, D]) 
s = torch.cuda.Stream(device=dev)
s.wait_stream(torch.cuda.current_stream())
x = torch.rand(input_shape, dtype=torch.float32, device=dev)
L = nn.Linear(D,D, bias=None).to(dev)
x.requires_grad = True

def loss(x):
    x = F.linear(x, L.weight)
    x = F.softplus(x)
    x = F.linear(x, L.weight)
    return x.sum()

with torch.cuda.device(0): # this clause should not have effect, the issue occurs with or without it
    print("Warmup 3 batches")
    with torch.cuda.stream(s): # warmup on a side stream, why?
        for i in range(3):
            l = loss(x)
            l.backward()
    torch.cuda.current_stream().wait_stream(s)
    
    print("Capturing")       
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        l = loss(x)
        l.backward()
    print("Replay 3 batches")
    for i in range(3):
        g.replay()
        
    print(f'Test passed')

```

Output:
```
Warmup 3 batches
Capturing
Traceback (most recent call last):
  File "/home.nfs/shekhole/work/quant/demo/MWE.py", line 35, in <module>
    l.backward()
  File "/mnt/appl/software/PyTorch/1.11.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/torch/_tensor.py", line 363, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/mnt/appl/software/PyTorch/1.11.0-foss-2021a-CUDA-11.3.1/lib/python3.9/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: CUDA error: operation not permitted when stream is capturing
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
...
```

The workaround is therefore rather simple: to use CUDA_VISIBLE_DEVICES if we do not need more than 1 GPU or to enclose all the execution into `with torch.cuda.device(dev)'. But it appears to be a bug nevertheless. 


### Versions

At this moment I was able to test only on two servers with older GPUs, where the maximal version of pytorch supporting these GPUs is 1.11. I will try to find a way to run it in a newer setup. 

```
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (GCC) 10.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.28

Python version: 3.9.5 (default, Nov 29 2021, 15:17:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-4.19.0-21-amd64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 11.3.109
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: NVIDIA GeForce GTX 1080 Ti
GPU 1: NVIDIA GeForce GTX 1080 Ti
GPU 2: NVIDIA GeForce GTX 1080 Ti
GPU 3: NVIDIA GeForce GTX 1080 Ti
GPU 4: NVIDIA GeForce GTX 1080 Ti
GPU 5: NVIDIA GeForce GTX 1080 Ti
GPU 6: NVIDIA GeForce GTX 1080 Ti
GPU 7: NVIDIA GeForce GTX 1080 Ti

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] numpy==1.20.3
[pip3] torch==1.11.0
[pip3] torchvision==0.12.0
[conda] Could not collect
```
```
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (GCC) 10.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.28

Python version: 3.9.5 (default, Nov 29 2021, 15:17:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-4.19.0-18-amd64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 11.3.109
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: NVIDIA GeForce GTX 1080
GPU 1: NVIDIA GeForce GTX 1080

Nvidia driver version: 470.103.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] numpy==1.20.3
[pip3] torch==1.11.0
[pip3] torchvision==0.12.0
[conda] Could not collect
```