diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b50e17e5aa..fb5e9b5137 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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 diff --git a/studio/backend/main.py b/studio/backend/main.py index dc3d802b78..2b34b11fef 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -83,9 +83,9 @@ if sys.platform == "win32": # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ # bitsandbytes derives the rocm.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" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 6bc94e3d0f..7d1ae535fd 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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.""" diff --git a/unsloth/save.py b/unsloth/save.py index 4ff56cc477..629cbb9548 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -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: