Yesterday I upgraded to the latest version of bumblebee only to learn EXLA took on the latest XLA which dropped support for CUDA 11.2.
Because I was running Pop!OS I assumed it would be something I could `apt install` and call it good. Unfortunately it was more work than my first go round so I wanted to detail it here for anyone who might follow. This is loosely based off the older CUDA 11 install instructions I found but updated for ubuntu 22.04.
Before you get started downloading anything, setup a developer account with nvidia so you can download cuDNN in a few minutes.
$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
$ sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
$ sudo add-apt-repository "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
$ sudo apt update
$ sudo apt install cuda-toolkit-12-2
Next download cuDNN from nvidia and install it.
$ sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.4.25_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.4.25/cudnn-local-72322D7F-keyring.gpg /usr/share/keyrings/
$ sudo apt update
$ sudo apt install libcudnn8=8.9.4.25-1+cuda12.2
Finally, set a few env variables and the correct XLA_TARGET for EXLA
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
export CUDA_HOME=/usr/local/cuda
export PATH="/usr/local/cuda/bin:$PATH"
export XLA_TARGET=cuda120
export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda
Open a new terminal session and verify CUDA 12 is setup correctly by running this command
nvcc -V
To upgrade from 8.9 download cuDNN 9.1 from nvidia. Before you install it remove libcudnn8 as shown below.
$ sudo apt remove libcudnn8
$ sudo dpkg -i cudnn-local-repo-ubuntu2204-9.1.1_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2204-9.1.1/cudnn-local-AD7F4AC5-keyring.gpg /usr/share/keyrings/
$ sudo apt update
$ sudo apt install libcudnn9-cuda-12=9.1.1.17-1
$ sudo apt install cuda-toolkit-12-6
$ sudo apt update
After you upgrade to 9.1.1 if you ever want to downgrade to 8.9 follow these instructions.
$ sudo apt remove libcudnn9-cuda-12
$ sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-08A7D361-keyring.gpg /usr/share/keyrings/
$ sudo apt update
$ sudo apt install libcudnn8=8.9.7.29-1+cuda12.2
The only difference between 8.9 and 9.1 is a small tweak to the XLA_TARGET in your zshrc. For Nx 7 you would normally use cuda120 but after upgrading to Nx 9 and CUDNN 9.1.1 you need cuda12 instead
XLA_TARGET=cuda12
