diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 4e9ac4ff83..75f9d56fa4 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1248,7 +1248,13 @@ def run_training_process( sys.meta_path.append(_StubSubpackageFinder()) - if sys.platform == "win32": + # Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao + # is real and shadowing it breaks torchao-based quantization paths. + # HIP_PATH / ROCM_PATH are set by the AMD HIP SDK installer on ROCm machines. + _is_win32_rocm = sys.platform == "win32" and bool( + os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH") + ) + if _is_win32_rocm: # Seed torchao top-level + key submodules; the finder handles the rest. for _tao_name in ( "torchao", diff --git a/studio/backend/main.py b/studio/backend/main.py index 85e79069ed..b61e192cfc 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -54,7 +54,11 @@ if sys.platform == "win32": # this the server process crashes with "Configured ROCm binary not found". # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION # before any import that pulls in bitsandbytes (mirrors worker.py logic). - if "BNB_ROCM_VERSION" not in os.environ: + # Guard: only set on ROCm hosts (HIP_PATH/ROCM_PATH present) -- setting + # BNB_ROCM_VERSION on a Windows CUDA machine makes bitsandbytes look for a + # ROCm DLL that doesn't exist and fail to initialise the CUDA backend. + _is_rocm_host = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")) + if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ: import glob as _glob _bnb_rocm_ver = None diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index bd6589dd20..6929e5eab7 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -228,10 +228,24 @@ def _detect_rocm_version() -> tuple[int, int] | None: def _detect_windows_gfx_arch() -> str | None: - """Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None.""" + """Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None. + + Resolves hipinfo via PATH first, then HIP_PATH\\bin and ROCM_PATH\\bin as + fallbacks -- the AMD HIP SDK installer sets these env vars but does not + always add the bin dir to the system PATH. + """ import re hipinfo = shutil.which("hipinfo") + if not hipinfo: + # Fallback: AMD HIP SDK sets HIP_PATH / ROCM_PATH even when bin isn't on PATH + for _env_var in ("HIP_PATH", "ROCM_PATH"): + _root = os.environ.get(_env_var) + if _root: + _candidate = os.path.join(_root, "bin", "hipinfo.exe") + if os.path.isfile(_candidate): + hipinfo = _candidate + break if not hipinfo: return None try: