一键构建 pytorch cpp lib 前端和 wheel

系统环境：ubuntu 22.04

V100 GPU

cuda 12.8.2

cudnn 9.10.2.21

TensorRT 10.9.0.34

1. 使用构建脚本

构建 develop，或者 wheel。

第一次构建：

bash 复制代码

BUILD_MODE=wheel ./meme_build.sh

如果构建过一次之后，依赖包可以跳过，使用 SKIP_APT=1

bash 复制代码

BUILD_MODE=wheel  SKIP_APT=1 ./meme_build.sh

验证 triton：

bash 复制代码

cd ~
python3 -c "import triton; print(triton.__version__)"
# 期望输出: 3.4.0

python3 -c "import torch; f = torch.compile(lambda x: x * 2); print(f(torch.randn(8, device='cuda')))"

bash 复制代码

cd ~
python3 -c "import torch; print(torch.__file__)"
# 期望: /home/ruler/.local/lib/python3.10/site-packages/torch/__init__.py
# 或者 /usr/lib/python3/dist-packages/... 之类, 总之不能是 pytorch 源码目录

其中，BUILD_TRITON=1 忽略，triton 已经从 torch 2.4.0时，移出了 torch 源码库。

v2.8.0 的 make triton 实际执行的是 scripts/install_triton_wheel.sh，它做的事就是：

bash 复制代码

pip install pytorch-triton==3.4.0+git11ec6354 \
    --index-url https://download.pytorch.org/whl/nightly/

也就是下载 PyTorch 官方 CI 预先构建好的 wheel，本地一行 Triton 代码都没有编译。这背后的原因前面说过：Triton 早已不是 PyTorch 的子模块，PyTorch 把它改成了外部版本化依赖------由 PyTorch 的 CI 从 triton-lang/triton 的某个钉死 commit（这里就是 git11ec6354）构建 wheel，发布到自己的 nightly 源，make triton 只是帮你把钉死的那个版本装上。

2. 构建脚本

bash 复制代码

#!/usr/bin/env bash
# =============================================================================
# PyTorch v2.8.0 源码构建脚本
# 目标环境: V100 (sm_70) / CUDA 12.8.2 / cuDNN 9.10.2.21 (与CUDA同目录)
#          Ubuntu 22.04 / Python 3.10 (系统Python, 无需conda)
# 依据: pytorch/pytorch v2.8.0 tag README.md "From Source" 流程
#
# 用法:
#   chmod +x build_pytorch_2.8.0_v100.sh
#   ./build_pytorch_2.8.0_v100.sh              # 默认构建 wheel 并 pip 安装
#   BUILD_MODE=develop ./build_pytorch_2.8.0_v100.sh   # README 原版 develop 模式
#   SKIP_APT=1 ./build_pytorch_2.8.0_v100.sh   # 跳过系统依赖安装
# =============================================================================
set -euo pipefail

# ----------------------------- 可配置参数 ------------------------------------
PYTORCH_VERSION="v2.8.0"
BUILD_MODE="${BUILD_MODE:-wheel}"        # wheel | develop
SKIP_APT="${SKIP_APT:-0}"                # 1 = 跳过 apt 安装
BUILD_TRITON="${BUILD_TRITON:-0}"        # 1 = 构建 triton (torch.compile 用)
MAX_JOBS="${MAX_JOBS:-$(nproc)}"

CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
WORK_DIR="${WORK_DIR:-$HOME/ex_holoscan/pytorch_src}"
# -----------------------------------------------------------------------------

log() { echo -e "\n\033[1;32m=== $* ===\033[0m"; }
die() { echo -e "\033[1;31m[ERROR] $*\033[0m" >&2; exit 1; }

# ============================ 1/7 系统依赖 ====================================
if [[ "${SKIP_APT}" != "1" ]]; then
    log "1/7 安装系统依赖 (内部使用 sudo)"
    sudo apt update
    sudo apt install -y build-essential cmake ninja-build git python3-dev \
        python3-pip libopenblas-dev libnuma-dev
else
    log "1/7 跳过系统依赖安装 (SKIP_APT=1)"
fi

