Jax doesn't see my GPU, even though Pytorch does

Question

Jax doesn't see my GPU, even though Pytorch does

asemic-horizon opened this issue 4 years ago · comments

Jax sounds like an impressive project, thanks for working on it.

That said: on Ubuntu 18.04, this happens

➜  python
Python 3.6.9 (default, Oct  8 2020, 12:12:24) 
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.

>>> import torch, jax; print(torch.cuda.is_available()); print(jax.devices())
True
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

I first tried to pip install jax and got various errors; the error messages said it was common with old versions of pip that didn't support newer kinds of wheel and directed me to upgrade pip, which I did (from 9 to 20). Now Jax seems to be installed (at least various numpy-compatible functions do), but not to the point where it appears to see my GPU, a laptop "Geforce" card by Nvidia.

I'm not sure what systems diagnostics I can bring to help. This is the name of my card:

✗  lspci | grep NVIDIA
01:00.0 3D controller: NVIDIA Corporation GP108M [GeForce MX150] (rev a1)

As best as I understand it, these are drivers:

✗  lsmod | grep nvidia
nvidia_uvm            970752  0
nvidia_drm             53248  3
nvidia_modeset       1212416  2 nvidia_drm
nvidia              27643904  103 nvidia_uvm,nvidia_modeset
drm_kms_helper        184320  2 nvidia_drm,i915
drm                   491520  8 drm_kms_helper,nvidia_drm,i915

Thanks for reading this anyway.

8bitmp3 · Answer 1 · Sun Dec 20 2020 23:26:51 GMT+0800 (China Standard Time)

Hello 👋 Just to confirm - did you follow the Linux-specific installation instructions from the README? Also, have you tried installing JAX in a separate virtual environment that excludes PyTorch? 🤷‍♀️

https://github.com/google/jax#installation

On Linux, it is often necessary to first update pip to a version that supports manylinux2010 wheels.

If you want to install JAX with both CPU and GPU support, using existing CUDA and CUDNN7 installations on your machine (for example, preinstalled on your cloud VM), you can run
pip install --upgrade pip
pip install --upgrade jax jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html
The jaxlib version must correspond to the version of the existing CUDA installation you want to use, with cuda110 for CUDA 11.0, cuda102 for CUDA 10.2, and cuda101 for CUDA 10.1. You can find your CUDA version with: install path:
nvcc --version
Note that some GPU functionality expects the CUDA installation to be at /usr/local/cuda-X.X, where X.X should be replaced with the CUDA version number (e.g. cuda-10.2). If CUDA is installed elsewhere on your system, you can either create a symlink:
sudo ln -s /path/to/cuda /usr/local/cuda-X.X
Or set the following environment variable before importing JAX:
XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda

Gert-Jan Both · Answer 2 · Mon Dec 21 2020 19:46:06 GMT+0800 (China Standard Time)

I had the same issue, but managed to solve it. It seems pytorch bundles its own cuda, so that's why you don't have to install it separately but it sees your gpu and nvidia-smi works. Installling cuda for your GPU following these instructions solved the issue for me: https://developer.nvidia.com/cuda-downloads

8bitmp3 · Answer 3 · Mon Dec 21 2020 23:24:55 GMT+0800 (China Standard Time)

@GJBoth That's awesome. Can you confirm the following (and please correct me if I'm wrong):

Having JAX in a separate env should technically help identify if JAX can detect your CUDA (and not PyTorch's bundled one)
The current instructions assume that you've taken care of your CUDA installation (see extract below) but maybe it would help to nudge the users to go to https://developer.nvidia.com/cuda-downloads and install CUDA, if they haven't already.

"... using existing CUDA and CUDNN7 installations on your machine (for example, preinstalled on your cloud VM)..."
...
"Note that some GPU functionality expects the CUDA installation to be at /usr/local/cuda-X.X, where X.X should be replaced with the CUDA version number (e.g. cuda-10.2). If CUDA is installed elsewhere on your system..."

WDYT? @asemic-horizon @GJBoth

Gert-Jan Both · Answer 4 · Mon Dec 21 2020 23:32:29 GMT+0800 (China Standard Time)

I have no experience installing cuda in a specific env. It seems that the symbolic link wouldn't work, see this thread:
https://discuss.pytorch.org/t/where-is-cudatoolkit-path-when-installed-via-conda/47791/5

I think many new jax users will come from pytorch, so adding a nudge a la 'If you're coming from pytorch, make sure to install cuda separately, if you haven't yet.') Two more observations: I could run nvidia-smi, but not nvcc, so this might be a nice check to see if you have pytorch cuda or systemwide. Furthermore, jupyter notebooks tend to die silently with these issues, so running things as a script gives you much more info.

