Make GPU-JAX work when CUDA is not available

[carlosgmartin]

I have a Slurm cluster with both GPU and non-GPU nodes.

CPU-JAX (obtained via python3 -m pip install jax) works on both types of nodes, but does not use the GPUs on the GPU nodes.

GPU-JAX (obtained via python3 -m pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html) works on the GPU nodes (and uses their GPUs), but fails on the non-GPU nodes with the following error:

$ python3 -c "import jax"
2023-01-30 16:02:12.481821: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11/lib64:
2023-01-30 16:02:13.027993: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11/lib64:
2023-01-30 16:02:13.036866: W external/org_tensorflow/tensorflow/tsl/platform/default/dso_loader.cc:67] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11/lib64:

This is because the CUDA folder does not exist on the non-GPU nodes. What's the recommended way to address this? Can GPU-JAX be configured to ignore the above error and proceed with the CPUs only, when CUDA is not available?

[clemisch]

Have you tried setting the platform? Maybe that avoids the folder check.

import jax
jax.config.update("jax_platform_name", "cpu")

You would have to do this at the very tippy-top of your code, so that no other JAX computation happens before it.

[carlosgmartin]

@clemisch The error occurs as soon as import jax.

[clemisch]

Ah right, I should have seen that from your code...

[jakevdp]

You could set the jax_platforms configuration using an environment variable to force JAX to use CPU:

$ JAX_PLATFORMS=cpu python -c "import jax"

or in a script:

import os
os.environ['JAX_PLATFORMS'] = 'cpu'

import jax