# gcc 版本检查: README 要求 Linux 上 gcc >= 9.4.0
GCC_VER="$(gcc -dumpversion)"
log "gcc 版本: ${GCC_VER} (README 要求 >= 9.4.0)"

# ============================ 2/7 CUDA 环境 ===================================
log "2/7 配置 CUDA 环境"
[[ -x "${CUDA_HOME}/bin/nvcc" ]] || die "未找到 nvcc: ${CUDA_HOME}/bin/nvcc，请检查 CUDA_HOME"

export PATH="${CUDA_HOME}/bin:${PATH}"
# 用 :- 兜底, 避免 set -u 下 LD_LIBRARY_PATH 未定义时报 unbound variable
export LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-}"

# V100 = Volta = sm_70; 只编译目标架构, 大幅缩短构建时间
export TORCH_CUDA_ARCH_LIST="7.0"

nvcc --version | grep "release" || true
echo "TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}"

# ============================ 3/7 获取源码 ====================================
log "3/7 获取 PyTorch 源码 (${PYTORCH_VERSION})"
mkdir -p "${WORK_DIR}"
cd "${WORK_DIR}"

if [[ -d pytorch/.git ]]; then
    echo "复用已有源码目录: ${WORK_DIR}/pytorch"
    cd pytorch
    git fetch --tags origin || true
else
    git clone https://github.com/pytorch/pytorch
    cd pytorch
fi

# 关键: 必须先 checkout tag, 再更新 submodule, 保证子模块版本与 tag 匹配
git checkout "${PYTORCH_VERSION}"
git submodule sync
git submodule update --init --recursive

# ============================ 4/7 Python 构建依赖 =============================
log "4/7 安装 Python 构建依赖"
python3 -m pip install --upgrade pip
python3 -m pip install -r requirements.txt
python3 -m pip install mkl-static mkl-include

# ============================ 5/7 Triton (可选) ===============================
if [[ "${BUILD_TRITON}" == "1" ]]; then
    log "5/7 构建 Triton (torch.compile/inductor 用)"
    make triton
else
    log "5/7 跳过 Triton (BUILD_TRITON=${BUILD_TRITON})"
fi

# ============================ 6/7 构建并安装 ==================================
log "6/7 构建 PyTorch (BUILD_MODE=${BUILD_MODE}, MAX_JOBS=${MAX_JOBS})"
export MAX_JOBS

case "${BUILD_MODE}" in
    develop)
        # README 官方命令: 开发模式原地安装
        python3 setup.py develop
        ;;
    wheel)
        # 构建 wheel: wheel 内 torch/lib + torch/include 即为 libtorch,
        # 可供 holoscan-sdk 链接 (Torch_DIR=<site-packages>/torch/share/cmake/Torch)
        python3 setup.py bdist_wheel
        WHEEL="$(ls -t dist/torch-2.8.0*.whl | head -n1)"
        [[ -n "${WHEEL}" ]] || die "未找到构建产物 dist/torch-2.8.0*.whl"
        python3 -m pip install --force-reinstall "${WHEEL}"
        ;;
    *)
        die "未知 BUILD_MODE=${BUILD_MODE} (可选: wheel | develop)"
        ;;
esac

# ============================ 7/7 验证 ========================================
log "7/7 验证安装"
# 切出源码目录再 import, 避免从源码目录误导入
cd /
python3 - <<'EOF'
import torch
print("torch version :", torch.__version__)
print("built with CUDA:", torch.version.cuda)
print("cuDNN version :", torch.backends.cudnn.version())
print("cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("device 0      :", torch.cuda.get_device_name(0))
    print("capability    :", torch.cuda.get_device_capability(0))  # 期望 (7, 0)
    x = torch.randn(1024, 1024, device="cuda")
    print("matmul check  :", (x @ x).sum().item() != 0)
EOF

log "构建完成 ✅"
echo "源码目录 : ${WORK_DIR}/pytorch"
[[ "${BUILD_MODE}" == "wheel" ]] && echo "wheel 产物: ${WORK_DIR}/pytorch/dist/"
echo "提示: TensorRT (/usr/local/TensorRT-10.9.0.34/) 不参与 PyTorch 本体构建, 留给 holoscan 使用即可"