Julia GPU Programming with WSL2

A walkthrough on how to get started with Julia GPU programming under WSL2 on Windows.

Julia GPU Programming with WSL2

I use Windows for gaming. It's been a long time since I've last done any serious development work on my Windows machine and yet I still spent a good chunk of money on building out a beefy machine for my efforts to learn machine learning. It has taken me a few months to finally sit down and get this machine ready for anything other than gaming.

With the recent announcement of GUI support for WSLg I got really excited to try out WSL and see how good the GPU support actually is, but that's not the main reason. I've been shying away from developing on Windows because I'm used to a *NIX environment. WSL gives you that, but up until recently you wouldn't have been able to interact with any GPU - and this all changed with this announcement!

You can watch the video below to see what's coming for WSL2.

The first half of this article will show you how to get everything set up and in the second half we'll set up CUDA.jl in Julia. Since the latter part is about CUDA I'll assume that you have a compatible nVidia GPU.

Installation

Here's a summary of what we need to go through to prepare our environment:

  • Update Windows 10 to latest release on the dev channel
  • Install nVidia CUDA drivers
  • Install Ubuntu 20.04 in WSL2
  • Install Linux CUDA packages
  • 🎉

Windows 10 Insider Preview

At the time of writing all of the features are only available through the Windows Insider Program. The Windows Insider Program allows you to receive new Windows features before the hit the main update line. The program is split into three channels: Dev, Beta and Release Preview.

To receive the update with WSLg and GPU support we will need to switch to the dev channel.

Note: The dev channel comes with some rough edges and potential for system instability. Be mindful of this when you switch and make sure you have backups!

After installing all downloaded updates you should end up with OS Build 21364 or higher. You can check your OS Build by running winver in PowerShell/cmd.

Screenshot of
Run winver using the Windows Run command

With this all set we can hop on to install the latest WSL2 compatible CUDA drivers.

CUDA drivers

NVIDIA are providing special CUDA drivers for Windows 10 WSL. The link below will take you to the download page. It's required to sign up for the NVIDIA Developer Program, which is free.

GPU in Windows Subsystem for Linux (WSL)
CUDA on Windows Subsystem for Linux (WSL) - Public Preview Microsoft Windows is a ubiquitous platform for enterprise, business, and personal computing systems. However, industry AI tools, models, frameworks, and libraries are predominantly available on Linux OS. Now all users of AI - whether they ar…

Follow the setup wizard (I chose the express installation which keeps existing settings in place).

Note: There's also a documentation page provided by NVIDIA around setting up your GPU for WSL. I found the docs to be outdated and not working on my machine.

Installing Ubuntu 20.04 LTS with WSL2

Before we can proceed with installing Ubuntu I advise to update the WSL kernel by running:

wsl --update

In case you're like me and don't enjoy the default terminal Windows comes with I suggest you install Windows Terminal from the Microsoft Store. This terminal comes is a lot more pleasant to use than either cmd or the PowerShell terminal ever were.

It's also a good idea to set WSL to default to version 2:

wsl --set-default-version 2

Finally let's install Ubuntu with:

wsl --install --distribution Ubuntu-20.04

In case you were wondering what other distributions are available you can simply run: wsl --list --online.

Ubuntu 20.04 LTS

With Ubuntu installed we have a couple of final steps. First we will add an upstream repo to apt for getting the latest CUDA builds directly from NVIDIA:

sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
Add CUDA apt repository for Ubuntu 20.04

We also need to add NVIDIA's GPG key for the apt repo:

sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub 
Add nVidia gpg key

And finally to make sure that we prefer the packages provided by NVIDIA over packages in mainline Ubuntu we need to pin the apt repo:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
Pin file to prioritize CUDA apt repository

With this out of the way, we're ready to install the CUDA drivers inside our WSL Ubuntu installation:

 sudo apt update && sudo apt install -y cuda-drivers

In case setting up WSL2 with GPU support was all you wanted to do - we're done!

Julia + WSL2 + Cuda.jl

Ok, so as promised in the headline we're now going to install Julia inside our Ubuntu 20.04 and set up Cuda.jl.

Install Julia

At the time of writing Julia 1.6.1 is the latest version available (make sure to check for updates on Julia's download page).

First let's fetch the latest Julia tarball:

wget https://julialang-s3.julialang.org/bin/linux/x64/1.6/julia-1.6.1-linux-x86_64.tar.gz

Extract the .tar.gz:

tar -xvzf julia-1.6.1-linux-x86_64.tar.gz

Move the extracted folder to /opt:

sudo mv -r julia-1.6.1 /opt/

Finally, create a symbolic link to julia inside the /usr/local/bin folder:

sudo ln -s /opt/julia-1.6.1/bin/julia /usr/local/bin/julia

You may pick a different target directory for your installation of Julia or use a version manager like asdf-vm.

Install Cuda.jl

At this point simply run julia in your terminal and you should be dropped into the Julia REPL. I assume you've worked with Julia before and know how to operate it's package manager.

Hit ] to enter pkg mode and install CUDA with:

activate --temp
add CUDA

From here hit backspace and import CUDA. CUDA.jl provides a useful function called functional which will confirm that we've done everything right (well, that's the hope at least, right?).

using CUDA
CUDA.functional()
Ensuring Julia CUDA.jl is functional inside WSL2

You can additionally run CUDA.versioninfo() to get a more detailed breakdown of the supported features on your GPU.

At this point you should have a working installation with WSL2, Ubuntu 20.04, Julia and CUDA.jl. In case you're new to CUDA.jl I suggest you follow the excellent introduction to GPU programming by JuliaGPU or jump in at the deep end with FluxML's GPU support.


If you like articles like this one, please consider subscribing to my free newsletter where at least once a week I send out my latest work covering Julia, Python, Machine Learning and other tech.

You can also follow me on Twitter.