Clear ALLOW_BITSANDBYTES when the bitsandbytes native kernels are not real (#7578)
* Clear ALLOW_BITSANDBYTES when the bitsandbytes native kernels are not real From bitsandbytes 0.46 a wheel whose native library never loaded still imports and resolves every ctypes handle: BNBNativeLibrary.__getattr__ returns a throw_on_call closure, and a dead library is replaced wholesale by ErrorHandlerMockBNBNativeLibrary, which does the same for every name. Nothing raises while kernels/utils.py binds them at module scope, so device_type.py's guarded import sees a healthy wheel, ALLOW_BITSANDBYTES stays true, loader.py forwards the default load_in_4bit=True and the run dies inside a kernel instead of degrading to 16bit. Probe the handles the kernels actually bind and clear the flags when they are not native. A real handle is a ctypes function pointer and carries restype; a deferred failure is a Python function and does not. Scoped to the capability flags on purpose. The module stays bound and get_ptr keeps pointing at bitsandbytes, because these shapes import perfectly well and treating them as absent would disable a wheel whose Python side works - a CPU-only install is exactly that shape. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Only clear the flags when the native library is dead, not partially exporting ALLOW_BITSANDBYTES gates 8bit as well as 4bit - loader.py:505-510 clears both - so failing the check on one missing 4bit symbol would silently downgrade an otherwise valid LLM.int8 request to 16bit. A library that exports some of these handles is alive; only one where none of them is a ctypes function pointer is dead, which is the CPU-only and ErrorHandlerMockBNBNativeLibrary case this exists for. A genuinely missing symbol raises where kernels/utils.py binds it, so it is a crash no capability flag can rescue and not something to trade 8bit for. * Gate the bitsandbytes ctypes binds on the same verdict as the flags Clearing ALLOW_BITSANDBYTES is not enough on its own. kernels/utils.py guarded the bnb.functional.lib.* binds on `bnb is None` alone, so an importable but dead wheel still reached them at module scope: bitsandbytes 0.45.5, the floor in pyproject.toml, sets functional.lib = None when the native library fails to load, and None.cdequantize_blockwise_fp32 raises right there. That kills import unsloth outright instead of degrading to 16bit, which is the fallback the cleared flag exists to reach. Reuse native_kernels_ready so the bind path and the flag path agree, and take the _bnb_required branch when they say the library is dead. Touches only the guard expression, not the binds themselves. * Tighten the comments on the bitsandbytes kernel readiness probe * Require every probed handle, and license the module Apache like the rest of unsloth The readiness verdict now gates the module-scope ctypes binds as well as the flags, so "at least one handle is native" is no longer the right question. A library that resolves one symbol and not another passed the probe and then raised AttributeError at the bind the probe exists to prevent. Require all of them. That costs 8bit in the partial case, since ALLOW_BITSANDBYTES gates both, but a wheel missing a symbol is a shape no flag can make safe and refusing it beats crashing on it. Flipped the test that encoded the old behaviour and added the more realistic shape: the library loaded, one symbol is still a deferred-failure closure. LICENSE:190 assigns files under unsloth/* to Apache 2.0, and 87 of the 90 modules there carry that header, so use it here rather than AGPL. * State the all-handles rule once instead of three times --------- Co-authored-by: danielhanchen <unslothai@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
bc07d3a2df
commit
f44379d9e8
4 changed files with 285 additions and 7 deletions
174
tests/python/test_bitsandbytes_kernel_readiness.py
Normal file
174
tests/python/test_bitsandbytes_kernel_readiness.py
Normal file
|
|
@ -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)
|
||||
96
unsloth/bnb_availability.py
Normal file
96
unsloth/bnb_availability.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -252,7 +253,10 @@ 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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue