Repository: pytorch
Issue: Setting the cuda device when using start_processes in Jupyter on Ampere leads to CUDA reinitialization error
Labels: module: multiprocessing,module: cuda,triaged

### 🐛 Describe the bug

When using a non-amp GPU this issue does not arise (tested environment for success was two Tesla T4's out of GCP). However when testing on  an amp-enabled GPU system I cannot start a multi-gpu processes and it will lead to a re-initialization error. Tested environment for this was two RTX A4000 from runpod.io

To recreate, run the following code in a Jupyter Notebook cell:

```python
import torch
import os
from contextlib import contextmanager
from torch.multiprocessing import start_processes

@contextmanager
def patch_environment(**kwargs):
    """
    A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting.
    Will convert the values in `kwargs` to strings and upper-case all the keys.
    """
    for key, value in kwargs.items():
        os.environ[key.upper()] = str(value)

    yield

    for key in kwargs:
        del os.environ[key.upper()]

def training_loop():
    device = torch.device("cuda", int(os.environ.get("LOCAL_RANK", -1)))
    torch.cuda.set_device(device)
    print(device)

def launch_func(index, *args):
    with patch_environment(
        local_rank=str(index),
        rank=str(index)
    ):
        training_loop()

with patch_environment(
    world_size=2,
    master_addr="127.0.0.1",
):
    start_processes(launch_func, nprocs=2, start_method="fork")
```

Desired effect is I should print the current device that was set, instead I receive the following trace:

```bash
---------------------------------------------------------------------------
ProcessRaisedException                    Traceback (most recent call last)
/tmp/ipykernel_14417/2989602212.py in <module>
      4     master_addr="127.0.0.1",
      5 ):
----> 6     start_processes(launch_func, nprocs=2, start_method="fork")

/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    196 
    197     # Loop on join until it returns True or raises an exception.
--> 198     while not context.join():
    199         pass
    200 

/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    158         msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
    159         msg += original_trace
--> 160         raise ProcessRaisedException(msg, error_index, failed_process.pid)
    161 
    162 

ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/tmp/ipykernel_14417/1434743469.py", line 6, in launch_func
    training_loop()
  File "/tmp/ipykernel_14417/3423764464.py", line 3, in training_loop
    torch.cuda.set_device(device)
  File "/opt/conda/lib/python3.7/site-packages/torch/cuda/__init__.py", line 314, in set_device
    torch._C._cuda_setDevice(device)
  File "/opt/conda/lib/python3.7/site-packages/torch/cuda/__init__.py", line 208, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. To use CUDA with "
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
```

### Versions

PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.7.13 (default, Mar 29 2022, 02:18:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.0-47-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA RTX A4000
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.21.5
[pip3] torch==1.12.0
[pip3] torchelastic==0.2.0
[pip3] torchtext==0.13.0
[pip3] torchvision==0.13.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] numpy                     1.21.5           py37he7a7128_2  
[conda] numpy-base                1.21.5           py37hf524024_2  
[conda] pytorch                   1.12.0          py3.7_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchelastic              0.2.0                    pypi_0    pypi
[conda] torchtext                 0.13.0                     py37    pytorch
[conda] torchvision               0.13.0               py37_cu113    pytorch


cc @VitalyFedyunin @ngimel