Repository: pytorch
Issue: WSL - Unexpected OOM error upon initialising torch.cuda
Labels: module: cuda,module: memory usage,triaged,module: wsl

### 🐛 Describe the bug

# Issue
The following error occurs in one form or another whenever most PyTorch initialisation methods are called:
```py
    File "/tmp/pip-build-env-rrigorsf/overlay/lib/python3.11/site-packages/torch/cuda/__init__.py", line 439, in get_device_capability
      prop = get_device_properties(device)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/tmp/pip-build-env-rrigorsf/overlay/lib/python3.11/site-packages/torch/cuda/__init__.py", line 453, in get_device_properties
      _lazy_init()  # will define _get_device_properties
      ^^^^^^^^^^^^
    File "/tmp/pip-build-env-rrigorsf/overlay/lib/python3.11/site-packages/torch/cuda/__init__.py", line 302, in _lazy_init
      torch._C._cuda_init()
  RuntimeError: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 2: out of memory
```
# Reproducing
The issue is easily reproduced by any of the following:
```py
Python 3.11.9 (main, Apr  6 2024, 17:59:24) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.init()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/txin/.local/lib/python3.11/site-packages/torch/cuda/__init__.py", line 269, in init
    _lazy_init()
  File "/home/txin/.local/lib/python3.11/site-packages/torch/cuda/__init__.py", line 302, in _lazy_init
    torch._C._cuda_init()
RuntimeError: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 2: out of memory
>>> torch.cuda.is_available()
/home/txin/.local/lib/python3.11/site-packages/torch/cuda/__init__.py:141: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 2: out of memory (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)
  return torch._C._cuda_getDeviceCount() > 0
False
```
# Attempted workarounds
I have tried installing in many environments over the last year, none with much success. As long as I am using WSL2, changing none of the following made any difference:
- Linux distributions (Ubuntu vs Debian etc)
- NVIDIA driver versions (543 vs 551 etc)
- Python versions (3.9.10 vs 3.10.12 etc)
- PyTorch versions (2.0.1 vs 2.1.2 etc)
- CUDA versions (11.8 vs 12.1 etc)
- CUDA toolkit versions (although I don't remember the previous versions I attempted)

Ultimately, I have decided to post an issue here, as I have yet to see an official fix for this issue at any point in time. It is the main problem preventing me from working with a multitude of other programs/frameworks.
# Known workarounds
The following modifications have been known to "fix" (or more accurately, bypass this issue, allowing `torch.cuda.init()` and related functions to be called without erroring):
- Setting environment variable `PYTORCH_NVML_BASED_CUDA_CHECK=1` before running the script
  - Update: This no longer appears to work. It appears either the error is occurring somewhere else, or the flag is no longer being used
- Calling `torch.cuda.device_count()` immediately after torch is imported in the script
- Directly modifying `torch/cuda/__init__.py` to include a call to `device_count()` immediately after the function is defined
### Unfortunately, this is not a sustainable fix, as it does not work for installing libraries that depend on specific versions of one another, as if they include torch it will be reinstalled from downloaded source/binaries and installed in a separate temporary environment, and will not permit such modifications during the install.

### Versions

```py
python collect_env.py
Collecting environment information...
/home/txin/.local/lib/python3.11/site-packages/torch/cuda/__init__.py:141: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 2: out of memory (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)
  return torch._C._cuda_getDeviceCount() > 0
PyTorch version: 2.2.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.9 (main, Apr  6 2024, 17:59:24) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 4070
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3060
GPU 5: NVIDIA GeForce RTX 3060

Nvidia driver version: 552.12
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):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen 5 7600 6-Core Processor
CPU family:                         25
Model:                              97
Thread(s) per core:                 2
Core(s) per socket:                 6
Socket(s):                          1
Stepping:                           2
BogoMIPS:                           7599.87
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 tsc_reliable nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero xsaveerptr arat npt nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm
Virtualization:                     AMD-V
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          192 KiB (6 instances)
L1i cache:                          192 KiB (6 instances)
L2 cache:                           6 MiB (6 instances)
L3 cache:                           32 MiB (1 instance)
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 Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
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 conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.0
[pip3] triton==2.2.0
[conda] Could not collect
```

cc @ptrblck