ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(studio/rocm): gate ROCm-only side-effects on active torch runtime

Browse source 
Address five edge cases flagged during PR review:

1. studio/backend/main.py: BNB_ROCM_VERSION was set whenever HIP_PATH or
   ROCM_PATH was present in the environment. A Windows CUDA user who once
   installed the HIP SDK and reverted to a CUDA torch wheel still has those
   env vars set, so bitsandbytes would try to load libbitsandbytes_rocm72.dll
   against a CUDA torch and crash. Now probe torch.version.hip inside the
   env-var guard (worker.py already does this).

2. studio/backend/main.py: os.add_dll_directory returned handles were
   discarded. Per CPython docs, the directory leaves the DLL search list when
   the handle is garbage collected. Retain handles in module-level
   _ROCM_DLL_HANDLES list so they survive process lifetime.

3. studio/install_python_stack.py: _install_bnb_windows_rocm() returned None
   regardless of pip_install_try outcome, and the caller flipped
   _rocm_windows_torch_installed to True unconditionally. On a failed BNB
   install the post-install "manual install may be required" warning was
   suppressed and the user was misled. Helper now returns bool; caller gates
   on it.

4. studio/install_python_stack.py: _detect_windows_gfx_arch returned the raw
   capture group, so mixed-case hipinfo output ("Gfx1151") missed the
   lowercase keys in _GFX_TO_AMD_INDEX_ARCH and silently fell back to CPU
   torch. Lowercase the token.

5. studio/install_python_stack.py: UNSLOTH_ROCM_TORCH_INSTALLED=1 early-
   return trusted the env var even when the venv was wiped between runs.
   Subprocess-probe torch importability first; fall through to the full
   install path if the probe fails.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py
(adds one new test for case 5 fall-through).
This commit is contained in:
Daniel Han 2026-05-19 07:59:55 +00:00
commit 0c2020d5a1
3 changed files with 94 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

22
studio/backend/main.py
View file

@ -17,6 +17,10 @@ os.environ["PYTHONWARNINGS"] = "ignore"
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
if sys.platform == "win32":
# Retained at module scope -- os.add_dll_directory returns a handle that
# removes the search-path entry when garbage collected.
_ROCM_DLL_HANDLES: list = []
def _add_rocm_dll_dirs() -> None:
candidates = []
# 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer
@ -40,7 +44,7 @@ if sys.platform == "win32":
for _d in candidates:
if os.path.isdir(_d):
try:
os.add_dll_directory(_d)
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
except (OSError, AttributeError):
pass
@ -54,10 +58,18 @@ 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).
# 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"))
# Gate on the active torch runtime, not env-var presence -- HIP_PATH /
# ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
# CUDA torch wheel, and setting BNB_ROCM_VERSION there makes bitsandbytes
# look for a ROCm DLL that doesn't exist and crash the CUDA backend.
_is_rocm_host = False
if os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"):
try:
import torch as _torch_probe
_is_rocm_host = bool(getattr(_torch_probe.version, "hip", None))
del _torch_probe
except Exception:
pass
if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ:
import glob as _glob

62
studio/install_python_stack.py
View file

