Merge remote-tracking branch 'origin/main' into r7552

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)

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(

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

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

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 = (