diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index da7eee86d2..8d154d19db 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1321,19 +1321,45 @@ def run_training_process( os.environ["TORCHDYNAMO_DISABLE"] = "1" logger.info("Windows ROCm: torch.compile (dynamo) disabled") - # Force BNB to load libbitsandbytes_rocm72.dll regardless of the - # HIP version that torch reports. As of torch==2.11.0+rocm7.13.0 - # (AMD index, May 2026) torch.version.hip returns "7.13", which - # makes BNB look for rocm713.dll — a file our AMD Windows prerelease - # wheel does not ship. The wheel only ships rocm72.dll, so we pin - # BNB_ROCM_VERSION="72" here. Callers may override by setting the - # variable before launching the worker. + # BNB auto-detects the HIP version from torch.version.hip and uses + # it to choose which DLL to load (e.g. "7.13" → rocm713.dll). + # AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its + # version suffix does not always match the torch HIP version (e.g. + # torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still + # ships rocm72.dll). We detect the actual DLL name from the installed + # package and override BNB's auto-detection. "72" is a safe fallback + # if detection fails. Callers may override by pre-setting the var. if "BNB_ROCM_VERSION" not in os.environ: - os.environ["BNB_ROCM_VERSION"] = "72" + _bnb_rocm_ver = None + try: + import glob as _glob + import importlib.util as _ilu + import re as _re + + _bnb_spec = _ilu.find_spec("bitsandbytes") + if _bnb_spec and _bnb_spec.submodule_search_locations: + for _pkg_dir in _bnb_spec.submodule_search_locations: + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): + _m = _re.search( + r"libbitsandbytes_rocm(\d+)\.dll", + os.path.basename(_dll), + ) + if _m: + _bnb_rocm_ver = _m.group(1) + break + if _bnb_rocm_ver: + break + except Exception: + pass + _bnb_rocm_ver = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver logger.info( - "Windows ROCm: set BNB_ROCM_VERSION=72 " - "(AMD Windows BNB wheel ships rocm72.dll; " - "overrides auto-detection from torch.version.hip)" + "Windows ROCm: set BNB_ROCM_VERSION=%s " + "(detected from installed BNB wheel; " + "overrides torch.version.hip auto-detection)", + _bnb_rocm_ver, ) # Patch _grouped_mm CUDA dispatch with a safe Python mm fallback. diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3df1756fbf..bd6589dd20 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -106,12 +106,11 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "download/continuous-release_main/" "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" ), - # Windows ROCm wheel — ships libbitsandbytes_rocm72.dll only. - # As of torch==2.11.0+rocm7.13.0 (AMD index, May 2026), BNB auto-detects - # HIP version as "7.13" and looks for rocm713.dll — which does not exist. - # BNB_ROCM_VERSION=72 must be set in the environment before importing bnb - # to force it to load rocm72.dll. Set in worker.py (training subprocess) - # and in _install_bnb_windows_rocm() (install subprocess). + # Windows ROCm wheel — ships libbitsandbytes_rocm{VER}.dll. + # BNB auto-detects HIP version from torch.version.hip, which does not always + # match the DLL suffix in this prerelease wheel (e.g. torch 7.13 with a rocm72 + # DLL). We scan the installed wheel for the actual DLL name and set + # BNB_ROCM_VERSION accordingly in _install_bnb_windows_rocm() and worker.py. "win_amd64": ( "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" "download/continuous-release_main/" @@ -259,6 +258,29 @@ def _windows_rocm_index_url(gfx_arch: str | None) -> str | None: return f"{_ROCM_WINDOWS_INDEX_BASE}/{arch_family}/" +def _detect_bnb_rocm_dll_ver() -> str | None: + """Scan the installed bitsandbytes package for libbitsandbytes_rocm{VER}.dll. + + Returns the version suffix string (e.g. ``"72"``, ``"713"``) or ``None`` + if bitsandbytes is not installed or no ROCm DLL is found. Does NOT import + bitsandbytes — uses importlib.util.find_spec so it is safe to call before + BNB is imported. + """ + import glob + import importlib.util + import re + + spec = importlib.util.find_spec("bitsandbytes") + if spec is None or not spec.submodule_search_locations: + return None + for pkg_dir in spec.submodule_search_locations: + for dll in glob.glob(os.path.join(pkg_dir, "libbitsandbytes_rocm*.dll")): + m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll)) + if m: + return m.group(1) + return None + + def _has_rocm_gpu() -> bool: """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" import re @@ -360,10 +382,6 @@ def _install_bnb_windows_rocm() -> None: _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") if _bnb_win_url is None: return - # Pin BNB_ROCM_VERSION=72 in this process now so that any post-install - # import of bitsandbytes (e.g. health-checks) loads the correct DLL. - # The worker subprocess sets this independently in worker.py section 1f. - os.environ.setdefault("BNB_ROCM_VERSION", "72") _prev = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1" try: @@ -380,6 +398,13 @@ def _install_bnb_windows_rocm() -> None: os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) else: os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _prev + # 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). + # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run). + if "BNB_ROCM_VERSION" not in os.environ: + _ver = _detect_bnb_rocm_dll_ver() or "72" + os.environ["BNB_ROCM_VERSION"] = _ver def _ensure_rocm_torch() -> None: diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 96744a0043..dae2c2ff67 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1674,18 +1674,37 @@ class TestInstallBnbWindowsRocm: stack_mod._install_bnb_windows_rocm() mock_pip.assert_not_called() - def test_sets_bnb_rocm_version_72(self): - """BNB_ROCM_VERSION must be set to '72' before install. - - As of torch==2.11.0+rocm7.13.0 (AMD index, May 2026), BNB auto-detects - HIP 7.13 and looks for rocm713.dll — which the prerelease wheel does not - ship. _install_bnb_windows_rocm() must pin BNB_ROCM_VERSION=72 so that - bitsandbytes loads libbitsandbytes_rocm72.dll instead. - """ + def test_sets_bnb_rocm_version_from_detected_dll(self): + """BNB_ROCM_VERSION is set from the DLL detected after install.""" with patch.dict(os.environ, {}, clear = False): os.environ.pop("BNB_ROCM_VERSION", None) with patch.object(stack_mod, "pip_install_try", return_value = True): - stack_mod._install_bnb_windows_rocm() + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72" + ): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "72" + + def test_sets_bnb_rocm_version_from_newer_dll(self): + """If AMD ships a newer DLL (e.g. rocm713.dll), that version is used.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713" + ): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "713" + + def test_falls_back_to_72_when_detection_fails(self): + """Falls back to '72' when DLL detection returns None.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None + ): + stack_mod._install_bnb_windows_rocm() assert os.environ.get("BNB_ROCM_VERSION") == "72" def test_does_not_override_existing_bnb_rocm_version(self): @@ -1696,6 +1715,47 @@ class TestInstallBnbWindowsRocm: assert os.environ.get("BNB_ROCM_VERSION") == "60" +class TestDetectBnbRocmDllVer: + """Unit tests for _detect_bnb_rocm_dll_ver().""" + + def test_returns_none_when_bnb_not_installed(self): + """Returns None if bitsandbytes is not importable.""" + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = None): + assert stack_mod._detect_bnb_rocm_dll_ver() is None + + def test_detects_rocm72_dll(self, tmp_path): + """Returns '72' when libbitsandbytes_rocm72.dll is present.""" + (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() == "72" + + def test_detects_rocm713_dll(self, tmp_path): + """Returns '713' when libbitsandbytes_rocm713.dll is present.""" + (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() == "713" + + def test_returns_none_when_only_cuda_dlls(self, tmp_path): + """Returns None when only CUDA DLLs are present (no ROCm DLL).""" + (tmp_path / "libbitsandbytes_cuda121.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() is None + + # ============================================================================= # TEST: install_python_stack.py -- UNSLOTH_ROCM_TORCH_INSTALLED early-return path # ============================================================================= @@ -1786,16 +1846,19 @@ class TestWorkerWindowsRocmPatches: assert "TORCHDYNAMO_DISABLE" in source def test_bnb_rocm_version_set_on_windows_rocm(self): - """worker.py must pin BNB_ROCM_VERSION=72 in the Windows ROCm section. + """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section. - As of torch==2.11.0+rocm7.13.0, BNB auto-detects HIP 7.13 and looks for - rocm713.dll, which the AMD prerelease wheel does not ship. The worker - must force BNB_ROCM_VERSION=72 before any ML library is imported so that - bitsandbytes loads libbitsandbytes_rocm72.dll. + BNB auto-detects HIP version from torch.version.hip, which can mismatch + the DLL suffix in the AMD prerelease wheel. The worker must detect the + actual DLL suffix and override BNB's auto-detection before ML imports. """ source = _WORKER_PATH.read_text(encoding = "utf-8") + # Env var must be set assert "BNB_ROCM_VERSION" in source - assert '"72"' in source + # Detection helper must be used + assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source + # "72" must appear as the safe fallback + assert '"72"' in source or "'72'" in source def test_bnb_rocm_version_set_before_ml_imports(self): """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports."""