Auto-set BNB_ROCM_VERSION from the installed wheel on Windows + ROCm (#5986)

* Auto-set BNB_ROCM_VERSION from the installed wheel on Windows + ROCm

bitsandbytes derives its ROCm backend DLL name from `torch.version.hip`.
AMD's Windows bitsandbytes prerelease wheel ships a single
`libbitsandbytes_rocm<NN>.dll` whose suffix does not always match the torch
HIP version: e.g. `torch==2.11.0+rocm7.13.0` reports HIP 7.13, so bitsandbytes
looks for `libbitsandbytes_rocm713.dll`, but the wheel only ships
`libbitsandbytes_rocm72.dll`. The names disagree, the native library fails to
load, and every 4-bit / 8-bit path breaks for users running `import unsloth`
directly (Unsloth Studio already works around this in its worker).

Detect the suffix from the actually-installed wheel and pin BNB_ROCM_VERSION
before bitsandbytes is first imported (unsloth_zoo.device_type imports it during
`from .models import *`), so the correct backend loads. This is precisely the
override bitsandbytes itself recommends when the build/runtime ROCm versions
differ.

Strict no-op unless ALL of: running on Windows, a ROCm torch build, the var is
unset, and a `libbitsandbytes_rocm*.dll` is actually installed. Linux ROCm is
untouched (its multi-backend bitsandbytes resolves the backend correctly from
torch.version.hip). Honors a user-provided BNB_ROCM_VERSION and an explicit
opt-out (UNSLOTH_SKIP_BNB_ROCM_VERSION=1).

Verified on an AMD Radeon 8060S (gfx1151, Strix Halo) Windows 11 + ROCm box:
`import unsloth` now auto-sets BNB_ROCM_VERSION=72 and a native 4-bit
quantize/dequantize roundtrip succeeds with the var unset; previously it failed
to load `libbitsandbytes_rocm713.dll`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Gate BNB_ROCM_VERSION on the actual torch build, not runtime hints

_is_rocm_torch_build() falls back to environment and filesystem hints
(HIP_PATH, ROCM_PATH, ...) that are routinely present on Windows boxes
with the AMD HIP SDK installed but a CUDA or CPU torch. If such a box
also has a bitsandbytes wheel that ships a rocm DLL (AMD's Windows
prerelease wheel ships rocm72 alongside all the cuda DLLs), setting
BNB_ROCM_VERSION makes bitsandbytes raise at import on its CUDA build
and `import unsloth` breaks.

Add _is_hip_torch_build(): wheel version tag first (no torch import),
then torch.version.hip for untagged custom/source HIP builds, and use
it as the gate. The broader hint-based helper keeps its other callers.

Verified on a gfx1151 Windows box: True on the ROCm venv
(2.11.0+rocm7.13.0), False on a torch-less interpreter; 4 new unit
tests including the HIP-SDK-on-CUDA-box false-positive regression.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Relocate wiring so this PR composes with the bnb arch-detection PR

Both this PR and the fix_bitsandbytes_rocm_arch_detection PR anchored
their _gpu_init.py wiring and import_fixes.py additions on the same
configure_amdgpu_asic_id_table_path lines, so whichever merged second
hit a textual conflict in both files (verified by merging both onto
main in each order).

Move maybe_set_windows_rocm_bnb_version's wiring to a self-contained
block after the import-order warning (still before `import
unsloth_zoo`, which is what pulls in bitsandbytes on ROCm) and append
the helpers at the end of import_fixes.py. The hunks no longer
overlap, so the two PRs merge cleanly in either order. No behavior
change: the env var only needs to be set before bitsandbytes is first
imported, and it still is.

* Redetect sitecustomize-seeded BNB_ROCM_VERSION for PR #5986

After #6048, every Studio venv process starts with BNB_ROCM_VERSION seeded
by the managed sitecustomize.py block, which made this gate a no-op inside
Studio venvs and blind to wheel updates. Treat values marked
UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize as redetectable defaults,
stamp redetected values as detected, and keep the seeded value when no DLL
is found. Explicit caller values still win and the opt-out is unchanged.
Also merges latest main.

* Make BNB_ROCM_VERSION opt-out drop the sitecustomize-seeded default for PR #5986

UNSLOTH_SKIP_BNB_ROCM_VERSION=1 previously no-opped the helper but left a
sitecustomize-seeded BNB_ROCM_VERSION in the environment, so bitsandbytes
still consumed the override the user disabled. The opt-out now removes
values carrying the sitecustomize source marker; explicit user values have
no marker and are untouched. Adds tests for the opt-out paths and the
empty-string edge.

* Tighten comments and docstrings for PR #5986

Comment-only pass: shorten verbose docstrings on the internal helpers,
collapse multi-line inline comments, and drop wording that restates the
code. Verified code-identical via the comment_tools.py AST signature
check (3/3 files unchanged).

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

tests/test_windows_rocm_bnb_version.py

@@ -0,0 +1,251 @@
+# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+
+"""Tests for ``maybe_set_windows_rocm_bnb_version`` (unsloth/import_fixes.py).
+
+The module is loaded in isolation (stdlib + packaging only), so no torch /
+GPU is required and unsloth's GPU init never runs."""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import types
+from pathlib import Path
+
+import pytest
+
+_IMPORT_FIXES_PATH = Path(__file__).resolve().parent.parent / "unsloth" / "import_fixes.py"
+
+
+def _load_import_fixes():
+    spec = importlib.util.spec_from_file_location(
+        "unsloth_import_fixes_under_test", _IMPORT_FIXES_PATH
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture()
+def import_fixes():
+    return _load_import_fixes()
+
+
+@pytest.fixture()
+def clean_env(monkeypatch):
+    """Unset the env vars and remove them afterwards (the function writes
+    os.environ directly, which monkeypatch does not auto-revert)."""
+    for var in (
+        "BNB_ROCM_VERSION",
+        "UNSLOTH_SKIP_BNB_ROCM_VERSION",
+        "UNSLOTH_BNB_ROCM_VERSION_SOURCE",
+    ):
+        monkeypatch.delenv(var, raising = False)
+    yield monkeypatch
+    for var in (
+        "BNB_ROCM_VERSION",
+        "UNSLOTH_SKIP_BNB_ROCM_VERSION",
+        "UNSLOTH_BNB_ROCM_VERSION_SOURCE",
+    ):
+        os.environ.pop(var, None)
+
+
+def _force(import_fixes, monkeypatch, *, win, rocm, detected):
+    monkeypatch.setattr(import_fixes.sys, "platform", "win32" if win else "linux")
+    monkeypatch.setattr(import_fixes, "_is_hip_torch_build", lambda: rocm)
+    monkeypatch.setattr(import_fixes, "_detect_installed_bnb_rocm_version", lambda: detected)
+
+
+# ---------------------------------------------------------------------------
+# _detect_installed_bnb_rocm_version
+# ---------------------------------------------------------------------------
+
+
+def test_detect_picks_highest_rocm_suffix(import_fixes, tmp_path, monkeypatch):
+    pkg = tmp_path / "bitsandbytes"
+    pkg.mkdir()
+    for name in (
+        "libbitsandbytes_rocm72.dll",
+        "libbitsandbytes_rocm713.dll",  # numerically highest -> should win
+        "libbitsandbytes_cpu.dll",
+        "__init__.py",
+    ):
+        (pkg / name).write_text("")
+    fake_spec = types.SimpleNamespace(submodule_search_locations = [str(pkg)])
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: fake_spec)
+    assert import_fixes._detect_installed_bnb_rocm_version() == "713"
+
+
+def test_detect_none_when_only_non_rocm_dlls(import_fixes, tmp_path, monkeypatch):
+    pkg = tmp_path / "bitsandbytes"
+    pkg.mkdir()
+    (pkg / "libbitsandbytes_cpu.dll").write_text("")
+    (pkg / "libbitsandbytes_cuda124.dll").write_text("")
+    fake_spec = types.SimpleNamespace(submodule_search_locations = [str(pkg)])
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: fake_spec)
+    assert import_fixes._detect_installed_bnb_rocm_version() is None
+
+
+def test_detect_none_when_bnb_absent(import_fixes, monkeypatch):
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: None)
+    assert import_fixes._detect_installed_bnb_rocm_version() is None
+
+
+# ---------------------------------------------------------------------------
+# maybe_set_windows_rocm_bnb_version
+# ---------------------------------------------------------------------------
+
+
+def test_sets_bnb_version_on_windows_rocm(import_fixes, clean_env):
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() == "72"
+    assert os.environ["BNB_ROCM_VERSION"] == "72"
+    assert os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] == "detected"
+
+
+def test_noop_off_windows(import_fixes, clean_env):
+    # Linux ROCm resolves its backend correctly from torch.version.hip.
+    _force(import_fixes, clean_env, win = False, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert "BNB_ROCM_VERSION" not in os.environ
+
+
+def test_noop_when_not_rocm_torch(import_fixes, clean_env):
+    _force(import_fixes, clean_env, win = True, rocm = False, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert "BNB_ROCM_VERSION" not in os.environ
+
+
+def test_noop_when_no_rocm_dll_installed(import_fixes, clean_env):
+    # Never force a ROCm backend name when no ROCm DLL ships (avoid breaking a
+    # non-ROCm bitsandbytes that happens to sit next to a ROCm torch build).
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = None)
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert "BNB_ROCM_VERSION" not in os.environ
+
+
+def test_respects_user_provided_value(import_fixes, clean_env):
+    clean_env.setenv("BNB_ROCM_VERSION", "999")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert os.environ["BNB_ROCM_VERSION"] == "999"
+
+
+def test_explicit_opt_out(import_fixes, clean_env):
+    clean_env.setenv("UNSLOTH_SKIP_BNB_ROCM_VERSION", "1")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert "BNB_ROCM_VERSION" not in os.environ
+
+
+def test_redetects_sitecustomize_seeded_default(import_fixes, clean_env):
+    # Studio's installer persists a default via the venv sitecustomize.py; the
+    # wheel may have changed since, so the seeded value must be redetected.
+    clean_env.setenv("BNB_ROCM_VERSION", "72")
+    clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "713")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() == "713"
+    assert os.environ["BNB_ROCM_VERSION"] == "713"
+    assert os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] == "detected"
+
+
+def test_sitecustomize_default_kept_when_no_dll_found(import_fixes, clean_env):
+    # A failed redetect must not discard the seeded value.
+    clean_env.setenv("BNB_ROCM_VERSION", "72")
+    clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = None)
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert os.environ["BNB_ROCM_VERSION"] == "72"
+    assert os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] == "sitecustomize"
+
+
+def test_user_value_with_non_sitecustomize_marker_untouched(import_fixes, clean_env):
+    # Only the sitecustomize marker makes a value redetectable.
+    clean_env.setenv("BNB_ROCM_VERSION", "999")
+    clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "detected")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert os.environ["BNB_ROCM_VERSION"] == "999"
+
+
+def test_opt_out_unseats_sitecustomize_seeded_value(import_fixes, clean_env):
+    # The opt-out must also drop a default our own sitecustomize block seeded,
+    # so bitsandbytes never sees the override the user disabled.
+    clean_env.setenv("BNB_ROCM_VERSION", "72")
+    clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
+    clean_env.setenv("UNSLOTH_SKIP_BNB_ROCM_VERSION", "1")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "713")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert "BNB_ROCM_VERSION" not in os.environ
+    assert "UNSLOTH_BNB_ROCM_VERSION_SOURCE" not in os.environ
+
+
+def test_opt_out_keeps_explicit_user_value(import_fixes, clean_env):
+    # Opt-out must never remove a value the user set themselves (no marker).
+    clean_env.setenv("BNB_ROCM_VERSION", "999")
+    clean_env.setenv("UNSLOTH_SKIP_BNB_ROCM_VERSION", "1")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert os.environ["BNB_ROCM_VERSION"] == "999"
+
+
+def test_empty_string_value_without_marker_is_respected(import_fixes, clean_env):
+    # "" counts as present: without the sitecustomize marker it is not ours
+    # to overwrite.
+    clean_env.setenv("BNB_ROCM_VERSION", "")
+    _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
+    assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
+    assert os.environ["BNB_ROCM_VERSION"] == ""
+
+
+# ---------------------------------------------------------------------------
+# _is_hip_torch_build (the strict gate -- regression for the HIP-SDK-on-a-
+# CUDA-box false positive: env hints like HIP_PATH must NOT count)
+# ---------------------------------------------------------------------------
+
+
+def _fake_torch(hip):
+    return types.SimpleNamespace(version = types.SimpleNamespace(hip = hip))
+
+
+def test_hip_build_true_from_wheel_tag(import_fixes, monkeypatch):
+    monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.11.0+rocm7.13.0")
+    assert import_fixes._is_hip_torch_build() is True
+
+
+def test_hip_build_true_from_torch_version_hip(import_fixes, monkeypatch):
+    # Custom/source HIP build without the +rocm tag.
+    monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.11.0")
+    monkeypatch.setitem(__import__("sys").modules, "torch", _fake_torch("7.2.0"))
+    assert import_fixes._is_hip_torch_build() is True
+
+
+def test_hip_build_false_for_cuda_torch_despite_rocm_env_hints(import_fixes, monkeypatch):
+    """HIP SDK env vars set but CUDA torch: the strict gate must say False,
+    otherwise BNB_ROCM_VERSION gets set and CUDA bitsandbytes raises."""
+    monkeypatch.setenv("HIP_PATH", r"C:\Program Files\AMD\ROCm\6.2")
+    monkeypatch.setenv("ROCM_PATH", r"C:\Program Files\AMD\ROCm\6.2")
+    monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.9.0+cu126")
+    monkeypatch.setitem(__import__("sys").modules, "torch", _fake_torch(None))
+    assert import_fixes._is_hip_torch_build() is False
+
+
+def test_hip_build_false_when_torch_absent(import_fixes, monkeypatch):
+    def _raise(name):
+        raise Exception("no torch dist")
+
+    monkeypatch.setattr(import_fixes, "importlib_version", _raise)
+    monkeypatch.setitem(__import__("sys").modules, "torch", None)
+    assert import_fixes._is_hip_torch_build() is False

