XLA error when running on GPU

[mehdiataei]

Description

On version 0.4.3 I receive the following error when I run any jitted function on the GPU.

2023-02-15 05:44:23.038563: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2023-02-15 05:44:23.039220: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr 
*** Begin stack trace ***

The problem goes away with jaxlib 0.4.2.

Ubuntu 22.04, CUDA 11.8, CuDNN 8.6
GPU: A6000 (I have dual-GPU config, but I only utilize one of them when using JAX)

What jax/jaxlib version are you using?

jax 0.4.3/jaxlib 0.4.3

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.10

NVIDIA GPU info

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 520.61.05    Driver Version: 520.61.05    CUDA Version: 11.8     |
|-------------------------------+----------------------+----------------------+
| 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  NVIDIA RTX A6000    Off  | 00000000:41:00.0 Off |                  Off |
| 30%   45C    P0    80W / 300W |      0MiB / 49140MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA RTX A4000    Off  | 00000000:61:00.0 Off |                  Off |
| 34%   56C    P0    38W / 140W |      0MiB / 16376MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

[nouiz]

How did you install jaxlib and cudnn?
You could have a mismatch of version. Maybe you use a jaxlib that expect a different cudnn version?

[hawkinsp]

I would have expected CuDNN 8.6 to work. Can you try a newer CuDNN, e.g., 8.8?

[mehdiataei]

OK this is fixed with CuDNN 8.8 👍 Not sure if there was an issue with my cuDNN installation, so maybe 8.6 also works if I reinstall it. I may try it a bit later.

[mjsML]

@mehdiataei is the issue fixed so we can close the bug?