@ -259,7 +259,10 @@ def _detect_windows_gfx_arch() -> str | None:
return None
text = result.stdout.decode(errors = "replace")
m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
return m.group(1).strip() if m else None
# Lowercase the captured token -- some hipinfo builds emit "Gfx1151"
# which would miss the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and
# silently fall back to CPU torch.
return m.group(1).strip().lower() if m else None
except Exception:
return None
@ -384,8 +387,8 @@ def _detect_amd_gfx_codes() -> list[str]:
_rocm_windows_torch_installed: bool = False
def _install_bnb_windows_rocm() -> None:
"""Install the AMD Windows BNB prerelease wheel.
def _install_bnb_windows_rocm() -> bool:
"""Install the AMD Windows BNB prerelease wheel. Returns True on success.
The continuous-release wheel is intentionally mismatched: the filename
encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the
@ -395,11 +398,11 @@ def _install_bnb_windows_rocm() -> None:
"""
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is None:
return
return False
_prev = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK")
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1"
try:
pip_install_try(
_ok = pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
@ -412,6 +415,8 @@ def _install_bnb_windows_rocm() -> None:
os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None)
else:
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _prev
if not _ok:
return False
# After install: detect the actual ROCm DLL suffix from the wheel so any
# post-install BNB import in this process loads the correct DLL.
# The worker subprocess does the same detection independently (worker.py §1f).
@ -419,6 +424,7 @@ def _install_bnb_windows_rocm() -> None:
if "BNB_ROCM_VERSION" not in os.environ:
_ver = _detect_bnb_rocm_dll_ver() or "72"
os.environ["BNB_ROCM_VERSION"] = _ver
return True
def _ensure_rocm_torch() -> None:
@ -431,14 +437,38 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# setup.ps1 sets this when it already installed AMD wheels; skip the probe.
# setup.ps1 sets this when it already installed AMD wheels; skip the probe
# only when torch is actually importable as ROCm. If the venv was wiped
# between runs, the stale env-var would suppress a needed reinstall.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_rocm_windows_torch_installed = True
# setup.ps1 already installed ROCm torch, but we still need to install
# the AMD Windows BNB wheel here — the PyPI bitsandbytes wheel ships
# only CUDA DLLs and will fail to load on ROCm (no libbitsandbytes_rocm72.dll).
_install_bnb_windows_rocm()
return
_torch_ok = False
try:
_probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"import sys; "
"sys.exit(0 if (hip or 'rocm' in torch.__version__.lower()) else 1)"
),
],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
timeout = 30,
)
_torch_ok = _probe.returncode == 0
except (OSError, subprocess.TimeoutExpired):
pass
if _torch_ok:
_rocm_windows_torch_installed = True
# setup.ps1 already installed ROCm torch, but we still need to install
# the AMD Windows BNB wheel here -- the PyPI bitsandbytes wheel ships
# only CUDA DLLs and will fail to load on ROCm.
_install_bnb_windows_rocm()
return
# torch was wiped between runs; fall through to the full install path
if IS_MACOS:
return
@ -488,11 +518,13 @@ def _ensure_rocm_torch() -> None:
"torchaudio",
constrain = False,
)
# Always install AMD Windows bitsandbytes — the PyPI wheel ships only
# Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
# CUDA DLLs and will fail to load on ROCm. Install even when torch was
# already a ROCm build so that `studio update` repairs a broken bnb.
_install_bnb_windows_rocm()
_rocm_windows_torch_installed = True
# Only flip the success flag when the install actually succeeds; otherwise
# the post-install "manual install may be required" warning is suppressed.
if _install_bnb_windows_rocm():
_rocm_windows_torch_installed = True
return
# ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──

33
tests/studio/install/test_rocm_support.py
View file

@ -1788,11 +1788,19 @@ class TestDetectBnbRocmDllVer:
class TestRocmTorchInstalledEnvVar:
"""Verify UNSLOTH_ROCM_TORCH_INSTALLED=1 skips main install but still installs BNB."""
@staticmethod
def _ok_torch_probe(*a, **kw):
# subprocess.run probe returns 0 when torch imports as ROCm
rv = MagicMock()
rv.returncode = 0
return rv
@patch.object(stack_mod, "_install_bnb_windows_rocm")
@patch.object(stack_mod, "pip_install")
def test_env_var_skips_main_pip_install(self, mock_pip, mock_bnb):
"""UNSLOTH_ROCM_TORCH_INSTALLED=1 should not trigger torch pip_install."""
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}):
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \
patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe):
stack_mod._ensure_rocm_torch()
mock_pip.assert_not_called()
@ -1800,7 +1808,8 @@ class TestRocmTorchInstalledEnvVar:
@patch.object(stack_mod, "pip_install")
def test_env_var_calls_bnb_install(self, mock_pip, mock_bnb):
"""UNSLOTH_ROCM_TORCH_INSTALLED=1 should still call _install_bnb_windows_rocm."""
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}):
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \
patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe):
stack_mod._ensure_rocm_torch()
mock_bnb.assert_called_once()
@ -1809,10 +1818,28 @@ class TestRocmTorchInstalledEnvVar:
def test_env_var_sets_rocm_windows_flag(self, mock_pip, mock_bnb):
"""UNSLOTH_ROCM_TORCH_INSTALLED=1 should set _rocm_windows_torch_installed."""
stack_mod._rocm_windows_torch_installed = False
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}):
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \
patch.object(stack_mod.subprocess, "run", side_effect=self._ok_torch_probe):
stack_mod._ensure_rocm_torch()
assert stack_mod._rocm_windows_torch_installed is True
@patch.object(stack_mod, "_install_bnb_windows_rocm")
@patch.object(stack_mod, "pip_install")
def test_env_var_falls_through_when_torch_missing(self, mock_pip, mock_bnb):
"""If the venv was wiped between runs, the stale env-var must not suppress reinstall."""
stack_mod._rocm_windows_torch_installed = False
def _bad_probe(*a, **kw):
rv = MagicMock()
rv.returncode = 1
return rv
with patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), \
patch.object(stack_mod.subprocess, "run", side_effect=_bad_probe), \
patch.object(stack_mod, "IS_WINDOWS", False), \
patch.object(stack_mod, "IS_MACOS", True):
stack_mod._ensure_rocm_torch()
# macOS branch is the next exit -- but the point is the early-return did NOT fire.
mock_bnb.assert_not_called()
# =============================================================================
# TEST: worker.py -- Windows ROCm patches (source-level checks)