Setup PyTorch with CUDA in WSL-Ubuntu on Windows 11

CUDA-WSL.md

Setup PyTorch with CUDA in WSL-Ubuntu on Windows 11

On Windows 11:

1. Download and install Nvidia GPU Driver for Windows: https://www.nvidia.com/download/index.aspx
2. Download and install CUDA Toolkit for Windows: https://developer.nvidia.com/cuda-downloads

NOTE: Nvidia Official doc says, the GPU Driver includes "CUDA Support", which is very confusing. It doesn't mean CUDA is included in Driver. What included is only some support files, not CUDA itself. You still need to install CUDA separately.

On WSL-Ubuntu:

1. DO NOT Install Nvidia GPU Driver for Linux. The Windows Nvidia driver is already injected into WSL, if you install driver again in WSL, it will override the Windows one.
2. If you just need to run existing CUDA app, no need to install CUDA, it should already working.
3. If you need to compile new CUDA app in WSL-Ubuntu, then you need to install CUDA, but the sepcial WSL-Ubuntu distribution: https://developer.nvidia.com/cuda-downloads This one won't override the injected CUDA from Windows host machine.

To check CUDA is working in WSL-Ubuntu, run this script:

#!/usr/bin/env python3

import os
import torch

x = torch.rand(5, 3)
print(x)

b = torch.cuda.is_available()
print(f'cuda available: {b}')

n = torch.cuda.device_count()
print(f'device count: {n}')

n = torch.cuda.get_device_name(0)
print(f'device 0 name: {n}')

To check GPU is working with Docker in WSL-Ubuntu:

docker run --rm -it --gpus=all nvcr.io/nvidia/k8s/cuda-sample:nbody nbody -gpu -benchmark

FAQ:

I want to install CUDA 12.1, which is listed on PyTorch website. But the above Nvidia website only shows 12.4 for me ?

CUDA 12.4 should be compatible with 12.1, and supported by PyTorch too.

If you do care, search for "CUDA 12.1" in google, you will find it. In the generated shell commands, you will need to change:

sudo apt install cuda   # this will install cuda-12-4 still

to

sudo apt install cuda-12-1

And you can check with:

/usr/local/cuda-12.1/bin/nvcc --version

If I don't install Nvidia driver in WSL-Ubuntu, how can I use nvidia-smi ?

Use the one injected into WSL: /usr/lib/wsl/lib/nvidia-smi To make it easier to user, you can either:

• add /usr/lib/wsl/lib to your $PATH,or
• link above nvidia-smi into dir such as /usr/local/bin

re: https://docs.nvidia.com/cuda/wsl-user-guide/index.html