Jongrae.K · Answer 5 · Sun Mar 21 2021 16:21:01 GMT+0800 (China Standard Time)

I had the same problem with jax not recognizing GPU. I did the following two steps:

pip install --upgrade pip
pip install --upgrade jax jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

and make a link as follows:

ln -s /usr/lib/nvidia-cuda-toolkit /usr/local/cuda-10.1

After this, jax still didn't recognize GPU. Then, I did the following steps hinted from the warning message in jax about GPU:

cd /usr/lib/nvidia-cuda-toolkit
mkdir nvvm
cd nvvm
sudo ln -s /usr/lib/nvidia-cuda-toolkit/libdevice libdevice

You would need to use "sudo" for the above steps. After these, jax recognises my GPU.

Mohammed Al-Rawi · Answer 6 · Sat Apr 17 2021 05:58:15 GMT+0800 (China Standard Time)

sudo ln -s /usr/lib/nvidia-cuda-toolkit/libdevice libdevice

I am remotely connected to a slurm cluster and do not have sudo rights. In fact, I do not even have permission to make a symbolic link.
Plus, my environment has no GPU, the GPU is assigned via the sbatch job file using the command "#SBATCH --gres=gpu:1" .
This is way too complicated. Yet, PyTorch seems to work perfectly well.

yongjiezhu · Answer 7 · Tue Apr 27 2021 14:03:45 GMT+0800 (China Standard Time)

sudo ln -s /usr/lib/nvidia-cuda-toolkit/libdevice libdevice

I am remotely connected to a slurm cluster and do not have sudo rights. In fact, I do not even have permission to make a symbolic link.
Plus, my environment has no GPU, the GPU is assigned via the sbatch job file using the command "#SBATCH --gres=gpu:1" .
This is way too complicated. Yet, PyTorch seems to work perfectly well.

Hi, I have the same problem with remotely connected to a slurm cluster. Do you solve this issue? How to solve it?

Mohammed Al-Rawi · Answer 8 · Tue Apr 27 2021 18:07:28 GMT+0800 (China Standard Time)

Do you solve this issue?

Unfortunately not. I don't have sudo control over the cluster and this makes it hard. The best way for JAX is that they ship a cuda bundle with the installation, similar to PyTorch.

Peter Hawkins · Answer 9 · Thu Apr 29 2021 00:17:22 GMT+0800 (China Standard Time)

FYI #6581 bundles libdevice.10.bc with jaxlib wheels, which hopefully will help avoid this particular problem. If you are feeling motivated, you could try patching in that PR and building a jaxlib from source to see if it fixes your problems.

josephrocca · Answer 10 · Fri Sep 03 2021 14:43:18 GMT+0800 (China Standard Time)

Similar to myjr52, I was able to solve this simply by replacing this:

pip install --upgrade jax jaxlib

