Repository: pytorch
Issue: Incorrect GPU management and deadlocks without torch.cuda.set_device
Labels: oncall: distributed,triaged,actionable

## Issue description

When I am running distributed and I simply set `CUDA_VISIBLE_DEVICES` in each rank:
- Running `torch.distributed.barrier()` makes rank 1 occupy GPU memory on the GPU of rank 0 (even though that GPU shouldn't be visible due to `CUDA_VISIBLE_DEVICES`)
- Running FSDP forward deadlocks

Both issues go away if I set `torch.cuda.set_device`

## Code example

```python
import os
import time
import torch
import torch.distributed.fsdp
import functools

from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy


def train() -> None:

    torch.distributed.init_process_group('nccl')

    local_rank = torch.distributed.get_rank()

    # DOESN'T WORK
    os.environ['CUDA_VISIBLE_DEVICES'] = str(local_rank)
    torch.cuda.empty_cache()
    
    # torch.cuda.set_device(local_rank)  # FIXES PROBLEMS

    print("BEFORE", torch.distributed.get_rank())
    time.sleep(5)

    torch.distributed.barrier()

    print("AFTER", torch.distributed.get_rank())
    time.sleep(15)

    module = torch.nn.Sequential(
        torch.nn.Conv2d(10, 10, 3),
        torch.nn.Conv2d(10, 10, 3),
        torch.nn.Conv2d(10, 10, 3),
        torch.nn.Conv2d(10, 10, 3),
    )
    module.cuda()

    my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=100
    )

    fsdp_module = torch.distributed.fsdp.FullyShardedDataParallel(
        module,
        auto_wrap_policy=my_auto_wrap_policy,
        use_orig_params=True,
    )

    print("FORWARD", torch.distributed.get_rank())

    out = fsdp_module(torch.zeros([10, 10, 900, 900]))

    print("END", torch.distributed.get_rank())
    time.sleep(15)

    torch.distributed.destroy_process_group()


if __name__ == "__main__":
    train()
```

I run this with
```
torchrun --standalone --nnodes 1 --nproc-per-node 2 --rdzv-backend=c10d --rdzv-endpoint=localhost:0 test_train.py
```
As a result, after the call to `torch.distributed.barrier()`, `nvidia-smi` reports:
```
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A   1528872      C   ...mamba/envs/debug-env/bin/python3.10        404MiB |
|    0   N/A  N/A   1528873      C   ...mamba/envs/debug-env/bin/python3.10        262MiB |
|    1   N/A  N/A   1528873      C   ...mamba/envs/debug-env/bin/python3.10        404MiB |
+-----------------------------------------------------------------------------------------+
```
Rank 1 has occupied memory of the GPU of rank 0, which shouldn't be happening.

The process also hang at the call to `FSDP.forward`, even though both ranks reach that stage. At that point, `nvidia-smi` reports even more unbalanced memory allocation
```
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A   1540224      C   ...mamba/envs/debug-env/bin/python3.10       1024MiB |
|    0   N/A  N/A   1540225      C   ...mamba/envs/debug-env/bin/python3.10        596MiB |
|    1   N/A  N/A   1540225      C   ...mamba/envs/debug-env/bin/python3.10        400MiB |
+-----------------------------------------------------------------------------------------+
```

If I simply set `torch.cuda.set_device(local_rank)`, everything works properly. Note that I don't even need to set `CUDA_VISIBLE_DEVICES`. For comparison, after the call to `torch.distributed.barrier`, `nvidia-smi` reports
```
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A   1529231      C   ...mamba/envs/debug-env/bin/python3.10        404MiB |
|    1   N/A  N/A   1529232      C   ...mamba/envs/debug-env/bin/python3.10        404MiB |
+-----------------------------------------------------------------------------------------+
```
and at the end of the program it reports
```
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A   1532552      C   ...mamba/envs/debug-env/bin/python3.10       2012MiB |
|    1   N/A  N/A   1532553      C   ...mamba/envs/debug-env/bin/python3.10       2012MiB |
+-----------------------------------------------------------------------------------------+
```
Equal memory usage on both GPUs at both stages, as expected

## System Info
```
Collecting environment information...
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux trixie/sid (x86_64)
GCC version: (Debian 14.2.0-1) 14.2.0
Clang version: Could not collect
CMake version: version 3.29.6
Libc version: glibc-2.39

Python version: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:45:18) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-6.9.12-1-insait-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000

Nvidia driver version: 550.90.07
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               96
On-line CPU(s) list:                  0-95
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 74F3 24-Core Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   24
Socket(s):                            2
Stepping:                             1
Frequency boost:                      enabled
CPU(s) scaling MHz:                   45%
CPU max MHz:                          4037.5000
CPU min MHz:                          1500.0000
BogoMIPS:                             6400.62
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm debug_swap
Virtualization:                       AMD-V
L1d cache:                            1.5 MiB (48 instances)
L1i cache:                            1.5 MiB (48 instances)
L2 cache:                             24 MiB (48 instances)
L3 cache:                             512 MiB (16 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-23,48-71
NUMA node1 CPU(s):                    24-47,72-95
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Vulnerable
Vulnerability Spec store bypass:      Vulnerable
Vulnerability Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.4.0+cu121
[pip3] triton==3.0.0
[conda] Could not collect
```
The problem also happens with torch-2.2.0 and torch-2.3.0

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o