Repository: pytorch
Issue: Repro for non-deterministic "operation not permitted when stream is capturing" crash
Labels: oncall: distributed

### 🐛 Describe the bug

Running on 2.1.0, seems to introduce a strange non-deterministic "operation not permitted when stream is capturing" error on the following code. Bisected to [this commit](https://github.com/pytorch/pytorch/commit/f89af601833f5e755a10596733b0834f6161e206) as the one that triggers it:

```
import os
import time
import torch.distributed
import torch.multiprocessing

def run_processes():
    mp = torch.multiprocessing.get_context(method="forkserver")
    processes = [mp.Process(target=func, args=[rank]) for rank in range(2)]
    for p in processes:
        p.start()
    for p in processes:
        p.join()

def func(local_rank):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(local_rank)
    torch.distributed.init_process_group(
        "nccl",
        init_method=f"tcp://127.0.0.1:2379",
        rank=local_rank,
        world_size=2,
    )

    input = torch.ones((1,), device="cuda")
    output = torch.ones((2, 1), device="cuda")
    group = torch.distributed.new_group([0, 1], backend="nccl")

    torch.distributed.all_gather_into_tensor(output, input, group=group)
    for i in range(10000):
        print(i)
        with torch.cuda.graph(torch.cuda.CUDAGraph()):
            torch.distributed.all_gather_into_tensor(output, input, group=group)
            time.sleep(1.0) # Adding this makes it crash. Why??

if __name__ == '__main__':
    run_processes()
```

A few other things of note:
* While the time.sleep() seems silly, this is my minimal repro from a larger project that does not do silly sleeps.
* The crash requires more than one rank, and requires cuda graph captures. Without those, there is no crash.
* Also strangely, using the default process group makes the crash go away. It only crashes if I create my own (identical) process group.
* The crash only occurs with the time.sleep(0.3) inside the cuda graph capture. Removing the sleep, or moving the sleep outside the graph capture, makes the problem go away.
* The crash is non-deterministic, and crashes at a different number of iterations each time I run this. I think I've gotten the repro such that it crashes on the second iteration, but with other variations (ie the gc.collect one below, or different sleep amounts), I've seen it take ~500 iterations to crash.
* The crash _also_ occurs if I do not have the time.sleep(), but I instead remove the gc.collect() inside torch.cuda.graph.__enter__. That makes it seem like garbage collections are causing the issue when they nondeterministically run, but if I remove gc.collect() _and disable garbage collection for the duration of the cuda graph capture,_ it still crashes! My theory is that it's the _time_ spent waiting for gc.collect() that is important (rather than the fact that it is garbage collecting), as replacing gc.collect() with a time.sleep() seems to have a similar effect without the GC.
* The stack traces printed are very unhelpful at figuring out _which_ operation is causing the issue, and in a real-world project the crashes seem to occur at random places in the python code tree. It took me a bit to track this down to the minimal repro I'm sharing here, due to the difficulty of interpreting the stack traces.

The full error returned is:
```
Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7f2e6a312617 in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7f2e6a2cd98d in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7f2e6a3cd9f8 in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10d::ProcessGroupNCCL::WorkNCCL::finishedGPUExecutionInternal() const + 0x80 (0x7f2df5e3caf0 in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #4: c10d::ProcessGroupNCCL::WorkNCCL::isCompleted() + 0x58 (0x7f2df5e40918 in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #5: c10d::ProcessGroupNCCL::workCleanupLoop() + 0x24b (0x7f2df5e5715b in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #6: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x78 (0x7f2df5e57468 in /home/user/.pyenv/versions/3.11.5/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so)
frame #7: <unknown function> + 0xdc253 (0x7f2e3a6b0253 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #8: <unknown function> + 0x94b43 (0x7f2e76cfab43 in /lib/x86_64-linux-gnu/libc.so.6)
frame #9: <unknown function> + 0x126a00 (0x7f2e76d8ca00 in /lib/x86_64-linux-gnu/libc.so.6)
```

Is this possibly related to what https://github.com/pytorch/pytorch/pull/105182 is trying to fix? I am dumping a repro on stock 2.1.0 before going deeper

### 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 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.0
Libc version: glibc-2.35

Python version: 3.11.5 (main, Sep 26 2023, 02:55:14) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.4.209-116.367.amzn2.x86_64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.57.02
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:                   46 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          96
On-line CPU(s) list:             0-4,6-10,12-16,18-22,24-28,30-34,36-40,42-46,48-52,54-58,60-64,66-70,72-76,78-82,84-88,90-94
Off-line CPU(s) list:            5,11,17,23,29,35,41,47,53,59,65,71,77,83,89,95
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
CPU family:                      6
Model:                           85
Thread(s) per core:              2
Core(s) per socket:              20
Socket(s):                       2
Stepping:                        7
BogoMIPS:                        6000.01
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.3 MiB (40 instances)
L1i cache:                       1.3 MiB (40 instances)
L2 cache:                        40 MiB (40 instances)
L3 cache:                        71.5 MiB (2 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-4,6-10,12-16,18-22,48-52,54-58,60-64,66-70
NUMA node1 CPU(s):               24-28,30-34,36-40,42-46,72-76,78-82,84-88,90-94
Vulnerability Itlb multihit:     KVM: Vulnerable
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] flake8-2020==1.8.1
[pip3] flake8-bugbear==23.7.10
[pip3] flake8-datetimez==20.10.0
[pip3] flake8-noqa==1.3.2
[pip3] flake8-pie==0.16.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pytest-style==1.7.2
[pip3] flake8-trio==23.2.5
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] torch==2.1.0
[pip3] torchtyping==0.1.4
[pip3] torchvision==0.16.0+75046bd.post10
[pip3] triton==2.1.0
[pip3] triton-ops==0.0.0
[conda] Could not collect

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