jax-ml / jax-ml/jax

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

Open
#29,068 2 comments 1 reaction 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.