[NVIDIA] Add opt in CUDA availability check using NVML#10933
Conversation
warrendeng
changed the title
| return [libdevice_dir, *libcuda_dirs()] | ||
|
|
||
|
|
||
| # Adapted from PyTorch's NVML-based CUDA availability check: |
There was a problem hiding this comment.
these helper functions are equivalent to how PyTorch uses nvml
25a55dc to
0af5c91
Compare
| def is_active(): | ||
| if os.getenv("TRITON_NVML_BASED_CUDA_CHECK") == "1": | ||
| count = _device_count_nvml() | ||
| if count >= 0: |
There was a problem hiding this comment.
if env var is not set or if _device_count_nvml returns -1 (inconclusive) we fall back to old path
|
|
||
| def test_cuda_driver_is_active_with_nvml(monkeypatch): | ||
| cuda_calls = [] | ||
| monkeypatch.setenv("TRITON_NVML_BASED_CUDA_CHECK", "1") |
There was a problem hiding this comment.
Remove this. You can test this with knobs infrastructure. You should add this to the Cuda knobs.
|
|
||
| @staticmethod | ||
| def is_active(): | ||
| if os.getenv("TRITON_NVML_BASED_CUDA_CHECK") == "1": |
There was a problem hiding this comment.
do we need an env var at all? Can we check nvml and if not there fall back? Wouldn't nvml always be there?
|
|
||
| # Adapted from PyTorch's NVML-based CUDA availability check: | ||
| # https://github.com/pytorch/pytorch/blob/70d99e998b4955e0049d13a98d77ae1b14db1f45/torch/cuda/__init__.py | ||
| def _parse_visible_devices(): |
There was a problem hiding this comment.
this looks quite over-complicated, so we really need all this?
CudaDriver.is_active() calls cuInit(), which initializes CUDA driver state in a parent process and can prevent a subsequently forked child from initializing CUDA. Add an opt-in NVML probe so callers that need to fork can check NVIDIA device availability without initializing CUDA when NVML returns a conclusive result. Keep the existing CUDA probe as the default and as the fallback for inconclusive NVML checks.
0af5c91 to
7620712
Compare
njriasan
left a comment
njriasan
left a comment
There was a problem hiding this comment.
let's focus on removing all this string parsing.
|
|
||
| @staticmethod | ||
| def is_active(): | ||
| if os.getenv("TRITON_NVML_BASED_CUDA_CHECK") == "1": |
| return 0 | ||
| try: | ||
| if isinstance(visible_devices[0], str): | ||
| if visible_devices[0].startswith("MIG-"): |
There was a problem hiding this comment.
Is the PyTorch version in C++? We should do this in native code as it could get expensive.
|
|
||
| # Adapted from PyTorch's NVML-based CUDA availability check: | ||
| # https://github.com/pytorch/pytorch/blob/70d99e998b4955e0049d13a98d77ae1b14db1f45/torch/cuda/__init__.py | ||
| def _parse_visible_devices(): |
There was a problem hiding this comment.
Why do we need this much string parsing? Is there not an Nvidia API to get the first device? How do we find a device normally?
3 participants
Context
Triton #9578 changed
CudaDriver.is_active()implementation from usingtorch.cuda.is_available()tocuInit(0). Downstream customers such as vLLM and PyTorch may callCudaDriver.is_activein a parent process before forking a child that later initializes CUDA/callsCudaDriver.is_active()again. BecausecuInit(0)poisons the CUDA state of any forked children, the change causes child processes to unable to initialize CUDA.Examples of this are in vllm-project/vllm#46996 and pytorch/pytorch#189287.
Details
To mitigate the issue, we add an additional path for checking NVIDIA driver status that queries NVML without having to initialize CUDA. We gate this path with env var
TRITON_NVML_BASED_CUDA_CHECK=1and set the var to 0 by default. If this new path fails or if the env var is not set, we fall back to the existingcuInit(0)path.The default behavior remains unchanged and we continue to preserve Triton's torch-independent runtime path.
Specifically for vLLM, the env var
PYTORCH_NVML_BASED_CUDA_CHECKis set to 1 by default. When this is set,torch.cuda.is_available()actually uses NVML instead ofcuInit(0)to check driver status. Thus for vLLM, setting our new env varTRITON_NVML_BASED_CUDA_CHECK=1should have restore behavior prior to #9578 with no other effect.Test Plan
python -m pytest -s --tb=short python/test/unit/runtime/test_driver.py python/test/unit/runtime/test_no_torch_dispatch.pyfork():fork(): child CUDA initialization passedpre-commit run --from-ref origin/main --to-ref HEAD