Which CUDA versions does JAX support?

[cool-RR]

I bought a new RTX 5080 and I'm trying to run JAX on it. It's not working, i.e. I'm getting this:

(viola_env) ➜  viola git:(development) uv pip install "jax[cuda12]" && python -c "import jax; print(jax.devices())"
Using Python 3.13.1 environment at: /home/ramrachum/.venvs/viola_env
Audited 1 package in 1ms
[CpuDevice(id=0)]

I wonder whether it's because I have the wrong CUDA version. I initially tried CUDA 12.8 which didn't work. Then I tried CUDA 12.3 which didn't work either. Is there a different CUDA version I should use?

Update:

I tried building JAX, and got these wheels:

-rw-rw-r-- 1 ramrachum ramrachum  97M Mar 12 18:55 jax_cuda12_pjrt-0.5.0.dev20250312-py3-none-manylinux2014_x86_64.whl
-rw-rw-r-- 1 ramrachum ramrachum 9.5M Mar 12 18:49 jax_cuda12_plugin-0.5.0.dev20250312-cp313-cp313-manylinux2014_x86_64.whl
-rw-rw-r-- 1 ramrachum ramrachum  94M Mar 12 18:49 jaxlib-0.5.0.dev20250312-cp313-cp313-manylinux2014_x86_64.whl

I installed them, but still JAX doesn't use the GPU:

(viola_env) ➜  dist git:((jax-v0.5.0)) python -c "import jax; print(jax.devices())"
[CpuDevice(id=0)]

[hawkinsp]

The released version should work. For a Blackwell GPU, you probably need CUDA 12.8 at a minimum.

What goes wrong when you try with the released version?

[cool-RR]

(viola_env) ➜  Desktop export TF_CPP_MIN_LOG_LEVEL=0
(viola_env) ➜  Desktop uv pip install "jax[cuda12]" && TF_CPP_MIN_LOG_LEVEL=0 python -c "import jax; print(jax.devices())"
Using Python 3.13.1 environment at: /home/ramrachum/.venvs/viola_env
Audited 1 package in 1ms
[CpuDevice(id=0)]


(viola_env) ➜  Desktop uv pip freeze | grep -E "jax|nvidia|cuda"
Using Python 3.13.1 environment at: /home/ramrachum/.venvs/viola_env
cupy-cuda12x==13.4.0
jax==0.5.0
jax-cuda12-pjrt==0.5.0
jax-cuda12-plugin==0.5.0
jaxlib==0.5.0
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvcc-cu12==12.8.93
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-cusparselt-cu12==0.6.2
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127

Full git freeze output: https://gist.github.com/cool-RR/91188b755fc993999ad8eda5e0daf957

[hawkinsp]

It looks like you have CUDA 12.4 installed: nvidia-cuda-runtime-cu12==12.4.127

Can you try in a fresh virtual environment? (It would probably also work to update all those packages to their newest versions, but clean venv is probably easiest)

[hawkinsp]

(And just to clarify: JAX is totally fine with that version of CUDA. But your GPU is not.)

[cool-RR]

Tried it, it still fails:

➜  viola git:(development) uv venv ~/.venvs/fresh_env
Using CPython 3.13.1
Creating virtual environment at: /home/ramrachum/.venvs/fresh_env
Activate with: source /home/ramrachum/.venvs/fresh_env/bin/activate
➜  viola git:(development) . ~/.venvs/fresh_env/bin/activate
(fresh_env) ➜  viola git:(development) uv pip install "jax[cuda12]" && TF_CPP_MIN_LOG_LEVEL=0 python -c "import jax; print(jax.devices())"
Using Python 3.13.1 environment at: /home/ramrachum/.venvs/fresh_env
Resolved 18 packages in 550ms
Installed 18 packages in 48ms
 + jax==0.5.2
 + jax-cuda12-pjrt==0.5.1
 + jax-cuda12-plugin==0.5.1
 + jaxlib==0.5.1
 + ml-dtypes==0.5.1
 + numpy==2.2.3
 + nvidia-cublas-cu12==12.8.4.1
 + nvidia-cuda-cupti-cu12==12.8.90
 + nvidia-cuda-nvcc-cu12==12.8.93
 + nvidia-cuda-runtime-cu12==12.8.90
 + nvidia-cudnn-cu12==9.8.0.87
 + nvidia-cufft-cu12==11.3.3.83
 + nvidia-cusolver-cu12==11.7.3.90
 + nvidia-cusparse-cu12==12.5.8.93
 + nvidia-nccl-cu12==2.25.1
 + nvidia-nvjitlink-cu12==12.8.93
 + opt-einsum==3.4.0
 + scipy==1.15.2
[CpuDevice(id=0)]
(fresh_env) ➜  viola git:(development) uv pip freeze | grep -E "jax|nvidia|cuda"
Using Python 3.13.1 environment at: /home/ramrachum/.venvs/fresh_env
jax==0.5.2
jax-cuda12-pjrt==0.5.1
jax-cuda12-plugin==0.5.1
jaxlib==0.5.1
nvidia-cublas-cu12==12.8.4.1
nvidia-cuda-cupti-cu12==12.8.90
nvidia-cuda-nvcc-cu12==12.8.93
nvidia-cuda-runtime-cu12==12.8.90
nvidia-cudnn-cu12==9.8.0.87
nvidia-cufft-cu12==11.3.3.83
nvidia-cusolver-cu12==11.7.3.90
nvidia-cusparse-cu12==12.5.8.93
nvidia-nccl-cu12==2.25.1
nvidia-nvjitlink-cu12==12.8.93

[cool-RR]

One thing I'm noticing with all of the people who are (generously) helping me is that JAX is mostly treated as a black box. I probably know the least here about how JAX works, so I may be wrong. I assume that JAX has some logic of "Try to use the GPU, if it fails for some reason, use the CPU". If we could know in which step JAX failed to use the GPU, we may know what the problem is.

[andportnoy]

What does nvidia-smi show?

[andportnoy]

Let's try to isolate environment issues. Can you try running docker run --gpus all ghcr.io/nvidia/jax:jax-2025-03-12 python -c 'import jax; print(jax.devices())'? This prints [CudaDevice(id=0)] for me on an 5080 machine, I'm curious what you see on your side.

[cool-RR]

I stripped away the download output to save space here.

(fresh_env) ➜  ~ docker run --gpus all ghcr.io/nvidia/jax:jax-2025-03-12 python -c 'import jax; print(jax.devices())'
docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]]

Run 'docker run --help' for more information

[cool-RR]

Should I install the nvidia docker stuff as shown here?

[andportnoy]

Yes, please follow the official instructions here.

[cool-RR]

I went ahead and installed it. Now it runs and I believe I got the output you expected:

(viola_env) ➜  viola git:(development) docker run --gpus all ghcr.io/nvidia/jax:jax-2025-03-12 python -c 'import jax; print(jax.devices())'

==========
== CUDA ==
==========

NVIDIA Release  (build )
CUDA Version 12.8.0.038
Container image Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
                                                                                                                                                                                                                 Various files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

This container image and its contents are governed by the NVIDIA Deep Learning Container License.
By pulling and using the container, you accept the terms and conditions of this license:
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license
                                                                                                                                                                                                                 NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be
   insufficient for CUDA.  NVIDIA recommends the use of the following flags:
   docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ...

[CudaDevice(id=0)]

[cool-RR]

I had export JAX_PLATFORMS=cpu in one of my dotfiles 😩

Thank you everyone for helping.