使用 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 -sf /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/conda create -n torch-builder python=3.12
   conda activate torch-builder
   pip install cmake==3.18.4 mkl-static mkl-include
   pip install -r requirements.txt
   export TORCH_CUDA_ARCH_LIST="8.0;8.9"
   export _GLIBCXX_USE_CXX11_ABI=1
   export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
   

3. 编译：

   python setup.py develop
   

参考：pytorch/pytorch#v2.4.0