Installing `jax[cuda12]==0.6.1` leads to a RuntimeError because the cuSPARSE library cannot be found

[maxencefaldor]

Description

Installing jax[cuda12]==0.6.1 leads to a RuntimeError because the cuSPARSE library cannot be found, causing JAX to fall back to CPU execution despite the presence of CUDA-enabled jaxlib components.

Here is how to reproduce:

$ uv venv
Using CPython 3.13.3
Creating virtual environment at: .venv
Activate with: source .venv/bin/activate

$ source .venv/bin/activate

(.venv) $ uv pip install -U "jax[cuda12]"
Using Python 3.13.3 environment at: .venv
Resolved 19 packages in 304ms
Prepared 19 packages in 37ms
Installed 19 packages in 6.00s
 + 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

(.venv) $ python
Python 3.13.3 (main, Apr  9 2025, 04:03:52) [Clang 20.1.0 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.numpy.zeros(3)
ERROR:2025-05-22 12:51:37,974:jax._src.xla_bridge:444: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File ".venv/lib/python3.13/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 ".venv/lib/python3.13/site-packages/jax/_src/xla_bridge.py", line 442, in discover_pjrt_plugins
    plugin_module.initialize()
    ~~~~~~~~~~~~~~~~~~~~~~~~^^
  File ".venv/lib/python3.13/site-packages/jax_plugins/xla_cuda12/__init__.py", line 230, in initialize
    _check_cuda_versions(raise_on_first_error=True)
    ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.13/site-packages/jax_plugins/xla_cuda12/__init__.py", line 199, in _check_cuda_versions
    _version_check("cuSPARSE", cuda_versions.cusparse_get_version,
    ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                   cuda_versions.cusparse_build_version,
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                   # Ignore patch versions.
                   ^^^^^^^^^^^^^^^^^^^^^^^^
                   scale_for_comparison=100,
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
                   min_supported_version=12100)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.13/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-22 12:51:38,032: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.
Array([0., 0., 0.], dtype=float32)

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
jax:    0.6.1
jaxlib: 0.6.1
numpy:  2.2.6
python: 3.13.3 (main, Apr  9 2025, 04:03:52) [Clang 20.1.0 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='***', release='5.15.0-1018-gcp-tcpx', version='#18-Ubuntu SMP Fri Jul 26 14:21:24 UTC 2024', machine='x86_64')

$ nvidia-smi
Thu May 22 12:57:55 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| 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  NVIDIA H100 80GB HBM3          On  | 00000000:04:00.0 Off |                    0 |
| N/A   42C    P0             73W / 700W |  4MiB / 81559MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    No running processes found |
+---------------------------------------------------------------------------------------+

[ybaturina]

hi @maxencefaldor may I ask you to run the following and then attach the printed log here?

LD_DEBUG=files,libs python -c "import jax; jax.numpy.zeros(3)"

[maxencefaldor]

Here it is

log.txt

[ybaturina]

This is what's happening:

trying file=/home/project/._venv/lib/python3.13/site-packages/jax_cuda12_plugin/../nvidia/cusparse/lib/libcusparse.so.12
   1382674:	
   1382674:	file=libcusparse.so.12 [0];  generating link map
   1382674:	  dynamic: 0x000014e28ef10a90  base: 0x000014e271e3f000   size: 0x000000001d2042d0
   1382674:	    entry: 0x000014e271eab480  phdr: 0x000014e271e3f040  phnum:                  9
   1382674:	
   1382674:	
   1382674:	file=libnvJitLink.so.12 [0];  needed by /home/project/._venv/lib/python3.13/site-packages/jax_cuda12_plugin/../nvidia/cusparse/lib/libcusparse.so.12 [0]
   1382674:	find library=libnvJitLink.so.12 [0]; searching
   1382674:	 search path=/home/common_modules/nccl/nccl_2.18.3-1+cuda12.1_x86_64/lib:/home/common_modules/cudnn/cudnn-8.9.7.29_cuda12/lib:/home/common_modules/cuda/cuda-12.1/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64		(LD_LIBRARY_PATH)
   1382674:	  trying file=/home/common_modules/nccl/nccl_2.18.3-1+cuda12.1_x86_64/lib/libnvJitLink.so.12
   1382674:	  trying file=/home/common_modules/cudnn/cudnn-8.9.7.29_cuda12/lib/libnvJitLink.so.12
   1382674:	  trying file=/home/common_modules/cuda/cuda-12.1/lib64/libnvJitLink.so.12
   1382674:	
   1382674:	file=libnvJitLink.so.12 [0];  generating link map

The libcusparse.12 is found properly, but libnvJitLink.so.12 is needed by libcusparse.so.12, and libnvJitLink.so.12 is found in a system-wide CUDA installation with the version 12.1.

Let me check the LD_DEBUG log on my machine to see where libnvJitLink.so.12 is found.

[ybaturina]

Would you unset LD_LIBRARY_PATH please and try again?

[maxencefaldor]

It works when I unset LD_LIBRARY_PATH. Thank you for solving the issue.

Should I always unset LD_LIBRARY_PATH before using JAX? I didn't have this issue before, and I am wondering if something is wrong with how my system is configured, or if something changed in JAX?

[ybaturina]

I think there is incompatibility between libcusparse version required by jax wheels and libnvjitlink version available in your system.

The ELF files in jax wheels have the RPATH tag set in NEEDED section (you can check it via readelf --dynamic <path to ELF file>). dlopen() searches RPATH first (if set), then LD_LIBRARY_PATH, then RUNPATH (see https://man7.org/linux/man-pages/man3/dlopen.3.html)

What happened here is that the libraries required by jax wheels were found in correct location (via RPATH pointing to NVIDIA pypi wheels installation). However, libcusparse.so.12 (which is not controlled by JAX) uses RUNPATH, and when dlopen() was searching for libnvJitLink.so.12 location, RUNPATH value got overriden by LD_LIBRARY_PATH on your system.

So the libcusparse version was correct, and libnvjitlink version was old and incompatible with libcusparse.

So yes, I suggest not using LD_LIBRARY_PATH if you use JAX.