diff --git a/bitsandbytes/cextension.py b/bitsandbytes/cextension.py index e234f20d3..714a44ddf 100644 --- a/bitsandbytes/cextension.py +++ b/bitsandbytes/cextension.py @@ -19,6 +19,59 @@ logger = logging.getLogger(__name__) +def _is_missing_dependency_error(error_msg: str) -> bool: + """True when loading failed because a *dependency* of the native library is missing. + + Windows surfaces this two ways. A raw ERROR_MOD_NOT_FOUND arrives as + "[WinError 126] The specified module could not be found", but ctypes' own loader + raises FileNotFoundError with ``winerror`` unset and the text "Could not find module + '' (or one of its dependencies)". Both mean the same thing -- the ROCm/CUDA + runtime is not on the loader search path -- so both must reach the same advice. + """ + lowered = error_msg.lower() + return ( + "winerror 126" in lowered + or "specified module could not be found" in lowered + or "could not find module" in lowered + or "(or one of its dependencies)" in lowered + ) + + +def _dependency_search_path_hint(backend: str, major_version: str) -> str: + """Platform-appropriate instructions for making the runtime libraries findable.""" + if os.name == "nt": + if backend == "ROCm": + return ( + f"{backend} {major_version}.x runtime libraries were not found on the DLL search path.\n\n" + f"To fix this, make sure that:\n" + f"1. You have installed the {backend} {major_version}.x SDK (the ROCm/HIP SDK installer,\n" + f" or the `rocm[libraries,devel]` wheels from the AMD nightly index)\n" + f"2. The runtime DLLs (amdhip64*.dll, rocblas.dll, hipblas.dll) are in the SDK's `bin`\n" + f" directory, which bitsandbytes registers with os.add_dll_directory\n\n" + f"You can point bitsandbytes at the SDK root explicitly:\n" + f" set ROCM_PATH=C:\\path\\to\\rocm-{major_version}.x\n\n" + ) + return ( + f"{backend} {major_version}.x runtime DLLs (cudart64*.dll, cublas64*.dll) were not found on\n" + f"the DLL search path.\n\n" + f"To fix this, make sure that:\n" + f"1. You have installed the CUDA {major_version}.x toolkit\n" + f"2. Its `bin` directory is on PATH, or CUDA_PATH is set to the toolkit root\n\n" + f"You can check what Windows can see with:\n" + f" where cudart64*.dll\n\n" + ) + + lib_dir = "lib64" if not HIP_ENVIRONMENT else "lib" + return ( + f"{backend} {major_version}.x runtime libraries were not found in the LD_LIBRARY_PATH.\n\n" + f"To fix this, make sure that:\n" + f"1. You have installed {backend} {major_version}.x toolkit on your system\n" + f"2. The {backend} runtime libraries are in your LD_LIBRARY_PATH\n\n" + f"You can add them with (and persist the change by adding the line to your .bashrc):\n" + f" export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/{backend.lower()}-{major_version}.x/{lib_dir}\n\n" + ) + + def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path: """ Get the path to the best matching CUDA/ROCm BNB native library for the given specs. @@ -40,8 +93,18 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path: override_value = os.environ.get(override_var) if override_value is not None: + # Accept the dotted spelling users naturally reach for ('12.4', '7.2') as well + # as the compact one that appears in the library filename ('124', '72'). + override_value = override_value.strip() if not override_value.isdigit(): - raise RuntimeError(f"{override_var}={override_value!r}: value must be digits only (e.g. '124' for 12.4).") + dotted = override_value.split(".") + if len(dotted) >= 2 and all(part.isdigit() for part in dotted[:2]): + override_value = "".join(dotted[:2]) + else: + raise RuntimeError( + f"{override_var}={override_value!r}: value must be digits only " + f"(e.g. '124' for 12.4), or a dotted version (e.g. '12.4')." + ) library_name = f"libbitsandbytes_{prefix}{override_value}{DYNAMIC_LIBRARY_SUFFIX}" logger.warning( f"WARNING: {override_var}={override_value} environment variable detected; loading {library_name}.\n" @@ -96,8 +159,11 @@ def __getattr__(self, name): def throw_on_call(*args, **kwargs): raise RuntimeError( - f"Method '{name}' not available in CPU-only version of bitsandbytes.\n" - "Reinstall with GPU support or use CUDA-enabled hardware." + f"Method '{name}' is not exported by the loaded bitsandbytes native library " + f"(backend: {BNB_BACKEND}).\n" + "This usually means a CPU-only library was loaded, or the prebuilt binary for " + "your backend was built without this kernel. Run `python -m bitsandbytes` for " + "the backend, versions and library path that were actually resolved." ) return throw_on_call @@ -197,8 +263,9 @@ def __init__(self, error_msg: str): user_version_str = f"{user_version[0]}.{user_version[1]}" if user_version else "unknown" self.requested_version = parse_cuda_version(override_value) if override_value else user_version_str - # Pre-generate the error message based on error type - if "cannot open shared object file" in error_msg: + # Pre-generate the error message based on error type. A missing *dependency* + # surfaces as a dlerror() string on Linux and as WinError 126 on Windows. + if "cannot open shared object file" in error_msg or _is_missing_dependency_error(error_msg): self.formatted_error = self._format_dependency_error() else: # lib loading errors self.formatted_error = self._format_lib_error_message( @@ -217,8 +284,9 @@ def _format_lib_error_message( ) -> str: """Format detailed error message for library loading failures""" analysis = "" - no_cpu_lib_found = "libbitsandbytes_cpu.so: cannot open" in original_error - no_cuda_lib_found = f"{BNB_BACKEND} binary not found" in original_error + no_cpu_lib_found = f"libbitsandbytes_cpu{DYNAMIC_LIBRARY_SUFFIX}" in original_error + # get_native_library() raises "No compatible {backend} binary found at ". + no_cuda_lib_found = f"{BNB_BACKEND} binary found at" in original_error if no_cpu_lib_found: analysis = "\n🚨 Failed to load CPU-only bitsandbytes library 🚨\n\n" @@ -302,13 +370,7 @@ def _format_dependency_error(self) -> str: return ( f"\n🚨 {BNB_BACKEND} SETUP ERROR: Missing dependency: {missing_lib} 🚨\n\n" - f"{BNB_BACKEND} {cuda_major_version}.x runtime libraries were not found in the LD_LIBRARY_PATH.\n\n" - f"To fix this, make sure that:\n" - f"1. You have installed {BNB_BACKEND} {cuda_major_version}.x toolkit on your system\n" - f"2. The {BNB_BACKEND} runtime libraries are in your LD_LIBRARY_PATH\n\n" - f"You can add them with (and persist the change by adding the line to your .bashrc):\n" - f" export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/{BNB_BACKEND.lower()}-{cuda_major_version}.x/" - f"{'lib64' if not HIP_ENVIRONMENT else 'lib'}\n\n" + f"{_dependency_search_path_hint(BNB_BACKEND, cuda_major_version)}" f"Original error: {self.error_msg}\n\n" f"🔍 Run this command for detailed diagnostics:\n" f"python -m bitsandbytes\n\n" @@ -345,6 +407,48 @@ def get_xpu_bnb_library_path() -> Path: return PACKAGE_DIR / f"libbitsandbytes_xpu{DYNAMIC_LIBRARY_SUFFIX}" +# Handles returned by os.add_dll_directory(). They must stay alive for the process +# lifetime: releasing a handle removes the directory from the search path again. +_DLL_DIR_HANDLES: list = [] + + +def _register_windows_dll_directories() -> None: + """Register the directories holding the CUDA/ROCm runtime DLLs on Windows. + + ctypes loads with LOAD_LIBRARY_SEARCH_DEFAULT_DIRS, which searches the + application directory, System32 and explicitly registered directories -- but + *not* PATH and not the loaded DLL's own directory. Without this, a Windows ROCm + install whose amdhip64*.dll / rocblas.dll live in the SDK's `bin` directory can + only be resolved by accident, e.g. when torch happens to have already loaded a + module with the same base name. + """ + if os.name != "nt" or _DLL_DIR_HANDLES: + return + + candidates = [] + for var in ("ROCM_PATH", "HIP_PATH", "ROCM_HOME", "CONDA_PREFIX"): + root = os.environ.get(var) + if root: + candidates += [Path(root) / "bin", Path(root) / "lib"] + # torch ships the runtime it was built against, which on a ROCm build is a + # reliable source of amdhip64*.dll. + candidates.append(Path(torch.__file__).parent / "lib") + candidates += [Path(entry) for entry in os.environ.get("PATH", "").split(os.pathsep) if entry] + + seen = set() + for directory in candidates: + try: + if directory in seen or not directory.is_dir(): + continue + seen.add(directory) + if not any(directory.glob("amdhip64*.dll")) and not any(directory.glob("cudart64*.dll")): + continue + _DLL_DIR_HANDLES.append(os.add_dll_directory(str(directory))) + logger.debug(f"Registered DLL search directory: {directory}") + except OSError as exc: + logger.warning(f"Could not register DLL search directory {directory}: {exc}") + + def get_native_library() -> BNBNativeLibrary: """ Load CUDA library XOR CPU, as the latter contains a subset of symbols of the former. @@ -365,6 +469,8 @@ def get_native_library() -> BNBNativeLibrary: logger.debug(f"Loading bitsandbytes native library from: {binary_path}") + _register_windows_dll_directories() + # Try to load the library - any errors will propagate up dll = ct.cdll.LoadLibrary(str(binary_path)) diff --git a/bitsandbytes/cuda_specs.py b/bitsandbytes/cuda_specs.py index 25ce3cd1e..bc7912365 100644 --- a/bitsandbytes/cuda_specs.py +++ b/bitsandbytes/cuda_specs.py @@ -1,9 +1,6 @@ import dataclasses from functools import lru_cache import logging -import platform -import re -import subprocess from typing import Optional import torch @@ -17,7 +14,13 @@ class CUDASpecs: @property def has_imma(self) -> bool: - return torch.version.hip or self.highest_compute_capability >= (7, 5) + if torch.version.hip: + # ROCm: int8 matrix cores (MFMA) start at CDNA1 (gfx908). The GCN5 parts + # (gfx900/gfx906/gfx90c) have none. This previously returned + # `torch.version.hip` -- a version *string* -- which is always truthy, so + # every AMD GPU was reported as IMMA-capable. + return not get_rocm_gpu_arch().startswith(("gfx900", "gfx906", "gfx90c")) + return self.highest_compute_capability >= (7, 5) def get_compute_capabilities() -> list[tuple[int, int]]: @@ -79,33 +82,21 @@ def get_cuda_specs() -> Optional[CUDASpecs]: return None -def get_rocm_gpu_arch() -> str: - """Get ROCm GPU architecture.""" +def get_rocm_gpu_arch(device_index: int = 0) -> str: + """Get the gfx architecture of a ROCm device, e.g. 'gfx1201'. + + Reads the device property torch already exposes rather than shelling out to + hipinfo.exe/rocminfo: the subprocess is not on PATH for pip-installed ROCm SDKs + (TheRock), its output format differs per platform, and the result is available + directly from the runtime. + """ logger = logging.getLogger(__name__) + if not torch.version.hip or not torch.cuda.is_available(): + return "unknown" try: - if torch.version.hip: - # On Windows, use hipinfo.exe; on Linux, use rocminfo - if platform.system() == "Windows": - cmd = ["hipinfo.exe"] - arch_pattern = r"gcnArchName:\s+gfx([a-zA-Z\d]+)" - else: - cmd = ["rocminfo"] - arch_pattern = r"Name:\s+gfx([a-zA-Z\d]+)" - - result = subprocess.run(cmd, capture_output=True, text=True) - match = re.search(arch_pattern, result.stdout) - if match: - return "gfx" + match.group(1) - else: - return "unknown" - else: - return "unknown" + # gcnArchName carries feature suffixes, e.g. 'gfx90a:sramecc+:xnack-'. + name = getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "") or "" + return name.split(":")[0] or "unknown" except Exception as e: logger.error(f"Could not detect ROCm GPU architecture: {e}") - if torch.cuda.is_available(): - logger.warning( - """ -ROCm GPU architecture detection failed despite ROCm being available. - """, - ) return "unknown" diff --git a/tests/test_cextension_errors.py b/tests/test_cextension_errors.py new file mode 100644 index 000000000..78c490f0e --- /dev/null +++ b/tests/test_cextension_errors.py @@ -0,0 +1,235 @@ +"""Regression tests for the native-library loader's error classification. + +These cover failures that only ever surfaced on Windows + ROCm: + +* the "no compatible binary" message is classified by a substring that the loader + never actually raised, so users were told to build a CPU-only library; +* a missing dependency on Windows is reported as ``[WinError 126]``, not as a + ``dlerror()`` string, so it never reached the dependency guidance; +* that guidance only ever described ``LD_LIBRARY_PATH`` and ``.bashrc``; +* ``os.add_dll_directory`` was never called, and ctypes does not search ``PATH``; +* ``has_imma`` returned the HIP version *string* rather than a bool; +* ``BNB_ROCM_VERSION=7.2`` (the natural spelling) was a hard error. + +Everything here is platform-independent: no GPU and no ROCm toolchain required. +""" + +from __future__ import annotations + +import pytest + +from bitsandbytes import cextension +from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX +from bitsandbytes.cuda_specs import CUDASpecs + + +def _spec() -> CUDASpecs: + """A spec that will not match any packaged library, so resolution is exercised.""" + return CUDASpecs( + cuda_version_string="72", + highest_compute_capability=(12, 1), + cuda_version_tuple=(7, 2), + ) + + +def _override_var() -> str: + return "BNB_ROCM_VERSION" if cextension.HIP_ENVIRONMENT else "BNB_CUDA_VERSION" + + +# -------------------------------------------------------------------------------------- +# Missing-dependency detection +# -------------------------------------------------------------------------------------- + + +def test_windows_module_not_found_is_a_missing_dependency(): + assert cextension._is_missing_dependency_error("[WinError 126] The specified module could not be found") + + +def test_ctypes_file_not_found_is_also_a_missing_dependency(): + """The form ctypes actually raises. + + Reproduced on Windows by copying the built extension somewhere without the ROCm + DLLs reachable: ``ct.cdll.LoadLibrary`` raises FileNotFoundError with ``winerror`` + unset (None) and this exact text. The "[WinError 126]" form does not appear, so + matching only on that string would miss the failure this whole path exists for. + """ + real = ( + "Could not find module 'C:\\Temp\\libbitsandbytes_rocm72.dll' " + "(or one of its dependencies). Try using the full path with constructor syntax." + ) + assert cextension._is_missing_dependency_error(real) + + +def test_unrelated_error_is_not_a_missing_dependency(): + assert not cextension._is_missing_dependency_error("something else went wrong") + + +def test_dependency_hint_is_platform_specific(monkeypatch): + monkeypatch.setattr(cextension.os, "name", "nt") + windows_hint = cextension._dependency_search_path_hint("ROCm", "7") + assert "DLL search path" in windows_hint + assert "ROCM_PATH" in windows_hint + assert "LD_LIBRARY_PATH" not in windows_hint + + monkeypatch.setattr(cextension.os, "name", "posix") + monkeypatch.setattr(cextension, "HIP_ENVIRONMENT", True) + posix_hint = cextension._dependency_search_path_hint("ROCm", "7") + assert "LD_LIBRARY_PATH" in posix_hint + assert "ROCM_PATH" not in posix_hint + + +# -------------------------------------------------------------------------------------- +# "No compatible binary" must reach the branch that explains it +# -------------------------------------------------------------------------------------- + + +def test_missing_binary_raises_the_message_the_formatter_keys_on(monkeypatch, tmp_path): + monkeypatch.setattr(cextension, "get_cuda_specs", _spec) + monkeypatch.setattr(cextension, "PACKAGE_DIR", tmp_path) + monkeypatch.setattr(cextension, "_find_cuda_libs", lambda prefix, is_hip: {}) + + with pytest.raises(RuntimeError) as excinfo: + cextension.get_native_library() + + # This is the exact coupling that was broken: the formatter matched + # "{backend} binary not found" while the loader raised "{backend} binary found at". + assert f"{cextension.BNB_BACKEND} binary found at" in str(excinfo.value) + + +def test_missing_binary_does_not_advise_a_cpu_only_build(): + message = ( + f"Original error: No compatible {cextension.BNB_BACKEND} binary found at " + f"C:/pkg/libbitsandbytes_rocm72{DYNAMIC_LIBRARY_SUFFIX}\n" + ) + formatted = cextension.ErrorHandlerMockBNBNativeLibrary(message).formatted_error + assert "COMPILE FROM SOURCE for CPU-only" not in formatted + + +def test_cpu_library_probe_uses_the_platform_suffix(): + # The old probe hardcoded ".so", so it never matched on Windows. Exercise the + # formatter directly: a real CPU-load failure also contains a dlerror() string, + # which routes to the dependency formatter before this branch is reached. + lib = cextension.ErrorHandlerMockBNBNativeLibrary("boom") + formatted = lib._format_lib_error_message( + available_versions=[], + user_cuda_version="7.2", + original_error=f"Original error: libbitsandbytes_cpu{DYNAMIC_LIBRARY_SUFFIX}: failed to load\n", + ) + assert "Failed to load CPU-only bitsandbytes library" in formatted + + +# -------------------------------------------------------------------------------------- +# Windows DLL search path +# -------------------------------------------------------------------------------------- + + +def test_dll_registration_is_a_noop_off_windows(monkeypatch): + monkeypatch.setattr(cextension, "_DLL_DIR_HANDLES", []) + monkeypatch.setattr(cextension.os, "name", "posix") + cextension._register_windows_dll_directories() + assert cextension._DLL_DIR_HANDLES == [] + + +def test_dll_registration_finds_the_rocm_runtime(monkeypatch, tmp_path): + rocm_root = tmp_path / "rocm" + (rocm_root / "bin").mkdir(parents=True, exist_ok=True) + (rocm_root / "bin" / "amdhip64_7.dll").write_bytes(b"") + + registered: list[str] = [] + monkeypatch.setattr(cextension, "_DLL_DIR_HANDLES", []) + monkeypatch.setattr(cextension.os, "name", "nt") + monkeypatch.setattr( + cextension.os, "add_dll_directory", lambda path: registered.append(path) or object(), raising=False + ) + monkeypatch.setattr(cextension.torch, "__file__", str(tmp_path / "torch" / "__init__.py"), raising=False) + monkeypatch.setenv("ROCM_PATH", str(rocm_root)) + for var in ("HIP_PATH", "ROCM_HOME", "CONDA_PREFIX"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PATH", "") + + cextension._register_windows_dll_directories() + + assert registered == [str(rocm_root / "bin")] + assert len(cextension._DLL_DIR_HANDLES) == 1 + + +def test_dll_registration_skips_directories_without_a_runtime(monkeypatch, tmp_path): + (tmp_path / "bin").mkdir(exist_ok=True) + registered: list[str] = [] + monkeypatch.setattr(cextension, "_DLL_DIR_HANDLES", []) + monkeypatch.setattr(cextension.os, "name", "nt") + monkeypatch.setattr( + cextension.os, "add_dll_directory", lambda path: registered.append(path) or object(), raising=False + ) + monkeypatch.setattr(cextension.torch, "__file__", str(tmp_path / "torch" / "__init__.py"), raising=False) + monkeypatch.setenv("ROCM_PATH", str(tmp_path)) + for var in ("HIP_PATH", "ROCM_HOME", "CONDA_PREFIX"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("PATH", "") + + cextension._register_windows_dll_directories() + + assert registered == [] + + +# -------------------------------------------------------------------------------------- +# BNB_ROCM_VERSION / BNB_CUDA_VERSION spelling +# -------------------------------------------------------------------------------------- + + +def test_dotted_override_is_accepted(monkeypatch): + monkeypatch.setenv(_override_var(), "7.2") + path = cextension.get_cuda_bnb_library_path(_spec()) + prefix = "rocm" if cextension.HIP_ENVIRONMENT else "cuda" + assert path.name == f"libbitsandbytes_{prefix}72{DYNAMIC_LIBRARY_SUFFIX}" + + +def test_compact_override_is_still_accepted(monkeypatch): + monkeypatch.setenv(_override_var(), "72") + path = cextension.get_cuda_bnb_library_path(_spec()) + prefix = "rocm" if cextension.HIP_ENVIRONMENT else "cuda" + assert path.name == f"libbitsandbytes_{prefix}72{DYNAMIC_LIBRARY_SUFFIX}" + + +def test_garbage_override_still_raises(monkeypatch): + monkeypatch.setenv(_override_var(), "seventy") + with pytest.raises(RuntimeError, match="digits only"): + cextension.get_cuda_bnb_library_path(_spec()) + + +# -------------------------------------------------------------------------------------- +# ROCm capability detection +# -------------------------------------------------------------------------------------- + + +class _FakeProps: + gcnArchName = "gfx90a:sramecc+:xnack-" + + +def test_rocm_arch_uses_device_properties(monkeypatch): + from bitsandbytes import cuda_specs + + monkeypatch.setattr(cuda_specs.torch.version, "hip", "7.2.0", raising=False) + monkeypatch.setattr(cuda_specs.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(cuda_specs.torch.cuda, "get_device_properties", lambda index=0: _FakeProps()) + assert cuda_specs.get_rocm_gpu_arch() == "gfx90a" + + +def test_rocm_arch_is_unknown_without_hip(monkeypatch): + from bitsandbytes import cuda_specs + + monkeypatch.setattr(cuda_specs.torch.version, "hip", None, raising=False) + assert cuda_specs.get_rocm_gpu_arch() == "unknown" + + +@pytest.mark.parametrize("arch,expected", [("gfx1201", True), ("gfx942", True), ("gfx906", False), ("gfx900", False)]) +def test_has_imma_returns_a_bool(monkeypatch, arch, expected): + from bitsandbytes import cuda_specs + + monkeypatch.setattr(cuda_specs.torch.version, "hip", "7.2.0", raising=False) + monkeypatch.setattr(cuda_specs, "get_rocm_gpu_arch", lambda *args, **kwargs: arch) + + result = _spec().has_imma + + assert isinstance(result, bool), "has_imma must not leak the HIP version string" + assert result is expected