GPUs Not Recognized

[chelseas]

I just downloaded jax_verify And I'm trying to run the demo scripts but there seems to be an issue recognizing my GPUs:

(jax_verify) chelseas@server:~/jax_verify$ python3 examples/run_boundprop.py --boundprop_method=backward_crown_bound_propagation
I1223 17:36:04.620655 139962850735168 xla_bridge.py:355] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I1223 17:36:04.620828 139962850735168 xla_bridge.py:355] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1223 17:36:04.620904 139962850735168 xla_bridge.py:355] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1223 17:36:04.621421 139962850735168 xla_bridge.py:355] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
I1223 17:36:04.621530 139962850735168 xla_bridge.py:355] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
W1223 17:36:04.621601 139962850735168 xla_bridge.py:362] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Any help?

My GPUs:

(base) chelseas@server:~/$ nvidia-smi
Fri Dec 23 17:43:24 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   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  Tesla K40c          Off  | 00000000:02:00.0 Off |                    0 |
| 23%   29C    P8    21W / 235W |      3MiB / 11441MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla K40c          Off  | 00000000:03:00.0 Off |                    0 |
| 23%   39C    P8    22W / 235W |      3MiB / 11441MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce ...  Off  | 00000000:05:00.0 Off |                  N/A |
|  0%   46C    P8    14W / 151W |   1682MiB /  8117MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    2   N/A  N/A   2514150      C   python                            839MiB |
|    2   N/A  N/A   2544384      C   python                            839MiB |
+-----------------------------------------------------------------------------+

[jakevdp]

Did you install JAX with GPU support? See https://github.com/google/jax#pip-installation-gpu-cuda for instructions, and make sure that the jaxlib version you specify matches your local CUDA and CUDNN version.

[chelseas]

Thanks, I just installed with pip install jax initially, so I guess I did not. I instead installed with:

conda install jax cuda-nvcc -c conda-forge -c nvidia

but it's still giving me the same result :/

[jakevdp]

The "Unable to initialize cuda" error looks like it comes from the CPU version of jaxlib. Make certain you install a gpu-compabile version of jaxlib, following the instructions I linked to above. Note that if you already have a cpu version of jaxlib installed, you may have to uninstall it before attempting to install the GPU-specific wheel.

I'm not certain how GPU-compatible jaxlib is installed via conda, so for troubleshooting using conda instead of pip I'd suggest asking on conda-specific channels.

[Xiang-Pan]

I create a new env with CUDA/CUDNN available, but I can not find the GPU. Is there any way to debug it?

pip install "jax[cuda]" jaxlib -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --force

Also tried

pip install jaxlib==0.3.7+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_releases.html --force

[jakevdp]

Hard to say without more information. One idea: is your CUDA bundled with PyTorch? If so see #5231

[Xiang-Pan]

Same