diff --git a/install.sh b/install.sh index 9df0bb7347..56f6e8e9de 100755 --- a/install.sh +++ b/install.sh @@ -1110,7 +1110,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN" else 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 \ + 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" # AMD ROCm: install bitsandbytes with AMD support case "$TORCH_INDEX_URL" in diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ff3ef92842..40364f765e 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1311,17 +1311,20 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if cuda_visible is not None: - # "" means zero GPUs, "0" means 1, "0,1,2" means 3 - cuda_visible = cuda_visible.strip() - if cuda_visible == "" or cuda_visible == "-1": + # Use _get_parent_visible_gpu_spec() which already handles + # HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm. + visible_spec = _get_parent_visible_gpu_spec() + if visible_spec["raw"] is not None: + raw = visible_spec["raw"].strip() + if raw == "" or raw == "-1": _visible_gpu_count = 0 + elif visible_spec["numeric_ids"] is not None: + _visible_gpu_count = len(visible_spec["numeric_ids"]) else: - _visible_gpu_count = len([x for x in cuda_visible.split(",") if x.strip()]) + _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # CUDA_VISIBLE_DEVICES not set -- try torch, fall back to physical count + # No visibility env var set -- try torch, fall back to physical count try: import torch diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index da2f47a04a..423d8f6fd3 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1466,13 +1466,8 @@ def detect_host() -> HostInfo: if _marker in _result.stdout.lower(): has_rocm = True break - # Fallback: HIP runtime DLL indicates a working HIP installation - if not has_rocm and any( - Path(d).joinpath("amdhip64.dll").exists() - for d in os.environ.get("PATH", "").split(os.pathsep) - if d - ): - has_rocm = True + # Note: amdhip64.dll presence alone is NOT treated as GPU evidence + # since the HIP SDK can be installed without an AMD GPU. return HostInfo( system = system, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9f3c3ddd04..d7f61e5580 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -187,8 +187,8 @@ def _ensure_rocm_torch() -> None: "--force-reinstall", "--no-cache-dir", "torch>=2.4,<2.11.0", - "torchvision", - "torchaudio", + "torchvision<0.26.0", + "torchaudio<2.11.0", "--index-url", index_url, constrain = False,