fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: build torchvision too (transformers Gemma 4 processor imports it)

Browse Source
This commit is contained in:
fabolous005
2026-10-27 16:04:22 +01:00
parent 341ad52baa
commit 7be39c9287
1 changed files with 13 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/build.sh
+13 -1
View File
@@ -4,14 +4,16 @@
# 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
# torchvision torchvision 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
# usage: scripts/build.sh STEP... Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VISION_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}
VISION_REF=${PRAIN_VISION_REF:-v0.28.0}
VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
PYTHON=${PRAIN_PYTHON:-python3.13}
@@ -90,6 +92,16 @@ torch() {
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())'
}
# transformers' Gemma 4 processor imports torchvision, even for text-only serving.
torchvision() {
[ -d build/vision/.git ] || git clone --depth 1 --branch "$VISION_REF" \
https://github.com/pytorch/vision build/vision
(cd build/vision && env FORCE_CUDA=1 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" 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/torchvision-*.whl
env -u LD_LIBRARY_PATH python -c 'import torch, torchvision; print(torchvision.__version__, torch.ops.torchvision.nms(torch.rand(4, 4, device="cuda"), torch.rand(4, device="cuda"), 0.5).device)'
}
vllm() {
local src=build/vllm site
site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')