unsloth/_gpu_init.py

@@ -96,6 +96,13 @@ if already_imported:
     )
 del already_imported, critical_modules
 
+# Pin BNB_ROCM_VERSION before bitsandbytes is first imported (`import
+# unsloth_zoo` below pulls it in on ROCm hosts).
+from .import_fixes import maybe_set_windows_rocm_bnb_version
+
+maybe_set_windows_rocm_bnb_version()
+del maybe_set_windows_rocm_bnb_version
+
 # Multi-GPU is not yet supported (beta available on request).
 
 # Fixes https://github.com/unslothai/unsloth/issues/1266

unsloth/import_fixes.py

@@ -2162,3 +2162,91 @@ def disable_broken_causal_conv1d():
         "Unsloth: Detected broken causal_conv1d binary; "
         "disabling causal_conv1d fast path and continuing import."
     )
+
+
+_BNB_ROCM_DLL_RE = re.compile(r"libbitsandbytes_rocm(\d+)\.dll", re.IGNORECASE)
+
+
+def _is_hip_torch_build():
+    """True only when torch itself is a HIP/ROCm build. Env hints (HIP_PATH
+    etc.) do not count: CUDA bitsandbytes raises at import when the ROCm
+    override is set. Wheel tag first (no torch import); torch.version.hip
+    fallback for source builds."""
+    try:
+        if "rocm" in str(importlib_version("torch")).lower():
+            return True
+    except Exception:
+        pass
+    try:
+        import torch
+        return bool(getattr(torch.version, "hip", None))
+    except Exception:
+        return False
+
+
+def _detect_installed_bnb_rocm_version():
+    """Highest installed ``libbitsandbytes_rocm<NN>.dll`` suffix ("72", "713")
+    or ``None``. Listing order is unordered, so take the numeric max."""
+    try:
+        spec = importlib.util.find_spec("bitsandbytes")
+    except Exception:
+        return None
+    if spec is None or not spec.submodule_search_locations:
+        return None
+
+    suffixes = []
+    for pkg_dir in spec.submodule_search_locations:
+        try:
+            entries = os.listdir(pkg_dir)
+        except Exception:
+            continue
+        for entry in entries:
+            match = _BNB_ROCM_DLL_RE.fullmatch(entry)
+            if match is not None:
+                suffixes.append(match.group(1))
+    if not suffixes:
+        return None
+    return max(suffixes, key = lambda value: int(value))
+
+
+def maybe_set_windows_rocm_bnb_version():
+    """Pin ``BNB_ROCM_VERSION`` from the installed wheel on Windows + ROCm torch.
+
+    AMD's Windows wheel ships one ``libbitsandbytes_rocm<NN>.dll`` whose
+    suffix can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll),
+    breaking the native 4-bit/8-bit paths. Pin the installed suffix before
+    bitsandbytes is first imported.
+
+    No-op unless ALL of: Windows, a real HIP torch build (env hints like
+    HIP_PATH do not count), a ROCm DLL installed, and no explicit user value.
+    Linux is untouched. Values seeded by Studio's venv sitecustomize.py
+    (marked ``UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize``) are
+    redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
+    opts out and drops a seeded default. Returns the value set, else None.
+    """
+    if sys.platform != "win32":
+        return None
+    if os.environ.get("UNSLOTH_SKIP_BNB_ROCM_VERSION") == "1":
+        # Real opt-out: drop our seeded default (marker present); explicit
+        # user values carry no marker and are kept.
+        if os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize":
+            os.environ.pop("BNB_ROCM_VERSION", None)
+            os.environ.pop("UNSLOTH_BNB_ROCM_VERSION_SOURCE", None)
+        return None
+    if "BNB_ROCM_VERSION" in os.environ and (
+        os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") != "sitecustomize"
+    ):
+        return None
+    if not _is_hip_torch_build():
+        return None
+    version = _detect_installed_bnb_rocm_version()
+    if version is None:
+        return None
+    os.environ["BNB_ROCM_VERSION"] = version
+    os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
+    if UNSLOTH_ENABLE_LOGGING:
+        logger.info(
+            f"Unsloth: set BNB_ROCM_VERSION={version} "
+            "(detected from the installed bitsandbytes ROCm wheel on Windows)."
+        )
+    return version