d75e765a6b* 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>
251 lines
11 KiB
Python
251 lines
11 KiB
Python
# 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
|