Question on installing Jax using CUDA #20678
Answered
by
jakevdp
pengzhi1998
asked this question in
Q&A
-
|
Dear all, After installing JAX, I found there were some problems to run it with GPU. My environment is: I tried to install it with command: and found the output of the installation: which indicates it has been installed successfully. Based on this |
Beta Was this translation helpful? Give feedback.
Answered by
jakevdp
Apr 10, 2024
Replies: 3 comments 12 replies
-
|
The key line in your output is this:
You should try For more information on GPU-compatible JAX installation, see Installation: NVIDIA GPU. |
Beta Was this translation helpful? Give feedback.
0 replies
-
|
Thank you so much for your quick reply! If I need to install with Besides, I tried to run: which is also based on this advice. However, I got the error of: |
Beta Was this translation helpful? Give feedback.
11 replies
-
|
What GPU hardware are you using? This error might come up if you are using unsupported GPU hardware, such as a Kepler GPU. Given that you have such an old CUDA version installed, I think it's possible that you also have older, unsupported hardware. (the |
Beta Was this translation helpful? Give feedback.
-
|
My GPU is |
Beta Was this translation helpful? Give feedback.
-
|
OK, I think the hardware is compatible. Next guess would be that your CUDNN version is incompatible with your CUDA version. Given that you have CUDA 10 installed locally, and CUDA 11 or 12 installed via pip, it's likely that there's cross-talk between the two installations. I'd make sure that your system paths are setup such that you're only loading one set of CUDA packages at runtime. Alternatively, you could try wiping the CUDA 10 installation from your system, and installing CUDA 11 or CUDA 12 locally, so you can run jax against these and avoid having multiple CUDA versions on your system. |
Beta Was this translation helpful? Give feedback.
Answer selected by
pengzhi1998
-
|
Thank you so much for your reply and continuous help!
it also reported: And I tried: but got a warning: if using
but got: during installation. Seems like it searches through all JAX versions but none of them could provide an installation with CUDA? Does it mean in these cases, cuda11 and 12 are all not installed for JAX actually? Could this issue be related to my Python version (3.8) being outdated? Though I need this version for another critical dependency.
|
Beta Was this translation helpful? Give feedback.
-
|
The latest JAX doesn't support cuda 11 anymore (https://jax.readthedocs.io/en/latest/installation.html#pip-installation-gpu-cuda-installed-via-pip-easier). |
Beta Was this translation helpful? Give feedback.
-
|
Oh ok. Thank you for pointing this out! It seems a bit tricky to install it on my laptop then. |
Beta Was this translation helpful? Give feedback.
-
|
JAX 0.4.13 does support the I believe you can fix this by updating to the most recent |
Beta Was this translation helpful? Give feedback.
-
|
Hi, import jax.numpy as jp
jp.zeros((1,))raised the following error 2025-01-28 11:34:00.709230: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:439] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2025-01-28 11:34:00.709269: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:443] Memory usage: 734789632 bytes free, 8361738240 bytes total.However jp.array([1.]).device()returns My environment: Ubuntu 20.04, python 3.8, GPU Quadro RTX 4000 and no local cuda toolkit installed, LD_LIBRARY_PATH is not empty but no cuda related path included. I have pytorch with cuda 11.8 installed I installed jax with pip install "jax[cuda11_pip]==0.4.13" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlcuda related wheels in my env: pip list | grep nvidia
nvidia-cublas-cu11 11.11.3.6
nvidia-cuda-cupti-cu11 11.8.87
nvidia-cuda-nvcc-cu11 2022.5.4
nvidia-cuda-nvcc-cu117 11.7.64
nvidia-cuda-nvrtc-cu11 11.8.89
nvidia-cuda-runtime-cu11 11.8.89
nvidia-cudnn-cu11 9.1.0.70
nvidia-cufft-cu11 10.9.0.58
nvidia-curand-cu11 10.3.0.86
nvidia-cusolver-cu11 11.4.1.48
nvidia-cusparse-cu11 11.7.5.86
nvidia-nccl-cu11 2.20.5
nvidia-nvtx-cu11 11.8.86
nvidia-pyindex 1.0.9 |
Beta Was this translation helpful? Give feedback.
1 reply
-
|
This specific version of jax requires cudnn<9, downgrading it solves the problem: # for example
pip install nvidia-cudnn-cu11==8.9.6.50If you are using together with pytorch, it would have to be downgraded as well; The lastest compatible pytorch version I found is 2.3.1: pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118This will automatically install |
Beta Was this translation helpful? Give feedback.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
OK, I think the hardware is compatible. Next guess would be that your CUDNN version is incompatible with your CUDA version. Given that you have CUDA 10 installed locally, and CUDA 11 or 12 installed via pip, it's likely that there's cross-talk between the two installations. I'd make sure that your system paths are setup such that you're only loading one set of CUDA packages at runtime.
Alternatively, you could try wiping the CUDA 10 installation from your system, and installing CUDA 11 or CUDA 12 locally, so you can run jax against these and avoid having multiple CUDA versions on your system.