Issue installing older version of JAX with CUDA11

[seishiroono]

Description

I am trying to JAX version 0.4.25 (or 0.4.24) with CUDA 11. I followed the recommended command here:

Since it is an older version of JAX, I think I have to introduce CUDA environment by myself.
Then, I performed the following commands. (I could not find CuDNN 8.6 in conda)

conda install nvidia/label/cuda-11.8.0::cuda-toolkit
conda install conda-forge::cudnn=8.8.0
pip install jax==0.4.25 jaxlib==0.4.25+cuda11_cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Error

When I did the following, I found that my GPUs were not recognized.

>>> import jax
>>> jax.devices()

CUDA backend failed to initialize: Unable to load CUDA. Is it installed? (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

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

jax:    0.4.25
jaxlib: 0.4.25
numpy:  1.26.4
Python 3.11.7 (main, Dec 15 2023, 18:12:31) [GCC 11.2.0] on linux

$ nvcc -V
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| 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 P100-PCIE...  Off  | 00000000:37:00.0 Off |                    0 |
| N/A   29C    P0    31W / 250W |  16131MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla P100-PCIE...  Off  | 00000000:86:00.0 Off |                    0 |
| N/A   30C    P0    32W / 250W |    534MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      2856      G   /usr/bin/X                         63MiB |
|    0   N/A  N/A      3007      G   /usr/bin/gnome-shell               15MiB |
|    0   N/A  N/A      8408      C   ...vs/jax_py311_4/bin/python     3471MiB |
|    0   N/A  N/A     21256      C   .../envs/jax_test/bin/python    12579MiB |
|    1   N/A  N/A      2856      G   /usr/bin/X                         22MiB |
|    1   N/A  N/A      8408      C   ...vs/jax_py311_4/bin/python      255MiB |
|    1   N/A  N/A     21256      C   .../envs/jax_test/bin/python      255MiB |
+-----------------------------------------------------------------------------+

[bantin]

I've also been unable to install jax with CUDA 11 since they dropped support for it. I have a different error when I try to create an array:

RuntimeError: Unable to initialize backend 'cuda': INVALID_ARGUMENT: Mismatched PJRT plugin PJRT API version (0.51) and framework PJRT API version 0.42). (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)

If a member of the jax team could tell us how to install for CUDA 11, that would be very helpful...

[jakevdp]

This looks similar to #20148 – is it possible that you are using JAX in a runtime where another framework has already claimed the GPU? Or perhaps you are attempting to use JAX in multiple processes?

[bantin]

Really appreciate the response @jakevdp! Unfortunately I don't think either of these are the issue on my end. I'm attempting to install in a new conda environment and nothing else is running that could be claiming the gpu.

[bantin]

Ah, I think I got it figured out. I had previously tried installing jax with pip via pip install "jax[cuda12]" which didn't work on my system, but introduced some cuda files which were messing things up. Things appear to be working now after running the 3 commands suggested by the OP.

Thanks again.

[seishiroono]

@bantin Do you mean that the three commands I shared work for your case? If possible, could you share your environment with me?

[bantin]

yeah, that worked for me in a fresh conda env.

[selamw1]

Hi @seishiroono

Thanks for the report. JAX now supports CUDA 12.1 or newer only. Support for CUDA 11.8 has been dropped.
For your reference, you can find more details in the official documentation: https://jax.readthedocs.io/en/latest/changelog.html#:~:text=jaxlib%200.4.26%20(April,supports%20NumPy%202.0.

[jakevdp]

I think it should still be possible to use CUDA 11 by installing older JAX versions – that looks to be what this question is about.

[jakevdp]

@seishiroono I'm not sure what the root cause is here, but I suspect it may be one of two things (1) perhaps the cudnn version mismatch is causing a problem (2) perhaps you have a different CUDA version installed on your system somewhere that is being found at runtime before your conda-installed version

[seishiroono]

@jakevdp Again, thanks for your help.

I have a question about the version of CUDA. I checked conda list, and then I found that many of the packages installed by conda install nvidia/label/cuda-11.8.0::cuda-toolkit are version 12.4. Do you think that this is related to what you told me?

