Loaded runtime CuDNN library: 8.0.5 but source was compiled with: 8.1.0.

When I run the train_test.py within the ViT project (https://github.com/google-research/vision_transformer/blob/master/vit_jax/train_test.py, which is based on JAX) I got the bug and the error looks like this:

Convolution performance may be suboptimal.
2021-04-20 23:10:23.094910: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc:352] Loaded runtime CuDNN library: 8.0.5 but source was compiled with: 8.1.0.  CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library.  If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.

And it ends up with: RuntimeError: Unimplemented: DNN library is not found.

I use CUDA 11.2 and cudnn 8.0.5. I install the jax with: pip install --upgrade jax jaxlib==0.1.65+cuda112 -f https://storage.googleapis.com/jax-releases/jax_releases.html

But this version only supports the cudnn 8.1.0 not cudnn 8.0.5.

Is there a jaxlib version which also supports CUDA 11.2 and cudnn 8.0.5?

Thanks in advance.

[hawkinsp]

Actually, I have another thing you can try. It turns out that CUDA 11.1 wheels are actually compatible with CUDA 11.2, and they are built with CUDNN 8.0. So you could install the cuda111 wheel, and provided you set an environment variable telling it where the CUDA 11.2 installation is (XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda/) then I think everything should work!

(We’re thinking about doing that instead of shipping CUDA 11.2 wheels at all, now that backward compatibility between minor CUDA versions is a thing.)

[hawkinsp]

In our wheel build, we pin the oldest release of CuDNN that Nvidia provides for a given CUDA version: https://cs.opensource.google/jax/jax/+/master:build/install_cuda.sh;drc=55c75c8ca30be89f1c6b9d19214cb23f429c6f47;l=32

We build in a Ubuntu 16.04 docker container, and for Ubuntu 16.04 and CUDA 11.2, the oldest wheel appears to be cudnn 8.1: https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/

So there’s no easy fix to this.