diff --git a/install.sh b/install.sh index e295c70a2b..8a0326b49d 100755 --- a/install.sh +++ b/install.sh @@ -1132,9 +1132,11 @@ _pick_radeon_wheel() { base = $NF sub(/[?#].*/, "", base) # strip query / fragment prefix = pkg "-" - suffix = "-" tag "-" tag "-linux_x86_64.whl" + # Match cpXY-cpXY or cpXY-abi3 with any linux x86_64 platform tag + # (linux_x86_64, manylinux_2_28_x86_64, manylinux2014_x86_64, etc.) if (substr(base, 1, length(prefix)) == prefix && - substr(base, length(base) - length(suffix) + 1) == suffix) + index(base, "-" tag "-") > 0 && + match(base, /x86_64\.whl$/)) print $0 }' \ | sort -V \ @@ -1247,22 +1249,26 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) && _tv_arg="$_tv_whl" _ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) && _ta_arg="$_ta_whl" _tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) && _tri_arg="$_tri_whl" + # Build install args; skip empty _tri_arg to avoid passing "" to uv + _radeon_pkgs="$_torch_arg $_tv_arg $_ta_arg" + [ -n "$_tri_arg" ] && _radeon_pkgs="$_tri_arg $_radeon_pkgs" run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ - "$_tri_arg" "$_torch_arg" "$_tv_arg" "$_ta_arg" + $_radeon_pkgs substep "installing bitsandbytes for AMD Radeon..." run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" \ "bitsandbytes>=0.49.1" else - substep "[WARN] Radeon repo unavailable; falling back to CPU-only PyTorch" "$C_WARN" + substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ - "torch>=2.4,<2.11.0" "torchvision<0.26.0" "torchaudio<2.11.0" \ - --index-url "${TORCH_INDEX_URL%/*}/cpu" + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" fi else - substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to CPU-only PyTorch" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" "torchvision<0.26.0" "torchaudio<2.11.0" \ - --index-url "${TORCH_INDEX_URL%/*}/cpu" + substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." @@ -1277,7 +1283,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then esac fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch - substep "installing unsloth (this may take a few minutes)..." + substep "installing unsloth (this may take a few minutes)..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. @@ -1448,4 +1454,4 @@ else substep "source ${VENV_DIR}/bin/activate" substep "unsloth studio -H 0.0.0.0 -p 8888" echo "" -fi \ No newline at end of file +fi diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index b9b61cdcfe..17d731d7b2 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -5,11 +5,9 @@ Hardware detection and GPU utilities """ +from . import hardware as _hardware from .hardware import ( DeviceType, - DEVICE, - CHAT_ONLY, - IS_ROCM, detect_hardware, get_device, is_apple_silicon, @@ -83,3 +81,11 @@ __all__ = [ "extract_arch_config", "estimate_training_vram", ] + + +def __getattr__(name: str): + """Resolve mutable module-level flags (DEVICE, CHAT_ONLY, IS_ROCM) at access + time so callers always see the current value after detect_hardware() runs.""" + if name in {"DEVICE", "CHAT_ONLY", "IS_ROCM"}: + return getattr(_hardware, name) + raise AttributeError(name) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index d71f1cd494..c3a0536a56 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -82,13 +82,19 @@ def _parse_memory_mb(value: Any) -> Optional[float]: if num is None: return None - # Explicit unit conversion - if "gib" in unit or "gb" in unit: + # Explicit unit conversion -- distinguish binary (GiB) from SI (GB) + if "gib" in unit: return num * 1024 - if "mib" in unit or "mb" in unit: + if "gb" in unit: + return num * 1000 + if "mib" in unit: return num - if "kib" in unit or "kb" in unit: + if "mb" in unit: + return num + if "kib" in unit: return num / 1024 + if "kb" in unit: + return num / 1000 if unit and ( "b" in unit and "g" not in unit and "m" not in unit and "k" not in unit ): diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 40364f765e..edeada0190 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1356,8 +1356,13 @@ def apply_gpu_ids(gpu_ids) -> None: value = str(gpu_ids) os.environ["CUDA_VISIBLE_DEVICES"] = value + # Keep ROCm visibility env vars in sync so _get_parent_visible_gpu_spec() + # picks up the narrowed set on AMD systems. + if IS_ROCM: + os.environ["HIP_VISIBLE_DEVICES"] = value + os.environ["ROCR_VISIBLE_DEVICES"] = value _visible_gpu_count = None - logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) + logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm=%s)", value, IS_ROCM) def get_device_map( diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d7f61e5580..ef3e511319 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -59,6 +59,25 @@ def _detect_rocm_version() -> tuple[int, int] | None: except Exception: pass + # Try amd-smi version (outputs "... | ROCm version: X.Y.Z") + amd_smi = shutil.which("amd-smi") + if amd_smi: + try: + result = subprocess.run( + [amd_smi, "version"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 5, + ) + if result.returncode == 0: + import re + m = re.search(r"ROCm version:\s*(\d+)\.(\d+)", result.stdout) + if m: + return int(m.group(1)), int(m.group(2)) + except Exception: + pass + # Try hipconfig --version (outputs bare version like "6.3.21234.2") hipconfig = shutil.which("hipconfig") if hipconfig: