diff --git a/scripts/build.sh b/scripts/build.sh index 8845f3b..90539e2 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -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"])')