# Name                    Version                   Build  Channel
cuda-cccl                 12.4.127             h06a4308_2  
cuda-cccl_linux-64        12.4.127             h06a4308_2  
cuda-command-line-tools   12.4.1               h06a4308_1  
cuda-compiler             11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-cudart               12.4.127             h99ab3db_0  
cuda-cudart-dev           12.4.127             h99ab3db_0  
cuda-cudart-dev_linux-64  12.4.127             hd681fbe_0  
cuda-cudart-static        12.4.127             h99ab3db_0  
cuda-cudart-static_linux-64 12.4.127             hd681fbe_0  
cuda-cudart_linux-64      12.4.127             hd681fbe_0  
cuda-cuobjdump            12.4.127             h6a678d5_1  
cuda-cupti                12.4.127             h6a678d5_1  
cuda-cupti-dev            12.4.127             h6a678d5_1  
cuda-cuxxfilt             12.4.127             h6a678d5_1  
cuda-documentation        11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-driver-dev           12.4.127             h99ab3db_0  
cuda-driver-dev_linux-64  12.4.127             hd681fbe_0  
cuda-gdb                  12.4.127             h122497a_1  
cuda-libraries            12.4.1               h06a4308_1  
cuda-libraries-dev        12.4.1               h06a4308_1  
cuda-nsight               12.4.127             h06a4308_1  
cuda-nvcc                 11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-nvdisasm             12.4.127             h6a678d5_1  
cuda-nvml-dev             12.4.127             h6a678d5_1  
cuda-nvprof               12.4.127             h6a678d5_1  
cuda-nvprune              12.4.127             h6a678d5_1  
cuda-nvrtc                12.4.127             h99ab3db_1  
cuda-nvrtc-dev            12.4.127             h99ab3db_1  
cuda-nvtx                 12.4.127             h6a678d5_1  
cuda-nvvp                 12.4.127             h6a678d5_1  
cuda-opencl               12.4.127             h6a678d5_0  
cuda-opencl-dev           12.4.127             h6a678d5_0  
cuda-profiler-api         12.4.127             h06a4308_1  
cuda-sanitizer-api        12.4.127             h99ab3db_1  
cuda-toolkit              11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-tools                12.4.1               h06a4308_1  
cuda-version              12.4                 hbda6634_3  
cuda-visual-tools         12.4.1               h06a4308_1  
cudnn                     8.8.0.121            h264754d_4    conda-forge

[jakevdp]

Yes, that looks like it could be problematic. I'm not familiar with the conda cuda installation command, but I would probably start by figuring out how to get a working CUDA 11 installation.

[bantin]

@seishiroono That screenshot above look similar to the problem I was having. If you make a new conda environment (so that there are no Cuda 12 packages present)) and run only the three commands you listed at the top of this thread, what happens?

[seishiroono]

@bantin Thanks for your comment. Actually, this is exactly what I did in the original post. Then, conda automatically installed some packages whose versions are 12.4...

[seishiroono]

@jakevdp Let me update the situation. When I use the command conda install cuda -c nvidia/label/cuda-11.8.0, I get CUDA 11.8 correctly. Then, my GPUs are recognized, but I get the following error. While it looks like CUDA 11.8.0 mismatches with the CUDA driver, I previously used JAX 0.4.24 with Driver Version: 510.47.03 CUDA Version: 11.6. Could you help me figure out what kinds of problems there are?

>>> import jax
>>> import jax.numpy as jnp
>>> jax.devices()
[cuda(id=0), cuda(id=1)]

>>>a = jnp.array(1.)
>>>b = jnp.array(2.)
>>>a+b
E0530 17:42:45.526298   23236 pjrt_stream_executor_client.cc:2804] Execution of replica 0 failed: INTERNAL: Failed to load in-memory CUBIN (compiled for a different GPU?).: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[5], line 1
----> 1 a+b

