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

Studio: require an installed ROCm DLL before forcing BNB_ROCM_VERSION; drop shadowing shutil imports in save.py (#6194)

Browse source 
* Require a found ROCm DLL before forcing BNB_ROCM_VERSION in Studio paths

main.py previously set BNB_ROCM_VERSION=72 whenever HIP_PATH or ROCM_PATH
was set, and the training worker fell back to a blind 72 when DLL
detection found nothing. On a Windows machine with the AMD HIP SDK
installed but CUDA or CPU torch, that forces a ROCm backend onto a
non-ROCm bitsandbytes wheel, which raises at import. Both paths now only
write the override when a libbitsandbytes_rocm DLL actually exists (or a
seeded value is already present), matching the strict gates in
unsloth/import_fixes.py.

Also removes four redundant local import shutil statements in
unsloth/save.py that shadow the module-level import, the same pattern
that caused the UnboundLocalError fixed in #6149.

* Worker: gate the BNB override on a found ROCm DLL, preserving seeded marker

Review follow-ups: track _found_rocm_bnb in the worker like main.py so a
ROCm DLL with an unparsable name still gets the seeded or 72 fallback,
and skip the env write entirely when no DLL exists so a seeded value
keeps its sitecustomize marker and stays redetectable by later import
fixes.
This commit is contained in:
Daniel Han 2026-06-11 06:52:38 -07:00 • committed by GitHub
commit 73973435ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 45 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -2143,14 +2143,14 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB
# wheel may ship a DLL whose suffix doesn't match. Detect the actual
# DLL name and override; "72" is a safe fallback. Values seeded by
# the installer are redetectable defaults, while caller overrides
# remain authoritative.
# DLL name and override. Values seeded by the installer are
# redetectable defaults, while caller overrides remain authoritative.
if (
"BNB_ROCM_VERSION" not in os.environ
or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"
):
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
import glob as _glob
import importlib.util as _ilu
@ -2163,6 +2163,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_found_rocm_bnb = True
_m = _re.search(
r"libbitsandbytes_rocm(\d+)\.dll",
os.path.basename(_dll),
@ -2174,15 +2175,20 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
pass
_bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Only when a ROCm bnb DLL actually exists (mirrors main.py):
# without one the seeded value and its marker stay untouched,
# so later import fixes can still redetect or opt out. DLL
# with unparsable name -> seeded value or "72".
if _found_rocm_bnb:
_bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Parse HIP version for the kernel-fix gate below, falling back to
# the rocm version embedded in torch.__version__ when version.hip is

15
studio/backend/main.py
View file

@ -83,9 +83,9 @@ if sys.platform == "win32":
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
# bitsandbytes derives the rocm<ver>.dll name from torch.version.hip, but the
# wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not
# found") without this. Detect the shipped DLL and fall back to "72" (mirrors
# worker.py). Gate on the rocm bnb DLL / HIP_PATH rather than torch.version.hip
# to avoid importing torch on every Windows host.
# found") without this. Detect the shipped DLL (mirrors worker.py); gate on
# the rocm bnb DLL rather than torch.version.hip to avoid importing torch on
# every Windows host.
# Values seeded by the installer's sitecustomize.py are redetectable
# defaults; explicit caller values remain authoritative.
if (
@ -95,7 +95,6 @@ if sys.platform == "win32":
import glob as _glob
import logging as _logging
_hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"))
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
@ -118,11 +117,13 @@ if sys.platform == "win32":
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
except Exception as _e:
_logging.getLogger(__name__).warning(
"Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
"Windows ROCm: BNB DLL detection failed (%s); leaving BNB_ROCM_VERSION as is",
_e,
)
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72")
if _found_rocm_bnb or _hip_env:
# Only when a ROCm bnb DLL actually exists: HIP_PATH/ROCM_PATH alone
# (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a
# non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72".
if _found_rocm_bnb:
_bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"

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

@ -2241,7 +2241,22 @@ class TestRuntimeBnbRocmSourceGuards:
"""A failed redetect must not downgrade a persisted suffix to '72'."""
for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
source = path.read_text(encoding = "utf-8")
assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name
assert (
'_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source
), path.name
def test_main_requires_found_rocm_dll(self):
"""HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force
a ROCm backend onto a non-ROCm bitsandbytes."""
source = self._MAIN_PATH.read_text(encoding = "utf-8")
assert "if _found_rocm_bnb:" in source
assert "_hip_env" not in source
def test_worker_requires_found_rocm_dll(self):
"""No DLL found: the worker must not write any override or touch the
seeded marker (later import fixes must still see sitecustomize)."""
source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8")
assert "if _found_rocm_bnb:" in source
class TestDetectBnbRocmDllVer:
@ -2443,8 +2458,9 @@ class TestWorkerWindowsRocmPatches:
assert "BNB_ROCM_VERSION" 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
# Falls back to the seeded value, never a blind "72" (which would
# force a ROCm backend onto a non-ROCm bitsandbytes wheel)
assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' 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."""

4
unsloth/save.py
View file

@ -1092,7 +1092,6 @@ def unsloth_save_model(
gc.collect()
# Remove temporary location
import shutil
shutil.rmtree(temporary_location, ignore_errors = True)
@ -1224,7 +1223,6 @@ def install_llama_cpp_old(version = -10):
for i in range(30):
print(f"**[WARNING]** Deleting llama.cpp directory... {30-i} seconds left.")
time.sleep(1)
import shutil
shutil.rmtree("llama.cpp", ignore_errors = True)
@ -2494,7 +2492,6 @@ def unsloth_push_to_hub_gguf(
except Exception as e:
if cleanup_temp:
import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
try:
shutil.rmtree(d)
@ -2681,7 +2678,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
# Clean up temporary directory
if cleanup_temp:
print("Unsloth: Cleaning up temporary files...")
import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
if os.path.exists(d):
try: