An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
- Dominant language
- Python
- Stars
- 36.3k
- Forks
- 3.8k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 399
Description
### Description
Using a GCP VM of type N1 with Tesla T1 GPU.
I updated jax to 6.0.1 and getting error with recognizing the GPU.
Consider the following reproduction:
```bash
myname@mymachine:~$ uv init jaxerror
Initialized project `jaxerror` at `/home/myname/jaxerror`
myname@mymachine:~$ cd jaxerror/
myname@mymachine:~/jaxerror$ uv add jax[cuda12]
Using CPython 3.12.8
Creating virtual environment at: .venv
Resolved 20 packages in 428ms
Installed 19 packages in 41ms
+ jax==0.6.1
+ jax-cuda12-pjrt==0.6.1
+ jax-cuda12-plugin==0.6.1
+ jaxlib==0.6.1
+ ml-dtypes==0.5.1
+ numpy==2.2.6
+ nvidia-cublas-cu12==12.8.4.1
+ nvidia-cuda-cupti-cu12==12.9.19
+ nvidia-cuda-nvcc-cu12==12.9.41
+ nvidia-cuda-runtime-cu12==12.9.37
+ nvidia-cudnn-cu12==9.10.1.4
+ nvidia-cufft-cu12==11.4.0.6
+ nvidia-cusolver-cu12==11.7.4.40
+ nvidia-cusparse-cu12==12.5.9.5
+ nvidia-nccl-cu12==2.26.5
+ nvidia-nvjitlink-cu12==12.9.41
+ nvidia-nvshmem-cu12==3.2.5
+ opt-einsum==3.4.0
+ scipy==1.15.3
myname@mymachine:~/jaxerror$ vi main.py
myname@mymachine:~/jaxerror$ uv run main.py
ERROR:2025-05-28 18:06:26,600:jax._src.xla_bridge:444: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
File "/home/myname/jaxerror/.venv/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 135, in _version_check
version = get_version()
^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:81: operation cusparseGetProperty(MAJOR_VERSION, &major) failed: The cuSPARSE library was not found.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/myname/jaxerror/.venv/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 442, in discover_pjrt_plugins
plugin_module.initialize()
File "/home/myname/jaxerror/.venv/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 230, in initialize
_check_cuda_versions(raise_on_first_error=True)
File "/home/myname/jaxerror/.venv/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 199, in _check_cuda_versions
_version_check("cuSPARSE", cuda_versions.cusparse_get_version,
File "/home/myname/jaxerror/.venv/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 139, in _version_check
raise RuntimeError(err_msg) from e
RuntimeError: Unable to load cuSPARSE. Is it installed?
WARNING:2025-05-28 18:06:26,604:jax._src.xla_bridge:791: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
[0. 1.05 2.1 3.1499999 4.2 ]
```
### System info (python version, jaxlib version, accelerator, etc.)
```
jax: 0.6.1
jaxlib: 0.6.1
numpy: 2.2.6
python: 3.12.8 (main, Dec 19 2024, 14:33:20) [Clang 18.1.8 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='gal-test', release='5.10.0-34-cloud-amd64', version='#1 SMP Debian 5.10.234-1 (2025-02-24)', machine='x86_64')
$ nvidia-smi
Wed May 28 18:08:43 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 Tesla T4 On | 00000000:00:04.0 Off | 0 |
| N/A 41C P8 15W / 70W | 3MiB / 15360MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| No running processes found |
+-----------------------------------------------------------------------------------------+
```
Contributor guide
Research direction
Reproduce the failure with the shown uv environment and inspect jax_plugins/xla_cuda12/__init__.py, especially initialize and _check_cuda_versions, along with jax/_src/xla_bridge.py. Confirm how the cuSPARSE lookup behaves with the reported driver and installed packages; done means the CUDA plugin initializes and JAX reports the Tesla T4 instead of falling back to CPU.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- linux, python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100