diff --git a/tests/python/test_bitsandbytes_kernel_readiness.py b/tests/python/test_bitsandbytes_kernel_readiness.py new file mode 100644 index 0000000000..db6ec74e57 --- /dev/null +++ b/tests/python/test_bitsandbytes_kernel_readiness.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""`ALLOW_BITSANDBYTES` must follow the kernels, not the mere presence of the module. + +From bitsandbytes 0.46 a wheel whose native library never loaded still imports and +resolves every ctypes handle to a `throw_on_call` closure, so a probe made of attribute +reads alone sees a healthy wheel, the loader selects a 4bit checkpoint, and the failure +lands inside a kernel mid-run instead of degrading to 16bit. +""" + +from __future__ import annotations + +import ast +import importlib.util +import types +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_probe(): + """Import by path, not as ``unsloth.bnb_availability``, which would run the package + __init__ and pull in torch. Works only because the module is a leaf - the property + that lets device_type.py, imported very early, use it without a cycle.""" + path = REPO_ROOT / "unsloth" / "bnb_availability.py" + spec = importlib.util.spec_from_file_location("_unsloth_bnb_availability", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _fake_bnb(lib): + functional = types.ModuleType("bitsandbytes.functional") + functional.get_ptr = lambda tensor: None + functional.lib = lib + bnb = types.ModuleType("bitsandbytes") + bnb.__version__ = "0.50.0" + bnb.functional = functional + return bnb + + +class _DeferredFailureLib: + """What bitsandbytes >= 0.46 hands back when the native library is dead.""" + + def __getattr__(self, name): + def throw_on_call(*args, **kwargs): + raise RuntimeError(f"Method '{name}' not available in CPU-only version") + + return throw_on_call + + +class _RealHandleLib: + """ctypes caches the function object on first lookup; its handles carry restype.""" + + def __getattr__(self, name): + def handle(*args, **kwargs): + return None + + handle.restype = None + setattr(self, name, handle) + return handle + + +def test_probe_covers_every_module_scope_ctypes_bind(): + """A probe that misses one of the import-time binds lets a dead wheel through.""" + tree = ast.parse((REPO_ROOT / "unsloth" / "kernels" / "utils.py").read_text(encoding = "utf-8")) + bound = { + node.attr + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Attribute) + and node.value.attr == "lib" + } + probe = _load_probe() + xpu = set(probe.bitsandbytes_symbols("xpu")) + cuda = set(probe.bitsandbytes_symbols("cuda")) + assert bound == xpu | cuda, f"probe and module-scope binds differ: {bound ^ (xpu | cuda)}" + # xpu probes the gemv pair, every other device the naive gemm pair, never both. + assert xpu - cuda and cuda - xpu, "the device split collapsed" + + +def test_a_deferred_failure_handle_is_not_ready(): + probe = _load_probe() + bnb = _fake_bnb(_DeferredFailureLib()) + for device in ("cuda", "xpu"): + assert probe.native_kernels_ready(bnb, device) is False, device + + +def test_a_real_ctypes_handle_is_ready(): + probe = _load_probe() + bnb = _fake_bnb(_RealHandleLib()) + for device in ("cuda", "xpu"): + assert probe.native_kernels_ready(bnb, device) is True, device + + +def test_a_lib_that_never_loaded_is_not_ready(): + """bitsandbytes 0.45.5, the floor in pyproject.toml, sets ``functional.lib = None``.""" + probe = _load_probe() + assert probe.native_kernels_ready(_fake_bnb(None), "cuda") is False + + +def test_a_partially_exporting_library_is_not_ready(): + """One resolvable symbol is not enough: the same verdict gates the module-scope + binds, so a partial library would pass here and raise `AttributeError` at the bind.""" + + class _MissingOne(_RealHandleLib): + def __getattr__(self, name): + if name == "cgemm_4bit_inference_naive_bf16": + raise AttributeError(name) + return super().__getattr__(name) + + probe = _load_probe() + assert probe.native_kernels_ready(_fake_bnb(_MissingOne()), "cuda") is False + + +def test_one_dead_handle_among_live_ones_is_not_ready(): + """The realistic partial shape: the library loaded but one symbol is a closure.""" + + class _OneDeferred(_RealHandleLib): + def __getattr__(self, name): + if name == "cdequantize_blockwise_bf16_nf4": + return lambda *a, **k: None + return super().__getattr__(name) + + probe = _load_probe() + assert probe.native_kernels_ready(_fake_bnb(_OneDeferred()), "cuda") is False + + +def test_absent_bitsandbytes_is_not_ready(): + probe = _load_probe() + assert probe.native_kernels_ready(None, "cuda") is False + + +def test_device_type_gates_the_flags_on_the_kernels(): + """The flags must follow ``native_kernels_ready``, not the bare import.""" + head = (REPO_ROOT / "unsloth" / "device_type.py").read_text(encoding = "utf-8") + head = head.split('if DEVICE_TYPE == "hip":')[0] + assert "import bitsandbytes as _bnb_probe" in head + assert 'find_spec("bitsandbytes")' not in head, "find_spec cannot see a broken wheel" + assert "native_kernels_ready(_bnb_probe, DEVICE_TYPE)" in head + assert ( + head.count("ALLOW_BITSANDBYTES = False") >= 2 + ), "both the failed-import path and the dead-kernels path must clear the flag" + + +def test_the_ctypes_binds_are_gated_on_the_same_verdict(): + """Clearing the flag is not enough on its own: ``bnb is None`` alone let an + importable-but-dead wheel reach the binds, and 0.45.5 sets ``functional.lib = None`` + on a native-load failure, so they killed ``import unsloth`` outright instead of + degrading to 16bit.""" + source = (REPO_ROOT / "unsloth" / "kernels" / "utils.py").read_text(encoding = "utf-8") + assert "from ..bnb_availability import native_kernels_ready" in source + assert ( + "if bnb is None or not native_kernels_ready(bnb, DEVICE_TYPE):" in source + ), "the ctypes bind block must take the _bnb_required branch on a dead library too" + guarded = source.split("if bnb is None or not native_kernels_ready(bnb, DEVICE_TYPE):")[1] + assert "bnb.functional.lib" in guarded, "the binds must sit under that guard" + + +def test_the_kernel_check_reads_the_submodule_not_the_parent_attribute(): + """A part-initialised bitsandbytes leaves the parent without ``functional`` while + the submodule stays in sys.modules, which ``import bitsandbytes.functional`` reads + directly.""" + probe = _load_probe() + bnb = types.ModuleType("bitsandbytes") # zombie: parent has no `functional` + bnb.__version__ = "0.50.0" + import sys + + real = sys.modules.get("bitsandbytes.functional") + if real is None: + return # bitsandbytes not importable here; the fallback has nothing to read + # Falls back to the cached submodule instead of raising on the missing attribute. + assert probe.native_kernels_ready(bnb, "cuda") in (True, False) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 682f3ae6c6..7e8f9ced46 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -303,13 +303,19 @@ if DEVICE_TYPE == "cuda": # Try loading bitsandbytes and triton try: import bitsandbytes as bnb + + # Bind the submodule by name: a half-imported bitsandbytes leaves the parent + # without a `functional` attribute, which would otherwise be misreported below + # as a CUDA linking failure. See unsloth/kernels/utils.py. + import bitsandbytes.functional as bnb_functional except: print( "Unsloth: `bitsandbytes` is not installed - 4bit QLoRA unallowed, but 16bit and full finetuning works!" ) bnb = None + bnb_functional = None try: - cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 + cdequantize_blockwise_fp32 = bnb_functional.lib.cdequantize_blockwise_fp32 libcuda_dirs() except: if hasattr(os, "geteuid") and os.geteuid() == 0: @@ -351,7 +357,7 @@ if DEVICE_TYPE == "cuda": pass else: from triton.common.build import libcuda_dirs - cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 + cdequantize_blockwise_fp32 = bnb_functional.lib.cdequantize_blockwise_fp32 libcuda_dirs() except: warnings.warn( diff --git a/unsloth/bnb_availability.py b/unsloth/bnb_availability.py new file mode 100644 index 0000000000..9d14bbb0f3 --- /dev/null +++ b/unsloth/bnb_availability.py @@ -0,0 +1,96 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Can bitsandbytes actually run a 4bit kernel here? A successful import does not say. + +From 0.46 a wheel whose native library never loaded still imports and hands back a +`throw_on_call` closure for every symbol, so attribute reads alone see a healthy wheel, +`ALLOW_BITSANDBYTES` stays true and 4bit dies inside a kernel instead of falling back to +16bit up front. A real handle is a ctypes function pointer and carries `restype`; a +deferred failure is a plain Python function and does not. That is the whole test, applied +to every probed handle: the same verdict gates the module-scope binds in kernels/utils.py, +where one bad symbol is the crash this exists to prevent. + +Decides the capability flags only, never importability - a CPU-only install is exactly +this shape and its Python side works. A leaf module: imports nothing from unsloth +(device_type.py imports it very early, so anything else is a cycle) and takes the +device type as an argument. +""" + +__all__ = [ + "bitsandbytes_symbols", + "check_native_kernels", + "native_kernels_ready", +] + +# The ctypes handles kernels/utils.py binds at module scope; a test asserts they match. +_C_SYMBOLS = ( + "cdequantize_blockwise_fp32", + "cdequantize_blockwise_fp16_nf4", + "cdequantize_blockwise_bf16_nf4", +) +# 4bit inference is a gemv on xpu and a naive gemm elsewhere; probing the wrong pair +# would write off a perfectly good wheel. +_C_SYMBOLS_XPU = ( + "cgemv_4bit_inference_fp16", + "cgemv_4bit_inference_bf16", +) +_C_SYMBOLS_GEMM = ( + "cgemm_4bit_inference_naive_fp16", + "cgemm_4bit_inference_naive_bf16", +) + + +def bitsandbytes_symbols(device_type): + """Names kernels/utils.py reads off `bitsandbytes.functional.lib`.""" + tail = _C_SYMBOLS_XPU if device_type == "xpu" else _C_SYMBOLS_GEMM + return _C_SYMBOLS + tail + + +def check_native_kernels(bnb, device_type): + """Raise unless every handle kernels/utils.py is about to bind is a real kernel. + + All of them: one that resolves here but not at the bind gives back the AttributeError + this prevents. Partial export costs 8bit too (`ALLOW_BITSANDBYTES` gates both), but a + wheel missing a symbol is a shape no flag makes safe. Safe to repeat - ctypes caches + each handle on first lookup, so these are the ones bound later. + """ + if bnb is None: + raise ImportError("Unsloth: `bitsandbytes` is not installed.") + functional = getattr(bnb, "functional", None) + if functional is None: + # A part-initialised bitsandbytes leaves the parent without the attribute while + # the submodule stays in sys.modules, which `import x.y as z` reads directly. + import bitsandbytes.functional as functional + + lib = functional.lib + if lib is None: + # 0.45.5, the floor in pyproject.toml, on a native-load failure. + raise AttributeError("Unsloth: `bitsandbytes.functional.lib` is None.") + for symbol in bitsandbytes_symbols(device_type): + handle = getattr(lib, symbol) # AttributeError here is itself a failed check + if not hasattr(handle, "restype"): + raise AttributeError( + f"Unsloth: `bitsandbytes.functional.lib.{symbol}` is not a native " + "function pointer - the bitsandbytes native library did not load." + ) + + +def native_kernels_ready(bnb, device_type): + """Is the bitsandbytes native library alive? Gates the flags, never the import.""" + try: + check_native_kernels(bnb, device_type) + except Exception: + return False + return True diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 058e166b08..968062c7c1 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -27,6 +27,7 @@ import functools import inspect import os from unsloth_zoo.utils import Version +from .bnb_availability import native_kernels_ready def is_mlx_available(): @@ -117,17 +118,20 @@ DEVICE_COUNT: int = get_device_count() ALLOW_PREQUANTIZED_MODELS: bool = True # HSA_STATUS_ERROR_EXCEPTION checks - sometimes AMD fails for BnB ALLOW_BITSANDBYTES: bool = True -# Unusable bitsandbytes on any backend, not just hip: clear the flags the loader -# reads before it selects a 4bit checkpoint. Same guarded import the fallbacks in -# _gpu_init.py and kernels/utils.py use rather than a find_spec probe, so an -# installed-but-broken wheel (missing .so, wrong ROCm/CUDA build) is treated as -# unavailable by all three, not only by the ones that import it. +# Unusable bitsandbytes on any backend, not just hip: clear the flags the loader reads +# before it picks a 4bit checkpoint. A guarded import, not find_spec, since importable +# is not usable - from 0.46 a dead native library still resolves every ctypes handle to +# a closure that raises only when called, so 4bit would die mid-run, not fall back here. try: import bitsandbytes as _bnb_probe - del _bnb_probe except Exception: ALLOW_PREQUANTIZED_MODELS = False ALLOW_BITSANDBYTES = False +else: + if not native_kernels_ready(_bnb_probe, DEVICE_TYPE): + ALLOW_PREQUANTIZED_MODELS = False + ALLOW_BITSANDBYTES = False + del _bnb_probe # gfx906 (MI50 / Radeon VII / Vega 20): Dynamo/Inductor codegen is broken on this # legacy GCN arch (ROCm dropped it after 6.3) - compiled graphs crash or miscompile # while the eager path trains fine. Default compile off; setdefault so a user diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 2118e65aef..839eb9db84 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -29,6 +29,7 @@ from ..device_type import ( DEVICE_COUNT, ALLOW_PREQUANTIZED_MODELS, ) +from ..bnb_availability import native_kernels_ready from .fp8 import weight_dequant, fp8_linear import functools @@ -135,11 +136,18 @@ def calculate_settings( HAS_CUDA_STREAM = False try: import bitsandbytes as bnb + + # If an earlier `import bitsandbytes` died inside __init__, CPython evicts only + # the parent from sys.modules and keeps its submodules, so this retry re-executes + # __init__ without rebinding `bnb.functional`. `import x.y as z` reads sys.modules + # directly and survives that, plain attribute access does not. + import bitsandbytes.functional as bnb_functional except Exception: # device_type.py already degrades to 16bit/full finetuning when bnb is missing # (e.g. gfx906, whose generic wheel has no kernels). Keep the import working and # fail only if a 4bit path is actually entered. bnb = None + bnb_functional = None def _bnb_required(*args, **kwargs): @@ -152,7 +160,7 @@ def _bnb_required(*args, **kwargs): if bnb is not None: # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3") - get_ptr = bnb.functional.get_ptr + get_ptr = bnb_functional.get_ptr else: get_ptr = _bnb_required @@ -252,25 +260,28 @@ else: # Bitsandbytes operations ctypes_c_int = ctypes.c_int ctypes_c_int32 = ctypes.c_int32 -if bnb is None: +# Same verdict device_type.py used to clear ALLOW_BITSANDBYTES, applied to the binds +# themselves. 0.45.5 leaves `functional.lib = None` when the native library fails to +# load, so these lookups would kill `import unsloth` instead of degrading to 16bit. +if bnb is None or not native_kernels_ready(bnb, DEVICE_TYPE): cdequantize_blockwise_fp32 = _bnb_required cdequantize_blockwise_fp16_nf4 = _bnb_required cdequantize_blockwise_bf16_nf4 = _bnb_required cgemm_4bit_inference_naive_fp16 = _bnb_required cgemm_4bit_inference_naive_bf16 = _bnb_required else: - cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 - cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4 - cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 + cdequantize_blockwise_fp32 = bnb_functional.lib.cdequantize_blockwise_fp32 + cdequantize_blockwise_fp16_nf4 = bnb_functional.lib.cdequantize_blockwise_fp16_nf4 + cdequantize_blockwise_bf16_nf4 = bnb_functional.lib.cdequantize_blockwise_bf16_nf4 if DEVICE_TYPE == "xpu": # https://github.com/bitsandbytes-foundation/bitsandbytes/blob/c3b8de268fdb55a88f92feada23fc811a1e6877a/bitsandbytes/backends/xpu/ops.py#L115 # for xpu, inference gemv using above link - cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemv_4bit_inference_fp16 - cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemv_4bit_inference_bf16 + cgemm_4bit_inference_naive_fp16 = bnb_functional.lib.cgemv_4bit_inference_fp16 + cgemm_4bit_inference_naive_bf16 = bnb_functional.lib.cgemv_4bit_inference_bf16 else: - cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 - cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 + cgemm_4bit_inference_naive_fp16 = bnb_functional.lib.cgemm_4bit_inference_naive_fp16 + cgemm_4bit_inference_naive_bf16 = bnb_functional.lib.cgemm_4bit_inference_naive_bf16 torch_device_stream = (