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 <danielhanchen@gmail.com>
This commit is contained in:
Eda Z
Daniel Han
parent
66f250a614
commit
7fa0192bb0
2 changed files with 169 additions and 4 deletions
54
install.sh
54
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)..."
|
||||
|
|
|
|||
|
|
@ -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" <<MOCK
|
||||
#!/bin/sh
|
||||
cat <<AMD_OUT
|
||||
AMDSMI Tool: 25.0.1+2b74356 | AMDSMI Library version: 25.0.1.0 | ROCm version: $1
|
||||
AMD_OUT
|
||||
MOCK
|
||||
chmod +x "$_dir/amd-smi"
|
||||
echo "$_dir"
|
||||
}
|
||||
|
||||
# Build a minimal tools directory with symlinks to essential commands
|
||||
# (uname, grep, head, etc.) but WITHOUT nvidia-smi.
|
||||
# (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi.
|
||||
_TOOLS_DIR=$(mktemp -d)
|
||||
for _cmd in uname grep sed head sh bash cat; do
|
||||
for _cmd in uname grep sed head sh bash cat awk printf; do
|
||||
_real=$(command -v "$_cmd" 2>/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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue