使用 CUDA 12.9 编译 PyTorch 2.4.0

最近跑的一个项目需要 torch==2.4.0，但是 GPU 需要 CUDA 12.9，PyTorch 官方这个配置的预编译包，因此需要手动编译。

• 操作系统：Ubuntu 24.04 LTS
• GPU：NVIDIA RTX PRO 6000 (SM 12.0, CUDA 12.9)

1. 获取源码：

   git clone -b v2.4.0 --depth 1 https://github.com/pytorch/pytorch
   cd pytorch
   git submodule sync
   git submodule update --init --recursive --depth 1 --progress
   

2. 安装编译环境：

   sudo apt install libnvtoolsext1
   sudo ln -s /usr/lib/x86_64-linux-gnu/libnvToolsExt.so.1 /usr/lib/x86_64-linux-gnu/libnvToolsExt.so
   sudo ln -s /usr/lib/x86_64-linux-gnu/libnvToolsExt.so.1 /usr/local/cuda-12.9/lib64/libnvToolsExt.so
   sudo ln -s /usr/lib/x86_64-linux-gnu/libnvToolsExt.so.1.0.0 /usr/local/cuda-12.9/lib64/libnvToolsExt.so.1
   cd pytorch
   sudo cp third_party/nccl/nccl/src/include/nvtx3/nvToolsExt.h /usr/local/cuda-12.9/include/
   sudo cp third_party/nccl/nccl/src/include/nvtx3/nvToolsExtCuda*.h /usr/local/cuda-12.9/include/
   export CPATH="$CUDA_HOME/targets/x86_64-linux/include/nvtx3"
   export 
   
   conda create -n torch-builder python=3.12
   conda activate torch-builder
   pip install -U cmake~=3.18.0 mkl-static mkl-include
   pip install -r requirements.txt
   

3. 安装 triton：

   curl -sSL https://raw.githubusercontent.com/pytorch/pytorch/refs/heads/main/.ci/docker/common/install_magma.sh | sudo bash -s -- 12.9  # 修改 CUDA 版本
   make triton
   

4. 编译安装 torch：

   export TORCH_CUDA_ARCH_LIST="12.0"
   export _GLIBCXX_USE_CXX11_ABI=1
   export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
   python setup.py develop
   

参考：pytorch/pytorch#v2.4.0