Running PyTorch and Triton on the RTX 5080
I was beyond stoked at the opportunity of getting my hands on a new RTX 5080 to speed up my machine learning developments! Unfortunately, as soon as I connected the new GPU to my workstation I quickly realised that things did not run quite as smoothly as I hoped for. Drivers needed to be re-installed, and machine learning libraries had to be re-downloaded and even built from source. I thought I’d write up a quick how-to for people to potentially save them some hours and to just get going - I’m pretty sure in a few weeks most of these issues won’t be issues anymore with new software releases.
Installing compatible drivers
I initially just updated my drivers to the nvidia-driver-570 (as recommended by “ubuntu-drivers devices”) but I was greeted with a lovely “No devices were found” whenever I ran “nvidia-smi”. After some trial-and-error, I figured that I needed to install the NVIDIA open kernel modules and not the proprietary ones:
sudo apt install -y nvidia-driver-570-open
[Maybe make sure to uninstall all previous drivers if this does not work.]
Installing PyTorch
So, this wasn’t as bad as I thought - all works by downloading the nightly build (I’m assuming the official release will come soon):
pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
Incase this is important, I’m using Python 3.10 and confirm all works here.
Building Triton
Okay, so this one was a bit of a pain. The official Triton GitHub page suggests building PyTorch from source (which is not needed - you can download the nightly build) and build Triton from source (this is needed). I followed the instructions here from their GitHub page:
git clone https://github.com/triton-lang/triton.git
cd triton
pip install ninja cmake wheel pybind11 # build-time dependencies
pip install -e python
However, I have to add that I found that the installation got stuck on “running build_ext” - but it turns out that the Triton builder is just downloading something at this step, so just give it a good amount of time (like 10 to 30 minutes). Then I got a couple of errors like missing “Zlib” so I went ahead and did a “sudo apt install -y zlib1g-dev”. Also had a missing “GLIBCXX” error so I did a “conda install -c conda-forge libstdcxx-ng”).
All seemed to work, but then when I ran my Triton kernels I was hit with a “@triton.jit not being found” error. Turns out there is a simple fix to this of just building Triton with a “python setup.py install” as mentioned in this issue.