Repository: pytorch
Issue: `torch.distributed.destroy_process_group()` hangs after CUDA graph capture of NCCL operations
Labels: oncall: distributed,module: c10d

### 🐛 Describe the bug

When I use the CUDA graphs API to capture execution of a model that performs NCCL collectives as part of that execution, it puts the NCCL process group into a state where it cannot cleanly exit -- a call to `destroy_process_group()` hangs. Squinting at a profile, it looks like it is a hang in `ncclCommDestroy`. Here is a fairly small reproducer:

```
import os

import torch
import torch.distributed as dist

class Model(torch.nn.Module):
    def __init__(self, dim, nlayers):
        super().__init__()
        self.layers = torch.nn.ModuleList()
        for _ in range(nlayers):
            self.layers.append(torch.nn.Linear(dim, dim, bias=False))

    @torch.inference_mode()
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
            dist.all_reduce(x)
        return x

def main():
    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    my_device = f"cuda:{rank}"
    torch.cuda.set_device(my_device)

    model = Model(2048, 12)
    model.cuda()

    g = torch.cuda.CUDAGraph()
    static_input = torch.empty((1024, 2048), dtype=torch.float32, device=my_device)
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    # Warmup
    with torch.cuda.stream(s):
        for _ in range(11):
            static_output = model(static_input)
    torch.cuda.current_stream().wait_stream(s)
    with torch.cuda.graph(g):
        static_output = model(static_input)

    torch.cuda.synchronize()
    print("Calling destroy_process_group()...")
    dist.destroy_process_group()
    print("Done.")

if __name__ == "__main__":
    main()
```
Assuming this file is `test.py`, I ran it with:
```
torchrun --standalone --nproc-per-node 2 test.py
```
and you see both processes hang with `Calling destroy_process_group()...`.

A few more notes:
- It doesn't appear to matter whether you every replay the graph (I don't in the example above)
- But if you remove the graph capture itself, it exits cleanly
- And if you remove the `dist.all_reduce()` call inside the model definition, it also exits cleanly

For long-running training jobs, this isn't much of an issue, but for various testing and inference workloads, it is useful to have short-lived process groups (usually spawned with `torch.multiprocessing`), and it's important to clean those up without needing to tear down the parent process as well.

### Versions

Collecting environment information...
PyTorch version: 2.1.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.4
Libc version: glibc-2.31

Python version: 3.9.18 (main, Sep 11 2023, 13:41:44)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-75-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A5000
GPU 1: NVIDIA RTX A5000
GPU 2: NVIDIA RTX A5000
GPU 3: NVIDIA RTX A5000

Nvidia driver version: 535.54.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
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
Byte Order:                      Little Endian
Address sizes:                   48 bits physical, 48 bits virtual
CPU(s):                          64
On-line CPU(s) list:             0-63
Thread(s) per core:              1
Core(s) per socket:              64
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      25
Model:                           1
Model name:                      AMD EPYC 7513 32-Core Processor
Stepping:                        1
CPU MHz:                         2599.998
BogoMIPS:                        5199.99
Virtualization:                  AMD-V
L1d cache:                       4 MiB
L1i cache:                       4 MiB
L2 cache:                        32 MiB
L3 cache:                        16 MiB
NUMA node0 CPU(s):               0-63
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
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 rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr wbnoinvd arat npt lbrv nrip_save tsc_scale vmcb_clean pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid arch_capabilities

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.0
[pip3] torch==2.1.0+cu121
[pip3] torch-tb-profiler==0.4.1
[pip3] torchaudio==2.1.0+cu121
[pip3] torchvision==0.16.0+cu121
[pip3] triton==2.1.0
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] torch                     2.1.0+cu121              pypi_0    pypi
[conda] torch-tb-profiler         0.4.1                    pypi_0    pypi
[conda] torchaudio                2.1.0+cu121              pypi_0    pypi
[conda] torchvision               0.16.0+cu121             pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225