Overview

This page documents how to set up WSL in order to run torch.compile on Windows. torch.compile relies on triton which only supports Linux, so WSL (Windows Linux Subsystem) is used to get full support of torch.compile. Why use torch.compile?

torch.compile is the latest method to speed up your PyTorch code! torch.compile makes PyTorch code run faster by JIT-compiling PyTorch code into optimized kernels, all while requiring minimal code changes.

In short, torch.compile speeds up training. Note that the first iteration of training will be slow as the model will be compiled, but subsequent iterations should be faster compared to not using torch.compile. The amount of time to compile the model in the first iteration, as well as the total amount of speedup after compiling, both depend on the specific architecture. Heavy transformer architectures can take several minutes to compile.

Instructions

1. We will use the default WSL settings and install Ubuntu. Install WSL by opening a CMD or PowerShell prompt and running the command wsl --install, restart the PC when prompted and then finish set up by creating your Ubuntu username and password.
2. After installation you should have an Ubuntu command prompt open, or you can open it by selecting Ubuntu from your start menu. Install Python 3.12 with these commands:

   sudo add-apt-repository ppa:deadsnakes/ppa
   sudo apt update
   sudo apt install python3.12-full -y
   

3. Install build essential:

   sudo apt-get update
   sudo apt-get install build-essential
   

4. Install Python 3 development headers:

   sudo apt-get install python3.12-dev 
   

5. Clone traiNNer-redux, and install its dependencies:

   git clone https://github.com/the-database/traiNNer-redux.git
   cd traiNNer-redux
   chmod +x install.sh && ./install.sh