Verifying CUDA and CUDNN installation paths for JAX

[simon-bachhuber]

Jax seems to work with only the nvidia driver installed.

import jax
# gpu-device
jax.numpy.array([1,]).device()

However, now how can i know whether or not the system installation of CUDA / CUDNN and the paths are set correctly?

A related question is: If CUDA is required, is CUDNN required?
Another related question is: If CUDA / CUDNN is installed using wheels, how can i pin the version of jax / jaxlib?

[hawkinsp]

Correct. If you are using the pip wheel installation of CUDA, JAX only needs you to install the NVIDIA driver (https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier).

The jax wheel prefers to use the pip-installed copy of CUDA/CUDNN if it is present. If it is not present, JAX will look for a system copy of CUDA/CUDNN. This is done primarily by adding a relative rpath that points to the pip-installed CUDA libraries; if they are not present, then your LD_LIBRARY_PATH will be searched, which should locate a self-installed CUDA/CUDNN.

Both CUDA and CUDNN are required. It is permissible to, say, install CUDA on your system yourself and CUDNN using a pip wheel, for example, provided the versions are compatible: it should work but we don't test it ourselves.

I don't entirely follow the question about version pinning. Can you clarify? You can pin the version of jax/jaxlib and all the CUDA wheels if you like, provided they are compatible.

[simon-bachhuber]

Regarding the pinning of jax jaxlib version, what i mean is that the following is not possible (but would be quite helpful)
pip install --upgrade "jax[cuda11_pip]==0.4.2" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"

A possible workaround that one might think would work is to use
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"

At this point the environment works.

jax                      0.4.8
jaxlib                   0.4.7+cuda11.cudnn86
ml-dtypes                0.0.4
numpy                    1.24.2
nvidia-cublas-cu11       11.11.3.6
nvidia-cuda-nvcc-cu11    11.8.89
nvidia-cuda-runtime-cu11 11.8.89
nvidia-cudnn-cu11        8.8.1.3
nvidia-cufft-cu11        10.9.0.58
nvidia-cusolver-cu11     11.4.1.48
nvidia-cusparse-cu11     11.7.5.86
opt-einsum               3.3.0
pip                      23.0.1
scipy                    1.10.1
setuptools               67.6.1
wheel                    0.40.0

However, i would like a lower version of jax.

So let's uninstall using pip uninstall -y jax jaxlib.

ml-dtypes                0.0.4
numpy                    1.24.2
nvidia-cublas-cu11       11.11.3.6
nvidia-cuda-nvcc-cu11    11.8.89
nvidia-cuda-runtime-cu11 11.8.89
nvidia-cudnn-cu11        8.8.1.3
nvidia-cufft-cu11        10.9.0.58
nvidia-cusolver-cu11     11.4.1.48
nvidia-cusparse-cu11     11.7.5.86
opt-einsum               3.3.0
pip                      23.0.1
scipy                    1.10.1
setuptools               67.6.1
wheel                    0.40.0

and re-install using the desired versions from the .whl hosted here, and with pip install jax==0.4.2.

jax                      0.4.2
jaxlib                   0.4.2+cuda11.cudnn86
ml-dtypes                0.0.4
numpy                    1.24.2
nvidia-cublas-cu11       11.11.3.6
nvidia-cuda-nvcc-cu11    11.8.89
nvidia-cuda-runtime-cu11 11.8.89
nvidia-cudnn-cu11        8.8.1.3
nvidia-cufft-cu11        10.9.0.58
nvidia-cusolver-cu11     11.4.1.48
nvidia-cusparse-cu11     11.7.5.86
opt-einsum               3.3.0
pip                      23.0.1
scipy                    1.10.1
setuptools               67.6.1
wheel                    0.40.0

However, this then produces

>>> import jax
2023-04-05 16:05:42.459035: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1/lib64
2023-04-05 16:05:42.521737: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1/lib64
2023-04-05 16:05:42.522559: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1/lib64

Note the cuda-12.1 is the system-wide installation.

[simon-bachhuber]

To expand on this a little more:

The following snippet doesn't actually confirm anything. It also works with only CUDA.

import jax
# gpu-device
jax.numpy.array([1,]).device()

However, this snippet only works with CUDA & CUDNN

import jax
# gpu-device
m = jax.numpy.array([1,])
# matrix multiplication will fail
m@m

[cdelavegamartin]

We are finding a similar issue where in our cluster max available cudnn is 8.9.4.
We can specify the jaxlib version in our pyproject.toml, and rely on the system cudnn and cuda,

dependencies = [
    "jax==0.4.28",
    "jaxlib @ https://storage.googleapis.com/jax-releases/cuda12/jaxlib-0.4.28+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl",
    "flax==0.10.1",
    "numpy==1.26.4",
    "optax==0.2.4",
    "orbax-checkpoint==0.6.4",
    "wandb==0.18.7",
    "matplotlib==3.8.3",
    "scipy==1.12.0",
    "scikit-fem==10.0.1",
    "hydra-core==1.3.2",
    "PyDMD==2024.12.1",
    "einops==0.7.0",
    "hydra-joblib-launcher==1.2.0",
    "numba==0.60.0",
    "diffrax==0.6.0",
    "pandas==2.2.3",
    "tqdm==4.67.1"
]

But we can't figure out how to install and use the nvidia pip wheels with it correctly instead.
Our interest in this is that we are ok relying on the system versions for ourselves while developing in the cluster, but once the tool goes public, we don't want to require people to manually install an specific CUDA and cuDNN version in their machines.
And we need to make sure strict reproducibility, so we'd want to pin that version of jax and jaxlib together with CUDA and cuDNN for everyone.

Any solutions?
@rodrigodzf