Charles Herring 14 March 2024
Pytorch

Updated 5/27/2024 to support PyTorch 2.4, Rust and CUDA 12.5

Background

I do my development work on VSCode on a Windows 11 Pro workstation that has 2 NVIDIA GPUs. I want to use the newest features from PyTorch and the NVIDIA CUDA toolkit while having bare metal access. Also I want to compile with access to emerging libraries that add features and speed to generative artificial intelligence (genai) development. This document outlines the steps I took to create an environment that has the following characteristics:

  • Python 3.10
  • NVIDIA CUDA Toolkit 12.5
  • NVIDIA cuDNN 9.0
  • PyTorch 2.4
  • Flash Attention support

Install Anaconda

Almost all AI development work is currently done in Python. Anaconda is the standard way for managing Python environments. In this set up we will create an environment called ai. Anaconda will allow for other Python environments that do not conflict with the fragile and complicated `ai` build for Windows. It is a straightforward installer: Free Download | Anaconda

Install Visual Studio Community 2022

To compile much of the code, Visual Studio Community 2022 Build Tools are needed. You want to include C++ and Python modules when you install. I installed all modules. You can grab it here: https://visualstudio.microsoft.com/downloads/. This is not included in VSCode and must be installed separately.

Install NVIDIA CUDA Toolkit 12.5

Installing the NVIDIA CUDA Toolkit is straightforward. Just make sure you have uninstalled any previous versions before installing it. It can be grabbed here: https://developer.nvidia.com/cuda-downloads

Install NVIDA Nsight NVTX (via CUDA 11.8)

The newer toolkit is failing to install the required Nsight NVTX package. Compiling will throw an error message stating “nvToolsExt” cannot be found. To fix this:

  1. Download CUDA Toolkit 11.8: https://developer.nvidia.com/cuda-11-8-0-download-archive
  2. Perform a Custom install
  3. Deselect all features except “Nsight NVTX” and install.

A huge shout out to Te93 in figuring that out: https://discuss.pytorch.org/t/failed-to-find-nvtoolsext/179635/2

Install NVIDIA cuDNN Library

The NVIDIA CUDA® Deep Neural Network library (cuDNN) is used in several PyTorch operations. It can be downloaded here: https://developer.nvidia.com/cudnn-downloads. The installer is straightforward.

Git for Windows

Install Git for grabbing PyTorch code: https://git-scm.com/download/win

Rust for Windows

Newer versions of PyTorch now require Rust to build: https://www.rust-lang.org/tools/install

Compiling PyTorch from Source

Compiling PyTorch will maximize compatibility with your hardware. It will also allow access to the newest features in PyTorch and CUDA toolkit. It takes a lot of time and processing power to complete. It is an advanced exercise and all kind of things can go wrong. The official documentation on accomplishing it can be found here: https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source

Steps

Detailed steps are laid out above. These are the steps I carried out:

  • Open Anaconda Powershell Prompt (use Run as Administrator)
  • Set-Item -Path Env:CUDA_VERSION -Value "12.5" //Define the CUDA Version installed
  • Set-Item -Path Env:USE_CUDA -Value "1" //Force the use of CUDA
  • Set-Item -Path Env:USE_FLASH_ATTENTION -Value "1" //Enable support for Flash Attention
  • git clone --recursive https://github.com/pytorch/pytorch
  • cd pytorch
  • conda create -n ai python=3.10 //Create a conda environment to install all the goodness
  • conda activate ai //Switch to the environment
  • conda install cmake ninja //Install basic compiler components
  • pip install -r requirements.txt //Install PyTorch required Python components
  • conda install intel::mkl-static intel::mkl-include //Intel libraries
  • conda install -c conda-forge libuv=1.39 //Downgrade to support PyTorch
  • python setup.py develop //Execute the compile action
  • python setup.py install //Install the package

Troubleshooting Compile

Here are a few steps I took along the way:

  1. conda uninstall cmake ninja intel::mkl-static intel::mkl-include
  2. conda clean –all
  3. Uninstall anaconda
  4. Uninstall Visual Studio (all versions)
  5. Uninstall CUDA and cuDNN  (all versions)
  6. Reboot
  7. Reinstall anaconda
  8. Reinstall VS Build Tools
  9. Reinstall CUDA and cuDNN
  10. Reboot
  11. Retry compile steps

Conclusion

When the compilation is completed, I can change my terminal interpreter to the `conda ai` environment to use the libraries. For guidance on using anaconda environments in VSCode see: https://code.visualstudio.com/docs/python/environments.

I hope this saves someone from the same frustrations I had. Best of luck!

Tags