File [~/miniconda3/envs/jax_test2/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:264](http://127.0.0.1:8888/lab/tree/Interpolation/ai2kitaev/plot/~/miniconda3/envs/jax_test2/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py#line=263), in _defer_to_unrecognized_arg.<locals>.deferring_binary_op(self, other)
    262 args = (other, self) if swap else (self, other)
    263 if isinstance(other, _accepted_binop_types):
--> 264   return binary_op(*args)
    265 # Note: don't use isinstance here, because we don't want to raise for
    266 # subclasses, e.g. NamedTuple objects that may override operators.
    267 if type(other) in _rejected_binop_types:

    [... skipping hidden 10 frame]

File [~/miniconda3/envs/jax_test2/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py:1209](http://127.0.0.1:8888/lab/tree/Interpolation/ai2kitaev/plot/~/miniconda3/envs/jax_test2/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py#line=1208), in ExecuteReplicated.__call__(self, *args)
   1207   self._handle_token_bufs(result_token_bufs, sharded_runtime_token)
   1208 else:
-> 1209   results = self.xla_executable.execute_sharded(input_bufs)
   1210 if dispatch.needs_check_special():
   1211   out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: INTERNAL: Failed to load in-memory CUBIN (compiled for a different GPU?).: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid

The following is part of conda list.

cuda                      11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-cccl                 11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-command-line-tools   11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-compiler             11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-cudart               11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-cudart-dev           11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-cuobjdump            11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-cupti                11.8.87                       0    nvidia/label/cuda-11.8.0
cuda-cuxxfilt             11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-demo-suite           11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-documentation        11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-driver-dev           11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-gdb                  11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-libraries            11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-libraries-dev        11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-memcheck             11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nsight               11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nsight-compute       11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-nvcc                 11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-nvdisasm             11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nvml-dev             11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nvprof               11.8.87                       0    nvidia/label/cuda-11.8.0
cuda-nvprune              11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nvrtc                11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-nvrtc-dev            11.8.89                       0    nvidia/label/cuda-11.8.0
cuda-nvtx                 11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-nvvp                 11.8.87                       0    nvidia/label/cuda-11.8.0
cuda-profiler-api         11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-runtime              11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-sanitizer-api        11.8.86                       0    nvidia/label/cuda-11.8.0
cuda-toolkit              11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-tools                11.8.0                        0    nvidia/label/cuda-11.8.0
cuda-version              12.4                 hbda6634_3  
cuda-visual-tools         11.8.0                        0    nvidia/label/cuda-11.8.0
cudnn                     8.8.0.121            h264754d_4    conda-forge
jax                       0.4.25                   pypi_0    pypi
jaxlib                    0.4.25+cuda11.cudnn86          pypi_0    pypi
libcublas                 11.11.3.6                     0    nvidia/label/cuda-11.8.0
libcublas-dev             11.11.3.6                     0    nvidia/label/cuda-11.8.0
libcufft                  10.9.0.58                     0    nvidia/label/cuda-11.8.0
libcufft-dev              10.9.0.58                     0    nvidia/label/cuda-11.8.0
libcufile                 1.4.0.31                      0    nvidia/label/cuda-11.8.0
libcufile-dev             1.4.0.31                      0    nvidia/label/cuda-11.8.0
libcurand                 10.3.0.86                     0    nvidia/label/cuda-11.8.0
libcurand-dev             10.3.0.86                     0    nvidia/label/cuda-11.8.0
libcusolver               11.4.1.48                     0    nvidia/label/cuda-11.8.0
libcusolver-dev           11.4.1.48                     0    nvidia/label/cuda-11.8.0
libcusparse               11.7.5.86                     0    nvidia/label/cuda-11.8.0
libcusparse-dev           11.7.5.86                     0    nvidia/label/cuda-11.8.0
libffi                    3.4.4                h6a678d5_1  
libgcc-ng                 13.2.0               h77fa898_7    conda-forge
libgomp                   13.2.0               h77fa898_7    conda-forge
libnpp                    11.8.0.86                     0    nvidia/label/cuda-11.8.0
libnpp-dev                11.8.0.86                     0    nvidia/label/cuda-11.8.0
libnvjpeg                 11.9.0.86                     0    nvidia/label/cuda-11.8.0
libnvjpeg-dev             11.9.0.86                     0    nvidia/label/cuda-11.8.0
libstdcxx-ng              13.2.0               hc0a3c3a_7    conda-forge
libuuid                   1.41.5               h5eee18b_0  
libzlib                   1.2.13               h4ab18f5_6    conda-forge