Using GPUs with Jax

[akhilnadigatla]

Hello!

I am trying to get Jax running on my GPUs, but face these error messages:

I0311 16:52:36.027944 140563802497728 xla_bridge.py:247] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
2022-03-11 16:52:36.058206: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
I0311 16:52:36.058603 140563802497728 xla_bridge.py:247] Unable to initialize backend 'gpu': FAILED_PRECONDITION: No visible GPU devices.
I0311 16:52:36.059262 140563802497728 xla_bridge.py:247] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
W0311 16:52:36.059466 140563802497728 xla_bridge.py:252] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

I have looked for solutions to the problems, and made sure all drivers/CUDA version/cuDNN version are as expected. I am working on Ubuntu 18.04 and the outputs of some key commands are:

>> nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.103.01   Driver Version: 470.103.01   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| 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  Quadro RTX 8000     Off  | 00000000:15:00.0 Off |                  Off |
| 61%   80C    P2   248W / 260W |  32490MiB / 48601MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Quadro RTX 8000     Off  | 00000000:2D:00.0 Off |                  Off |
| 76%   86C    P2   261W / 260W |   7237MiB / 48600MiB |     99%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

>> nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Wed_Jun__2_19:15:15_PDT_2021
Cuda compilation tools, release 11.4, V11.4.48
Build cuda_11.4.r11.4/compiler.30033411_0

>> cat /usr/local/cuda/include/cudnn_version.h
...
/**
 * \file: The master cuDNN version file.
 */

#ifndef CUDNN_VERSION_H_
#define CUDNN_VERSION_H_

#define CUDNN_MAJOR 8
#define CUDNN_MINOR 2
#define CUDNN_PATCHLEVEL 4

#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

#endif /* CUDNN_VERSION_H */

All the path variables (to my knowledge) have been set as expected:

export PATH=/usr/local/cuda-11.4/bin:${PATH}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:${LD_LIBRARY_PATH}
export CUDA_HOME=/usr/local/cuda

And Tensorflow does appear to recognize the existence of two GPUs:

>> python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')]

I am not sure what could be the issue here. Any help would be greatly appreciated! Thank you.

[YouJiacheng]

Do you install correct version jaxlib?
Can you uninstall jax and jaxlib then reinstall with pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html?

[akhilnadigatla]

Thank you for the suggestions! I did install the correct version of jaxlib initially:

>> pip freeze | grep jaxlib
jaxlib==0.3.0+cuda11.cudnn82

[YouJiacheng]

@akhilnadigatla Maybe you need to create a issue?

[akhilnadigatla]

Yeah, I created #9859. Hopefully, it gets sorted there.