prain: build scripts (ROCm SDK venv, PyTorch and vLLM from source with clang)
This commit is contained in:
Executable
+121
@@ -0,0 +1,121 @@
|
||||
#!/bin/bash
|
||||
# Build prain's stack on the GPU host, with clang, tuned for this machine:
|
||||
# wrappers clang/clang++ wrappers in build/bin (libstdc++ + libgcc, lld)
|
||||
# sdk venv + AMD's ROCm SDK wheels for $PRAIN_GPU_ARCH
|
||||
# fetch clone PyTorch and vLLM into build/
|
||||
# torch PyTorch from source -> build/wheels, installed into the venv
|
||||
# vllm vLLM from source (adds $PRAIN_GPU_ARCH to its HIP arch list), installed editable
|
||||
# prain this package, installed editable, and its test suite
|
||||
# usage: scripts/build.sh STEP... Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VLLM_REF
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
ROCM_INDEX=${PRAIN_ROCM_INDEX:-https://stable.repo.amd.com/rocm/whl-next/}
|
||||
ROCM_VERSION=${PRAIN_ROCM_VERSION:-10.1.0}
|
||||
TORCH_REF=${PRAIN_TORCH_REF:-v2.13.0}
|
||||
VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
|
||||
PYTHON=${PRAIN_PYTHON:-python3.13}
|
||||
|
||||
wrappers() {
|
||||
mkdir -p build/bin
|
||||
local gccdir
|
||||
gccdir=$(gcc -print-search-dirs | sed -n 's/^install: //p')
|
||||
for c in clang clang++; do
|
||||
cat > build/bin/$c <<EOF
|
||||
#!/bin/sh
|
||||
case "\$*" in
|
||||
# clang 23 spins for hours on PyTorch's baseline (non-AVX) linear algebra kernel with
|
||||
# -march=native. This CPU runs the AVX2 variant, so the baseline one is built with -O1.
|
||||
*LinearAlgebraKernel.cpp.DEFAULT.cpp*) set -- "\$@" -O1 ;;
|
||||
esac
|
||||
exec $(command -v $c) --no-default-config --gcc-install-dir=${gccdir%/} -stdlib=libstdc++ -rtlib=libgcc -unwindlib=libgcc -fuse-ld=lld -Wno-unused-command-line-argument "\$@"
|
||||
EOF
|
||||
chmod +x build/bin/$c
|
||||
done
|
||||
printf '#include <iostream>\nint main() { std::cout << "ok" << std::endl; }\n' > build/bin/probe.cpp
|
||||
build/bin/clang++ -march=native -O2 build/bin/probe.cpp -o build/bin/probe
|
||||
build/bin/probe
|
||||
ldd build/bin/probe | grep -E "c\+\+|gcc|unwind"
|
||||
}
|
||||
|
||||
sdk() {
|
||||
[ -x .venv/bin/python ] || $PYTHON -m venv .venv
|
||||
.venv/bin/python -m pip install -U pip wheel 'setuptools<80' 'cmake<4' ninja
|
||||
.venv/bin/python -m pip install --index-url "$ROCM_INDEX" \
|
||||
"rocm[libraries,devel]==$ROCM_VERSION" "rocm-sdk-device-$PRAIN_GPU_ARCH==$ROCM_VERSION"
|
||||
.venv/bin/rocm-sdk init
|
||||
.venv/bin/rocm-sdk version
|
||||
# libstdc++ 16's <format> has [[__gnu__::__noinline__]], and HIP's host_defines.h empties
|
||||
# __noinline__ for host compilers: parse <format> first, like Gentoo's hip ebuild does
|
||||
# (https://github.com/ROCm/rocm-systems/issues/1874).
|
||||
python - "$(echo .venv/lib/python3*/site-packages)/_rocm_sdk_devel/include/hip/amd_detail/host_defines.h" <<'EOF'
|
||||
import sys
|
||||
path = sys.argv[1]
|
||||
text = open(path).read()
|
||||
fix = ("#if defined(__cplusplus) && defined(__has_include)\n# if __has_include(<__config>)\n"
|
||||
"# include <__config>\n# else\n# include <format>\n# endif\n#endif\n")
|
||||
if fix not in text:
|
||||
assert text.count("\n#define __noinline__\n") == 1, path
|
||||
open(path, "w").write(text.replace("\n#define __noinline__\n", "\n" + fix + "#define __noinline__\n"))
|
||||
print("host_defines.h: <format> before __noinline__")
|
||||
EOF
|
||||
}
|
||||
|
||||
fetch() {
|
||||
[ -d build/pytorch/.git ] || git clone --depth 1 --branch "$TORCH_REF" --recurse-submodules \
|
||||
--shallow-submodules -j 8 https://github.com/pytorch/pytorch build/pytorch
|
||||
[ -d build/vllm/.git ] || git clone --depth 1 --branch "$VLLM_REF" \
|
||||
https://github.com/vllm-project/vllm build/vllm
|
||||
}
|
||||
|
||||
# PyTorch preloads these SDK libraries on import (torch/_rocm_init.py), so no LD_LIBRARY_PATH.
|
||||
TORCH_PRELOADS='"amd_comgr", "amdhip64", "hiprtc", "roctx64", "hipblas", "hipblaslt", "hipfft",
|
||||
"hiprand", "hipsparse", "hipsparselt", "hipsolver", "miopen"'
|
||||
|
||||
torch() {
|
||||
local src=build/pytorch
|
||||
python -m pip install -r $src/requirements-build.txt filelock 'fsspec>=0.8.5' jinja2 \
|
||||
'networkx>=2.5.1' 'optree>=0.13.0' psutil 'sympy>=1.13.3'
|
||||
(cd $src && python tools/amd_build/build_amd.py > ../logs/hipify.log)
|
||||
printf 'def initialize():\n import rocm_sdk\n rocm_sdk.initialize_process(\n preload_shortnames=[%s],\n check_version="%s")\n' \
|
||||
"$TORCH_PRELOADS" "$ROCM_VERSION" > $src/torch/_rocm_init.py
|
||||
# Only what prain needs on one RDNA2 GPU: no CK/aotriton attention (not built for gfx103x),
|
||||
# no NCCL (vLLM falls back to gloo on one GPU), no profiler, no mobile/quantized CPU backends.
|
||||
(cd $src && env USE_ROCM=1 USE_CUDA=0 USE_XPU=0 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
|
||||
USE_ROCM_CK_GEMM=0 USE_ROCM_CK_SDPA=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 \
|
||||
USE_MSLK=0 USE_NCCL=0 USE_MPI=0 USE_KINETO=0 USE_FBGEMM=0 USE_NNPACK=0 \
|
||||
USE_PYTORCH_QNNPACK=0 USE_XNNPACK=0 USE_ITT=0 USE_MAGMA=0 BUILD_TEST=0 BLAS=OpenBLAS \
|
||||
MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v \
|
||||
-w "$PRAIN_ROOT/build/wheels" .)
|
||||
python -m pip install --no-deps --force-reinstall build/wheels/torch-*.whl
|
||||
env -u LD_LIBRARY_PATH python -c 'import torch; x = torch.ones(4, device="cuda"); print(torch.__version__, torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName, (x * 2).sum().item(), torch.backends.openmp.is_available(), torch.get_num_threads())'
|
||||
}
|
||||
|
||||
vllm() {
|
||||
local src=build/vllm site
|
||||
site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
|
||||
# vLLM lists gfx1030 but not its siblings gfx1031/gfx1032: add ours.
|
||||
grep -q "$PRAIN_GPU_ARCH" $src/CMakeLists.txt ||
|
||||
sed -i "s/;gfx1030;/;gfx1030;$PRAIN_GPU_ARCH;/" $src/CMakeLists.txt
|
||||
# Keep our torch and AMD's Triton for this ROCm release: pip fails rather than swap them.
|
||||
printf 'torch==%s\ntriton==3.8.0+git669b31ac.rocm%s\n' \
|
||||
"$(python -c 'import torch; print(torch.__version__)')" "$ROCM_VERSION" > build/constraints.txt
|
||||
python -m pip install -c build/constraints.txt --extra-index-url "$ROCM_INDEX" triton \
|
||||
-r $src/requirements/common.txt 'setuptools-scm>=9.2.0' 'setuptools-rust>=1.9.0' 'jinja2>=3.1.6'
|
||||
# The SDK's own amdsmi package (vLLM's ROCm platform imports it), not PyPI's.
|
||||
echo "$site/_rocm_sdk_core/share/amd_smi" > "$site/rocm-sdk-amdsmi.pth"
|
||||
(cd $src && env VLLM_TARGET_DEVICE=rocm PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
|
||||
MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip install --no-build-isolation --no-deps -v -e .)
|
||||
env -u LD_LIBRARY_PATH python -c 'import vllm, triton, amdsmi; print(vllm.__version__, triton.__version__)'
|
||||
}
|
||||
|
||||
prain() {
|
||||
python -m pip install --no-deps -e .
|
||||
python -m unittest discover -s tests -t .
|
||||
}
|
||||
|
||||
for step in "$@"; do
|
||||
. scripts/env.sh
|
||||
mkdir -p build/logs
|
||||
echo "== $step $(date -Is)"
|
||||
"$step" 2>&1 | tee "build/logs/$step.log"
|
||||
done
|
||||
@@ -0,0 +1,26 @@
|
||||
# prain environment: the project venv with AMD's ROCm SDK wheels (TheRock) and a clang host
|
||||
# toolchain tuned for this CPU. Source it from bash: . scripts/env.sh
|
||||
# Overrides: PRAIN_GPU_ARCH (default gfx1032), PRAIN_JOBS (parallel build jobs, default 16).
|
||||
PRAIN_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
|
||||
export PRAIN_ROOT
|
||||
export PRAIN_GPU_ARCH=${PRAIN_GPU_ARCH:-gfx1032}
|
||||
export PRAIN_JOBS=${PRAIN_JOBS:-16}
|
||||
export VIRTUAL_ENV=$PRAIN_ROOT/.venv
|
||||
case ":$PATH:" in *":$VIRTUAL_ENV/bin:"*) ;; *) PATH=$VIRTUAL_ENV/bin:$PATH ;; esac
|
||||
_site=$(echo "$VIRTUAL_ENV"/lib/python3*/site-packages)
|
||||
if [ -d "$_site/_rocm_sdk_devel" ]; then
|
||||
_dev=$_site/_rocm_sdk_devel
|
||||
_core=$_site/_rocm_sdk_core
|
||||
export ROCM_PATH=$_dev ROCM_HOME=$_dev HIP_PATH=$_dev CMAKE_PREFIX_PATH=$_dev
|
||||
export HIP_CLANG_PATH=$_dev/lib/llvm/bin HIP_DEVICE_LIB_PATH=$_core/lib/llvm/amdgcn/bitcode
|
||||
case ":${LD_LIBRARY_PATH:-}:" in *":$_dev/lib:"*) ;; *)
|
||||
export LD_LIBRARY_PATH=$_dev/lib:$_dev/lib/rocm_sysdeps/lib:$_core/lib:$_core/lib/rocm_sysdeps/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ;;
|
||||
esac
|
||||
case ":$PATH:" in *":$_dev/bin:"*) ;; *) PATH=$_dev/bin:$PATH ;; esac
|
||||
fi
|
||||
export PATH
|
||||
# Host C/C++: the system clang through wrappers (scripts/build.sh wrappers) that drop Gentoo's
|
||||
# libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries.
|
||||
export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++
|
||||
export CFLAGS="-march=native" CXXFLAGS="-march=native"
|
||||
unset _site _dev _core
|
||||
Reference in New Issue
Block a user