122 lines
6.5 KiB
Bash
Executable File
122 lines
6.5 KiB
Bash
Executable File
#!/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
|