with this (you'll have the change cuda111 based on your output of nvcc --version - mine is 11.1):

pip install --upgrade jax jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html

I didn't need to do any of the additional steps mentioned by myjr52.

Peng Ding · Answer 11 · Thu Oct 28 2021 03:43:22 GMT+0800 (China Standard Time)

Guess it's a bit late for this. But I got mine fixed by specifying the exact whl link found in the https://storage.googleapis.com/jax-releases/jax_releases.html. Just I need cuda 11.0.

The one I used was:

pip uninstall jax jaxlib -y
pip install https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.71+cuda110-cp38-none-manylinux2010_x86_64.whl

Jithendaraa Subramanian · Answer 12 · Sat Nov 13 2021 12:44:10 GMT+0800 (China Standard Time)

Do you solve this issue?

Unfortunately not. I don't have sudo control over the cluster and this makes it hard. The best way for JAX is that they ship a cuda bundle with the installation, similar to PyTorch.

@morawi Curious to know if this is solved for you yet since I'm going through the same thing with JAX on a slurm cluster

Mohammed Al-Rawi · Answer 13 · Sat Nov 13 2021 21:06:57 GMT+0800 (China Standard Time)

Do you solve this issue?

Unfortunately not. I don't have sudo control over the cluster and this makes it hard. The best way for JAX is that they ship a cuda bundle with the installation, similar to PyTorch.

@morawi Curious to know if this is solved for you yet since I'm going through the same thing with JAX on a slurm cluster

I just stopped using it.

Ritobrata Ghosh · Answer 14 · Sun Dec 26 2021 19:46:12 GMT+0800 (China Standard Time)

I have solved this problem very easily just following this issue, some googling and stumbling upon two SO questions, and the readme of this project.

I had nvidia drivers installed in my laptop through the Pop OS store, and I installed nvidia-cuda-toolkit through apt, and then installed PyTorch (earlier).

My cuda version is 11.2.

I did not have to do any other installation.

I upgraded pip.
I installed jax[cuda11] instead if just jax.
Followed other generic instructions from the ReadMe.
Created two symlinks- one for nvcc and another for cuda.

Then it started working great.

Juan Ignacio Folco · Answer 15 · Tue Sep 20 2022 02:52:42 GMT+0800 (China Standard Time)

If someone else stumbles into this, the CUDA wheel releases are now stored on https://storage.googleapis.com/jax-releases/jax_cuda_releases.html for some reason.

Francesco Saverio Zuppichini · Answer 16 · Mon Oct 10 2022 20:43:19 GMT+0800 (China Standard Time)

Same issue, using

FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
...
RUN python3 -m pipinstall "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
...

Jax doesn't see the gpu

import jax
print(jax.devices())

only cpu

Alessio Quaresima · Answer 17 · Wed Oct 12 2022 22:43:44 GMT+0800 (China Standard Time)

Hello,

same issue here! Torch can find my GPU, JAX does not!

nvcc --version                                                                             [16:16:43]
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Jun__8_16:49:14_PDT_2022
Cuda compilation tools, release 11.7, V11.7.99
Build cuda_11.7.r11.7/compiler.31442593_0

pacman -Q cudnn                                                                            [16:37:00]
cudnn 8.5.0.96-1

I followed the installation guide here:

When I run
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
It first asks me for the storage.googleapis.com user (WTF ??)
then I just press enter and it return an access error!

Collecting jaxlib==0.3.22+cuda11.cudnn82
User for storage.googleapis.com:   WARNING: 401 Error, Credentials not correct for https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22%2Bcuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl
  ERROR: HTTP error 401 while getting https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22%2Bcuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl (from https://storage.googleapis.com/jax-releases/jax_cuda_releases.html)
ERROR: Could not install requirement jaxlib==0.3.22+cuda11.cudnn82 from https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22%2Bcuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl (from jax[cuda]) because of HTTP error 401 Client Error: Unauthorized for url: https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22%2Bcuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl for URL https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22%2Bcuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl (from https://storage.googleapis.com/jax-releases/jax_cuda_releases.html)
FAIL: 1

If I run

pip install --upgrade "jax[cuda]" or pip install "jax[cuda11_cudnn82]"

Then it unrolls the version till the very first that has not cuda! (WTF^2)

....
  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
  Using cached jax-0.2.22-py3-none-any.whl
WARNING: jax 0.2.22 does not provide the extra 'cuda'
Installing collected packages: jax
  Attempting uninstall: jax
    Found existing installation: jax 0.3.22
    Uninstalling jax-0.3.22:
      Successfully uninstalled jax-0.3.22
Successfully installed jax-0.2.22

The pip install --upgrade "jax[all]" runs just fine: Successfully installed jax-0.3.22
but GPU access is not available (see topmost).

Thanks for the help

Alessio Quaresima · Answer 18 · Wed Oct 12 2022 22:56:46 GMT+0800 (China Standard Time)

Update:

If I download my appropriate version

[cuda11/jaxlib-0.3.22+cuda11.cudnn82-cp39-cp39-manylinux2014_x86_64.whl](https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.22+cuda11.cudnn82-cp39-cp39-manylinux2014_x86_64.whl)

and run:

pip install --upgrade "jax[cuda11_cudnn82]" -f ~/Downloads/jax

Here we go again:

[16:53:15]
Looking in links: /home/cocconat/Downloads/jax
Requirement already satisfied: jax[cuda11_cudnn82] in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (0.3.22)
Requirement already satisfied: absl-py in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (1.2.0)
Requirement already satisfied: opt-einsum in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (3.3.0)
Requirement already satisfied: numpy>=1.20 in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (1.23.3)
Requirement already satisfied: etils[epath] in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (0.8.0)
Requirement already satisfied: typing-extensions in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (4.3.0)
Requirement already satisfied: scipy>=1.5 in /home/cocconat/.virtualenvs/learning/lib/python3.10/site-packages (from jax[cuda11_cudnn82]) (1.9.1)
Collecting jax[cuda11_cudnn82]
  Using cached jax-0.3.22-py3-none-any.whl
  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 jax-0.3.17-py3-none-any.whl
  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)
^C  Preparing metadata (setup.py) ... canceled
ERROR: Operation cancelled by user
FAIL: 1

Peter Hawkins · Answer 19 · Wed Oct 12 2022 23:00:44 GMT+0800 (China Standard Time)

@aquaresima Please open a new issue, please don't ping long-closed issues.