Repository: pytorch
Issue: torch.compile/triton holding GIL during compilation and CompiledKernel call results in deadlocks during distributed training
Labels: triaged,oncall: pt2,module: inductor,upstream triton,module: startup-tracing-compile

### 🐛 Describe the bug

We are seeing our training processes get deadlocked as follows. There is a NCCL error causing the allreduce in the backward pass to get stuck. When this occurs, typically the GPU locks up until a [ncclCommAbort](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/comms.html#ncclcommabort) is called and even simple cuda calls like cudaLaunchKernel, cudaEventQuery etc. get stuck on the host side.  As a result, the main thread is stuck as follows:

```
Thread 432919 (idle): "MainThread"
    0x155555277197 (libc.so.6)
    pthread_cond_wait (libc.so.6)
    __gnu_cxx::__nothrow_wait_cv::wait (compatibility-condvar.cc:100)
    torch::autograd::ReadyQueue::pop (libtorch_cpu.so)
    torch::autograd::Engine::thread_main (libtorch_cpu.so)
    torch::autograd::Engine::execute_with_graph_task (libtorch_cpu.so)
    torch::autograd::python::PythonEngine::execute_with_graph_task (libtorch_python.so)
    torch::autograd::Engine::execute (libtorch_cpu.so)
    torch::autograd::python::PythonEngine::execute (libtorch_python.so)
    THPEngine_run_backward (libtorch_python.so)
    backward (torch/autograd/__init__.py:200)
    backward (torch/_tensor.py:503)
  ```
  
  Simultaneously, triton is stuck due to the above NCCL issue while trying to call a compiled kernel. The bad part here though is that triton is holding GIL (see active+gil in the trace below) and is stuck. As a result, the entire Python process freezes up and no other python threads are able to execute:

```
Thread 493983 (active+gil): "Dummy-20"
    sched_yield (libc.so.6)
    0x15550de3cbdb (libcuda.so.535.54.03)
    0x15550e10adb5 (libcuda.so.535.54.03)
    0x15550e10b32a (libcuda.so.535.54.03)
    0x15550df18454 (libcuda.so.535.54.03)
    0x15550ddf7836 (libcuda.so.535.54.03)
    0x15550ddf7f3d (libcuda.so.535.54.03)
    0x15550ddfaa27 (libcuda.so.535.54.03)
    0x15550dfe9b3e (libcuda.so.535.54.03)
    launch (triton_.so)
    launcher (<string>:6)
    run (torch/_inductor/triton_ops/autotune.py:190)
    call (coksmpcybzxz4bxfvl56kcd6k5czlqpovjyjwg4yry6do7p7hcim.py:871)
    run (torch/_inductor/compile_fx.py:248)
    _fn (torch/_dynamo/eval_frame.py:209)
    call_func_with_args (torch/_functorch/aot_autograd.py:1249)
    call_compiled_backward (torch/_functorch/aot_autograd.py:2341)
    backward (torch/_functorch/aot_autograd.py:2365)
```

Another example is that triton also get stuck during compilation as well due to a `torch.cuda.synchronize` getting stuck as NCCL kernels are stuck. Even in this case, triton is stuck while holding the GIL and results in the entire Python process to lock up:

```
Thread 1957540 (active+gil): "Dummy-20"
    synchronize (torch/cuda/__init__.py:688)
    _precompile_config (torch/_inductor/triton_ops/autotune.py:94)
    <listcomp> (torch/_inductor/triton_ops/autotune.py:70)
    precompile (torch/_inductor/triton_ops/autotune.py:69)
    _load_kernel (torch/_inductor/codecache.py:554)
    result (torch/_inductor/codecache.py:574)
    wait (torch/_inductor/codecache.py:715)
    <module> (cpb2lazn3hvywa5maow5ngritvyt6pwlcq6ela3fl5rbblyj636x.py:1578)
    load (torch/_inductor/codecache.py:528)
    compile_to_module (torch/_inductor/graph.py:575)
    time_wrapper (torch/_dynamo/utils.py:163)
    compile_to_fn (torch/_inductor/graph.py:586)
    compile_fx_inner (torch/_inductor/compile_fx.py:177)
    inner (contextlib.py:75)
    inner (torch/_inductor/debug.py:239)
    debug_wrapper (torch/_dynamo/debug_utils.py:595)
    bw_compiler (torch/_inductor/compile_fx.py:441)
    time_wrapper (torch/_dynamo/utils.py:163)
    _fn (torch/_dynamo/eval_frame.py:209)
    _wrapped_bw_compiler (torch/_dynamo/backends/common.py:38)
    call_compiled_backward (torch/_functorch/aot_autograd.py:2336)
    backward (torch/_functorch/aot_autograd.py:2365)
    apply (torch/autograd/function.py:274)
```

We have some custom logic in our application to recover from NCCL errors by calling `_abort` on the ProcessGroup here: https://github.com/pytorch/pytorch/blob/main/torch/csrc/distributed/c10d/init.cpp#L2226 (We set `NCCL_ASYNC_ERROR_HANDLING=0` and manage errors ourselves). This logic is usually able to recover from the NCCL errors. However, now since triton is holding GIL and is stuck, the entire Python process is frozen and none our custom logic in python can run and recover from the situation.


I'm wondering if it is possible to release GIL before entering C/C++ land (similar to other PyTorch operations). This would resolve the issue and allow other python threads in the application to recover from the situation.

### Error logs

_No response_

### Minified repro

_No response_

### Versions

PyTorch 2.0.1

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @muchulee8 @aakhundov @ColinPeppler @Xia-Weiwen @ngimel