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

fix: detect BNB ROCm DLL suffix dynamically instead of hardcoding '72'

Browse source 
BNB_ROCM_VERSION was pinned to '72' which works today (AMD wheel ships
rocm72.dll) but would break again if AMD ships a future wheel with a
different DLL suffix (e.g. rocm713.dll).

Add _detect_bnb_rocm_dll_ver() to install_python_stack.py: scans the
installed bitsandbytes package dir for libbitsandbytes_rocm{VER}.dll
using importlib.util.find_spec (no BNB import needed) and returns the
suffix.  '72' remains the fallback when detection fails.

Apply the same detection inline in worker.py section 1f.  Both paths
still respect a pre-set BNB_ROCM_VERSION (caller override wins).

Tests: +8 cases covering detection logic and fallback (147 passed, 2 skipped).
This commit is contained in:
LeoBorcherding 2026-05-15 14:05:04 -05:00
commit c55aaa5809
3 changed files with 149 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

48
studio/backend/core/training/worker.py
View file

@ -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.

45
studio/install_python_stack.py
View file

@ -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:

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

@ -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."""