Getting Bumblebee GPU Support Online in WSL2

Update 11/13/23: The ecosystem has caught up to CUDA 12 now, and this guide is out of date. The instructions at https://docs.nvidia.com/cuda/wsl-user-guide/index.html have worked fine for me recently.

Update 9/12/23: I ran through the guide again today on my workstation reinstall. I updated the asdf to the latest version, but ran into an issue after installing: syntax errors complaining about \r CRLFs in the asdf script. I ran “git config —global core.autocrlf input”, deleted the .asdf folder and recloned it, and everything worked fine.

At the time of this writing, if you follow the basic instructions you’ll find online from sources like Microsoft and Nvidia, you’ll end up with pain down the road. This is because Nvidia has released a new version of CUDA—12.0—and essentially nothing works with it yet. PyTorch and XLA both have issues open about this, and Pytorch has nightly builds with support, but overall the situation is “install CUDA 11.8 and wait if you need something that works now”.

Bumblebee with GPU support has a couple of dependencies from nvidia: CUDA and cuDNN.

I don’t know if this is the most efficient way to do this, but it works for me.

Pre-requisites:

• WSL installed, Ubuntu flavor.
• GeForce drivers installed (I used GeForce Experience to do this)

Install CUDA Toolkit

Source: CUDA Toolkit 11.8 Downloads | NVIDIA Developer

wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

Then add this to your .bashrc file:

export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}

export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64\
                         ${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

Install cuDNN

Setup repo

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-key del 7fa2af80

Install cuDNN packages

sudo apt-get install libcudnn8=8.9.2.*-1+cuda11.8
sudo apt-get install libcudnn8-dev=8.9.2.*-1+cuda11.8
sudo apt-get install libcudnn8-samples=8.9.2.*-1+cuda11.8

Your GPU should now be online.

Install Elixir Ecosystem

We’re going to use asdf to manage our erlang and elixir versions.

Install dependencies

asdf and elixir require a few packages that aren’t installed by default.

sudo apt install curl git libssl-dev automake autoconf libncurses5-dev unzip

Install asdf

git clone https://github.com/asdf-vm/asdf.git ~/.asdf --branch v0.13.0

For bash, edit your .bashrc and add these lines:

. "$HOME/.asdf/asdf.sh"
. "$HOME/.asdf/completions/asdf.bash"

Install erlang / elixir

asdf plugin-add erlang
asdf plugin-add elixir

asdf install erlang latest:26.0
asdf global erlang latest:26.0

asdf install elixir 1.15.0-otp-26
asdf global elixir 1.15.0-otp-26

Install Phoenix

mix archive.install hex phx_new

Install Livebook

mix do local.rebar --force, local.hex --force
mix escript.install hex livebook
asdf reshim elixir

Launch livebook

livebook server

Add environment variable XLA_TARGET=cuda118

Check that the GPU is recognized by XLA:

Mix.install([
  {:bumblebee, "~> 0.3.0"},
  {:nx, "~> 0.5.1"},
  {:exla, "~> 0.5.1"},
  {:kino, "~> 0.8.0"}
])

Nx.global_default_backend(EXLA.Backend)

EXLA.Client.get_supported_platforms()

You should see cuda on the list.

And that’s it. You’re ready to go. If you want to dive right in with Elixir Machine Learning on the GPU, copy over the livebook at https://github.com/elixir-nx/bumblebee/blob/main/notebooks/stable_diffusion.livemd and give it a try.