Lambda Labs "No GPU/TPU found, falling back to CPU."

[Aymericr]

My setup:

Ubuntu 20.04.5 LTS NVIDIA GPU A100 (lambdalabs) NVIDIA-SMI / Driver Version: 515.65.01 CUDA Version: 11.7 Python 3.10.11 (via Anaconda) pip install git+https://github.com/sanchit-gandhi/whisper-jax.git

When I run in a Jupyter Notebook:

from whisper_jax import FlaxWhisperPipline
import jax.numpy as jnp
import jax
jax.devices()

I get:

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

[EpicCodeWizard]

I'm facing this issue as well. Has a solution been found?

[Aymericr]

I was able to fix it. It was a Lambda Labs issue.

Here: https://docs.lambdalabs.com/cloud/cannot-find-nvidia-cudnn/

Software outside of Lambda Stack usually looks for the cuDNN library files in /usr/lib/x86_64-linux-gnu. However, on our instances, the cuDNN library files are in /usr/lib/python3/dist-packages/tensorflow.

Creating symbolic links, or “symlinks,” for the cuDNN library files might allow your program to find the cuDNN library on our instances.

Run the following command to create symlinks for the cuDNN library files:

for cudnn_so in /usr/lib/python3/dist-packages/tensorflow/libcudnn*; do
  sudo ln -s "$cudnn_so" /usr/lib/x86_64-linux-gnu/
done

Another alternative is to install cuDNN manually by signing up as a developer and downloading the Linux installer from the NVIDIA website (cumbersome).

Also, please note:

• I changed to python3.9.16 (Lambda labs currently defaults to 3.8)
• I had trouble running Whisper-JAX from a virtualenv or conda so I overwrote the default python:

  sudo apt-get update
  sudo apt install software-properties-common -y
  sudo add-apt-repository ppa:deadsnakes/ppa -y
  sudo apt install python3.9 -y
  sudo apt install python3-pip -y
  python3.9 --version

• make sure you follow the right instructions to install JAX for cuda: pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Cheers mate

[silvacarl2]

after doing this: pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

i now get this:

pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

nvcc --version nvcc: NVIDIA (R) Cuda compiler driver Copyright (c) 2005-2022 NVIDIA Corporation Built on Tue_May__3_18:49:52_PDT_2022 Cuda compilation tools, release 11.7, V11.7.64 Build cuda_11.7.r11.7/compiler.31294372_0