GH 200 - PyTorch - Triton setup

gistfile1.txt

#!/bin/bash

# Exit on any error
set -e

# Step 1: Create a new virtual environment
echo "Creating virtual environment 'myenv'..."
python3 -m venv myenv
source myenv/bin/activate
echo "Virtual environment activated."

# Step 1a: Install CUDA 12.8 based on system architecture
echo "Detecting system architecture..."
ARCH=$(uname -m)
if [ "$ARCH" = "x86_64" ]; then
    echo "Installing CUDA 12.8 for x86_64..."
    wget https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda_12.8.0_560.35.03_linux.run
    sudo sh cuda_12.8.0_560.35.03_linux.run --toolkit --silent --override
    rm cuda_12.8.0_560.35.03_linux.run
elif [ "$ARCH" = "aarch64" ]; then
    echo "Installing CUDA 12.8 for aarch64..."
    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-ubuntu2204.pin
    sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
    wget https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda-repo-ubuntu2204-12-8-local_12.8.0-570.86.10-1_arm64.deb
    sudo dpkg -i cuda-repo-ubuntu2204-12-8-local_12.8.0-570.86.10-1_arm64.deb
    sudo cp /var/cuda-repo-ubuntu2204-12-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
    sudo apt-get update
    sudo apt-get -y install cuda-toolkit-12-8
    rm cuda-repo-ubuntu2204-12-8-local_12.8.0-570.86.10-1_arm64.deb
else
    echo "Unsupported architecture: $ARCH"
    exit 1
fi
echo "CUDA 12.8 installation complete."

# Set CUDA environment variables
export PATH=/usr/local/cuda-12.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH

# Step 2: Build and install PyTorch from source
echo "Building PyTorch from source..."
git clone https://github.com/pytorch/pytorch
cd pytorch
git checkout v2.6.0-rc9
git submodule sync
git submodule update --init --recursive -j 8
pip install -r requirements.txt
pip install mkl-static mkl-include wheel
export CUDA_HOME=/usr/local/cuda-12.8
export CUDA_PATH=$CUDA_HOME
export TORCH_CUDA_ARCH_LIST="8.0 9.0"  # Adjust based on your GPU; e.g., Blackwell is 9.0
python setup.py develop
cd ..
echo "PyTorch installation complete."

# Step 3: Build and install Triton from source
echo "Building Triton from source..."
git clone https://github.com/triton-lang/triton.git
cd triton
pip install ninja cmake wheel pybind11
pip install -e python
cd ..
echo "Triton installation complete."

echo "Setup complete! Virtual environment 'myenv' is ready with CUDA 12.8, PyTorch, and Triton."