From 7fa0192bb0ce16d7120511e2450031fff6d912e8 Mon Sep 17 00:00:00 2001 From: Eda Z Date: Tue, 31 Mar 2026 08:37:33 +0000 Subject: [PATCH] Add ROCm detection to install.sh and expand shell tests Add AMD ROCm GPU detection to get_torch_index_url() in install.sh. When nvidia-smi is not found, probe for ROCm via amd-smi, /opt/rocm version file, hipconfig, dpkg-query, and rpm. Includes validation guard for malformed _rocm_tag, Debian epoch prefix stripping, ROCm 7.2+ cap to rocm7.1 index, bitsandbytes AMD install, and status messaging. Shell tests expanded to 23 cases. Co-authored-by: Daniel Han --- install.sh | 54 +++++++++++- tests/sh/test_get_torch_index_url.sh | 119 ++++++++++++++++++++++++++- 2 files changed, 169 insertions(+), 4 deletions(-) diff --git a/install.sh b/install.sh index 9ea80bc161..b8f49abf51 100755 --- a/install.sh +++ b/install.sh @@ -982,7 +982,44 @@ get_torch_index_url() { elif [ -x "/usr/bin/nvidia-smi" ]; then _smi="/usr/bin/nvidia-smi" fi - if [ -z "$_smi" ]; then echo "$_base/cpu"; return; fi + if [ -z "$_smi" ]; then + # No NVIDIA GPU -- check for AMD ROCm + _rocm_tag="" + _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ + amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[^0-9.]/, "", $2); split($2,a,"."); print "rocm"a[1]"."a[2]; ok=1; exit} END{exit !ok}'; } || \ + { [ -r /opt/rocm/.info/version ] && \ + awk -F. '{print "rocm"$1"."$2; exit}' /opt/rocm/.info/version; } || \ + { command -v hipconfig >/dev/null 2>&1 && \ + hipconfig --version 2>/dev/null | awk 'NR==1{split($1,a,"."); if(a[1]+0>0) print "rocm"a[1]"."a[2]}'; } || \ + { command -v dpkg-query >/dev/null 2>&1 && \ + ver="$(dpkg-query -W -f='${Version}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | sed 's/^[0-9]*://' | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; } || \ + { command -v rpm >/dev/null 2>&1 && \ + ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null + # Validate _rocm_tag: must match "rocmX.Y" with leading digits + case "$_rocm_tag" in + rocm[0-9]*.[0-9]*) : ;; # valid + *) _rocm_tag="" ;; # reject malformed (empty version, garbled output) + esac + if [ -n "$_rocm_tag" ]; then + # ROCm 7.2 only has torch 2.11.0 which exceeds current bounds (<2.11.0). + # Fall back to rocm7.1 index which has torch 2.10.0. + # TODO: uncomment the next line when torch upper bound is bumped to >=2.11.0 + # echo "$_base/$_rocm_tag"; return + case "$_rocm_tag" in + rocm7.2*|rocm7.3*|rocm7.4*|rocm7.5*|rocm8*|rocm9*) + echo "$_base/rocm7.1" ;; + *) + echo "$_base/$_rocm_tag" ;; + esac + return + fi + echo "$_base/cpu"; return + fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -1007,13 +1044,19 @@ case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then echo "" - echo " NOTE: No NVIDIA GPU detected (nvidia-smi not found)." + echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)." echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," echo " re-run with --no-torch for a faster, lighter install:" echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" + echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" echo "" fi ;; + */rocm*) + echo "" + echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)" + echo "" + ;; esac # ── Install unsloth directly into the venv (no activation needed) ── @@ -1051,6 +1094,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" + # AMD ROCm: install bitsandbytes with AMD support + case "$TORCH_INDEX_URL" in + */rocm*) + substep "installing bitsandbytes for AMD ROCm..." + run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" + ;; + esac fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch substep "installing unsloth (this may take a few minutes)..." diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 6387922712..81da79aa32 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -45,10 +45,23 @@ MOCK echo "$_dir" } +# Helper: create a mock amd-smi that prints a given ROCm version string +make_mock_amd_smi() { + _dir=$(mktemp -d) + cat > "$_dir/amd-smi" </dev/null || true) [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd" done @@ -119,6 +132,108 @@ _result=$(run_func "$_dir") assert_eq "unparseable -> cu126" "https://download.pytorch.org/whl/cu126" "$_result" rm -rf "$_dir" +# 9) ROCm 6.3 (no nvidia-smi) -> rocm6.3 +_dir=$(make_mock_amd_smi "6.3") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.3 -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result" +rm -rf "$_dir" + +# 10) ROCm 7.1 (no nvidia-smi) -> rocm7.1 +_dir=$(make_mock_amd_smi "7.1") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.1 (capped due to torch <2.11.0) +_dir=$(make_mock_amd_smi "7.2") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.2 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence +_cuda_dir=$(make_mock_smi "12.6") +_amd_dir=$(make_mock_amd_smi "6.3") +_combined_dir=$(mktemp -d) +ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi" +ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi" +_result=$(run_func "$_combined_dir") +assert_eq "CUDA+ROCm -> CUDA precedence" "https://download.pytorch.org/whl/cu126" "$_result" +rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir" + +# 13) No nvidia-smi, no amd-smi -> cpu (duplicate of test 1, confirms ROCm didn't break it) +_result=$(run_func "none") +assert_eq "no GPU -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" + +# 14) ROCm 6.1 (no nvidia-smi) -> rocm6.1 +_dir=$(make_mock_amd_smi "6.1") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.1 -> rocm6.1" "https://download.pytorch.org/whl/rocm6.1" "$_result" +rm -rf "$_dir" + +# 15) ROCm 6.4 (no nvidia-smi) -> rocm6.4 +_dir=$(make_mock_amd_smi "6.4") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.4 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result" +rm -rf "$_dir" + +# 16) ROCm 7.0 (no nvidia-smi) -> rocm7.0 +_dir=$(make_mock_amd_smi "7.0") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result" +rm -rf "$_dir" + +# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.1 (capped) +_dir=$(make_mock_amd_smi "8.0") +_result=$(run_func "$_dir") +assert_eq "ROCm 8.0 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 18) Malformed amd-smi output (empty version field) -> cpu +_dir=$(mktemp -d) +cat > "$_dir/amd-smi" <<'MOCK' +#!/bin/sh +echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: " +MOCK +chmod +x "$_dir/amd-smi" +_result=$(run_func "$_dir") +assert_eq "empty amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 19) amd-smi with "N/A" version -> cpu +_dir=$(mktemp -d) +cat > "$_dir/amd-smi" <<'MOCK' +#!/bin/sh +echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: N/A" +MOCK +chmod +x "$_dir/amd-smi" +_result=$(run_func "$_dir") +assert_eq "N/A amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 20) ROCm version with trailing text (e.g. "6.3.1-beta") -> rocm6.3 +_dir=$(make_mock_amd_smi "6.3.1-beta") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.3.1-beta -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result" +rm -rf "$_dir" + +# 22) CUDA 12.6 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "12.6") +_result=$(run_func "$_dir") +assert_eq "CUDA 12.6 regression -> cu126" "https://download.pytorch.org/whl/cu126" "$_result" +rm -rf "$_dir" + +# 23) CUDA 13.0 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "13.0") +_result=$(run_func "$_dir") +assert_eq "CUDA 13.0 regression -> cu130" "https://download.pytorch.org/whl/cu130" "$_result" +rm -rf "$_dir" + +# 24) CUDA 12.8 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir") +assert_eq "CUDA 12.8 regression -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR"