can't install newer jax versions with cuda support. #12941

erazor-de · 2022-10-23T06:54:58Z

erazor-de
Oct 23, 2022

I try to install jax with cuda support but I only get an old version installed.

I have cuda 11.5 installed:

> nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Nov_18_09:45:30_PST_2021
Cuda compilation tools, release 11.5, V11.5.119
Build cuda_11.5.r11.5/compiler.30672275_0

I have cuda installed in /usr/local/cuda-11.5 with links from /usr/local/cuda and /usr/local/cuda-11

Installed cuDNN 8.6 from cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz on top of this and can compile an example.

The libraries are found by ldconfig:

> ldconfig -p | grep libcudnn
        libcudnn_ops_train.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
        libcudnn_ops_train.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_ops_train.so
        libcudnn_ops_infer.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
        libcudnn_ops_infer.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_ops_infer.so
        libcudnn_cnn_train.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
        libcudnn_cnn_train.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_cnn_train.so
        libcudnn_cnn_infer.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
        libcudnn_cnn_infer.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_cnn_infer.so
        libcudnn_adv_train.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
        libcudnn_adv_train.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_adv_train.so
        libcudnn_adv_infer.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
        libcudnn_adv_infer.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn_adv_infer.so
        libcudnn.so.8 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn.so.8
        libcudnn.so (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libcudnn.so

I have a halfway recent driver:

> nvidia-smi
Sun Oct 23 08:33:28 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.48.07    Driver Version: 515.48.07    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| 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 GeForce ...  Off  | 00000000:2B:00.0  On |                  N/A |
| 30%   31C    P8    13W / 172W |    488MiB /  6144MiB |     14%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1447      G   /usr/libexec/Xorg                 177MiB |
|    0   N/A  N/A      1585      G   /usr/bin/gnome-shell               20MiB |
|    0   N/A  N/A      2144      G   ...n/opt/firefox/firefox-bin      152MiB |
|    0   N/A  N/A      4842      G   ...t/thunderbird/thunderbird      129MiB |
|    0   N/A  N/A      6108      G   ...tefan/opt/bin/wezterm-gui        5MiB |
+-----------------------------------------------------------------------------+

I have an actual pip:

> pip3 --version
pip 22.3 from /home/stefan/.local/lib/python3.9/site-packages/pip (python 3.9)

Yet when I try to issue the command

pip3 install --upgrade --user "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

I get this:

Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Requirement already satisfied: jax[cuda] in /home/stefan/.local/lib/python3.9/site-packages (0.2.22)
Collecting jax[cuda]
  Using cached jax-0.3.23-py3-none-any.whl
Requirement already satisfied: opt-einsum in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (3.3.0)
Requirement already satisfied: absl-py in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (1.3.0)
Requirement already satisfied: typing-extensions in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (4.4.0)
Requirement already satisfied: numpy>=1.20 in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (1.23.4)
Requirement already satisfied: etils[epath] in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (0.8.0)
Requirement already satisfied: scipy>=1.5 in /home/stefan/.local/lib/python3.9/site-packages (from jax[cuda]) (1.9.3)
  Using cached jax-0.3.22.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.21.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.20.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.19.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached https://storage.googleapis.com/jax-releases/jax/jax-0.3.18.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.17.tar.gz (1.1 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.16.tar.gz (1.0 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.15.tar.gz (1.0 MB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.14.tar.gz (990 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.13.tar.gz (951 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.12.tar.gz (947 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.11.tar.gz (947 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.10.tar.gz (939 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.9.tar.gz (937 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.8.tar.gz (935 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.7.tar.gz (944 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.6.tar.gz (936 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.5.tar.gz (946 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.4.tar.gz (924 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.3.tar.gz (924 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.2.tar.gz (926 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.1.tar.gz (912 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.3.0.tar.gz (896 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.2.28.tar.gz (887 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.2.27.tar.gz (873 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.2.26.tar.gz (850 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.2.25.tar.gz (786 kB)
  Preparing metadata (setup.py) ... done
  Using cached jax-0.2.24.tar.gz (786 kB)
  Preparing metadata (setup.py) ... done
WARNING: jax 0.2.22 does not provide the extra 'cuda'

I tried this with jax[cuda11_cudnn82] jax[cuda111] and jax[cuda11]
Why doesn't it install the newest jax version when using cuda?

Answered by erazor-de

Oct 23, 2022

So https://storage.googleapis.com/jax-releases/jax_cuda_releases.html is the right link for the cuda versions nowadays instead of https://storage.googleapis.com/jax-releases/jax_releases.html. So

pip3 install --upgrade --user "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

does the right thing.

View full answer

erazor-de · 2022-10-23T17:50:44Z

erazor-de
Oct 23, 2022
Author

One addition:
The html at address https://storage.googleapis.com/jax-releases/jax_releases.html does not list any cuda variants while the XML at https://storage.googleapis.com/jax-releases/ does.

0 replies

erazor-de · 2022-10-23T18:03:55Z

erazor-de
Oct 23, 2022
Author

So https://storage.googleapis.com/jax-releases/jax_cuda_releases.html is the right link for the cuda versions nowadays instead of https://storage.googleapis.com/jax-releases/jax_releases.html. So

pip3 install --upgrade --user "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

does the right thing.

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

can't install newer jax versions with cuda support. #12941

Uh oh!

{{title}}

Uh oh!

Replies: 2 comments

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

can't install newer jax versions with cuda support. #12941

Uh oh!

erazor-de Oct 23, 2022

Replies: 2 comments

Uh oh!

erazor-de Oct 23, 2022 Author

Uh oh!

erazor-de Oct 23, 2022 Author

erazor-de
Oct 23, 2022

erazor-de
Oct 23, 2022
Author

erazor-de
Oct 23, 2022
Author