diff --git a/studio/backend/main.py b/studio/backend/main.py index 0bca224f89..db7467dcb8 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1832,7 +1832,9 @@ async def get_gpu_visibility(current_subject: str = Depends(get_current_subject) @app.get("/api/system/hardware") def get_hardware_info( - include_details: bool = Query(False), current_subject: str = Depends(get_current_subject) + include_details: bool = Query(False), + include_accelerators: bool = Query(False), + current_subject: str = Depends(get_current_subject), ): """Return GPU name, total VRAM, and key ML package versions. @@ -1847,6 +1849,7 @@ def get_hardware_info( from utils.hardware import ( get_gpu_summary, get_package_versions, + get_accelerator_report, export_capability, video_capability, ) @@ -1871,6 +1874,14 @@ def get_hardware_info( for d in sorted(devices, key = lambda d: d.get("visible_ordinal", 0)) ] body["llama_cpp"] = get_installed_llama_version() + if include_accelerators: + # Whether the optimized kernels (xformers / flash-attn / torchao / bitsandbytes) + # are installed, import, and actually load. Behind its own flag rather than + # include_details: the detail path is also read by Export, Video and onboarding, + # and this one spawns an interpreter to do the imports somewhere they cannot hurt + # anything. Cached per process, so only the first Settings open pays. Additive: + # no existing key changes shape, and older clients ignore it. + body["accelerators"] = get_accelerator_report() return body diff --git a/studio/backend/tests/test_accelerator_probe_kernels.py b/studio/backend/tests/test_accelerator_probe_kernels.py new file mode 100644 index 0000000000..6a693c4d26 --- /dev/null +++ b/studio/backend/tests/test_accelerator_probe_kernels.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The probe child must answer "does this actually work", not "did it import". + +Three packages can import cleanly while the thing they exist for is absent, and all three +rendered as "Working" in Settings: + +* bitsandbytes >= 0.46 hands back a ``throw_on_call`` closure for every ctypes symbol when + its native library did not load, so 4-bit dies mid-run instead of falling back; +* torchao imports with its C++ extension "cleanly skipped" on a torch it has no build for + (the Python stack pins 0.17.0 on torch 2.10+cu130 for exactly that reason), leaving no + optimized quantization kernels; +* xformers loads its library on a GPU it ships no attention kernel for, and every call is + capability-rejected back to SDPA -- the degraded state this whole report exists to name. + +Hermetic: the packages are stubbed through ``sys.modules``, so none of this needs a GPU, a +real xformers, or a working bitsandbytes. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +import utils.hardware.accelerator_probe as probe + + +@pytest.fixture(autouse = True) +def _no_real_capability(monkeypatch): + # Never shell out to nvidia-smi from a unit test; each test says what the host is. + monkeypatch.delenv("UNSLOTH_PROBE_DEVICE_CC", raising = False) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ()) + + +def _op( + *, + minimum = None, + maximum = None, + operator = object(), +): + op = types.SimpleNamespace(OPERATOR = operator) + if minimum is not None: + op.CUDA_MINIMUM_COMPUTE_CAPABILITY = minimum + if maximum is not None: + op.CUDA_MAXIMUM_COMPUTE_CAPABILITY = maximum + return op + + +def _install_xformers(monkeypatch, ops): + xformers = types.ModuleType("xformers") + cpp_lib = types.ModuleType("xformers._cpp_lib") + cpp_lib._cpp_library_load_exception = None + ops_mod = types.ModuleType("xformers.ops") + fmha = types.ModuleType("xformers.ops.fmha") + fmha.ALL_FW_OPS = ops + ops_mod.fmha = fmha + xformers._cpp_lib = cpp_lib + xformers.ops = ops_mod + for name, module in ( + ("xformers", xformers), + ("xformers._cpp_lib", cpp_lib), + ("xformers.ops", ops_mod), + ("xformers.ops.fmha", fmha), + ): + monkeypatch.setitem(sys.modules, name, module) + + +def test_a_loaded_xformers_with_no_kernel_for_this_gpu_is_not_working(monkeypatch): + # sm_120 (RTX 50-series) against a build whose ops all cap at sm_90: the library loads, + # _cpp_library_load_exception is None, and attention silently runs on SDPA. + _install_xformers(monkeypatch, [_op(maximum = (9, 0)), _op(minimum = (10, 0), operator = None)]) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((12, 0),)) + + entry = probe.probe_xformers() + assert entry["imports"] is True + assert entry["runs"] is False + assert "no memory-efficient attention kernel" in entry["error"] + assert "12.0" in entry["error"] + + +def test_a_kernel_that_covers_this_gpu_is_unknown_not_working(monkeypatch): + """The op table admitting this GPU is not evidence the build ships a kernel image for it. + + CUDA_MINIMUM/MAXIMUM_COMPUTE_CAPABILITY are class constants describing what the op + supports in principle. A source build with TORCH_CUDA_ARCH_LIST set for other + architectures, or a wheel that dropped one, registers the very same op and fails the + first launch with "no kernel image is available". Establishing coverage needs a launch, + which this child does not do, so the honest answer is the one probe_flash_attn gives.""" + _install_xformers(monkeypatch, [_op(maximum = (9, 0)), _op(minimum = (10, 0))]) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((12, 0),)) + + entry = probe.probe_xformers() + assert entry["runs"] is None + assert "no kernel was launched" in entry["error"] + + +def test_an_unknown_capability_leaves_the_load_status_alone(monkeypatch): + # No nvidia-smi, no answer. "Cannot be checked" must not become "broken" -- nor "Working". + _install_xformers(monkeypatch, [_op(maximum = (9, 0))]) + + entry = probe.probe_xformers() + assert entry["runs"] is None + assert "compute capability could not be read" in entry["error"] + + +def test_an_unrecognised_op_table_leaves_the_load_status_alone(monkeypatch): + # A future xformers that renames ALL_FW_OPS must not be reported as broken. + _install_xformers(monkeypatch, []) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((12, 0),)) + + entry = probe.probe_xformers() + assert entry["runs"] is None + assert "could not be enumerated" in entry["error"] + + +def test_a_cpp_lib_that_raises_on_import_is_broken_not_unknown(monkeypatch): + """The parent adds a package to `degraded` on imports=False or runs=False only, so a + native load error raised by _cpp_lib itself showed as "Not checked" with no banner -- + on precisely the corrupt install this report exists to name.""" + import builtins + + _install_xformers(monkeypatch, [_op(minimum = (7, 0))]) + real_import = builtins.__import__ + + def _raise_for_cpp_lib(name, *args, **kwargs): + if name == "xformers._cpp_lib" or (name == "xformers" and "_cpp_lib" in (args[2] or ())): + raise OSError("libc10.so: undefined symbol") + return real_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "xformers._cpp_lib") + monkeypatch.setattr(builtins, "__import__", _raise_for_cpp_lib) + entry = probe.probe_xformers() + assert entry["imports"] is True + assert entry["runs"] is False + assert "undefined symbol" in entry["error"] + + +def test_a_layout_without_a_cpp_lib_at_all_is_still_unknown(monkeypatch): + # A future rename is not a dead install, so ModuleNotFoundError stays unknown. + import builtins + + _install_xformers(monkeypatch, [_op(minimum = (7, 0))]) + real_import = builtins.__import__ + + def _missing(name, *args, **kwargs): + if name == "xformers._cpp_lib" or (name == "xformers" and "_cpp_lib" in (args[2] or ())): + raise ModuleNotFoundError("No module named 'xformers._cpp_lib'") + return real_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "xformers._cpp_lib") + monkeypatch.setattr(builtins, "__import__", _missing) + assert probe.probe_xformers()["runs"] is None + + +def test_an_unresolvable_mask_does_not_fall_back_to_the_whole_box(monkeypatch): + """The parent answers None when it cannot resolve a numeric mask, and used to serialize + that as an empty override -- which this child reads as no override at all and answers + from nvidia-smi over every physical GPU, the exact verdict the parent declined to give.""" + import shutil as _shutil + import subprocess as _subprocess + + monkeypatch.setenv("UNSLOTH_PROBE_DEVICE_CC", probe._CC_UNKNOWN) + monkeypatch.setattr( + _shutil, "which", lambda name: pytest.fail("the child must not look for nvidia-smi") + ) + monkeypatch.setattr( + _subprocess, "run", lambda *a, **k: pytest.fail("the child must not ask nvidia-smi") + ) + assert probe._device_compute_capabilities() == () + + +def test_a_failed_library_load_still_wins(monkeypatch): + _install_xformers(monkeypatch, [_op(minimum = (7, 0))]) + sys.modules["xformers._cpp_lib"]._cpp_library_load_exception = OSError("undefined symbol") + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((12, 0),)) + + entry = probe.probe_xformers() + assert entry["runs"] is False and "undefined symbol" in entry["error"] + + +def test_every_visible_gpu_has_to_have_a_kernel(monkeypatch): + """CUDA_VISIBLE_DEVICES=0,1 across a mixed pair. The rank that lands on the sm_120 card + falls back to SDPA whatever the sm_90 card can do, so a verdict taken from the first + visible GPU is the same false all-clear in a smaller box.""" + _install_xformers(monkeypatch, [_op(maximum = (9, 0))]) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((9, 0), (12, 0))) + + entry = probe.probe_xformers() + assert entry["runs"] is False + assert "12.0" in entry["error"], "the report must name the GPU that is not covered" + + +def test_a_pair_the_build_covers_is_unknown_not_working(monkeypatch): + # Covering both cards clears the proven-broken verdict; it does not earn "Working". + _install_xformers(monkeypatch, [_op(minimum = (7, 0))]) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((9, 0), (12, 0))) + assert probe.probe_xformers()["runs"] is None + + +@pytest.mark.parametrize("capabilities", [((12, 0),), ((9, 0),), ((8, 6),), ()]) +def test_flash_attn_kernel_coverage_is_unknown_without_a_launch(monkeypatch, capabilities): + """Importing the extension is not the question, on any card. + + A build with no cubin or PTX image for this architecture imports fine and fails its + first launch with "no kernel image is available" -- true of source builds and of older + wheels, not only of the sm_100+ cards our installer refuses to fetch a wheel for. + flash-attn exposes no list of the architectures it was compiled for, and this child + never launches a kernel, so nothing here can establish support.""" + monkeypatch.setitem(sys.modules, "flash_attn", types.ModuleType("flash_attn")) + monkeypatch.setitem(sys.modules, "flash_attn.flash_attn_interface", types.ModuleType("iface")) + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: capabilities) + + entry = probe.probe_flash_attn() + assert entry["imports"] is True + assert entry["runs"] is None, "unknown, not a verdict either way" + assert "without launching one" in entry["error"] + + +def test_a_flash_attn_that_cannot_import_is_still_broken(monkeypatch): + # Unknown is for "cannot be established", never for "it raised". A None entry in + # sys.modules is how the import system spells a submodule that will not load. + monkeypatch.setitem(sys.modules, "flash_attn", types.ModuleType("flash_attn")) + monkeypatch.setitem(sys.modules, "flash_attn.flash_attn_interface", None) + + entry = probe.probe_flash_attn() + assert entry["imports"] is True + assert entry["runs"] is False and entry["error"] + + +def test_the_capability_is_parsed_from_the_override(monkeypatch): + # Reload past the autouse stub: this one is about the real reader. + import importlib + fresh = importlib.reload(probe) + try: + monkeypatch.setenv("UNSLOTH_PROBE_DEVICE_CC", "12.0") + assert fresh._device_compute_capabilities() == ((12, 0),) + # The parent sends every visible one, comma separated; the child cannot re-resolve + # the mask because the parent cleared it. + monkeypatch.setenv("UNSLOTH_PROBE_DEVICE_CC", "9.0,12.0") + assert fresh._device_compute_capabilities() == ((9, 0), (12, 0)) + monkeypatch.setenv("UNSLOTH_PROBE_DEVICE_CC", "not a capability") + assert fresh._device_compute_capabilities() == () + finally: + importlib.reload(probe) + + +def test_bitsandbytes_with_dead_native_handles_is_not_working(monkeypatch): + # The 0.46+ shape: every symbol resolves to a plain closure that raises when called, + # so attribute reads see a healthy wheel and 4-bit dies inside a kernel. + def throw_on_call(*args, **kwargs): + raise RuntimeError("native library not loaded") + + lib = types.SimpleNamespace( + cdequantize_blockwise_fp32 = throw_on_call, + cdequantize_blockwise_fp16_nf4 = throw_on_call, + cdequantize_blockwise_bf16_nf4 = throw_on_call, + cgemm_4bit_inference_naive_fp16 = throw_on_call, + cgemm_4bit_inference_naive_bf16 = throw_on_call, + ) + bnb = types.ModuleType("bitsandbytes") + functional = types.ModuleType("bitsandbytes.functional") + functional.lib = lib + bnb.functional = functional + monkeypatch.setitem(sys.modules, "bitsandbytes", bnb) + monkeypatch.setitem(sys.modules, "bitsandbytes.functional", functional) + + entry = probe.probe_bitsandbytes() + assert entry["imports"] is True + assert entry["runs"] is False + assert "native library did not load" in entry["error"] + + +def test_bitsandbytes_with_real_handles_is_working(monkeypatch): + real = types.SimpleNamespace(restype = None) + lib = types.SimpleNamespace( + cdequantize_blockwise_fp32 = real, + cdequantize_blockwise_fp16_nf4 = real, + cdequantize_blockwise_bf16_nf4 = real, + cgemm_4bit_inference_naive_fp16 = real, + cgemm_4bit_inference_naive_bf16 = real, + ) + bnb = types.ModuleType("bitsandbytes") + functional = types.ModuleType("bitsandbytes.functional") + functional.lib = lib + bnb.functional = functional + monkeypatch.setitem(sys.modules, "bitsandbytes", bnb) + monkeypatch.setitem(sys.modules, "bitsandbytes.functional", functional) + + entry = probe.probe_bitsandbytes() + assert entry["runs"] is True and entry["error"] is None + + +def test_the_bitsandbytes_check_is_the_one_the_loader_gates_on(): + # Loaded by path, not reimplemented: the report must not be able to say "Working" + # about a wheel unsloth's own loader has already written off. + module = probe._load_bnb_availability() + assert module is not None + assert hasattr(module, "check_native_kernels") + + +def test_torchao_without_native_operators_is_unknown_not_broken(monkeypatch): + """The dispatcher's table, NOT dir(torch.ops.torchao): touching that attribute creates an + empty _OpNamespace whose dir() already lists __name__, __spec__ and friends, so the + no-native-operators case -- the one this exists to catch -- read as healthy. + + And the verdict is UNKNOWN, not broken: the managed stack pins torchao 0.17 on torch + 2.10+cu130 knowing its extension is cleanly skipped, and torchao keeps working through + its Python fallbacks. Calling that degraded would put a destructive banner on the + standard configuration.""" + monkeypatch.setitem(sys.modules, "torchao", types.ModuleType("torchao")) + monkeypatch.setattr(probe, "_registered_ops", lambda namespace: 0) + + entry = probe.probe_torchao() + assert entry["imports"] is True + assert entry["runs"] is None + assert "registered no native operators" in entry["error"] + + +def test_torchao_with_native_operators_is_working(monkeypatch): + monkeypatch.setitem(sys.modules, "torchao", types.ModuleType("torchao")) + monkeypatch.setattr(probe, "_registered_ops", lambda namespace: 12) + + entry = probe.probe_torchao() + assert entry["runs"] is True and entry["error"] is None + + +def test_an_empty_ops_namespace_is_not_mistaken_for_a_loaded_extension(): + """Against the real torch in this environment: torch.ops. materialises on attribute + access, so the count has to come from the dispatcher.""" + pytest.importorskip("torch") + import torch + + assert probe._registered_ops("aten") > 0 + # A namespace nothing has registered under. Reading it through torch.ops would create it. + assert probe._registered_ops("unsloth_no_such_namespace") == 0 + assert len(dir(getattr(torch.ops, "unsloth_no_such_namespace"))) > 0 + + +def test_the_dispatch_table_uses_the_kernel_aware_probes(): + # The parent only ever calls through PROBES; a table entry left on the plain import is + # the same false all-clear with the checks sitting unused next to it. + assert probe.PROBES["bitsandbytes"] is probe.probe_bitsandbytes + assert probe.PROBES["torchao"] is probe.probe_torchao + assert probe.PROBES["xformers"] is probe.probe_xformers + + +def test_the_kernel_verdict_also_applies_on_the_register_extensions_path(monkeypatch): + # The older xformers layout has no _cpp_library_load_exception and is probed by + # re-registering the extensions. Same question, same answer required. + _install_xformers(monkeypatch, [_op(maximum = (9, 0))]) + cpp_lib = sys.modules["xformers._cpp_lib"] + del cpp_lib._cpp_library_load_exception + cpp_lib._register_extensions = lambda: None + monkeypatch.setattr(probe, "_device_compute_capabilities", lambda: ((12, 0),)) + + entry = probe.probe_xformers() + assert entry["runs"] is False and "no memory-efficient attention kernel" in entry["error"] diff --git a/studio/backend/tests/test_hardware_accelerator_report.py b/studio/backend/tests/test_hardware_accelerator_report.py new file mode 100644 index 0000000000..fbe9a9d6fc --- /dev/null +++ b/studio/backend/tests/test_hardware_accelerator_report.py @@ -0,0 +1,754 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the accelerator-stack health report (NVIDIA QA P0-1). + +The managed Windows xformers was built for torch 2.10.0+cu128 and Python 3.10.11 while +the app ran cu130 and Python 3.13.2, so its CUDA extensions never loaded and +memory-efficient attention silently went missing. ``get_package_versions()`` reported +nothing at all about it -- it covered only unsloth/torch/transformers plus torch's CUDA +version, so a mismatched wheel looked identical to a healthy one. + +These tests cover what the report must now answer: what Python is running, which +optimized kernels are installed, which of them actually load, and -- for xformers -- what +the wheel on disk was compiled against. Nothing here needs a GPU or a working xformers. +""" + +import ast +import json +import types +from importlib.metadata import PackageNotFoundError +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +@pytest.fixture(autouse = True) +def _clear_report_cache(): + # The report is cached for the process; a stale cache would leak between tests. + hw._accelerator_report_cache = None + yield + hw._accelerator_report_cache = None + + +@pytest.fixture +def on_accelerator(monkeypatch): + """Pin the host to a plain CUDA box so the probe is considered applicable. + + Without this every probe assertion below silently passes on a CPU-only runner by + never probing at all. + """ + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "IS_ROCM", False) + + +@pytest.fixture +def fake_probe(monkeypatch): + """Replace the probe subprocess. Returns a setter for its result. + + Nothing in this file may run the real probe: it spawns an interpreter that imports + the host's native wheels, which is slow, host-dependent, and on a genuinely broken + wheel can abort the child -- the very case these tests describe. + """ + calls = {"n": 0, "names": None} + + def install(results): + def run( + names, + timeout = 180, + status = None, + keep_cuda = False, + ): + calls["n"] += 1 + calls["names"] = list(names) + return results + + monkeypatch.setattr(hw, "_run_probe_subprocess", run) + return calls + + return install + + +def _cpp_lib( + torch = "2.10.0+cu128", + cuda = 1208, + python = "3.10.11", + hip = None, +): + """The real shape of an xformers wheel's cpp_lib.json (verified against 0.0.34).""" + return { + "version": {"cuda": cuda, "hip": hip, "torch": torch, "python": python}, + "env": {"XFORMERS_PACKAGE_FROM": "wheel-v0.0.34"}, + } + + +# -- get_package_versions stays additive ---------------------------------------------------------- + + +def test_package_versions_keeps_its_existing_keys(): + # Every existing consumer (About tab, training method auto-selection, the ROCm and + # XPU tests) reads these; the new entries must not disturb them. + versions = hw.get_package_versions() + for key in ("unsloth", "torch", "transformers", "cuda", "rocm"): + assert key in versions + # Flat str-or-None, as before: nothing nested crept into this dict. + for key, value in versions.items(): + assert value is None or isinstance(value, str), f"{key} is {type(value)}" + + +def test_package_versions_reports_python_and_the_accelerators(): + versions = hw.get_package_versions() + assert versions["python"] == __import__("platform").python_version() + for key in ("xformers", "flash_attn", "torchao", "bitsandbytes"): + assert key in versions, f"{key} missing from get_package_versions()" + + +def test_flash_attn_key_is_an_identifier(): + # The distribution is "flash-attn" but the JSON key must be usable as a property + # name on the frontend without bracket access. + assert "flash-attn" not in hw.get_package_versions() + + +# -- build metadata ------------------------------------------------------------------------------- + + +def test_built_for_is_read_from_cpp_lib_json(tmp_path, monkeypatch): + import importlib.util + import types + + package = tmp_path / "xformers" + package.mkdir() + (package / "cpp_lib.json").write_text(json.dumps(_cpp_lib()), encoding = "utf-8") + spec = types.SimpleNamespace( + submodule_search_locations = [str(package)], origin = str(package / "__init__.py") + ) + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: spec) + + assert hw._xformers_built_for() == { + "torch": "2.10.0+cu128", + # xformers stores major * 100 + minor; its own message prints the raw 1208. + "cuda": "12.8", + "hip": None, + "python": "3.10.11", + } + + +def test_built_for_is_none_without_xformers(monkeypatch): + import importlib.util + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: None) + assert hw._xformers_built_for() is None + + +def test_break_description_names_both_sides(monkeypatch): + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0+cu130", "cuda": "13.0", "python": "3.13.2"}, + ) + reason = hw._describe_xformers_break( + {"torch": "2.10.0+cu128", "cuda": "12.8", "hip": None, "python": "3.10.11"}, + "OSError: [WinError 126] The specified module could not be found", + ) + # The whole P0 in one line: what it was built for, what is actually running. + assert "2.10.0+cu128" in reason + assert "3.10.11" in reason + assert "2.10.0+cu130" in reason + assert "3.13.2" in reason + + +def test_break_description_falls_back_to_the_raw_error(): + # No cpp_lib.json (source install): the exception is all we have, and saying nothing + # would be worse. + reason = hw._describe_xformers_break(None, "ImportError: undefined symbol: _ZN3c10") + assert "undefined symbol" in reason + + +def test_a_matching_build_is_not_called_a_mismatch(monkeypatch): + # [WinError 126] with a perfectly matching wheel is the most common Windows xformers + # failure that is NOT a version mismatch (missing VC++ runtime or CUDA DLL). Claiming + # a mismatch whenever build metadata merely exists misdiagnoses it AND throws the real + # error away, leaving the user nothing to search for. + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0+cu130", "cuda": "13.0", "python": "3.13.2"}, + ) + reason = hw._describe_xformers_break( + {"torch": "2.10.0+cu130", "cuda": "13.0", "hip": None, "python": "3.10.11"}, + "OSError: [WinError 126] The specified module could not be found", + ) + assert reason == "OSError: [WinError 126] The specified module could not be found" + + +def test_a_python_only_difference_is_not_a_mismatch(monkeypatch): + # The wheels are abi3/none-tagged and _C loads through torch.ops.load_library, not the + # CPython ABI: this very repo runs a 3.10-built xformers on 3.13 with working kernels. + # Naming Python here sends people off to reinstall Python for nothing. + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.9.1+cu128", "cuda": "12.8", "python": "3.13.12"}, + ) + reason = hw._describe_xformers_break( + {"torch": "2.9.1+cu128", "cuda": "12.8", "hip": None, "python": "3.10.19"}, + "OSError: boom", + ) + assert reason == "OSError: boom" + + +def test_a_cuda_minor_difference_is_not_a_mismatch(monkeypatch): + # CUDA minor version compatibility: a cu126-built wheel loads against a cu128 torch. + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0", "cuda": "12.8", "python": "3.13.2"}, + ) + reason = hw._describe_xformers_break( + {"torch": None, "cuda": "12.6", "hip": None, "python": None}, "OSError: boom" + ) + assert reason == "OSError: boom" + + +# -- the report ----------------------------------------------------------------------------------- + + +def test_report_has_the_documented_shape(on_accelerator, fake_probe): + fake_probe( + { + name: {"imports": True, "runs": None, "error": None} + for name, _ in hw._ACCELERATOR_PACKAGES + } + ) + report = hw.get_accelerator_report(refresh = True) + assert report["python_version"] == __import__("platform").python_version() + assert set(report["packages"]) == {"xformers", "flash_attn", "torchao", "bitsandbytes"} + for name, entry in report["packages"].items(): + assert set(entry) >= {"version", "installed", "imports", "runs", "reason"}, name + # installed must follow the metadata, not the probe: that split is the whole + # point (a mismatched wheel is installed, imports, and does not work). + assert entry["installed"] is (entry["version"] is not None), name + assert isinstance(report["degraded"], list) + assert report["torch_version"] == hw._running_torch()["torch"] + # xformers is the one with a wheel-recorded build, so it carries built_for. + assert "built_for" in report["packages"]["xformers"] + + +def test_report_flags_an_installed_but_dead_package(monkeypatch, on_accelerator, fake_probe): + monkeypatch.setattr(hw, "pkg_version", lambda name: "0.0.34") + monkeypatch.setattr(hw, "_xformers_built_for", lambda: _cpp_lib()["version"] | {"cuda": "12.8"}) + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0+cu130", "cuda": "13.0", "python": "3.13.2"}, + ) + fake_probe( + { + "xformers": {"imports": True, "runs": False, "error": "OSError: [WinError 126]"}, + "flash_attn": {"imports": True, "runs": True, "error": None}, + "torchao": {"imports": True, "runs": None, "error": None}, + "bitsandbytes": {"imports": True, "runs": None, "error": None}, + } + ) + + report = hw.get_accelerator_report(refresh = True) + assert report["degraded"] == ["xformers"] + assert report["packages"]["xformers"]["runs"] is False + assert "2.10.0+cu128" in report["packages"]["xformers"]["reason"] + + +def test_only_the_installed_and_applicable_packages_are_probed( + monkeypatch, on_accelerator, fake_probe +): + # Probing something that is not installed wastes an import and can only produce a + # ModuleNotFoundError the metadata already told us about. + def only_xformers(name): + if name == "xformers": + return "0.0.34" + raise PackageNotFoundError(name) + + monkeypatch.setattr(hw, "pkg_version", only_xformers) + calls = fake_probe({"xformers": {"imports": True, "runs": True, "error": None}}) + hw.get_accelerator_report(refresh = True) + assert calls["names"] == ["xformers"] + + +def test_the_probe_is_one_child_per_visibility_group(monkeypatch, on_accelerator): + """The interpreter start and the torch import dominate, so one child per PACKAGE would + cost four times as much for the same answer. Two groups, though, not one: bitsandbytes + picks its native library from torch.cuda.is_available(), so with the GPUs hidden a healthy + CUDA install loads the CPU library and reports dead handles. It gets the app's own mask; + everything else keeps the hidden one, which is what stops a diagnostic taking VRAM from a + run in progress.""" + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.0") + seen: list[tuple] = [] + + def run( + names, + timeout = 180, + status = None, + keep_cuda = False, + ): + seen.append((tuple(names), keep_cuda)) + return {name: {"imports": True, "runs": None, "error": None} for name in names} + + monkeypatch.setattr(hw, "_run_probe_subprocess", run) + hw.get_accelerator_report(refresh = True) + + assert seen == [ + (("xformers", "flash_attn", "torchao"), False), + (("bitsandbytes",), True), + ] + + +def test_absent_package_is_not_degraded(monkeypatch, on_accelerator, fake_probe): + # "Not installed" is a normal configuration, not a broken one. Reporting it as + # degraded would make the UI banner permanent on every machine without flash-attn. + from importlib.metadata import PackageNotFoundError + + def missing(name): + raise PackageNotFoundError(name) + + monkeypatch.setattr(hw, "pkg_version", missing) + fake_probe({}) + report = hw.get_accelerator_report(refresh = True) + assert report["degraded"] == [] + for entry in report["packages"].values(): + assert entry["installed"] is False + assert entry["reason"] is None + + +def test_probe_can_be_skipped(monkeypatch, on_accelerator): + # Escape hatch for an install where importing a broken native wheel is worse than + # not knowing. Versions still report; nothing is claimed to be broken. + monkeypatch.setenv(hw._ACCELERATOR_PROBE_ENV, "1") + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + + def must_not_run(*args, **kwargs): + raise AssertionError("probe ran despite the skip flag") + + monkeypatch.setattr(hw, "_run_probe_subprocess", must_not_run) + + report = hw.get_accelerator_report(refresh = True) + assert report["probed"] is False + assert report["degraded"] == [] + assert report["packages"]["xformers"]["reason"] == "not probed" + + +def test_a_probe_that_cannot_answer_is_unknown_not_broken(monkeypatch, on_accelerator): + # The child timed out, crashed, or printed nothing. That says nothing about the + # packages, so it must not light the banner. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + monkeypatch.setattr( + hw, + "_run_probe_subprocess", + lambda names, timeout = 180, status = None, keep_cuda = False: None, + ) + + report = hw.get_accelerator_report(refresh = True) + assert report["probed"] is False + assert report["degraded"] == [] + assert report["packages"]["torchao"]["reason"] == "could not be checked" + + +def test_a_long_reason_is_capped_before_it_reaches_the_ui(monkeypatch, on_accelerator, fake_probe): + # `reason` lands in a settings row description AND in its aria-label, so an + # unbounded traceback string would be read out in full by a screen reader. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + fake_probe({"torchao": {"imports": False, "runs": None, "error": "x" * 5000}}) + report = hw.get_accelerator_report(refresh = True) + assert len(report["packages"]["torchao"]["reason"]) <= hw._MAX_REASON_CHARS + + +def test_report_is_cached(monkeypatch, on_accelerator, fake_probe): + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + calls = fake_probe( + { + name: {"imports": True, "runs": None, "error": None} + for name, _ in hw._ACCELERATOR_PACKAGES + } + ) + + hw.get_accelerator_report(refresh = True) + first = calls["n"] + hw.get_accelerator_report() + hw.get_accelerator_report() + assert calls["n"] == first, "the report spawned another probe on a cache hit" + + +def test_cached_report_is_not_shared_by_reference(on_accelerator, fake_probe): + # Callers mutate response bodies; the cache must not be editable through them. + fake_probe({"xformers": {"imports": True, "runs": True, "error": None}}) + first = hw.get_accelerator_report(refresh = True) + first["packages"]["xformers"]["version"] = "tampered" + assert hw.get_accelerator_report()["packages"]["xformers"]["version"] != "tampered" + + +# -- endpoint wiring (ast, so it runs without starting the app) ------------------------------------ + + +def test_hardware_endpoint_gates_the_report_behind_its_own_flag(): + source = (_BACKEND / "main.py").read_text(encoding = "utf-8") + node = next( + n + for n in ast.walk(ast.parse(source)) + if isinstance(n, ast.FunctionDef) and n.name == "get_hardware_info" + ) + body = ast.get_source_segment(source, node) + # Its own flag, not include_details: the detail path is also read by Export, Video and + # onboarding, and this one spawns an interpreter. + assert "include_accelerators" in body + before, marker, after = body.partition("if include_accelerators:") + assert marker, "include_accelerators branch went missing" + assert "get_accelerator_report()" in after + assert "get_accelerator_report()" not in before + assert '"accelerators"' in after + assert '"accelerators"' not in before + + +@pytest.mark.parametrize( + "device, is_rocm, label", + [ + (lambda: hw.DeviceType.MLX, False, "Apple Silicon"), + (lambda: hw.DeviceType.CPU, False, "CPU-only"), + ], +) +def test_a_host_these_kernels_cannot_run_on_is_never_degraded(monkeypatch, device, is_rocm, label): + # bitsandbytes is a dependency on every platform but the stock build is CUDA-only, so + # on these hosts its import failure is the design, not a broken acceleration stack. + # Probing there pins a permanent false banner to those installs, which is how a + # warning stops being read. + monkeypatch.setattr(hw, "get_device", device) + monkeypatch.setattr(hw, "IS_ROCM", is_rocm) + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + + def must_not_run(*args, **kwargs): + raise AssertionError(f"probed {label}, which cannot use these kernels") + + monkeypatch.setattr(hw, "_run_probe_subprocess", must_not_run) + + report = hw.get_accelerator_report(refresh = True) + assert report["probed"] is False + assert report["degraded"] == [] + assert report["packages"]["bitsandbytes"]["reason"] == "not used on this device" + + +def test_an_xpu_host_probes_the_one_package_that_runs_there(monkeypatch, fake_probe): + """bitsandbytes has an XPU backend, and the loader gates 4-bit on + native_kernels_ready(..., "xpu") with its own symbol set -- so a broken XPU wheel is + something training will reach, and calling it "not used on this device" suppressed the + warning until the first 4-bit load failed. flash-attn / xformers / torchao publish nothing + for XPU, so they stay out.""" + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.XPU) + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + calls = fake_probe( + {"bitsandbytes": {"imports": True, "runs": False, "error": "AttributeError: ..."}} + ) + + report = hw.get_accelerator_report(refresh = True) + assert calls["names"] == ["bitsandbytes"] + assert report["degraded"] == ["bitsandbytes"] + assert report["packages"]["xformers"]["reason"] == "not used on this device" + # ...and it is checked against the XPU symbol set, not the CUDA one. + assert hw._probe_device_type() == "xpu" + + +def test_a_rocm_host_probes_only_what_it_can_actually_load(monkeypatch, fake_probe): + """ROCm is the subtle one: those hosts report DeviceType.CUDA internally (there is + deliberately no DeviceType.ROCM), so a plain device check waves them through -- but + they are not empty either. Unsloth imports and enables flash-attn under + DEVICE_TYPE == "hip", so a broken ROCm FlashAttention is something training will + actually reach, and calling it "not used on this device" hides the one banner that + explains the crash. bitsandbytes is reached on every backend -- device_type.py imports it + and asks native_kernels_ready(..., DEVICE_TYPE) whether 4-bit loading is allowed -- so a + dead ROCm wheel silently costs quantized loading. xformers and torchao ship CUDA-only + builds, where an import failure is the design.""" + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + calls = fake_probe({"flash_attn": {"imports": False, "runs": None, "error": "OSError: boom"}}) + + report = hw.get_accelerator_report(refresh = True) + # Two children, because flash-attn needs the real device mask and bitsandbytes must not + # latch a CUDA context; the fixture records the last group. + # Two children: bitsandbytes is probed with the device mask intact, flash-attn without. + assert calls["n"] == 2 and calls["names"] == ["bitsandbytes"] + assert report["packages"]["bitsandbytes"]["reason"] != "not used on this device" + assert report["degraded"] == ["flash_attn"] + assert report["packages"]["xformers"]["reason"] == "not used on this device" + assert report["packages"]["torchao"]["reason"] == "not used on this device" + + +_SMI_ROWS = ( + "0, GPU-aaaaaaaa-1111-2222-3333-444444444444, 9.0\n" + "1, GPU-bbbbbbbb-5555-6666-7777-888888888888, 12.0\n" +) + + +def _fake_smi( + monkeypatch, + stdout = _SMI_ROWS, + returncode = 0, +): + monkeypatch.setattr(hw.shutil, "which", lambda name: "/usr/bin/nvidia-smi") + monkeypatch.setattr( + hw.subprocess, + "run", + lambda *a, **kw: types.SimpleNamespace(returncode = returncode, stdout = stdout), + ) + + +@pytest.mark.parametrize( + "mask, expected", + [ + # No mask: the whole box is visible, and one uncovered card in it is still an + # install this report has to call degraded. + (None, "9.0,12.0"), + ("1", "12.0"), + # BOTH, in mask order: a rank landing on the second card falls back to SDPA + # whatever the first can do, so a verdict from device 0 alone is a false all-clear. + ("1,0", "12.0,9.0"), + ("0,1", "9.0,12.0"), + ("GPU-bbbbbbbb-5555-6666-7777-888888888888", "12.0"), + ("GPU-bbbbbbbb", "12.0"), + ("7", None), + # CUDA stops at the first entry it cannot resolve, and so does this. + ("1,nonsense", "12.0"), + ], +) +def test_the_capability_follows_the_mask_the_app_runs_under(monkeypatch, mask, expected): + """nvidia-smi lists PHYSICAL devices and ignores CUDA_VISIBLE_DEVICES. On a mixed host + masked to the sm_120 card, reading row 0 evaluates the sm_90 op table and reports + xFormers as working while the GPU the app can actually use has no kernel -- exactly the + degraded state this report exists to surface. Resolved in the parent because the child + is started with the mask cleared.""" + _fake_smi(monkeypatch) + if mask is None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + else: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", mask) + assert hw._visible_compute_capability() == expected + + +def test_a_busy_gpu_does_not_turn_a_healthy_wheel_red(monkeypatch, on_accelerator, fake_probe): + """bitsandbytes keeps CUDA visible on purpose, so opening Settings while a trainer has + filled the GPU (or holds it in EXCLUSIVE_PROCESS) fails at context creation on a healthy + install -- and this answer is cached for the life of the backend, so the banner would sit + there until restart.""" + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + fake_probe( + { + "bitsandbytes": { + "imports": False, + "runs": None, + "error": "RuntimeError: CUDA error: all CUDA-capable devices are busy or unavailable", + } + } + ) + + report = hw.get_accelerator_report(refresh = True) + + assert "bitsandbytes" not in report["degraded"], "a busy GPU is not a broken wheel" + row = report["packages"]["bitsandbytes"] + assert row["probed"] is False and row["runs"] is None + assert "busy" in (row["reason"] or "").lower(), "and the row says why" + + +def test_a_real_import_failure_is_still_degraded(monkeypatch, on_accelerator, fake_probe): + # The classification must not swallow the case the report exists for. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + fake_probe( + {"bitsandbytes": {"imports": False, "runs": None, "error": "OSError: undefined symbol"}} + ) + + report = hw.get_accelerator_report(refresh = True) + + assert "bitsandbytes" in report["degraded"] + assert report["packages"]["bitsandbytes"]["probed"] is True + + +def test_an_unresolved_capability_reaches_the_child_as_unknown(monkeypatch): + """The child reads an EMPTY UNSLOTH_PROBE_DEVICE_CC as "no override" and falls back to + nvidia-smi over every physical GPU -- so serializing an unresolvable mask as "" handed it + the whole box, which is the verdict the resolver returned None to avoid.""" + seen = {} + + def _capture(argv, **kwargs): + seen.update(kwargs.get("env") or {}) + raise RuntimeError("stop here; the environment is what is under test") + + monkeypatch.setattr(hw, "_visible_compute_capability", lambda: None) + monkeypatch.setattr(hw.subprocess, "run", _capture) + hw._run_probe_subprocess(["xformers"], keep_cuda = False) + + assert seen.get("UNSLOTH_PROBE_DEVICE_CC") == hw._CC_UNKNOWN + assert seen.get("UNSLOTH_PROBE_DEVICE_CC") != "" + + +def test_a_numeric_mask_is_unknown_under_a_non_pci_device_order(monkeypatch): + """A numeric CUDA_VISIBLE_DEVICES entry is a CUDA ordinal in the CURRENT + CUDA_DEVICE_ORDER, while nvidia-smi's index column is the PCI_BUS_ID ordering. The + backend pins PCI_BUS_ID with setdefault, so a user who exported FASTEST_FIRST keeps it, + and there "1" is a speed rank we cannot invert without initialising CUDA. Answering + anyway hands the probe the wrong card.""" + _fake_smi(monkeypatch) + monkeypatch.setenv("CUDA_DEVICE_ORDER", "FASTEST_FIRST") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + assert hw._visible_compute_capability() is None + # A UUID names the same device in any ordering, so it still resolves. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-bbbbbbbb") + assert hw._visible_compute_capability() == "12.0" + # And the whole box is still readable when nothing is masked. + monkeypatch.delenv("CUDA_VISIBLE_DEVICES") + assert hw._visible_compute_capability() == "9.0,12.0" + + +def test_no_nvidia_smi_is_unknown_rather_than_a_guess(monkeypatch): + monkeypatch.setattr(hw.shutil, "which", lambda name: None) + assert hw._visible_compute_capability() is None + _fake_smi(monkeypatch, stdout = "", returncode = 9) + assert hw._visible_compute_capability() is None + + +def test_each_row_says_whether_it_was_probed(monkeypatch, fake_probe): + """Report-wide "probed" cannot answer for a row. The probe set is per device, so on this + ROCm host three of the four packages are installed and deliberately unprobed -- and with + only the global flag the About table rendered all three "Not loading", a red badge for + packages that are fine and that the banner does not even list.""" + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + fake_probe({"flash_attn": {"imports": True, "runs": True, "error": None}}) + + packages = hw.get_accelerator_report(refresh = True)["packages"] + assert packages["flash_attn"]["probed"] is True + for name in ("xformers", "torchao", "bitsandbytes"): + assert packages[name]["installed"] is True + assert packages[name]["probed"] is False + + +def test_a_cpu_host_is_distinguishable_from_an_opted_out_one(monkeypatch): + # Two different unknowns: "these kernels do not apply here" and "you told us not to + # look". Collapsing them would make the About tab lie about one of them. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setenv(hw._ACCELERATOR_PROBE_ENV, "1") + opted_out = hw.get_accelerator_report(refresh = True) + + monkeypatch.delenv(hw._ACCELERATOR_PROBE_ENV) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CPU) + no_gpu = hw.get_accelerator_report(refresh = True) + + assert opted_out["packages"]["torchao"]["reason"] == "not probed" + assert no_gpu["packages"]["torchao"]["reason"] == "not used on this device" + + +def test_a_child_that_dies_is_isolated_rather_than_erasing_the_report(monkeypatch, on_accelerator): + # A native wheel that ABORTS the interpreter instead of raising is a failure mode this + # probe is specifically built around -- and in the batch child it used to take the + # whole answer down with it: every package read "could not be checked", degraded came + # back empty, and Settings showed nothing at all for the one install that cannot load. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + seen = [] + + def run( + names, + timeout = 180, + status = None, + keep_cuda = False, + ): + seen.append(list(names)) + if "bitsandbytes" in names: + # bitsandbytes' child dies without answering. + if status is not None: + status["died"] = True + return None + return {name: {"imports": True, "runs": True, "error": None} for name in names} + + monkeypatch.setattr(hw, "_run_probe_subprocess", run) + + report = hw.get_accelerator_report(refresh = True) + assert report["probed"] is True + # The survivors keep their real answers... + assert report["packages"]["xformers"]["imports"] is True + assert report["packages"]["torchao"]["runs"] is True + # ...and the one that kills its own child is reported broken, not unknown. + assert report["degraded"] == ["bitsandbytes"] + assert report["packages"]["bitsandbytes"]["runs"] is False + assert "exited without answering" in report["packages"]["bitsandbytes"]["reason"] + # The hidden-GPU group answers in one child; the bitsandbytes group is alone and its own + # child already died, so that IS the diagnosis. No retry storm either way. + assert seen == [["xformers", "flash_attn", "torchao"], ["bitsandbytes"]] + + +def test_a_probe_that_never_ran_is_still_unknown(monkeypatch, on_accelerator): + # The other half of the same fence: a timeout or a missing script says nothing about + # any package, and re-probing one at a time would only repeat it. + monkeypatch.setattr(hw, "pkg_version", lambda name: "1.2.3") + calls = {"n": 0} + + def run( + names, + timeout = 180, + status = None, + keep_cuda = False, + ): + calls["n"] += 1 + return None # no status["died"]: the child never got to run + + monkeypatch.setattr(hw, "_run_probe_subprocess", run) + + report = hw.get_accelerator_report(refresh = True) + # One child per visibility group, and neither retried. + assert calls["n"] == 2 + assert report["probed"] is False and report["degraded"] == [] + assert report["packages"]["torchao"]["reason"] == "could not be checked" + + +def test_a_stable_abi_wheel_on_a_later_torch_is_not_called_a_mismatch(monkeypatch): + # xFormers 0.0.34+ targets the PyTorch stable ABI, so a 2.10-built wheel on 2.12 is + # the design. Diagnosing a torch mismatch there discards the REAL error (here a + # missing DLL, the most common Windows failure that is not a version problem) and + # blames a working pair. + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.12.1+cu128", "cuda": "12.8", "python": "3.13.2"}, + ) + error = "OSError: [WinError 126] The specified module could not be found" + reason = hw._describe_xformers_break( + {"torch": "2.10.0+cu128", "cuda": "12.8", "hip": None, "python": "3.10.11"}, error + ) + assert reason == error + + +def test_a_compatible_local_tag_is_not_a_torch_mismatch(monkeypatch): + # Same release, different CUDA MINOR: cu126 loads against cu128 fine, and the raw + # string compare claimed a torch mismatch before the CUDA-major check could speak. + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0+cu128", "cuda": "12.8", "python": "3.13.2"}, + ) + error = "OSError: [WinError 126] The specified module could not be found" + reason = hw._describe_xformers_break( + {"torch": "2.10.0+cu126", "cuda": "12.6", "hip": None, "python": "3.10.11"}, error + ) + assert reason == error + + +def test_the_stable_abi_guarantee_does_not_cover_going_backwards(monkeypatch): + monkeypatch.setattr( + hw, + "_running_torch", + lambda: {"torch": "2.10.0+cu128", "cuda": "12.8", "python": "3.13.2"}, + ) + reason = hw._describe_xformers_break( + {"torch": "2.12.0+cu128", "cuda": "12.8", "hip": None, "python": "3.10.11"}, None + ) + assert reason is not None and "2.12.0+cu128" in reason diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 4d06fb2073..40444cda19 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -18,6 +18,7 @@ from .hardware import ( log_gpu_memory, get_gpu_summary, get_package_versions, + get_accelerator_report, get_gpu_utilization, get_visible_gpu_utilization, get_backend_visible_gpu_info, @@ -95,6 +96,7 @@ __all__ = [ "log_gpu_memory", "get_gpu_summary", "get_package_versions", + "get_accelerator_report", "get_gpu_utilization", "get_visible_gpu_utilization", "get_backend_visible_gpu_info", diff --git a/studio/backend/utils/hardware/accelerator_probe.py b/studio/backend/utils/hardware/accelerator_probe.py new file mode 100644 index 0000000000..dca6ab9fd2 --- /dev/null +++ b/studio/backend/utils/hardware/accelerator_probe.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Import the optimized-kernel packages and report what happened, as a THROWAWAY process. + +Run as a script; prints one JSON object on stdout:: + + {"xformers": {"imports": true, "runs": false, "error": "OSError: ..."}, ...} + +This exists as a separate process, not a function call, for three reasons -- each of them +something that has already bitten this codebase: + +* A package whose ``__init__`` raises leaves every submodule it already executed behind in + ``sys.modules``. The next import re-runs ``__init__`` with those served from cache, so + attributes are never rebound and the package imports "successfully" while missing pieces. + See ``utils/torch_warmup.purge_partial_import`` and unslothai/unsloth#7580. A diagnostic + that poisons the import cache for the warm and for every later request is worse than no + diagnostic. +* ``import bitsandbytes`` creates a CUDA context. The backend deliberately never latches + one -- main.py pins CUDA_DEVICE_ORDER before any torch import, and the export planner + budgets from free VRAM read before a context exists -- so a diagnostic must not + permanently take several hundred MB off every later VRAM reading. +* A genuinely broken native wheel can abort the interpreter rather than raise (pybind11 + answers a duplicate type registration with ``std::terminate``). In a child that is just a + failed probe. In the server it is a dead app, on exactly the broken installs this is + meant to describe. + +Kept dependency-free (stdlib only, no studio imports) so the parent can run it with the +same interpreter and nothing else on the path. +""" + +import json +import os +import sys +from typing import Any, Dict + + +# Mirrors _CC_UNKNOWN in utils/hardware/hardware.py, which sets it. +_CC_UNKNOWN = "unknown" + + +def _error(exc: BaseException) -> str: + return f"{type(exc).__name__}: {exc}" + + +def probe_xformers() -> Dict[str, Any]: + """Does xformers import, and do its C++/CUDA extensions load? + + The extension load is what fails on an ABI mismatch: ``torch.ops.load_library`` raises + OSError, and xformers catches it and downgrades it to a logger warning, which is why + the package imports and reports a version while having no memory-efficient attention. + + Prefer the outcome ``xformers/_cpp_lib.py`` already recorded at its own import over + calling ``_register_extensions()`` again: re-calling it re-runs + ``os.add_dll_directory`` on Windows, and the private name is not guaranteed to exist in + every layout -- treating its absence as "broken" would put a red banner on a working + install. + """ + entry: Dict[str, Any] = {"imports": False, "runs": None, "error": None} + try: + import xformers # noqa: F401 + except BaseException as exc: + entry["error"] = _error(exc) + return entry + entry["imports"] = True + + try: + from xformers import _cpp_lib + except ModuleNotFoundError as exc: + # A layout with no _cpp_lib at all. Unknown, not broken: treating a future rename as a + # dead install would put a red banner on a working one. + entry["error"] = _error(exc) + return entry + except BaseException as exc: + # The module exists and raised: its native load failed, which IS the dead-kernel state + # this report exists to surface. Left as None it read "Not checked", stayed out of + # `degraded`, and showed no banner at all. + entry["runs"] = False + entry["error"] = _error(exc) + return entry + + if hasattr(_cpp_lib, "_cpp_library_load_exception"): + failure = _cpp_lib._cpp_library_load_exception + entry["runs"] = failure is None + if failure is not None: + entry["error"] = _error(failure) + return entry + return _with_kernel_verdict(entry) + + register = getattr(_cpp_lib, "_register_extensions", None) + if register is None: + # An xformers layout we do not recognise. Unknown is not broken. + return entry + try: + register() + entry["runs"] = True + except BaseException as exc: + entry["runs"] = False + entry["error"] = _error(exc) + return entry + return _with_kernel_verdict(entry) + + +def _with_kernel_verdict(entry: Dict[str, Any]) -> Dict[str, Any]: + """Settle a loaded xFormers into "does not run" or "cannot be established". + + The library loading is necessary, not sufficient: a build with no kernel for THIS GPU + (an sm_120 card against a wheel that ships none) loads fine, and every attention call + is then capability-rejected straight back to SDPA. That is precisely the degraded state + this report exists to surface, and it was rendering as "Working". + + Static, not a forward pass: this child runs with no visible GPU on purpose (a probe + must not latch a CUDA context or take VRAM from a run in progress), so the verdict + comes from the op table plus the capability read out of nvidia-smi, which needs no + context. That buys a sound NEGATIVE only: no op admitting this GPU proves attention + falls back, but an op that admits it proves nothing about the images the build was + compiled with. So every other exit here is unknown, never "Working" -- an install that + cannot be checked is neither broken nor verified. + """ + capabilities = _device_compute_capabilities() + if not capabilities: + return _unverified(entry, "this GPU's compute capability could not be read") + try: + from xformers.ops import fmha + except BaseException: + return _unverified(entry, "xformers' attention operators could not be enumerated") + ops = getattr(fmha, "ALL_FW_OPS", None) + if not ops: + return _unverified(entry, "xformers' attention operators could not be enumerated") + # ANY visible GPU without a kernel is a degraded install: with CUDA_VISIBLE_DEVICES=0,1 + # across a mixed pair, the rank that lands on the uncovered card falls back to SDPA + # whatever the other card can do. + for capability in capabilities: + if _has_usable_op(ops, capability): + continue + entry["runs"] = False + entry["error"] = ( + f"xformers loaded but ships no memory-efficient attention kernel for this GPU " + f"(compute capability {capability[0]}.{capability[1]}), so attention falls " + f"back to SDPA" + ) + return entry + # Coverage NOT established, only not-ruled-out. The capability bounds above are class + # constants describing what an op supports in principle; a wheel or source build compiled + # for other architectures still registers that op and still fails the first launch with + # "no kernel image is available". Establishing the compiled architectures needs a launch, + # which this child deliberately does not do (it holds no CUDA context). So the positive + # verdict becomes unknown, the same answer probe_flash_attn gives for the same build. + return _unverified(entry, "its operators admit this GPU, but no kernel was launched") + + +def _unverified(entry: Dict[str, Any], why: str) -> Dict[str, Any]: + """Kernel coverage could not be established. Not broken, not confirmed working.""" + entry["runs"] = None + entry["error"] = ( + f"xformers loaded, but whether its build ships a usable attention kernel for this " + f"GPU is unknown: {why}" + ) + return entry + + +def _has_usable_op(ops, capability) -> bool: + for op in ops: + minimum = getattr(op, "CUDA_MINIMUM_COMPUTE_CAPABILITY", None) + maximum = getattr(op, "CUDA_MAXIMUM_COMPUTE_CAPABILITY", None) + if getattr(op, "OPERATOR", False) is None: + continue # the build did not ship this op at all + if minimum is not None and capability < minimum: + continue + if maximum is not None and capability > maximum: + continue + return True + return False + + +def _device_compute_capability(): + """The FIRST visible compute capability as ``(major, minor)``, or None when unknown.""" + capabilities = _device_compute_capabilities() + return capabilities[0] if capabilities else None + + +def _device_compute_capabilities(): + """Every compute capability this process can use, as ``(major, minor)`` tuples. + + Through nvidia-smi rather than torch: reading it from torch initialises CUDA, and the + whole point of this child is that it never holds a context. UNSLOTH_PROBE_DEVICE_CC + carries the parent's already-mask-resolved answer (comma separated), which is also how + the tests drive it -- the child cannot resolve the mask itself, because the parent + cleared it. + """ + override = os.environ.get("UNSLOTH_PROBE_DEVICE_CC", "").strip() + if override == _CC_UNKNOWN: + # The parent could not resolve the mask (a numeric ordinal under a non-PCI_BUS_ID + # CUDA_DEVICE_ORDER). Falling back to nvidia-smi here would answer from every GPU on + # the box, which is exactly the verdict the parent declined to give. + return () + text = override + if not text: + try: + import shutil + import subprocess + + exe = shutil.which("nvidia-smi") + if not exe: + return None + result = subprocess.run( + [exe, "--query-gpu=compute_cap", "--format=csv,noheader"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 10, + ) + if result.returncode != 0: + return None + # Every row: with no mask in the environment the whole box is visible, and one + # uncovered card in it is still an install this report has to call degraded. + text = ",".join( + line.strip() for line in (result.stdout or "").splitlines() if line.strip() + ) + except BaseException: + return () + capabilities = [] + for part in str(text).split(","): + part = part.strip() + if not part: + continue + try: + major, _, minor = part.partition(".") + capabilities.append((int(major), int(minor or 0))) + except (AttributeError, ValueError): + return () + return tuple(capabilities) + + +def probe_flash_attn() -> Dict[str, Any]: + """flash-attn, forced past the lazy bit: the CUDA extension lives in the interface.""" + entry = probe_import("flash_attn") + if not entry["imports"]: + return entry + try: + import flash_attn.flash_attn_interface # noqa: F401 + except BaseException as exc: + entry["runs"] = False + entry["error"] = _error(exc) + return entry + # Importing the extension is not the question. A build with no cubin or PTX image for + # this architecture imports fine and fails its first launch with "no kernel image is + # available" -- true of source builds and of older wheels, not only of the sm_100+ cards + # our own installer refuses to fetch a prebuilt wheel for. flash-attn exposes no list of + # the architectures it was compiled for, and this child deliberately never launches a + # kernel, so support cannot be established here. Unknown, with the reason, is the honest + # answer; a failed import above is still broken. + entry["runs"] = None + entry["error"] = ( + "flash-attn imported; whether its kernels cover this GPU cannot be established " + "without launching one, which this probe does not do" + ) + return entry + + +def probe_bitsandbytes() -> Dict[str, Any]: + """bitsandbytes, past the import: are the ctypes kernel handles real? + + From 0.46 a wheel whose native library never loaded still imports cleanly and hands + back a ``throw_on_call`` closure for every symbol, so ``imports=True`` says nothing and + the row rendered as "Working" while the first 4-bit op was going to die mid-run. The + repository already answers this question in ``unsloth/bnb_availability.py``, a leaf + module that imports nothing from unsloth -- loaded here BY PATH so the verdict is the + same one the loader gates on, without dragging ``unsloth/__init__`` (and its patching) + into a diagnostic child. + """ + entry = probe_import("bitsandbytes") + if not entry["imports"]: + return entry + ready = _load_bnb_availability() + if ready is None: + # Cannot find the checker: leave runs unknown rather than inventing a verdict. + return entry + try: + import bitsandbytes + ready.check_native_kernels(bitsandbytes, _device_type()) + entry["runs"] = True + except BaseException as exc: + entry["runs"] = False + entry["error"] = _error(exc) + return entry + + +def _device_type() -> str: + """The device type ``bitsandbytes_symbols`` keys on. Only "xpu" differs.""" + return os.environ.get("UNSLOTH_PROBE_DEVICE_TYPE", "cuda").strip().lower() or "cuda" + + +def _bnb_availability_path(): + """Where ``unsloth/bnb_availability.py`` lives, or None. + + Two ways of asking, because either can come up empty. ``find_spec`` is the right answer + for an INSTALLED unsloth, but it needs the package to be importable from this child -- + which it is not when the backend runs from a checkout whose root is not on the child's + path, and the probe then silently reported bitsandbytes as unknown on every host. This + file's own location settles that case: studio/backend/utils/hardware -> repo root. + """ + candidates = [] + try: + import importlib.util + + spec = importlib.util.find_spec("unsloth") + origin = getattr(spec, "origin", None) if spec is not None else None + if origin: + candidates.append(os.path.join(os.path.dirname(origin), "bnb_availability.py")) + except BaseException: # noqa: BLE001 -- an unimportable unsloth is not an error here + pass + here = os.path.dirname(os.path.abspath(__file__)) + root = os.path.abspath(os.path.join(here, "..", "..", "..", "..")) + candidates.append(os.path.join(root, "unsloth", "bnb_availability.py")) + for path in candidates: + if os.path.isfile(path): + return path + return None + + +def _load_bnb_availability(): + """``unsloth.bnb_availability`` loaded standalone, or None when it cannot be found.""" + try: + import importlib.util + + path = _bnb_availability_path() + if path is None: + return None + leaf = importlib.util.spec_from_file_location("_unsloth_bnb_availability", path) + if leaf is None or leaf.loader is None: + return None + module = importlib.util.module_from_spec(leaf) + leaf.loader.exec_module(module) + return module + except BaseException: + return None + + +def probe_torchao() -> Dict[str, Any]: + """torchao, past the import: did its C++/CUDA extension actually load? + + A supported install can import torchao while its native operators are absent -- the + Python stack pins 0.17.0 on torch 2.10+cu130 precisely because its torch-2.11 extension + is "cleanly skipped" rather than crashed. Reporting THAT as "Working" is a false + all-clear; reporting it as broken is a false alarm on the managed stack, since torchao + keeps working through its supported fallbacks and only the optimized kernels are gone. + + So: unknown, with the reason. It shows in the About row without lighting the banner, + which is reserved for an install that genuinely cannot load. + """ + entry = probe_import("torchao") + if not entry["imports"]: + return entry + registered = _registered_ops("torchao") + if registered is None: + # Cannot tell (an unfamiliar torch): leave the import-only answer rather than invent + # a verdict. + return entry + if registered > 0: + entry["runs"] = True + return entry + entry["runs"] = None + entry["error"] = ( + "torchao imported but registered no native operators: its C++/CUDA extension was " + "skipped for this torch build, so the optimized quantization kernels are not " + "available and it falls back to its Python paths" + ) + return entry + + +def _registered_ops(namespace: str): + """How many operators ``namespace`` has registered with the dispatcher, or None. + + NOT ``dir(torch.ops.)``: touching that attribute CREATES an empty ``_OpNamespace``, + whose dir() is already non-empty (``__name__``, ``__spec__``, ...) before a single operator + exists -- so the no-native-operators case this exists to catch read as healthy. The + dispatcher's own table is the only thing that answers the question asked. + """ + try: + import torch + names = torch._C._dispatch_get_all_op_names() + except BaseException: # noqa: BLE001 — an unfamiliar torch means "cannot tell" + return None + prefix = f"{namespace}::" + return sum(1 for name in names if name.startswith(prefix)) + + +def probe_import(import_name: str) -> Dict[str, Any]: + """Plain import. An ABI mismatch in a native wheel surfaces here as an undefined symbol.""" + entry: Dict[str, Any] = {"imports": False, "runs": None, "error": None} + try: + __import__(import_name) + entry["imports"] = True + except BaseException as exc: + entry["error"] = _error(exc) + return entry + + +PROBES = { + "xformers": probe_xformers, + "flash_attn": probe_flash_attn, + "torchao": probe_torchao, + "bitsandbytes": probe_bitsandbytes, +} + + +def main(argv) -> int: + """Probe the names given on the command line (default: all of them).""" + wanted = [name for name in (argv or PROBES) if name in PROBES] + results = {} + for name in wanted: + try: + results[name] = PROBES[name]() + except BaseException as exc: + results[name] = {"imports": False, "runs": None, "error": _error(exc)} + # Some of these packages print to stdout on import, so the JSON goes out last behind a + # marker the parent can seek to rather than assuming it owns the stream. + sys.stdout.write("\n__UNSLOTH_ACCELERATOR_PROBE__" + json.dumps(results) + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 2e835c0d8a..73aca12da9 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -22,6 +22,7 @@ import glob import os import platform import re +import shutil import subprocess import sys import threading @@ -30,6 +31,7 @@ from contextlib import contextmanager from importlib.metadata import PackageNotFoundError, version as pkg_version import structlog from loggers import get_logger +from utils.child_stdio import utf8_child_env from enum import Enum from pathlib import Path from typing import Optional, Dict, Any @@ -871,15 +873,40 @@ def get_package_versions() -> Dict[str, Optional[str]]: Uses importlib.metadata (stdlib), no subprocess. CUDA version from torch.version.cuda. Returns dict keyed unsloth/torch/transformers/cuda; missing packages yield None. + + The accelerator packages (xformers, flash-attn, torchao, bitsandbytes) are here as + plain version strings only. Whether they actually *load* costs a real import, so it + lives in get_accelerator_report() behind the About/diagnostics detail path. """ - packages = ("unsloth", "torch", "transformers") + packages = ( + "unsloth", + "torch", + "transformers", + # Optimized kernels. A version string alone does not mean they work -- an + # xformers built for another torch reports its version happily and has no + # memory-efficient attention -- but "installed at all" is the first question. + "xformers", + "flash-attn", + "torchao", + "bitsandbytes", + ) versions: Dict[str, Optional[str]] = {} for name in packages: + # JSON key stays a valid identifier so clients can use dotted access: + # "flash-attn" is the distribution name, "flash_attn" is the import name. + key = name.replace("-", "_") try: - versions[name] = pkg_version(name) + versions[key] = pkg_version(name) except PackageNotFoundError: - versions[name] = None + versions[key] = None + except Exception as e: + logger.debug(f"Failed to read {name} version: {e}") + versions[key] = None + + # Which Python the app is running, so an About-tab screenshot answers the "built for + # 3.10, running 3.13" half of a mismatch report without a follow-up question. + versions["python"] = platform.python_version() # GPU runtime versions bundled with torch (CUDA, ROCm/HIP, Intel XPU) try: @@ -906,6 +933,660 @@ def get_package_versions() -> Dict[str, Optional[str]]: return versions +# ========== Accelerator stack health ========== +# +# NVIDIA QA P0-1: the managed Windows xformers was built for torch 2.10.0+cu128 and +# Python 3.10.11 while the app ran cu130 and Python 3.13.2, so its CUDA extensions never +# loaded and memory-efficient attention silently went missing. Nothing in the app said so: +# the package imported, reported a version, and quietly did nothing. get_package_versions +# above would have shown "xformers 0.0.34" and looked healthy. +# +# So report three separate things per package -- is it installed, does it import, do its +# kernels load -- plus what the wheel was compiled against, which is the part that +# actually names the fix. +# +# We do NOT block startup on this. Degraded attention is far better than a dead app: the +# whole failure mode is that training still works, just slower and hungrier. This is a +# report, and the UI is loud about it. + +# (import name, distribution name). The distribution name is what pip installed; the +# import name is what has to load. +_ACCELERATOR_PACKAGES = ( + ("xformers", "xformers"), + ("flash_attn", "flash-attn"), + ("torchao", "torchao"), + ("bitsandbytes", "bitsandbytes"), +) + +# What a ROCm host could actually load. flash-attn ships ROCm wheels and Unsloth enables it +# under DEVICE_TYPE == "hip"; xformers and torchao are CUDA-only builds there, so probing them +# would turn an expected import failure into a red banner. +# +# bitsandbytes belongs here too: unsloth/device_type.py imports it on EVERY backend and asks +# native_kernels_ready(..., DEVICE_TYPE) whether 4-bit and prequantized loading are allowed, so +# a ROCm wheel with a dead native library silently costs quantized loading. Skipping it labelled +# an installed package "not used on this device", which is the false all-clear this report is for. +_ROCM_ACCELERATOR_PACKAGES = frozenset({"flash_attn", "bitsandbytes"}) + +# And on Intel XPU. Same reasoning, other package: bitsandbytes has an XPU backend that the +# loader gates 4-bit on, while flash-attn / xformers / torchao publish nothing for it. +_XPU_ACCELERATOR_PACKAGES = frozenset({"bitsandbytes"}) + +# Importing native extension packages is the only honest way to know they load, but it is +# slow and, on a badly broken install, not risk-free. Computed at most once per process +# and skippable outright. +_ACCELERATOR_PROBE_ENV = "UNSLOTH_SKIP_ACCELERATOR_PROBE" +_accelerator_report_cache: Optional[Dict[str, Any]] = None +_accelerator_report_lock = threading.Lock() + + +def _xformers_build_metadata() -> Optional[Dict[str, Any]]: + """The installed xformers wheel's cpp_lib.json, without importing xformers. + + ``find_spec`` locates the package without executing ``xformers/__init__.py``, so this + costs no import, drags in no torch, and does not fire xformers' own warning. + + Deliberately a copy of ``unsloth.xformers_compat.xformers_build_metadata`` rather than + a call to it: the studio backend runs with ``studio/backend`` on sys.path, and + importing ``unsloth.xformers_compat`` would execute ``unsloth/__init__.py`` and pull in + torch -- which this module must work without (see the no-torch sandbox tests). Only + this file read is duplicated, never the version tables. + """ + try: + import importlib.util + spec = importlib.util.find_spec("xformers") + except Exception: + return None + if spec is None: + return None + locations = list(getattr(spec, "submodule_search_locations", None) or ()) + origin = getattr(spec, "origin", None) + if origin: + locations.append(os.path.dirname(origin)) + for location in locations: + path = os.path.join(location, "cpp_lib.json") + try: + import json + with open(path, "r", encoding = "utf-8") as handle: + metadata = json.load(handle) + except Exception: + continue + if isinstance(metadata, dict) and isinstance(metadata.get("version"), dict): + return metadata + return None + + +def _xformers_built_for() -> Optional[Dict[str, Optional[str]]]: + """cpp_lib.json -> {"torch", "cuda", "hip", "python"} strings, or None.""" + version_block = (_xformers_build_metadata() or {}).get("version") + if not isinstance(version_block, dict): + return None + cuda = version_block.get("cuda") + # xformers' setup.py stores major * 100 + minor: 1208 is CUDA 12.8, 1300 is 13.0. + # Its own error message prints the raw integer, which is why users misread it. + cuda_text = ( + f"{cuda // 100}.{cuda % 100}" + if isinstance(cuda, int) and not isinstance(cuda, bool) + else None + ) + return { + "torch": str(version_block["torch"]) if version_block.get("torch") else None, + "cuda": cuda_text, + "hip": str(version_block["hip"]) if version_block.get("hip") else None, + "python": str(version_block["python"]) if version_block.get("python") else None, + } + + +def _running_torch() -> Dict[str, Optional[str]]: + """{"torch", "cuda", "python"} for the interpreter answering the request.""" + try: + import torch + return { + "torch": getattr(torch, "__version__", None), + "cuda": getattr(torch.version, "cuda", None), + "python": platform.python_version(), + } + except Exception: + return {"torch": None, "cuda": None, "python": platform.python_version()} + + +# Longest reason we will put in a response. These are exception strings, and an unbounded +# one lands in a settings row's description AND in its aria-label. +_MAX_REASON_CHARS = 300 + + +# Failures that say "ask again later", not "this wheel is broken": the GPU was busy, full, or +# claimed by another process. Mirrors _INCONCLUSIVE_PROBE_ERRORS in +# unsloth/utils/attention_dispatch.py, which makes the same call for the in-process probe. +_INCONCLUSIVE_PROBE_ERRORS = ( + "out of memory", + "busy or unavailable", + "all cuda-capable devices are busy", + "no cuda-capable device", + "cuda_error_not_permitted", + "insufficient driver", + "initialization error", + "invalid device ordinal", + "invalid device id", + "currently in use", + "in use by another process", + "exclusive", +) + + +def _probe_was_inconclusive(error) -> bool: + """True when the probe failed for a reason that never tested the wheel.""" + if not error: + return False + text = str(error).lower() + return any(marker in text for marker in _INCONCLUSIVE_PROBE_ERRORS) + + +def _short_reason(reason: Optional[str]) -> Optional[str]: + if not reason: + return None + reason = " ".join(str(reason).split()) + if len(reason) <= _MAX_REASON_CHARS: + return reason + return reason[: _MAX_REASON_CHARS - 1].rstrip() + "…" + + +def _applicable_accelerator_packages() -> frozenset: + """Which of these packages this host could load at all. + + Everything here is a CUDA kernel package. Probing one on a host it was never going to + run on turns an expected outcome into a red banner, which is how a warning stops being + read: bitsandbytes is a dependency on every platform but the stock build is CUDA-only, + so on a Mac, a CPU-only box, an AMD ROCm host or an Intel XPU host its import failure + is the design, not a broken acceleration stack. + + ROCm is the subtle one -- those hosts report ``DeviceType.CUDA`` internally (there is + deliberately no DeviceType.ROCM), so a plain device check would wave them through. + They are not empty, though: Unsloth imports and enables flash-attn under + ``DEVICE_TYPE == "hip"`` (``unsloth/models/_utils.py``), so a broken ROCm + FlashAttention install is something training will actually try to use, and reporting + it as "not used on this device" suppresses the one banner that would explain the + crash. The stock bitsandbytes / xformers / torchao builds are CUDA-only, so those + stay out. + """ + try: + device = get_device() + if device is DeviceType.XPU: + # bitsandbytes ships an XPU backend and the loader gates 4-bit on + # native_kernels_ready(..., "xpu") with its own symbol set, so a broken XPU wheel + # is something training will reach. The other three are CUDA-only builds. + return _XPU_ACCELERATOR_PACKAGES + if device is not DeviceType.CUDA: + return frozenset() + if IS_ROCM: + return _ROCM_ACCELERATOR_PACKAGES + except Exception: + return frozenset() + return frozenset(name for name, _ in _ACCELERATOR_PACKAGES) + + +_PROBE_MARKER = "__UNSLOTH_ACCELERATOR_PROBE__" + + +def _run_probe_subprocess( + names, + timeout: int = 180, + status: Optional[Dict[str, Any]] = None, + keep_cuda: bool = False, +) -> Optional[Dict[str, Any]]: + """Run accelerator_probe.py in a throwaway interpreter. None if it could not answer. + + ``status`` is an optional dict the caller passes in to learn WHY a None came back: + ``status["died"]`` is True when a child ran and exited without printing its marker, + and absent when the probe never got to run (no script, launch error, timeout). The + two need different answers -- one package that kills the interpreter is a diagnosis, + a probe that could not start is an unknown -- and they are indistinguishable from the + return value alone. + + Out of process on purpose -- see that module's docstring. In short: a failed import + leaves half a package in sys.modules and poisons every later import of it, ``import + bitsandbytes`` permanently latches a CUDA context in whatever process runs it, and a + badly broken native wheel can abort the interpreter instead of raising. All three are + survivable in a child and fatal in the server. + """ + script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "accelerator_probe.py") + if not os.path.exists(script): + return None + try: + completed = subprocess.run( + [sys.executable, script, *names], + capture_output = True, + text = True, + # Native import errors quote DLL and file paths, which on Windows can carry + # non-ASCII; the ANSI codepage default would mangle them. errors="replace" so + # a stray byte degrades one character instead of losing the whole diagnosis. + encoding = "utf-8", + errors = "replace", + timeout = timeout, + # The child imports torch; keep it off the GPUs entirely. bitsandbytes and + # xformers both report their load status without one, and this way the probe + # cannot take VRAM from a run in progress. Through utf8_child_env because we + # decode as UTF-8 above and the child would otherwise emit the Windows ANSI + # code page -- which is precisely where the DLL paths this probe exists to + # report carry non-ASCII. + env = utf8_child_env( + { + **os.environ, + # Hidden for every probe that only needs to know whether a library LOADS, + # so the diagnostic cannot take VRAM from a run in progress. bitsandbytes + # is the exception (see _probe_packages): with no visible CUDA it picks its + # CPU library and a healthy install reports as broken. + **({} if keep_cuda else {"CUDA_VISIBLE_DEVICES": ""}), + "UNSLOTH_ALLOW_CPU": "1", + # Which ctypes symbol set bitsandbytes is checked against; only "xpu" + # differs, and the child cannot work it out with the GPUs hidden. + "UNSLOTH_PROBE_DEVICE_TYPE": _probe_device_type(), + # The capability of the GPU the APP can use. The child cannot read it: the + # mask is gone by the time it starts, and nvidia-smi's first row is a + # physical device this process may not even be allowed to touch. + # "unknown", never "": the child reads an EMPTY override as no + # override and falls back to nvidia-smi over the whole physical box, + # so an unresolvable mask (a user-preserved FASTEST_FIRST ordering) + # would be answered from every GPU present -- the false verdict the + # resolver returned None to avoid. + **( + {"UNSLOTH_PROBE_DEVICE_CC": (_visible_compute_capability() or _CC_UNKNOWN)} + if not keep_cuda + else {} + ), + } + ), + ) + except Exception as e: + logger.debug(f"Accelerator probe subprocess failed: {e}") + return None + # These packages print on import, so seek to the marker instead of assuming the JSON + # owns stdout. + _, marker, payload = (completed.stdout or "").rpartition(_PROBE_MARKER) + if not marker: + logger.debug( + f"Accelerator probe produced no result (rc={completed.returncode}): " + f"{(completed.stderr or '')[-400:]}" + ) + # It ran and did not answer: on a native wheel that aborts the interpreter this is + # the only trace the failure leaves. + if status is not None: + status["died"] = True + return None + try: + import json + results = json.loads(payload.strip().splitlines()[0]) + except Exception as e: + logger.debug(f"Accelerator probe output was not JSON: {e}") + return None + return results if isinstance(results, dict) else None + + +# Sentinel for "the capabilities could not be resolved", passed to the probe child in place of +# an empty string so it cannot read the absence as "no mask, read the whole box". +_CC_UNKNOWN = "unknown" + + +# The torch release xFormers moved to the PyTorch stable API/ABI at. Its v0.0.34 notes +# state that "binary builds targeting PyTorch 2.10+ will be compatible with any later +# version", so a 2.10-built wheel on 2.11 is the design. Mirrors _STABLE_ABI_TORCH_FLOOR +# in unsloth/xformers_compat.py, which makes the same call for the training-side message. +_STABLE_ABI_TORCH_FLOOR = (2, 10) + + +def _torch_release(version: Optional[str]) -> Optional[tuple]: + """'2.10.0+cu130' -> (2, 10, 0). None when there is no numeric release to read.""" + if not version: + return None + match = re.match(r"^(\d+(?:\.\d+)*)", str(version).strip().split("+", 1)[0]) + if match is None: + return None + try: + return tuple(int(part) for part in match.group(1).split(".")) + except ValueError: + return None + + +def _stable_abi_covers(built: Optional[tuple], running: Optional[tuple]) -> bool: + """Whether a wheel built for ``built`` is expected to load on ``running``. + + One-directional: 2.10-built on 2.11 is covered, 2.11-built on 2.10 is not, and below + the floor there is no guarantee to lean on. + """ + if built is None or running is None: + return False + return built[:2] >= _STABLE_ABI_TORCH_FLOOR and running >= built + + +def _probe_device_type() -> str: + try: + return "xpu" if get_device() is DeviceType.XPU else "cuda" + except Exception: + return "cuda" + + +# bitsandbytes decides which native library to load from torch.cuda.is_available(), so a child +# with the GPUs hidden picks the CPU one and a healthy CUDA install reports dead handles. It +# gets its own child, with the mask the app itself runs under. That child does latch a CUDA +# context, but it exits immediately and the SERVER never holds one -- which is what the hidden +# mask is protecting. +_CUDA_VISIBLE_PROBES = frozenset({"bitsandbytes"}) + + +def _visible_compute_capability() -> Optional[str]: + """Compute capabilities this process can actually use, comma-joined ("9.0,12.0"), or None. + + nvidia-smi lists PHYSICAL devices and ignores CUDA_VISIBLE_DEVICES, so the first row is not + necessarily the app's device: on a mixed host masked to an sm_120 card, reading row 0 (an + sm_90 part) evaluates the wrong op table and can report xFormers as working while the + visible GPU has no kernel. Resolved here, in the parent, because the child is started with + the mask cleared and cannot recover it. + + EVERY visible device, not the first: with CUDA_VISIBLE_DEVICES=0,1 across a mixed pair, a + rank landing on the second card falls back to SDPA while a verdict taken from the first + says "Working". A report that covers only one of the GPUs a run will use is the same false + all-clear in a smaller box. + """ + exe = shutil.which("nvidia-smi") + if not exe: + return None + try: + result = subprocess.run( + [exe, "--query-gpu=index,uuid,compute_cap", "--format=csv,noheader"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 10, + ) + except Exception: # noqa: BLE001 — no answer is "unknown", never a failure + return None + if result.returncode != 0: + return None + rows = [] + for line in (result.stdout or "").splitlines(): + parts = [part.strip() for part in line.split(",")] + if len(parts) >= 3: + rows.append((parts[0], parts[1], parts[2])) + if not rows: + return None + mask = (os.environ.get("CUDA_VISIBLE_DEVICES") or "").strip() + if not mask: + return ",".join(dict.fromkeys(capability for _i, _u, capability in rows)) + # A NUMERIC token is a CUDA ordinal in the current CUDA_DEVICE_ORDER, and nvidia-smi's + # index column is the PCI_BUS_ID ordering. main.py pins PCI_BUS_ID before torch loads, so + # the two agree on every ordinary host -- but it pins it with setdefault, so a user who + # exported FASTEST_FIRST keeps it, and there the ordinal is a speed rank we cannot invert + # without initialising CUDA. Answering anyway would hand the probe the wrong card: a false + # degraded warning, or worse, silence about an uncovered GPU. Unknown is the honest answer, + # and UUID tokens are unaffected because a UUID means the same device in any ordering. + ordinals_are_pci = (os.environ.get("CUDA_DEVICE_ORDER") or "").strip().upper() == "PCI_BUS_ID" + selected = [] + for token in mask.split(","): + token = token.strip() + if not token: + # CUDA stops at the first invalid entry, and so does this. + break + if token.isdigit() and not ordinals_are_pci: + return None + for index, uuid, capability in rows: + # The mask names either an ordinal or a GPU UUID, and a UUID may be abbreviated. + if token == index or uuid == token or uuid.startswith(token): + selected.append(capability) + break + else: + break + if not selected: + return None + return ",".join(dict.fromkeys(selected)) + + +def _probe_packages(names) -> Dict[str, Any]: + """Probe ``names``, isolating a package that KILLS the child rather than raising. + + One child for all of them is the fast path and the usual one. But aborting the + interpreter (pybind11 answers a duplicate type registration with ``std::terminate``) + is a failure mode this probe is specifically built to survive, and in the batch child + it takes the whole answer down with it: every installed package then reported "could + not be checked", ``degraded`` came back empty, and Settings showed no warning at all + for the one install that genuinely cannot load. So on a dead batch, re-probe one at a + time -- the survivors get their real answer, and the package whose own child dies is + reported as broken, which is what it is. + """ + names = list(names) + # Split off the packages that need the real device mask; each group gets its own child. + cuda_visible = [name for name in names if name in _CUDA_VISIBLE_PROBES] + hidden = [name for name in names if name not in _CUDA_VISIBLE_PROBES] + if cuda_visible and hidden: + results = _probe_packages_with(hidden, keep_cuda = False) + results.update(_probe_packages_with(cuda_visible, keep_cuda = True)) + return results + return _probe_packages_with(names, keep_cuda = bool(cuda_visible)) + + +def _probe_packages_with(names, *, keep_cuda: bool) -> Dict[str, Any]: + """One group of packages, under one CUDA visibility. Isolation rules as above.""" + if not names: + return {} + status: Dict[str, Any] = {} + results = _run_probe_subprocess(list(names), status = status, keep_cuda = keep_cuda) + if results is not None: + return results + if not status.get("died"): + # The probe never ran (no script, launch error, timeout). That says nothing about + # any package, and retrying per package would only repeat it. + return {} + if len(names) == 1: + # Its own child already died: that IS the diagnosis, not an unknown. + return { + names[0]: { + "imports": False, + "runs": False, + "error": "the probe process exited without answering (the import aborted " + "the interpreter)", + } + } + isolated: Dict[str, Any] = {} + for name in names: + isolated.update(_probe_packages_with([name], keep_cuda = keep_cuda)) + return isolated + + +def _describe_xformers_break( + built_for: Optional[Dict[str, Optional[str]]], error: Optional[str] +) -> Optional[str]: + """Name the mismatch when there is one, otherwise hand back the real error. + + Only claims a mismatch when the recorded build and the running runtime actually + differ. Asserting one whenever build metadata merely exists misdiagnoses the most + common Windows failure that is NOT a version mismatch -- a missing VC++ runtime or + CUDA DLL, which surfaces as ``[WinError 126]`` with a perfectly matching wheel. + + A Python difference alone is never a mismatch and is never named on its own: the + wheels are abi3/none-tagged and ``_C`` loads through ``torch.ops.load_library``, not + the CPython ABI, so 3.10-built kernels run fine on 3.13. Saying otherwise sends people + off to reinstall Python. + """ + running = _running_torch() + built_torch = (built_for or {}).get("torch") + built_cuda = (built_for or {}).get("cuda") + running_torch = running.get("torch") + running_cuda = running.get("cuda") + + mismatch = None + # Compare RELEASES, and only when the stable-ABI guarantee does not already cover the + # pair. Two things went wrong with the raw string compare: "2.10.0+cu126" vs + # "2.10.0+cu128" is the same release with a compatible local tag, and it claimed a + # torch mismatch before the CUDA-major check below could give the accurate answer; + # and a 2.10-built stable-ABI wheel on torch 2.11 is exactly what upstream says will + # work. In both cases a real, actionable error (a missing VC++ runtime, say) was + # thrown away and replaced with a confident wrong diagnosis. + built_release = _torch_release(built_torch) + running_release = _torch_release(running_torch) + if ( + built_release + and running_release + and built_release != running_release + and not _stable_abi_covers(built_release, running_release) + ): + mismatch = (f"torch {built_torch}", f"torch {running_torch}") + elif built_cuda and running_cuda and built_cuda.split(".")[0] != running_cuda.split(".")[0]: + # Majors only: CUDA minor version compatibility means a cu126 wheel loads fine + # against a cu128 torch, and flagging that would be crying wolf. + # + # Both torch builds are named alongside the CUDA versions. This is now the branch + # that reports the NVIDIA case (a cu128-built wheel beside a cu130 torch, same + # torch release on both sides), and "CUDA 12.8 vs CUDA 13.0" alone leaves the + # reader guessing which of their two torch installs is which. + mismatch = ( + f"torch {built_torch} (CUDA {built_cuda})" if built_torch else f"CUDA {built_cuda}", + f"torch {running_torch} (CUDA {running_cuda})" + if running_torch + else f"CUDA {running_cuda}", + ) + if mismatch is None: + return error + + built_text, running_text = mismatch + built_python = (built_for or {}).get("python") + if built_python: + built_text += f" / Python {built_python}" + if running.get("python"): + running_text += f" / Python {running['python']}" + return ( + f"xformers was built for {built_text} but this app runs {running_text}; " + f"its C++/CUDA extensions cannot load, so memory-efficient attention is off" + ) + + +def get_accelerator_report(refresh: bool = False) -> Dict[str, Any]: + """Whether each optimized-kernel package is installed, imports, and actually loads. + + Additive: nothing here changes the shape of get_package_versions() or of any existing + /api/system response field. Shape:: + + { + "python_version": "3.13.2", + "torch_version": "2.10.0+cu130", + "torch_cuda": "13.0", + "probed": true, + "packages": { + "xformers": {"version": "0.0.34", "installed": true, "imports": true, + "runs": false, "reason": "...", "built_for": {...}}, + ... + }, + "degraded": ["xformers"] + } + + ``runs`` is None where the question does not apply (package absent, or no separate + kernel-load step). Cached for the process: the probe spawns an interpreter and + re-running it per request would make the About tab expensive. Set + UNSLOTH_SKIP_ACCELERATOR_PROBE=1 to report versions only. + """ + global _accelerator_report_cache + + if not refresh and _accelerator_report_cache is not None: + return copy.deepcopy(_accelerator_report_cache) + + with _accelerator_report_lock: + if not refresh and _accelerator_report_cache is not None: + return copy.deepcopy(_accelerator_report_cache) + + disabled = os.environ.get(_ACCELERATOR_PROBE_ENV, "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + applicable = _applicable_accelerator_packages() + running = _running_torch() + packages: Dict[str, Any] = {} + degraded = [] + + versions: Dict[str, Optional[str]] = {} + for import_name, dist_name in _ACCELERATOR_PACKAGES: + try: + versions[import_name] = pkg_version(dist_name) + except PackageNotFoundError: + versions[import_name] = None + except Exception as e: + logger.debug(f"Failed to read {dist_name} version: {e}") + versions[import_name] = None + + # One child for all of them: the interpreter start and the torch import dominate, + # so four separate children would cost four times as much for the same answer. + to_probe = [ + name for name, version in versions.items() if version is not None and name in applicable + ] + results = {} if (disabled or not to_probe) else _probe_packages(to_probe) + probed = bool(results) + + for import_name, _ in _ACCELERATOR_PACKAGES: + version = versions[import_name] + entry: Dict[str, Any] = { + "version": version, + "installed": version is not None, + "imports": False, + # Per package, because the probe set is per device: a ROCm host never probes + # bitsandbytes, and one global flag would render that row "Not loading". + "probed": False, + "runs": None, + "reason": None, + } + if import_name == "xformers": + # Read straight off the wheel, so it is reported even when nothing is + # probed -- it costs one small file read and no import. + entry["built_for"] = _xformers_built_for() + + result = results.get(import_name) + if result is not None: + entry["probed"] = True + entry["imports"] = result.get("imports") is True + entry["runs"] = result.get("runs") if isinstance(result.get("runs"), bool) else None + error = result.get("error") + if _probe_was_inconclusive(error): + # The probe never got as far as testing the wheel. bitsandbytes keeps CUDA + # visible on purpose, so opening Settings while a trainer has filled the GPU + # (or holds it in EXCLUSIVE_PROCESS) fails at context creation on a perfectly + # healthy install -- and this answer is cached for the life of the backend, + # so a red banner would sit there until restart. Same classification the + # in-process xformers probe applies for the same reason. + # Reported as NOT probed, which is what happened: acceleratorHealth reads + # that as unknown and now renders the reason beside it, so the user sees + # why rather than a red row or a bare "Not checked". + entry["probed"] = False + entry["runs"] = None + entry["reason"] = _short_reason(error) + packages[import_name] = entry + continue + if import_name == "xformers": + error = _describe_xformers_break(entry.get("built_for"), error) + entry["reason"] = _short_reason(error) + if not entry["imports"] or entry["runs"] is False: + degraded.append(import_name) + elif version is not None: + # Three different unknowns, and the UI must not read any of them as broken. + if disabled: + entry["reason"] = "not probed" + elif import_name not in applicable: + entry["reason"] = "not used on this device" + else: + entry["reason"] = "could not be checked" + packages[import_name] = entry + + report = { + "python_version": platform.python_version(), + "torch_version": running.get("torch"), + "torch_cuda": running.get("cuda"), + "probed": probed, + "packages": packages, + "degraded": degraded, + } + _accelerator_report_cache = report + return copy.deepcopy(report) + + # ========== Torch-based GPU fallbacks (AMD ROCm, Intel XPU, nvidia-smi missing) ========== diff --git a/studio/frontend/src/features/settings/components/accelerator-status.tsx b/studio/frontend/src/features/settings/components/accelerator-status.tsx new file mode 100644 index 0000000000..cfb60462e6 --- /dev/null +++ b/studio/frontend/src/features/settings/components/accelerator-status.tsx @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** + * Optimized-kernel health, for the About tab and the Settings banner. + * + * NVIDIA QA P0-1: the managed Windows xFormers was built for PyTorch 2.10.0+cu128 and + * Python 3.10.11 while the app ran cu130 and Python 3.13.2, so its CUDA extensions never + * loaded and memory-efficient attention silently went missing. Nothing in the UI said so + * -- the About tab showed a version string, which a mismatched wheel reports just as + * happily as a working one. + * + * So the version alone is never the status here. The status is whether the kernels load, + * and when they do not, what the wheel was built for versus what is running -- the pair + * that names the fix. + */ + +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; +import { + type AcceleratorPackage, + type AcceleratorReport, + type Health, + acceleratorHealth, + acceleratorShowsReason, + hasDeadAccelerator, +} from "@/hooks/accelerator-report"; +import { useAcceleratorReport } from "@/hooks/use-accelerator-report"; +import { type TranslationKey, useT } from "@/i18n"; +import { cn } from "@/lib/utils"; +import { Alert01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { SettingsRow } from "./settings-row"; +import { SettingsSection } from "./settings-section"; + +// Import name -> display label. Not translated: these are package names. +const PACKAGE_LABELS: Record = { + xformers: "xFormers", + flash_attn: "FlashAttention", + torchao: "torchao", + bitsandbytes: "bitsandbytes", +}; + +function packageLabel(name: string): string { + return PACKAGE_LABELS[name] ?? name; +} + +const HEALTH_LABEL_KEYS: Record = { + working: "settings.about.accelerator.working", + broken: "settings.about.accelerator.notLoading", + absent: "settings.about.accelerator.notInstalled", + unknown: "settings.about.accelerator.notChecked", +}; + +const HEALTH_CLASSES: Record = { + // Deliberately not colour-only: each badge carries its own word, so the state + // survives a screenshot in greyscale and a colour-vision difference. + working: "text-muted-foreground", + broken: "text-destructive font-medium", + absent: "text-muted-foreground/70", + unknown: "text-muted-foreground/70", +}; + +/** "torch 2.10.0+cu128 / Python 3.10.11" from a wheel's recorded build. */ +function describeBuild(pkg: AcceleratorPackage): string | null { + const build = pkg.builtFor; + if (!build) return null; + const parts: string[] = []; + if (build.torch) { + parts.push(`torch ${build.torch}`); + } else if (build.cuda) { + parts.push(`CUDA ${build.cuda}`); + } else if (build.hip) { + parts.push(`ROCm ${build.hip}`); + } + if (build.python) parts.push(`Python ${build.python}`); + return parts.length > 0 ? parts.join(" / ") : null; +} + +function AcceleratorRow({ + pkg, + probed, +}: { + pkg: AcceleratorPackage; + probed: boolean; +}) { + const t = useT(); + const health = acceleratorHealth(pkg, probed); + const build = describeBuild(pkg); + // Explain a broken one, and an unknown one that came with a reason. On a healthy machine + // the build detail is noise, and the raw exception text is never the first thing to show. + // + // The unknown arm matters because several probes return `runs: null` DELIBERATELY, with an + // explanation: flash-attn imported but no kernel was launched, torchao registered no native + // operator. Without it the row read "Not checked" and threw the reason away, so a skipped + // native extension was indistinguishable from a probe that never ran at all. + const explained = acceleratorShowsReason(health, pkg.reason); + const detail = !explained + ? null + : health === "broken" && build + ? t("settings.about.accelerator.builtFor", { build }) + : pkg.reason; + + return ( + + + + {pkg.version ?? "—"} + + + {t(HEALTH_LABEL_KEYS[health])} + + + + ); +} + +export function AcceleratorSection() { + const t = useT(); + const report: AcceleratorReport | null = useAcceleratorReport(); + // Older backends do not send this at all; render nothing rather than an empty section. + if (!report || report.packages.length === 0) return null; + + return ( + + {report.packages.map((pkg) => ( + + ))} + + ); +} + +/** + * Shown on every Settings tab, not only the one a user would have to think to open. A + * dead kernel is silent by nature; this is the thing that surfaces it. + * + * Reads the hook itself rather than taking a prop so it can be mounted inside + * DialogContent, which Radix only renders while the dialog is open. That keeps the + * detail fetch -- and the native imports behind it -- off app startup. + */ +export function AcceleratorBanner() { + const t = useT(); + const report = useAcceleratorReport(); + if (!hasDeadAccelerator(report)) return null; + + const packages = (report?.degraded ?? []).map(packageLabel).join(", "); + return ( + + + {t("settings.about.accelerator.bannerTitle")} + + {t("settings.about.accelerator.bannerBody", { packages })} + + + ); +} diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 5e5ef9e645..ce1b223e0a 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -36,6 +36,7 @@ import { useRef, useState, } from "react"; +import { AcceleratorBanner } from "./components/accelerator-status"; import { SETTINGS_SEARCH_KEYWORDS, createSettingsSearchIndex, @@ -465,6 +466,10 @@ export function SettingsDialog() { ref={mainScrollRef} className="hover-scrollbar flex min-h-0 min-w-0 flex-1 flex-col overflow-y-auto p-6 [scrollbar-gutter:stable]" > + {/* Above every tab, not just About: an optimized kernel that is + installed and cannot load produces no other symptom than being + slower, so nothing else would ever tell the user. */} + {renderTab(panelTab)} diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx index 4670098322..93b103bcae 100644 --- a/studio/frontend/src/features/settings/tabs/about-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx @@ -18,6 +18,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; +import { AcceleratorSection } from "../components/accelerator-status"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -175,7 +176,9 @@ export function AboutTab() { - {hw.gpus.length > 0 || runtimes.length > 0 ? ( + {/* hw.python keeps the section alive on a CPU-only host, where there is no GPU + and no accelerator runtime but the Python version still belongs in a report. */} + {hw.gpus.length > 0 || runtimes.length > 0 || hw.python ? ( {hw.gpus.map((gpu, i) => ( ))} + {/* The other half of every "built for Python 3.10, running 3.13" report. It + was nowhere in the app, so bug reports never carried it. */} + {hw.python ? ( + + + {hw.python} + + + ) : null} ) : null} + + ; + return { + torch: (build.torch as string) ?? null, + cuda: (build.cuda as string) ?? null, + hip: (build.hip as string) ?? null, + python: (build.python as string) ?? null, + }; +} + +export function parseAcceleratorReport(raw: unknown): AcceleratorReport | null { + if (typeof raw !== "object" || raw === null) return null; + const report = raw as Record; + const rawPackages = (report.packages ?? {}) as Record< + string, + Record + >; + // Unknown names go last rather than being dropped: a backend that grows a package + // should not need a frontend release to show it. + const names = Object.keys(rawPackages).sort((a, b) => { + const ia = ACCELERATOR_ORDER.indexOf(a); + const ib = ACCELERATOR_ORDER.indexOf(b); + return ( + (ia < 0 ? ACCELERATOR_ORDER.length : ia) - + (ib < 0 ? ACCELERATOR_ORDER.length : ib) + ); + }); + return { + pythonVersion: (report.python_version as string) ?? null, + torchVersion: (report.torch_version as string) ?? null, + torchCuda: (report.torch_cuda as string) ?? null, + probed: report.probed !== false, + packages: names.map((name) => { + const entry = rawPackages[name] ?? {}; + return { + name, + version: (entry.version as string) ?? null, + installed: entry.installed === true, + probed: typeof entry.probed === "boolean" ? entry.probed : null, + imports: entry.imports === true, + runs: typeof entry.runs === "boolean" ? entry.runs : null, + reason: (entry.reason as string) ?? null, + builtFor: parseAcceleratorBuild(entry.built_for), + }; + }), + degraded: Array.isArray(report.degraded) + ? (report.degraded as string[]) + : [], + }; +} + +export type Health = "working" | "broken" | "absent" | "unknown"; + +/** + * Three questions collapse to one badge. `runs === false` and "imports but does not run" + * are the same thing to a user, and both mean the kernels are dead. + * + * `reportProbed` is only the fallback for a backend too old to answer per package. The + * probe set is per device -- a ROCm host never probes bitsandbytes -- so reading the + * report-wide flag onto every row renders a deliberately skipped package as "Not loading" + * the moment any other package is probed, and it is not in `degraded` either. + */ +export function acceleratorHealth( + pkg: AcceleratorPackage, + reportProbed: boolean, +): Health { + if (!pkg.installed) return "absent"; + if (!(pkg.probed ?? reportProbed)) return "unknown"; + if (!pkg.imports) return "broken"; + // `runs === null` from a probe that DID run means the kernel question could not be + // answered: an xformers layout with no recognised load record, a missing bitsandbytes + // checker, a torch with no dispatcher table, a flash-attn on a card no prebuilt wheel + // covers. Reading that as Working is the false all-clear this report exists to remove, + // and those packages are not in `degraded` either, so nothing else would say it. + if (pkg.runs === null) return "unknown"; + return pkg.runs ? "working" : "broken"; +} + +/** + * Does the About row owe the user the backend's reason text? + * + * A broken row always did. An unknown one has to as well, because most unknowns are + * DELIBERATE and carry an explanation the row was throwing away: flash-attn imported with no + * kernel launched, xformers registering an op it may have no image for, torchao with no + * native operator. Rendering only "Not checked" made a skipped native extension look exactly + * like a probe that never ran. A reasonless unknown (the probe genuinely did not run) still + * shows nothing, which is the honest answer there. + */ +export function acceleratorShowsReason( + health: Health, + reason: string | null | undefined, +): boolean { + if (health === "broken") return true; + return health === "unknown" && reason != null && reason !== ""; +} + +/** True when something is installed and cannot load -- the case worth a banner. */ +export function hasDeadAccelerator(report: AcceleratorReport | null): boolean { + return (report?.degraded.length ?? 0) > 0; +} diff --git a/studio/frontend/src/hooks/use-accelerator-report.ts b/studio/frontend/src/hooks/use-accelerator-report.ts new file mode 100644 index 0000000000..fc9177f038 --- /dev/null +++ b/studio/frontend/src/hooks/use-accelerator-report.ts @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { useEffect, useState } from "react"; +import { + type AcceleratorReport, + parseAcceleratorReport, +} from "./accelerator-report"; + +// Deliberately NOT part of useHardwareInfo. That hook's `include_details=true` response is +// read by Export, Video and onboarding, and the accelerator block costs a child interpreter +// on the backend to import the native packages somewhere they cannot poison anything. Only +// the Settings surfaces ask for it, and only while the Settings dialog is mounted. +const ENDPOINT = + "/api/system/hardware?include_details=true&include_accelerators=true"; + +// Module-level cache so the About section and the banner share one request. +let cached: AcceleratorReport | null = null; +let fetchPromise: Promise | null = null; + +export function invalidateAcceleratorReport() { + cached = null; + fetchPromise = null; +} + +async function fetchOnce(): Promise { + if (cached) return cached; + if (fetchPromise) return fetchPromise; + + fetchPromise = (async () => { + try { + const res = await authFetch(ENDPOINT); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + cached = parseAcceleratorReport(data?.accelerators); + return cached; + } catch { + // Reset so a later mount retries; a failed read is "unknown", and null renders as + // no section and no banner rather than as an all-clear. + fetchPromise = null; + return null; + } + })(); + + return fetchPromise; +} + +/** + * Optimized-kernel health, or null while unknown. + * + * Null covers all three of "still loading", "the read failed" and "this backend predates + * the field". None of them may render as healthy, and none of them may render a banner. + */ +export function useAcceleratorReport(): AcceleratorReport | null { + const [report, setReport] = useState(cached); + + useEffect(() => { + let cancelled = false; + fetchOnce().then((next) => { + if (!cancelled) setReport(next); + }); + return () => { + cancelled = true; + }; + }, []); + + return report; +} diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36f42f9405..09b2339ecd 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -4,6 +4,7 @@ import { authFetch } from "@/features/auth"; import { useEffect, useState } from "react"; + export interface GpuDevice { name: string | null; vramTotalGb: number | null; @@ -28,6 +29,9 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // The Python the backend runs on. Half of every "built for 3.10, running 3.13" report, + // and previously not shown anywhere in the app. + python: string | null; // Whether export can run here (true only on a supported accelerator), with a torch-aware // reason. `null` until the authoritative response lands, so callers don't briefly enable // export; `loaded` flips true once a real (non-error) response arrives. @@ -55,6 +59,7 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + python: null, exportSupported: null, exportUnsupportedReason: null, exportUnsupportedMessage: null, @@ -115,6 +120,7 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + python: data?.versions?.python ?? null, exportSupported: data?.export_supported ?? null, exportUnsupportedReason: data?.export_unsupported_reason ?? null, exportUnsupportedMessage: data?.export_unsupported_message ?? null, diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts index 6af9768606..ea9a9ee3b6 100644 --- a/studio/frontend/src/i18n/locales/ar.ts +++ b/studio/frontend/src/i18n/locales/ar.ts @@ -1083,6 +1083,20 @@ export const ar = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "النوى المُحسَّنة", + sectionDescription: + "نوى انتباه وتكميم أسرع. تُبنى كل نواة مقابل إصدار واحد من PyTorch، لذا عند اختلاف الإصدار لا يُحمَّل أي منها ويعود Unsloth بهدوء إلى المسارات الأبطأ.", + working: "تعمل", + notLoading: "لا تُحمَّل", + notInstalled: "غير مثبَّتة", + notChecked: "لم يتم الفحص", + builtFor: "مبنية من أجل {build}", + bannerTitle: "النوى المُحسَّنة مثبَّتة لكنها لا تُحمَّل", + bannerBody: + "لا يمكن تحميل {packages} هنا، لذا فإن نواته المُحسَّنة غير متاحة. تعرض القائمة أعلاه ما لا يزال يعمل.", + }, updates: "التحديث", help: "مساعدة", documentation: "الوثائق", diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts index 63de463195..052a46180d 100644 --- a/studio/frontend/src/i18n/locales/de.ts +++ b/studio/frontend/src/i18n/locales/de.ts @@ -1122,6 +1122,20 @@ export const de = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Optimierte Kernel", + sectionDescription: + "Schnellere Attention- und Quantisierungs-Kernel. Jeder wird gegen genau einen PyTorch-Build kompiliert, daher lädt bei abweichenden Versionen keiner davon und Unsloth fällt still auf langsamere Standardpfade zurück.", + working: "Funktioniert", + notLoading: "Lädt nicht", + notInstalled: "Nicht installiert", + notChecked: "Nicht geprüft", + builtFor: "Gebaut für {build}", + bannerTitle: "Optimierte Kernel sind installiert, laden aber nicht", + bannerBody: + "{packages} kann hier nicht geladen werden, daher stehen seine optimierten Kernel nicht zur Verfügung. Die Liste oben zeigt, was weiterhin funktioniert.", + }, updates: "Update", help: "Hilfe", documentation: "Dokumentation", diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 15e4eecce6..bd3420f760 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -1074,6 +1074,20 @@ export const en = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Optimized kernels", + sectionDescription: + "Faster attention and quantization kernels. Each is compiled against one PyTorch build, so a version mismatch loads nothing and quietly falls back to slower defaults.", + working: "Working", + notLoading: "Not loading", + notInstalled: "Not installed", + notChecked: "Not checked", + builtFor: "Built for {build}", + bannerTitle: "Optimized kernels are installed but not loading", + bannerBody: + "{packages} cannot load here, so its optimized kernels are unavailable. The list above shows what is still working.", + }, updates: "Update", help: "Help", documentation: "Documentation", diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts index c804a06237..e9f2e2e402 100644 --- a/studio/frontend/src/i18n/locales/es.ts +++ b/studio/frontend/src/i18n/locales/es.ts @@ -1114,6 +1114,20 @@ export const es = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Kernels optimizados", + sectionDescription: + "Kernels de atención y cuantización más rápidos. Cada uno se compila contra una única compilación de PyTorch, así que si las versiones no coinciden no se carga ninguno y Unsloth vuelve en silencio a las rutas más lentas.", + working: "Funcionando", + notLoading: "No se carga", + notInstalled: "No instalado", + notChecked: "Sin comprobar", + builtFor: "Compilado para {build}", + bannerTitle: "Los kernels optimizados están instalados pero no se cargan", + bannerBody: + "{packages} no se puede cargar aquí, así que sus kernels optimizados no están disponibles. La lista de arriba muestra lo que sigue funcionando.", + }, updates: "Actualización", help: "Ayuda", documentation: "Documentación", diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts index 64c6393f98..d3d3c9eaa8 100644 --- a/studio/frontend/src/i18n/locales/fr.ts +++ b/studio/frontend/src/i18n/locales/fr.ts @@ -1122,6 +1122,20 @@ export const fr = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Noyaux optimisés", + sectionDescription: + "Noyaux d'attention et de quantification plus rapides. Chacun est compilé pour une seule version de PyTorch : en cas d'écart de version, aucun ne se charge et Unsloth revient silencieusement aux chemins plus lents.", + working: "Fonctionne", + notLoading: "Ne se charge pas", + notInstalled: "Non installé", + notChecked: "Non vérifié", + builtFor: "Compilé pour {build}", + bannerTitle: "Des noyaux optimisés sont installés mais ne se chargent pas", + bannerBody: + "{packages} ne peut pas se charger ici, ses noyaux optimisés sont donc indisponibles. La liste ci-dessus indique ce qui fonctionne encore.", + }, updates: "Mise à jour", help: "Aide", documentation: "Documentation", diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts index 02168000a1..f5dbae45d9 100644 --- a/studio/frontend/src/i18n/locales/hi.ts +++ b/studio/frontend/src/i18n/locales/hi.ts @@ -1089,6 +1089,20 @@ export const hi = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "अनुकूलित कर्नेल", + sectionDescription: + "तेज़ अटेंशन और क्वांटाइज़ेशन कर्नेल। हर कर्नेल किसी एक PyTorch बिल्ड के लिए संकलित होता है, इसलिए संस्करण मेल न खाने पर कोई भी लोड नहीं होता और Unsloth चुपचाप धीमे डिफ़ॉल्ट पर लौट जाता है।", + working: "काम कर रहा है", + notLoading: "लोड नहीं हो रहा", + notInstalled: "इंस्टॉल नहीं है", + notChecked: "जाँचा नहीं गया", + builtFor: "{build} के लिए बनाया गया", + bannerTitle: "अनुकूलित कर्नेल इंस्टॉल हैं पर लोड नहीं हो रहे", + bannerBody: + "{packages} यहाँ लोड नहीं हो पा रहा, इसलिए इसके अनुकूलित कर्नेल उपलब्ध नहीं हैं। ऊपर दी गई सूची बताती है कि क्या अब भी काम कर रहा है।", + }, updates: "अपडेट", help: "सहायता", documentation: "दस्तावेज़", diff --git a/studio/frontend/src/i18n/locales/it.ts b/studio/frontend/src/i18n/locales/it.ts index 2e1c412fba..9bf5f69f92 100644 --- a/studio/frontend/src/i18n/locales/it.ts +++ b/studio/frontend/src/i18n/locales/it.ts @@ -1089,6 +1089,20 @@ export const it = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Kernel ottimizzati", + sectionDescription: + "Kernel di attenzione e quantizzazione più veloci. Ognuno è compilato per una sola build di PyTorch, quindi se le versioni non coincidono non ne viene caricato nessuno e Unsloth torna silenziosamente ai percorsi più lenti.", + working: "Funzionante", + notLoading: "Non si carica", + notInstalled: "Non installato", + notChecked: "Non verificato", + builtFor: "Compilato per {build}", + bannerTitle: "I kernel ottimizzati sono installati ma non si caricano", + bannerBody: + "{packages} non può essere caricato qui, quindi i suoi kernel ottimizzati non sono disponibili. L'elenco sopra mostra cosa funziona ancora.", + }, updates: "Aggiornamento", help: "Aiuto", documentation: "Documentazione", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index 97bda39a81..312bb19457 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -1058,6 +1058,20 @@ export const ja = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "最適化カーネル", + sectionDescription: + "高速なアテンションおよび量子化カーネルです。それぞれ特定の PyTorch ビルド向けにコンパイルされているため、バージョンが一致しないと 1 つも読み込まれず、Unsloth は静かに低速な既定の処理へ切り替わります。", + working: "動作中", + notLoading: "読み込めません", + notInstalled: "未インストール", + notChecked: "未確認", + builtFor: "{build} 向けにビルド", + bannerTitle: "最適化カーネルはインストール済みですが読み込めません", + bannerBody: + "{packages} をここで読み込めないため、その最適化カーネルは利用できません。上の一覧に、現在も動作しているものが表示されます。", + }, updates: "アップデート", help: "ヘルプ", documentation: "ドキュメント", diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts index 89051d7dfc..fb5ab0f130 100644 --- a/studio/frontend/src/i18n/locales/ko.ts +++ b/studio/frontend/src/i18n/locales/ko.ts @@ -1081,6 +1081,20 @@ export const ko = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "최적화 커널", + sectionDescription: + "더 빠른 어텐션 및 양자화 커널입니다. 각각 하나의 PyTorch 빌드에 맞춰 컴파일되므로 버전이 맞지 않으면 아무것도 로드되지 않고 Unsloth는 조용히 느린 기본 경로로 전환합니다.", + working: "정상 작동", + notLoading: "로드되지 않음", + notInstalled: "설치되지 않음", + notChecked: "확인하지 않음", + builtFor: "{build}용으로 빌드됨", + bannerTitle: "최적화 커널이 설치되어 있지만 로드되지 않습니다", + bannerBody: + "{packages}을(를) 여기서 로드할 수 없어 최적화 커널을 사용할 수 없습니다. 위 목록에서 여전히 작동하는 항목을 확인하세요.", + }, updates: "업데이트", help: "도움말", documentation: "문서", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index 96cefd9cdf..dc743933ce 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -1096,6 +1096,20 @@ export const ptBR = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Kernels otimizados", + sectionDescription: + "Kernels de atenção e quantização mais rápidos. Cada um é compilado para uma única build do PyTorch, então uma divergência de versão não carrega nenhum e o Unsloth volta silenciosamente aos caminhos mais lentos.", + working: "Funcionando", + notLoading: "Não carrega", + notInstalled: "Não instalado", + notChecked: "Não verificado", + builtFor: "Compilado para {build}", + bannerTitle: "Kernels otimizados estão instalados, mas não carregam", + bannerBody: + "{packages} não consegue carregar aqui, portanto seus kernels otimizados estão indisponíveis. A lista acima mostra o que continua funcionando.", + }, updates: "Atualização", help: "Ajuda", documentation: "Documentação", diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts index 4a2ac3a13e..6c49276f10 100644 --- a/studio/frontend/src/i18n/locales/ru.ts +++ b/studio/frontend/src/i18n/locales/ru.ts @@ -1094,6 +1094,20 @@ export const ru = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "Оптимизированные ядра", + sectionDescription: + "Быстрые ядра внимания и квантизации. Каждое собрано под одну сборку PyTorch, поэтому при несовпадении версий не загружается ни одно, и Unsloth молча переходит на более медленные пути.", + working: "Работает", + notLoading: "Не загружается", + notInstalled: "Не установлено", + notChecked: "Не проверено", + builtFor: "Собрано для {build}", + bannerTitle: "Оптимизированные ядра установлены, но не загружаются", + bannerBody: + "{packages} не удаётся загрузить здесь, поэтому его оптимизированные ядра недоступны. В списке выше показано, что продолжает работать.", + }, updates: "Обновление", help: "Справка", documentation: "Документация", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index cfe84b9af3..bf70821de4 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -1051,6 +1051,20 @@ export const zhCN = { cuda: "CUDA", rocm: "ROCm", xpu: "XPU", + python: "Python", + accelerator: { + sectionTitle: "优化内核", + sectionDescription: + "更快的注意力与量化内核。每个内核只针对一个 PyTorch 版本编译，版本不匹配时不会加载任何内核，Unsloth 会静默回退到较慢的默认实现。", + working: "正常", + notLoading: "无法加载", + notInstalled: "未安装", + notChecked: "未检查", + builtFor: "为 {build} 编译", + bannerTitle: "优化内核已安装但无法加载", + bannerBody: + "{packages} 无法在此加载，因此其优化内核不可用。上方列表显示仍在正常工作的组件。", + }, updates: "更新", help: "帮助", documentation: "文档", diff --git a/studio/frontend/tests/accelerator-report.test.ts b/studio/frontend/tests/accelerator-report.test.ts new file mode 100644 index 0000000000..14f972e5b4 --- /dev/null +++ b/studio/frontend/tests/accelerator-report.test.ts @@ -0,0 +1,270 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// NVIDIA QA P0-1: the managed Windows xFormers was built for PyTorch 2.10.0+cu128 and +// Python 3.10.11 while the app ran cu130 and Python 3.13.2, so its CUDA extensions never +// loaded. The About tab showed a version string, which a mismatched wheel reports exactly +// as happily as a working one -- so "installed", "imports" and "runs" have to survive the +// parse as three separate answers, and a backend that predates the field has to parse to +// null rather than to a false all-clear. + +import assert from "node:assert/strict"; +import test from "node:test"; + +import { registerBundlerResolver } from "./helpers/kit.ts"; + +registerBundlerResolver(); + +const { + parseAcceleratorReport, + hasDeadAccelerator, + acceleratorHealth, + acceleratorShowsReason, +} = await import("../src/hooks/accelerator-report.ts"); + +const BROKEN_XFORMERS = { + python_version: "3.13.2", + torch_version: "2.10.0+cu130", + torch_cuda: "13.0", + probed: true, + packages: { + bitsandbytes: { + version: "0.48.0", + installed: true, + imports: true, + runs: null, + reason: null, + }, + xformers: { + version: "0.0.34", + installed: true, + imports: true, + runs: false, + reason: "xformers was built for torch 2.10.0+cu128 ...", + built_for: { + torch: "2.10.0+cu128", + cuda: "12.8", + hip: null, + python: "3.10.11", + }, + }, + flash_attn: { + version: null, + installed: false, + imports: false, + runs: null, + reason: null, + }, + }, + degraded: ["xformers"], +}; + +test("an installed-but-dead package survives the parse as exactly that", () => { + const report = parseAcceleratorReport(BROKEN_XFORMERS); + assert.ok(report); + const xformers = report.packages.find((p) => p.name === "xformers"); + assert.ok(xformers); + // The three answers stay separate: collapsing them is how a version string alone + // came to stand in for "working". + assert.equal(xformers.installed, true); + assert.equal(xformers.imports, true); + assert.equal(xformers.runs, false); + assert.equal(xformers.builtFor?.torch, "2.10.0+cu128"); + assert.equal(xformers.builtFor?.python, "3.10.11"); +}); + +test("packages come back in a fixed order regardless of key order", () => { + const report = parseAcceleratorReport(BROKEN_XFORMERS); + assert.deepEqual( + report?.packages.map((p) => p.name), + ["xformers", "flash_attn", "bitsandbytes"], + ); +}); + +test("an unknown package is kept, and sorted after the known ones", () => { + // A backend that grows a package must not need a frontend release to show it, and it + // must not displace the ones the UI knows how to label. The single-package version of + // this test asserted an ordering it never exercised. + const report = parseAcceleratorReport({ + packages: { + somethingnew: { version: "1.0", installed: true, imports: true }, + bitsandbytes: { version: "0.48.0", installed: true, imports: true }, + xformers: { version: "0.0.34", installed: true, imports: true }, + }, + degraded: [], + }); + assert.deepEqual( + report?.packages.map((p) => p.name), + ["xformers", "bitsandbytes", "somethingnew"], + ); +}); + +test("a ROCm build reports hip, which the CUDA-only parse dropped", () => { + const report = parseAcceleratorReport({ + packages: { + xformers: { + version: "0.0.33", + installed: true, + imports: true, + runs: false, + built_for: { + torch: null, + cuda: null, + hip: "6.4.43483", + python: "3.10.11", + }, + }, + }, + degraded: ["xformers"], + }); + assert.equal(report?.packages[0].builtFor?.hip, "6.4.43483"); + assert.equal(report?.packages[0].builtFor?.cuda, null); +}); + +test("a missing runs is null, not false", () => { + // torchao and bitsandbytes have no separate kernel-load step, so `runs` is unknown + // rather than broken. Coercing it to false would light the banner on every machine. + const report = parseAcceleratorReport(BROKEN_XFORMERS); + const bnb = report?.packages.find((p) => p.name === "bitsandbytes"); + assert.equal(bnb?.runs, null); +}); + +test("an older backend that sends no report parses to null", () => { + // Not to an empty healthy report: "we cannot tell" must not render as "all fine". + assert.equal(parseAcceleratorReport(undefined), null); + assert.equal(parseAcceleratorReport(null), null); +}); + +test("probed defaults to true and only an explicit false turns it off", () => { + assert.equal(parseAcceleratorReport({ packages: {} })?.probed, true); + assert.equal( + parseAcceleratorReport({ packages: {}, probed: false })?.probed, + false, + ); +}); + +test("the banner fires only on a degraded package", () => { + assert.equal( + hasDeadAccelerator(parseAcceleratorReport(BROKEN_XFORMERS)), + true, + ); + assert.equal( + hasDeadAccelerator( + parseAcceleratorReport({ ...BROKEN_XFORMERS, degraded: [] }), + ), + false, + ); + // No report at all (older backend, or a fetch without include_details) must not + // produce a banner claiming something is broken. + assert.equal(hasDeadAccelerator(null), false); +}); + +test("a package the backend chose not to probe reads as unknown, not broken", () => { + // The probe set is per device: a ROCm host probes flash_attn and skips bitsandbytes, an + // Intel host does the reverse. Reading the report-wide flag onto every row rendered the + // skipped ones as "Not loading" as soon as one other row was probed -- a red badge for a + // package that is fine, and one the banner does not even list, since it is not degraded. + const report = parseAcceleratorReport({ + probed: true, + packages: { + xformers: { + version: "0.0.34", + installed: true, + probed: true, + imports: true, + runs: true, + }, + bitsandbytes: { + version: "0.48.0", + installed: true, + probed: false, + imports: false, + runs: null, + reason: "not used on this device", + }, + }, + degraded: [], + }); + assert.ok(report); + const [xformers, bnb] = report.packages; + assert.equal(acceleratorHealth(xformers, report.probed), "working"); + assert.equal(acceleratorHealth(bnb, report.probed), "unknown"); +}); + +test("a backend too old to answer per package falls back to the report flag", () => { + const older = parseAcceleratorReport({ + probed: true, + packages: { + xformers: { version: "0.0.34", installed: true, imports: true, runs: true }, + }, + degraded: [], + }); + assert.ok(older); + assert.equal(older.packages[0].probed, null); + assert.equal(acceleratorHealth(older.packages[0], older.probed), "working"); + assert.equal(acceleratorHealth(older.packages[0], false), "unknown"); +}); + +test("not installed stays not installed, probed or not", () => { + const report = parseAcceleratorReport(BROKEN_XFORMERS); + const flash = report?.packages.find((p) => p.name === "flash_attn"); + assert.ok(flash); + assert.equal(acceleratorHealth(flash, true), "absent"); + // And a probed package that is dead is still dead. + const xformers = report?.packages.find((p) => p.name === "xformers"); + assert.ok(xformers); + assert.equal(acceleratorHealth({ ...xformers, probed: true }, true), "broken"); +}); + +test("a probe that ran but could not decide is unknown, not working", () => { + // `runs === null` from a probe that DID run means the kernel question could not be + // answered: an xformers layout with no recognised load record, a missing bitsandbytes + // checker, a torch with no dispatcher table, a flash-attn on a card no prebuilt wheel + // covers. It rendered as Working -- the false all-clear this report exists to remove, + // and one nothing else corrects, since those packages are not in `degraded` either. + const report = parseAcceleratorReport({ + probed: true, + packages: { + flash_attn: { + version: "2.8.3", + installed: true, + probed: true, + imports: true, + runs: null, + reason: "no prebuilt wheel covers compute capability 12.0", + }, + xformers: { + version: "0.0.34", + installed: true, + probed: true, + imports: true, + runs: true, + }, + }, + degraded: [], + }); + assert.ok(report); + const [xformers, flash] = report.packages; + assert.equal(acceleratorHealth(xformers, report.probed), "working"); + assert.equal(acceleratorHealth(flash, report.probed), "unknown"); + // A failed import is still broken, whatever runs says. + assert.equal( + acceleratorHealth({ ...flash, imports: false }, report.probed), + "broken", + ); +}); + + +test("an unknown result keeps the reason the backend sent with it", () => { + // Most unknowns are deliberate and carry an explanation -- flash-attn imported with no + // kernel launched, xformers registering an op whose image may be missing, torchao with no + // native operator. The row showed "Not checked" and discarded all of it, so a skipped + // native extension looked identical to a probe that never ran. + assert.equal(acceleratorShowsReason("unknown", "no kernel was launched"), true); + assert.equal(acceleratorShowsReason("broken", null), true); + // A reasonless unknown really is "not checked"; there is nothing to say. + assert.equal(acceleratorShowsReason("unknown", null), false); + assert.equal(acceleratorShowsReason("unknown", ""), false); + assert.equal(acceleratorShowsReason("working", "ignored"), false); + assert.equal(acceleratorShowsReason("absent", "ignored"), false); +}); diff --git a/tests/utils/test_attention_masks.py b/tests/utils/test_attention_masks.py index 627da0552f..c9fe83151f 100644 --- a/tests/utils/test_attention_masks.py +++ b/tests/utils/test_attention_masks.py @@ -285,3 +285,145 @@ def test_run_attention_flash_varlen_receives_window_and_softcap(monkeypatch): """Unit tests for packed-attention mask helpers with sliding-window logic.""" + + +def test_run_attention_sdpa_windows_an_unpacked_unmasked_batch(monkeypatch): + """No packing, no padding mask: the case that had nothing to hang the window off. + + SDPA's ``is_causal`` is FULL causal -- it has no window -- so with neither the xformers + bias nor flash's ``window_size`` in play, a model whose config declares a sliding window + attended its entire causal history. That is reachable from a Mistral training step the + moment xFormers is disabled and FlashAttention is absent, which is precisely what the + kernel probe can now decide. + """ + captured = {} + + def _fake_sdpa(Q, K, V, **kwargs): + captured["mask"] = kwargs.get("attn_mask") + captured["is_causal"] = kwargs.get("is_causal") + return torch.zeros_like(Q) + + monkeypatch.setattr(attention_dispatch, "scaled_dot_product_attention", _fake_sdpa) + + config = attention_dispatch.AttentionConfig( + backend = attention_dispatch.SDPA, + n_kv_heads = 1, + n_groups = 1, + ) + context = attention_dispatch.AttentionContext( + bsz = 1, + q_len = 6, + kv_seq_len = 6, + n_heads = 1, + head_dim = 1, + requires_grad = True, + seq_info = None, + attention_mask = None, + causal_mask = None, + sliding_window = 3, + ) + Q = torch.zeros(1, 1, 6, 1) + + attention_dispatch.run_attention(config = config, context = context, Q = Q, K = Q, V = Q) + + mask = captured["mask"] + assert mask is not None, "a declared window must not fall through to plain is_causal" + assert captured["is_causal"] is False + assert mask.shape == (1, 1, 6, 6) + # Row 5 sees 3, 4, 5 and nothing older; the future stays masked either way. + assert [bool(v) for v in mask[0, 0, 5]] == [False, False, False, True, True, True] + + +def test_run_attention_sdpa_leaves_a_short_sequence_alone(monkeypatch): + # Shorter than the window: nothing to clamp, and the cheap is_causal path must survive. + captured = {} + monkeypatch.setattr( + attention_dispatch, + "scaled_dot_product_attention", + lambda Q, K, V, **kw: (captured.update(kw), torch.zeros_like(Q))[1], + ) + config = attention_dispatch.AttentionConfig( + backend = attention_dispatch.SDPA, n_kv_heads = 1, n_groups = 1 + ) + context = attention_dispatch.AttentionContext( + bsz = 1, + q_len = 4, + kv_seq_len = 4, + n_heads = 1, + head_dim = 1, + requires_grad = True, + seq_info = None, + attention_mask = None, + causal_mask = None, + sliding_window = 8, + ) + Q = torch.zeros(1, 1, 4, 1) + attention_dispatch.run_attention(config = config, context = context, Q = Q, K = Q, V = Q) + assert captured["attn_mask"] is None and captured["is_causal"] is True + + +def test_mistral_hands_the_dispatcher_its_configured_window(): + """The context Mistral builds omitted `sliding_window` entirely, so even a correct SDPA + window path had nothing to act on.""" + import ast + from pathlib import Path + + src = Path(attention_dispatch.__file__).resolve().parents[1] / "models" / "mistral.py" + tree = ast.parse(src.read_text(encoding = "utf-8")) + contexts = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "AttentionContext" + ] + assert contexts, "AttentionContext construction not found in mistral.py" + for call in contexts: + assert "sliding_window" in {kw.arg for kw in call.keywords} + + +def test_a_zero_configured_window_is_full_causal_not_a_blank_mask(): + """`sliding_window = 0` means "no local attention", the same as absent -- which is how + Mistral's own mask builders read it. Passing the 0 through makes the SDPA lower bound + `q_pos - (0 - 1)` sit above the causal upper bound, so every position is masked and the + layer returns nothing at all.""" + import ast + from pathlib import Path + + src = Path(attention_dispatch.__file__).resolve().parents[1] / "models" / "mistral.py" + text = src.read_text(encoding = "utf-8") + assert "isinstance(sw_cfg, int) and sw_cfg <= 0" in text, ( + "a non-positive configured window must be normalised before it reaches window_size " + "or the dispatcher" + ) + ast.parse(text) + + +def test_run_attention_sdpa_ignores_a_zero_window(monkeypatch): + # Belt and braces at the dispatcher: even handed a zero, it must not build a mask that + # hides everything. + captured = {} + monkeypatch.setattr( + attention_dispatch, + "scaled_dot_product_attention", + lambda Q, K, V, **kw: (captured.update(kw), torch.zeros_like(Q))[1], + ) + config = attention_dispatch.AttentionConfig( + backend = attention_dispatch.SDPA, n_kv_heads = 1, n_groups = 1 + ) + context = attention_dispatch.AttentionContext( + bsz = 1, + q_len = 4, + kv_seq_len = 4, + n_heads = 1, + head_dim = 1, + requires_grad = True, + seq_info = None, + attention_mask = None, + causal_mask = None, + sliding_window = 0, + ) + Q = torch.zeros(1, 1, 4, 1) + attention_dispatch.run_attention(config = config, context = context, Q = Q, K = Q, V = Q) + mask = captured["attn_mask"] + assert mask is None or bool(mask.any()), "a zero window must not mask everything" diff --git a/tests/utils/test_xformers_broken_warning.py b/tests/utils/test_xformers_broken_warning.py new file mode 100644 index 0000000000..ddf4feb8bc --- /dev/null +++ b/tests/utils/test_xformers_broken_warning.py @@ -0,0 +1,351 @@ +# 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 General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +"""A broken xformers must be LOUD on the default path (NVIDIA P0-1). + +Before this, ``unsloth/models/_utils.py`` silenced the xformers logger to ERROR right +before the import and then reported the failure only under ``UNSLOTH_ENABLE_LOGGING`` -- +so a wheel built for a different torch dropped the user to SDPA attention with no output +at all, which is how a cu128 wheel shipped beside a cu130 runtime unnoticed. + +Runs the real ``import unsloth`` in a subprocess against a stub xformers package that is +put first on ``sys.path``, so nothing here depends on the xformers actually installed -- +the stub ships its own ``.dist-info`` as well as its own ``cpp_lib.json``, because +``import unsloth`` reads both. + +The child also loads ``tests/conftest.py``, the same GPU-free harness pytest applies to +this directory, so the subprocess can finish the import on a host with no accelerator +exactly like the parent process can. See ``_import_unsloth_with``. +""" + +import importlib.util +import json +import os +import subprocess +import sys +import textwrap +import types +from pathlib import Path + +import pytest +import torch + + +# A stub that imports fine and reports a version, then fails exactly where the real one +# does: torch.ops.load_library inside _register_extensions. +_STUB_VERSION = "0.0.34" +_STUB_INIT = f"__version__ = {_STUB_VERSION!r}\n" +_STUB_CPP_LIB = textwrap.dedent( + """\ + def _register_extensions(): + raise OSError("[WinError 126] The specified module could not be found") + """ +) + + +def _write_stub_xformers(root: Path, built_torch: str, built_cuda: int) -> None: + package = root / "xformers" + package.mkdir() + (package / "__init__.py").write_text(_STUB_INIT, encoding = "utf-8") + (package / "_cpp_lib.py").write_text(_STUB_CPP_LIB, encoding = "utf-8") + (package / "cpp_lib.json").write_text( + json.dumps( + { + "version": { + "cuda": built_cuda, + "hip": None, + "torch": built_torch, + "python": "3.10.11", + }, + "env": {"XFORMERS_PACKAGE_FROM": f"wheel-v{_STUB_VERSION}"}, + } + ), + encoding = "utf-8", + ) + # A .dist-info beside the package, so the stub is a complete install rather than just + # an importable module. `import unsloth` asks importlib.metadata for the xformers + # version twice - unsloth/import_fixes.py:fix_xformers_performance_issue() and + # unsloth/models/_utils.py - and a module without metadata makes both of those read + # the HOST's xformers instead, which is precisely what this file's docstring promises + # not to do. On a host with no xformers at all there is nothing to read and + # importlib.metadata.version() raises PackageNotFoundError straight out of + # fix_xformers_performance_issue(), killing the child before it reaches the subject + # matter. Writing the metadata makes the stub self-sufficient in both cases. + dist_info = root / f"xformers-{_STUB_VERSION}.dist-info" + dist_info.mkdir() + (dist_info / "METADATA").write_text( + f"Metadata-Version: 2.1\nName: xformers\nVersion: {_STUB_VERSION}\n", + encoding = "utf-8", + ) + + +def _import_unsloth_with(stub_root, repo_root: Path, enable_logging: str) -> str: + """Import unsloth in a child, optionally with a stub xformers shadowing the real one.""" + inject = [f"sys.path.insert(0, {str(stub_root)!r})"] if stub_root is not None else [] + # Only meaningful with the stub: without one this asks the host for a version it may + # legitimately not have. + dist_probe = ( + [ + "from importlib.metadata import version as _dist_version", + 'print("XFORMERS_DIST_VERSION:", _dist_version("xformers"))', + ] + if stub_root is not None + else [] + ) + # tests/conftest.py is the GPU-free harness pytest already applies to everything under + # tests/: on a host with no accelerator it forces DEVICE_TYPE to "cuda" and stubs the + # torch.cuda probes unsloth fires at import time (get_device_capability, + # is_bf16_supported, mem_get_info). The parent pytest process gets it for free, a bare + # `python -c` child does not, and that asymmetry is the whole reason these two tests + # were the only red ones on the CPU-only CI runners: unsloth/_gpu_init.py calls + # torch.cuda.get_device_capability() at module scope and the child died on + # "Found no NVIDIA driver" before reaching any xformers code. Loaded by path under a + # private name so it neither shadows nor is shadowed by the injected stub, and on a + # real accelerator it is a no-op. + harness = repo_root / "tests" / "conftest.py" + assert harness.is_file(), f"GPU-free test harness missing at {harness}" + # Assembled line by line rather than through a dedented f-string: an interpolated + # multi-line fragment lands at column 0 and makes textwrap.dedent a no-op, which turns + # the whole child into an IndentationError. + code = "\n".join( + ["import importlib.util, sys"] + + inject + + [ + "_spec = importlib.util.spec_from_file_location(", + f" '_unsloth_gpu_free_harness', {str(harness)!r}", + ")", + "_harness = importlib.util.module_from_spec(_spec)", + "_spec.loader.exec_module(_harness)", + "import unsloth # noqa: F401", + "from unsloth.models._utils import XFORMERS_BROKEN_REASON, xformers", + 'print("REASON:", XFORMERS_BROKEN_REASON)', + 'print("XFORMERS_IS_NONE:", xformers is None)', + ] + + dist_probe + ) + result = subprocess.run( + [sys.executable, "-c", code], + cwd = str(repo_root), + capture_output = True, + text = True, + timeout = 900, + env = { + **os.environ, + "UNSLOTH_ENABLE_LOGGING": enable_logging, + # Pin one device rather than inheriting a list, so the child cannot pick a + # different GPU than the one this process was given (it never allocates on it - + # the stub xformers fails first). The harness above covers the case where the + # named device does not exist. + "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0], + }, + ) + assert result.returncode == 0, f"import unsloth failed:\n{result.stdout}\n{result.stderr}" + return result.stdout + result.stderr + + +@pytest.fixture(scope = "module") +def repo_root() -> Path: + return Path(__file__).resolve().parents[2] + + +def _mismatched_build(): + """A (torch, cuda-int) pair one CUDA major away from whatever this host runs. + + The host's torch is whatever the test env has, so pick the build relative to it + instead of hardcoding cu128/cu130 and hoping. + """ + running_major = int((torch.version.cuda or "12").split(".", 1)[0]) + built_cuda = 1300 if running_major != 13 else 1208 + built_torch = f"{torch.__version__.split('+')[0]}+cu{built_cuda // 100}{built_cuda % 100:d}" + return built_torch, built_cuda + + +@pytest.mark.skipif( + torch.version.cuda is None, + reason = "the mismatch is synthesized as a CUDA-major difference, which needs a CUDA " + "torch to compare against; on ROCm/XPU/CPU there is no running CUDA major", +) +def test_broken_xformers_warns_without_enable_logging(tmp_path, repo_root): + built_torch, built_cuda = _mismatched_build() + _write_stub_xformers(tmp_path, built_torch, built_cuda) + + output = _import_unsloth_with(tmp_path, repo_root, enable_logging = "0") + + # The stub really is what unsloth inspected, metadata included - not the host's + # xformers leaking in through importlib.metadata. + assert f"XFORMERS_DIST_VERSION: {_STUB_VERSION}" in output + # Loud on the default path: this is the whole point. + assert "Xformers is installed but its optimized kernels cannot load" in output + # Actionable: names what it was built for and what is actually running. + assert built_torch in output + assert "3.10.11" in output + assert torch.__version__ in output + # And it must actually fall back rather than half-work. + assert "XFORMERS_IS_NONE: True" in output + assert "REASON: None" not in output + + +def test_warning_is_printed_once(capsys): + # In-process, because a fresh subprocess reaches the announcement once no matter what: + # the subprocess version of this test passed with the once-guard deleted. + from unsloth.models import _utils + + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + try: + _utils._announce_xformers_breakage("first", build_mismatch = True) + _utils._announce_xformers_breakage("second", build_mismatch = True) + printed = capsys.readouterr().out + finally: + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + assert printed.count("Xformers is installed but its optimized kernels cannot load") == 1 + assert "second" not in printed + + +def test_a_live_flash_attention_is_not_reported_as_an_sdpa_fallback(capsys): + """select_attention_backend puts FlashAttention ahead of xformers, so on a dual install + a broken xformers costs nothing -- and "falling back to SDPA, it uses more memory" is a + performance warning about a run that is on the faster of the two kernels.""" + from unsloth.models import _utils + for flash, expected, forbidden in ( + (True, "FlashAttention is handling attention instead", "uses more memory"), + (False, "SDPA", "FlashAttention is handling attention instead"), + ): + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + original = _utils.HAS_FLASH_ATTENTION + try: + _utils.HAS_FLASH_ATTENTION = flash + _utils._announce_xformers_breakage("built for torch 2.9", build_mismatch = True) + printed = capsys.readouterr().out + finally: + _utils.HAS_FLASH_ATTENTION = original + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + assert expected in printed + assert forbidden not in printed + # The breakage itself is still reported either way; only the consequence changes. + assert "cannot load" in printed + + +def test_a_non_mismatch_failure_is_not_relabelled_as_a_build_mismatch(capsys): + # This arm also catches the sm_100/110/120 FA3 guard and the old-torch guards, whose + # messages are multi-line, fenced and already actionable. Reflowing one of those into + # "its optimized kernels cannot load ... install the matching build" states the wrong + # cause and mangles the text. + from unsloth.models import _utils + + original = ( + "Unsloth: Xformers 0.0.32.post2 has a broken FA3 dispatch on SM 12.0 GPUs.\n" + "```\npip install ninja\n```\n" + ) + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + try: + _utils._announce_xformers_breakage(original, build_mismatch = False) + printed = capsys.readouterr().out + finally: + _utils._XFORMERS_BREAKAGE_ANNOUNCED = False + assert "its optimized kernels cannot load" not in printed + assert "--force-reinstall" not in printed + # Verbatim: the fences and the newlines are part of a copy-pasteable instruction. + assert original in printed + + +@pytest.mark.skipif( + torch.version.cuda is None, reason = "needs a CUDA torch to have a matching xformers" +) +@pytest.mark.skipif( + importlib.util.find_spec("xformers") is None, + reason = "no xformers is installed on this host (importlib.util.find_spec('xformers') " + "is None), so there is no healthy install for `import unsloth` to stay quiet about; " + "the not-installed path is a different branch of unsloth/models/_utils.py and is " + "covered by test_xformers_compat.py", +) +def test_a_healthy_install_says_nothing(repo_root): + # The other half of "never cry wolf": on a machine where xformers works, importing + # unsloth must print none of this. Uses the real installed xformers, and skips when + # that one is itself broken -- then the warning is correct, not a false positive. + # The skipif above already established xformers IS installed, so xformers is None here + # can only mean the import of it failed, never that there was nothing to import. + output = _import_unsloth_with(None, repo_root, enable_logging = "0") + if "XFORMERS_IS_NONE: True" in output: + pytest.skip("this host's xformers is genuinely broken, so a warning is correct") + assert "Xformers is installed but" not in output + assert "REASON: None" in output + + +def test_the_fix_hint_names_the_index_not_just_the_version(monkeypatch): + # The reported case is a CUDA-family mismatch: torch 2.10.0+cu130 beside an xformers + # 0.0.34 built for cu128. All three families publish the SAME version string and PyPI + # carries only the cu128 flavour, so `pip install --force-reinstall xformers==0.0.34` + # reinstalls the identical broken wheel. The index is the part that repairs it. + from unsloth.models import _utils + + monkeypatch.setattr(_utils, "xformers_for_torch", lambda version: "0.0.34") + monkeypatch.setattr(_utils.torch, "version", types.SimpleNamespace(cuda = "13.0")) + hint = _utils._xformers_fix_hint() + assert "--index-url https://download.pytorch.org/whl/cu130" in hint + assert '"xformers==0.0.34"' in hint + + +def test_the_fix_hint_honours_a_configured_mirror(monkeypatch): + from unsloth.models import _utils + + monkeypatch.setattr(_utils, "xformers_for_torch", lambda version: "0.0.34") + monkeypatch.setattr(_utils.torch, "version", types.SimpleNamespace(cuda = "12.8")) + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", "https://mirror.example/whl/") + assert "--index-url https://mirror.example/whl/cu128" in _utils._xformers_fix_hint() + + +def test_the_fix_hint_redacts_a_credentialed_mirror(monkeypatch): + """UNSLOTH_PYTORCH_MIRROR can carry credentials or a signed-URL query, and this hint is + printed to stdout on the DEFAULT import path -- into logs, CI output and notebooks.""" + from unsloth.models import _utils + + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", "https://ci:s3cr3t@wheels.internal/simple") + index = _utils._xformers_torch_index_url() + if index is None: + pytest.skip("no CUDA torch, so there is no index to name") + assert "s3cr3t" not in index and "ci:" not in index + assert "***@wheels.internal" in index + + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", "https://wheels.internal/simple?token=abc123") + index = _utils._xformers_torch_index_url() + assert "abc123" not in index and "token" not in index + assert index.startswith("https://wheels.internal/simple/") + + # An ordinary mirror is passed through untouched, so the command stays copy-pasteable. + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", "https://wheels.internal/whl") + assert _utils._xformers_torch_index_url().startswith("https://wheels.internal/whl/") + + # An authority that cannot be parsed -- a bad IPv6 bracket, a non-numeric port -- is + # exactly the input whose secret this exists to withhold, so the fallback drops the URL + # instead of echoing it. The hint then omits --index-url rather than printing the token. + for malformed in ( + "https://ci:s3cr3t@[bad:ipv6/simple", + "https://wheels.internal:notaport/simple?token=abc123", + ): + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", malformed) + assert _utils._xformers_torch_index_url() is None + assert "s3cr3t" not in _utils._xformers_fix_hint() + assert "abc123" not in _utils._xformers_fix_hint() + + +def test_the_fix_hint_names_no_index_without_a_cuda_torch(monkeypatch): + # ROCm / XPU / CPU: there is no CUDA-matched xformers index to point at, and inventing + # one would send the user to a 404. + from unsloth.models import _utils + + monkeypatch.setattr(_utils, "xformers_for_torch", lambda version: "0.0.34") + monkeypatch.setattr(_utils.torch, "version", types.SimpleNamespace(cuda = None)) + hint = _utils._xformers_fix_hint() + assert "--index-url" not in hint and '"xformers==0.0.34"' in hint diff --git a/tests/utils/test_xformers_capability_gate.py b/tests/utils/test_xformers_capability_gate.py index 7514c623e3..7c2fe78c94 100644 --- a/tests/utils/test_xformers_capability_gate.py +++ b/tests/utils/test_xformers_capability_gate.py @@ -1,7 +1,33 @@ +# 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 General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + """Regression test for unslothai/unsloth#4631: xformers must not be blanket-disabled on sm_120 GPUs where its kernel actually runs (a ~57% attention-memory saving over the -SDPA packed-mask fallback). The gate now probes the real op instead of guessing by the -compute-capability major version.""" +SDPA packed-mask fallback). The gate probes the real op instead of guessing by the +compute-capability major version. + +Also covers NVIDIA QA P0-1: the probe used to be skipped entirely below sm_120, on the +assumption that xformers always works there. It does not when the wheel was built for a +different torch or CUDA major, and skipping the probe is what let a cu128-built managed +Windows package ship beside a cu130 runtime with its kernels silently dead.""" + +import contextlib +import os +import subprocess +import sys +from pathlib import Path import pytest import torch @@ -9,27 +35,41 @@ import unsloth # noqa: F401 from unsloth.utils import attention_dispatch as ad +_REPO_ROOT = Path(__file__).resolve().parents[2] + @pytest.mark.parametrize( "capability, probe_result, expect_disabled", [ - ((8, 9), None, False), # Ada: below sm_120, never probed, always kept - ((9, 0), None, False), # Hopper: below sm_120, kept - ((10, 0), None, False), # Blackwell B200 (sm_100): below sm_120, kept + # #4631: on sm_120 the answer must come from the real op, both ways. ((12, 0), True, False), # sm_120 where the kernel runs: keep xformers ((12, 0), False, True), # sm_120 where the kernel can't run: fall back to SDPA + # Every other capability is now probed too, and a working kernel is still kept: + # this is the half of #4631 that must not regress into "probe means disable". + ((7, 5), True, False), # Turing + ((8, 0), True, False), # Ampere + ((8, 9), True, False), # Ada + ((9, 0), True, False), # Hopper + ((10, 0), True, False), # Blackwell B200 (sm_100) + # P0-1: a mismatched build is dead on these too, and must be caught. + ((8, 9), False, True), + ((9, 0), False, True), + ((10, 0), False, True), ], ) def test_capability_gate(capability, probe_result, expect_disabled): + # `capability` is documentation now: the gate reads the kernel, not the number. It used + # to be read at the call site, which is what put an unguarded CUDA query at module scope. calls = {"n": 0} def probe(): calls["n"] += 1 return probe_result - assert ad._xformers_disabled_for_capability(capability, probe = probe) is expect_disabled - # Below sm_120 the probe must not run at all (no import-time kernel launch there). - assert calls["n"] == (0 if capability[0] < 12 else 1) + assert ad._xformers_disabled(probe = probe) is expect_disabled + # Exactly once, on every capability: the decision is the kernel's to make, and it is + # one tiny forward, so there is no reason to run it twice or to skip it. + assert calls["n"] == 1 @pytest.mark.skipif( @@ -50,14 +90,17 @@ def test_probe_shapes_are_valid_on_working_gpu(): @pytest.mark.parametrize( - "supports_bf16, expected_dtype", - [(True, torch.bfloat16), (False, torch.float16)], + "capability, expected_dtype", + [((8, 0), torch.bfloat16), ((7, 5), torch.float16)], ) -def test_probe_dtype_follows_bf16_support(monkeypatch, supports_bf16, expected_dtype): +def test_probe_dtype_follows_the_probed_device(monkeypatch, capability, expected_dtype): # Pre-Ampere GPUs (sm < 80: Turing/Volta, e.g. T4/V100) run xformers fine in # float16 but have no bfloat16 attention kernel, so a hardcoded bf16 probe would # raise there, get swallowed to False, and misreport a working xformers as broken. - # The probe must pick its dtype from SUPPORTS_BFLOAT16 (no Turing GPU needed here). + # + # Read off THE DEVICE BEING PROBED, not the module-level SUPPORTS_BFLOAT16, which + # describes device 0. On a mixed box where 0 is Ampere-or-newer and LOCAL_RANK selects + # a Turing card, the global says bf16 and this rank writes off a healthy install. captured = {} def fake_zeros( @@ -68,7 +111,9 @@ def test_probe_dtype_follows_bf16_support(monkeypatch, supports_bf16, expected_d captured["dtype"] = dtype raise RuntimeError("stop after capturing the probe dtype") - monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", supports_bf16) + # The opposite of the device's own answer, so a dtype taken from here fails the test. + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", capability[0] < 8) + monkeypatch.setattr(ad.torch.cuda, "get_device_capability", lambda index = None: capability) monkeypatch.setattr(ad.torch, "zeros", fake_zeros) ad._xformers_runs_on_device() # RuntimeError is swallowed; only the dtype matters assert captured["dtype"] is expected_dtype @@ -100,3 +145,304 @@ def test_probe_syncs_and_fails_on_deferred_async_error(monkeypatch): # Without the synchronize the stubbed op returns cleanly and the probe wrongly # reports True; the sync surfaces the deferred error so the probe returns False. assert ad._xformers_runs_on_device() is False + + +def test_probe_caches_the_failure_reason(monkeypatch): + # "xformers is off" with no reason is what made the mismatched Windows build so hard + # to diagnose. A failed probe must leave something a report can print. + monkeypatch.setattr(ad, "XFORMERS_PROBE_REASON", None) + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", True) + + def boom(*args, **kwargs): + raise RuntimeError("CUDA error: no kernel image is available for execution") + + monkeypatch.setattr(ad.torch, "zeros", boom) + assert ad._xformers_runs_on_device() is False + assert "no kernel image is available" in ad.XFORMERS_PROBE_REASON + assert ad.XFORMERS_PROBE_REASON.startswith("RuntimeError: ") + + +def test_probe_clears_a_stale_reason_on_success(monkeypatch): + # A later success must not leave the previous failure's reason behind, or a healthy + # install reports itself broken. + monkeypatch.setattr(ad, "XFORMERS_PROBE_REASON", "RuntimeError: stale") + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", True) + monkeypatch.setattr(ad.torch, "zeros", lambda *a, **k: object()) + monkeypatch.setattr( + ad, + "xformers", + type( + "X", + (), + { + "attn_bias": type( + "B", + (), + { + "BlockDiagonalCausalMask": type( + "M", (), {"from_seqlens": staticmethod(lambda seqlens: None)} + ) + }, + ) + }, + ), + ) + monkeypatch.setattr(ad, "xformers_attention", lambda *a, **k: None) + monkeypatch.setattr(ad.torch.cuda, "synchronize", lambda *a, **k: None) + assert ad._xformers_runs_on_device() is True + assert ad.XFORMERS_PROBE_REASON is None + + +def test_probe_never_raises_even_when_xformers_is_none(monkeypatch): + # Patch the globals the probe writes through, or its `global` assignment leaks the + # failure into every later test in the session. + monkeypatch.setattr(ad, "XFORMERS_PROBE_REASON", None) + monkeypatch.setattr(ad, "XFORMERS_PROBE_INCONCLUSIVE", False) + # The probe runs at import time on every CUDA capability now, so anything it touches + # being missing or broken must degrade to False, never to an ImportError at `import + # unsloth`. + monkeypatch.setattr(ad, "xformers", None) + monkeypatch.setattr(ad, "xformers_attention", None) + assert ad._xformers_runs_on_device() is False + assert ad.XFORMERS_PROBE_REASON + + +@pytest.mark.parametrize( + "message, inconclusive", + [ + ("CUDA out of memory. Tried to allocate 2.00 GiB", True), + ("CUDA error: all CUDA-capable devices are busy or unavailable", True), + ("CUDA error: no kernel image is available for execution on the device", False), + ("undefined symbol: _ZN3c105ErrorC1E", False), + ], +) +def test_a_busy_or_full_gpu_does_not_count_as_a_broken_build(monkeypatch, message, inconclusive): + # Device 0 being full, or claimed by another rank under EXCLUSIVE_PROCESS, says + # nothing about the wheel. Turning memory-efficient attention off for the whole + # process on that basis is a silent 2x memory regression caused by the probe itself. + monkeypatch.setattr(ad, "XFORMERS_PROBE_REASON", None) + monkeypatch.setattr(ad, "XFORMERS_PROBE_INCONCLUSIVE", False) + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", True) + + def boom(*args, **kwargs): + raise RuntimeError(message) + + monkeypatch.setattr(ad.torch, "zeros", boom) + assert ad._xformers_runs_on_device() is False + assert ad.XFORMERS_PROBE_INCONCLUSIVE is inconclusive + + +def test_the_probe_targets_this_rank_s_device(monkeypatch): + # Under torchrun each rank owns a different GPU, and on a mixed box device 0 is often + # the small display card. Probing 0 for everyone lets a wheel with no kernel for the + # weakest GPU disable xformers on the good ones. + captured = {} + + def fake_zeros( + *args, + device = None, + **kwargs, + ): + captured["device"] = device + raise RuntimeError("stop after capturing the device") + + # torch.cuda.device is stubbed so this runs on a host with any number of GPUs, and so + # the index the context is entered with can be asserted directly. + @contextlib.contextmanager + def fake_device(index): + captured["context"] = index + yield + + monkeypatch.setattr(ad, "XFORMERS_PROBE_REASON", None) + monkeypatch.setattr(ad, "XFORMERS_PROBE_INCONCLUSIVE", False) + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", True) + monkeypatch.setattr(ad, "_PROBE_DEVICE_INDEX", 3) + monkeypatch.setattr(ad.torch.cuda, "device", fake_device) + monkeypatch.setattr(ad.torch, "zeros", fake_zeros) + ad._xformers_runs_on_device() + assert captured["device"] == "cuda:3" + # The attn_bias must land there too. BlockDiagonalCausalMask builds its seqstart + # tensors on the CURRENT device, so without this context q went to cuda:N and the bias + # to cuda:0, and xformers rejected the pair -- every rank but 0 lost xformers on a + # healthy install, which is the exact silent downgrade this gate exists to prevent. + assert captured["context"] == 3 + + +def test_the_probe_device_is_clamped_to_a_device_that_exists(): + """LOCAL_RANK is a rank, not an index into the visible devices. + + Slurm with --gpus-per-task=1, and anything that narrows CUDA_VISIBLE_DEVICES per rank, + gives a rank one visible device while still exporting its global rank. accelerate and + transformers also use -1 for "not distributed". torch raises on an invalid ordinal and + the capability read is at module scope, so an unclamped index makes `import unsloth` + itself crash. + """ + visible = ad.torch.cuda.device_count() if ad.torch.cuda.is_available() else 0 + assert 0 <= ad._PROBE_DEVICE_INDEX < max(visible, 1) + if visible: + # Would raise "Invalid device id" for an out-of-range index. + ad.torch.cuda.get_device_capability(ad._PROBE_DEVICE_INDEX) + + +@pytest.mark.parametrize("local_rank", ["3", "-1", "07", "abc", ""]) +def test_import_survives_a_local_rank_that_names_no_visible_device(local_rank): + """The regression test the mocked one above cannot be: a real import, real env.""" + if not (ad.torch.cuda.is_available() and ad.torch.cuda.device_count() >= 1): + pytest.skip("needs at least one CUDA device") + env = { + **os.environ, + "LOCAL_RANK": local_rank, + # One visible device, so any rank above 0 is out of range. + "CUDA_VISIBLE_DEVICES": (os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]), + } + result = subprocess.run( + [ + sys.executable, + "-c", + "import unsloth;" + "from unsloth.utils import attention_dispatch as a;" + "print('IDX', a._PROBE_DEVICE_INDEX)", + ], + capture_output = True, + text = True, + timeout = 600, + env = env, + cwd = str(_REPO_ROOT), + ) + assert ( + result.returncode == 0 + ), f"import unsloth crashed with LOCAL_RANK={local_rank!r}:\n{result.stderr[-2000:]}" + assert "IDX 0" in result.stdout, result.stdout[-2000:] + + +def test_no_unguarded_cuda_query_runs_at_import(): + """A CUDA capability query is not safe to run at module scope. + + The driver refuses it when the device is busy, in exclusive-compute mode, or otherwise + unavailable, and at module scope that is `import unsloth` raising -- over a diagnostic + whose worst outcome is "keep xformers on and let the real forward decide". Every such + read has to sit inside a function that can answer "unknown".""" + import ast + + tree = ast.parse((_REPO_ROOT / "unsloth" / "utils" / "attention_dispatch.py").read_text()) + offenders = [] + for node in tree.body: + # Module scope only. A call inside a def runs when something calls it, and every + # such caller here is already guarded. + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + for child in ast.walk(node): + if not isinstance(child, ast.Call): + continue + if isinstance(child.func, ast.Attribute) and child.func.attr in ( + "get_device_capability", + "get_device_properties", + ): + offenders.append(ast.unparse(child)) + assert offenders == [], f"unguarded CUDA query at import: {offenders}" + + +def test_an_unavailable_device_leaves_the_fp32_answer_unknown(monkeypatch): + # Same query, same refusal. It must degrade to "cannot tell" rather than propagate. + monkeypatch.setattr(ad.torch.cuda, "is_available", lambda: True) + + def _refuse(index = None): + raise RuntimeError("CUDA error: device is currently in use by another process") + + monkeypatch.setattr(ad.torch.cuda, "get_device_capability", _refuse) + assert ad._probe_device_major() is None + + +def test_the_model_code_reads_the_probed_verdict_not_the_bare_import(): + """`HAS_XFORMERS = xformers is not None` recomputed in llama.py ignored the probe. + + Mistral answers "xFormers is on" by skipping the 4D sliding-window mask and letting the + xFormers bias carry the window. With the dispatcher already fallen back to SDPA, that + mask is the only thing making the window local, so every sequence longer than + config.sliding_window silently attended to the whole causal history.""" + import ast + + tree = ast.parse((_REPO_ROOT / "unsloth" / "models" / "llama.py").read_text()) + assigned = [ + target.id + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) and target.id == "HAS_XFORMERS" + ] + assert assigned == [], "llama.py recomputes HAS_XFORMERS instead of taking the probed one" + imported = [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and (node.module or "").endswith("attention_dispatch") + for alias in node.names + ] + assert "HAS_XFORMERS" in imported + + from unsloth.models import llama + from unsloth.utils import attention_dispatch + + assert llama.HAS_XFORMERS is attention_dispatch.HAS_XFORMERS + + +def test_an_exclusive_mode_refusal_keeps_xformers_on(monkeypatch): + """ "device is currently in use by another process" is the driver saying the GPU is + someone else's right now, not that the wheel is broken. Recorded as conclusive, it turned + a transient into a process-wide 2x memory regression -- caused by the diagnostic.""" + monkeypatch.setattr( + ad.torch.cuda, + "get_device_capability", + lambda index = None: (_ for _ in ()).throw( + RuntimeError("CUDA error: device is currently in use by another process") + ), + ) + assert ad._xformers_runs_on_device() is False + assert ( + ad.XFORMERS_PROBE_INCONCLUSIVE is True + ), "an inconclusive failure must leave xformers enabled for the real forward to decide" + + +def test_the_probed_device_follows_the_caller_selection(monkeypatch): + """A single-process app that calls torch.cuda.set_device(1) before importing us has said + where its work goes. Probing 0 anyway can disable xformers over a card nothing will + touch, and creates a context on it to do so. LOCAL_RANK still wins when it is usable.""" + monkeypatch.setattr(ad.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(ad.torch.cuda, "device_count", lambda: 4) + monkeypatch.setattr(ad.torch.cuda, "current_device", lambda: 1) + + monkeypatch.delenv("LOCAL_RANK", raising = False) + assert ad._resolve_probe_device_index() == 1 + + monkeypatch.setenv("LOCAL_RANK", "3") + assert ad._resolve_probe_device_index() == 3 + + # Out of range, the not-distributed sentinel, and junk all fall back to the selection. + for value in ("9", "-1", "not a rank", ""): + monkeypatch.setenv("LOCAL_RANK", value) + assert ad._resolve_probe_device_index() == 1 + + # And a current_device that raises still lands on 0 rather than propagating. + monkeypatch.setattr( + ad.torch.cuda, + "current_device", + lambda: (_ for _ in ()).throw(RuntimeError("no context")), + ) + assert ad._resolve_probe_device_index() == 0 + + +def test_the_inconclusive_branch_can_actually_read_the_logging_flag(): + """The inconclusive arm prints behind UNSLOTH_ENABLE_LOGGING at MODULE scope. + + attention_dispatch pulls _utils in with `import *`, and UNSLOTH_ENABLE_LOGGING is not in + _utils.__all__, so the name only exists here because it is imported explicitly. Without + that import a busy or out-of-memory GPU -- newly classified as inconclusive -- raises + NameError during `import unsloth` instead of keeping xformers on. + """ + from unsloth.models import _utils + + assert "UNSLOTH_ENABLE_LOGGING" not in getattr( + _utils, "__all__", () + ), "if the flag is exported, this explicit import can go -- but not before" + assert hasattr( + ad, "UNSLOTH_ENABLE_LOGGING" + ), "the inconclusive branch reads this name at import time" diff --git a/tests/utils/test_xformers_compat.py b/tests/utils/test_xformers_compat.py new file mode 100644 index 0000000000..b4bf555f3f --- /dev/null +++ b/tests/utils/test_xformers_compat.py @@ -0,0 +1,405 @@ +# 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 General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +"""Tests for unsloth.xformers_compat -- the offline torch <-> xformers ABI check behind +the NVIDIA P0-1 report ("xFormers was built for PyTorch 2.10.0+cu128 and Python 3.10.11 +while the app runs cu130 and Python 3.13.2, disabling its CUDA extensions"). + +Every version fact asserted here was read off the published wheels: +``pypi.org/pypi/xformers//json`` for the declared ``Requires-Dist``, and the +wheel's own ``xformers/cpp_lib.json`` for what it was actually compiled against. + +Pure stdlib module under test -- no torch, no xformers, no GPU needed. +""" + +import importlib.util +import json +import types + +import pytest + +from unsloth import xformers_compat as xc + + +# ---------------------------------------------------------------- version tables + + +@pytest.mark.parametrize( + "xformers_version, torch_release", + [ + # The old table stopped at torch 2.4; these are the releases it never covered. + ("0.0.33", "2.9.0"), + ("0.0.33.post1", "2.9.0"), + ("0.0.33.post2", "2.9.1"), # .post bumps the torch release, so it must be kept + ("0.0.34", "2.10.0"), + # 0.0.35 declares `torch>=2.10` but its _C is still one build against 2.10.0. + ("0.0.35", "2.10.0"), + ], +) +def test_expected_torch_covers_modern_releases(xformers_version, torch_release): + assert xc.expected_torch_for_xformers(xformers_version) == torch_release + + +def test_declared_pin_is_absent_for_the_range_release(): + # 0.0.35 must not appear in the `==` pin table: it declares a range, so a lookup + # there would report a pin the wheel never made. + assert "0.0.35" not in xc.XFORMERS_TORCH_PINS + assert xc.XFORMERS_BUILT_FOR_TORCH["0.0.35"] == "2.10.0" + + +@pytest.mark.parametrize( + "torch_version, xformers_version", + [ + ("2.9.0", "0.0.33.post1"), + ("2.9.1", "0.0.33.post2"), + ("2.9.1+cu128", "0.0.33.post2"), # local tag must not defeat the lookup + ("2.10.0", "0.0.34"), + ("2.10.0+cu130", "0.0.34"), + ], +) +def test_xformers_for_torch(torch_version, xformers_version): + assert xc.xformers_for_torch(torch_version) == xformers_version + + +@pytest.mark.parametrize("torch_version", ["2.11.0", "2.12.0", "2.13.0", "2.11.1+cu128"]) +def test_the_stable_abi_release_covers_a_torch_with_no_row(torch_version): + """0.0.35 is the answer above the floor, not "no release exists". + + xFormers moved to the PyTorch stable API/ABI in 0.0.34, whose notes state that binaries + targeting 2.10+ are "compatible with any later version" -- which is why + describe_xformers_mismatch accepts a 2.10-built wheel on 2.11. Returning None here made + the two halves disagree: the diagnosis said "your CUDA family is wrong" while the fix + hint said no release exists, so downgrade torch or build from source. The hint names the + CUDA index alongside the version, which is what repairs the family. + """ + assert xc.xformers_for_torch(torch_version) == "0.0.35" + + +@pytest.mark.parametrize("torch_version", ["2.9.5", "2.8.3", "2.11.0.dev20260101"]) +def test_a_torch_the_guarantee_does_not_reach_is_still_unknown(torch_version): + # Below the floor there is no stable ABI to lean on, and a pre-release is not something + # to make a compatibility promise about. Naming a version there sends the user back into + # the same mismatch. + assert xc.xformers_for_torch(torch_version) is None + + +def test_inverse_table_agrees_with_the_forward_tables(): + # TORCH_TO_XFORMERS is the inverse with posts winning; drift between the two is a + # silent way to recommend a wheel built for a different torch. + for torch_release, xformers_version in xc.TORCH_TO_XFORMERS.items(): + assert xc.expected_torch_for_xformers(xformers_version) == torch_release + + +# ---------------------------------------------------------------- version parsing + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("2.10.0+cu130", "2.10.0"), + ("2.10.0", "2.10.0"), + ("2.11.0.dev20260101+cu130", "2.11.0"), + ("2.10.0a0+gitabc123", "2.10.0"), + ("", None), + (None, None), + ("not-a-version", None), + ], +) +def test_normalize_release(raw, expected): + assert xc.normalize_release(raw) == expected + + +@pytest.mark.parametrize("version", ["0.0.35.dev1130", "0.0.36rc1", "0.0.36a0"]) +def test_a_pre_release_build_is_unknown_rather_than_folded_onto_its_release(version): + # The fourteen 0.0.35.devNNNN wheels on PyPI are built against torch nightlies, so + # answering "0.0.35, therefore torch 2.10.0" for one of them is confidently wrong. + assert xc.normalize_release_with_post(version) is None + assert xc.expected_torch_for_xformers(version) is None + + +def test_normalize_release_with_post_keeps_the_post(): + # 0.0.33 -> torch 2.9.0 but 0.0.33.post2 -> torch 2.9.1, so dropping the post here + # would silently answer for the wrong wheel. + assert xc.normalize_release_with_post("0.0.33.post2") == "0.0.33.post2" + assert xc.normalize_release("0.0.33.post2") == "0.0.33" + + +@pytest.mark.parametrize( + "torch_version, major", + [ + ("2.10.0+cu130", 13), + ("2.10.0+cu128", 12), + ("2.4.1+cu118", 11), + ("2.10.0", None), + ("2.10.0+cpu", None), + ("2.10.0+rocm6.4", None), + ], +) +def test_cuda_major_from_torch_version(torch_version, major): + assert xc.cuda_major_from_torch_version(torch_version) == major + + +@pytest.mark.parametrize( + "raw, formatted", + [(1208, "12.8"), (1300, "13.0"), (1126, "11.26"), (None, None), (True, None)], +) +def test_format_build_cuda(raw, formatted): + # cpp_lib.json stores major * 100 + minor. xformers' own exception prints the raw + # integer ("with CUDA 1208"), which is what makes its message unreadable. + assert xc.format_build_cuda(raw) == formatted + + +# ---------------------------------------------------------------- build metadata + + +def _cpp_lib( + torch = "2.10.0+cu128", + cuda = 1208, + python = "3.10.11", + hip = None, +): + """The real shape of an xformers wheel's cpp_lib.json (verified against 0.0.34).""" + return { + "version": {"cuda": cuda, "hip": hip, "torch": torch, "python": python}, + "env": {"XFORMERS_PACKAGE_FROM": "wheel-v0.0.34"}, + } + + +def test_build_metadata_is_read_from_disk_without_importing_xformers(tmp_path, monkeypatch): + package = tmp_path / "xformers" + package.mkdir() + (package / "cpp_lib.json").write_text(json.dumps(_cpp_lib()), encoding = "utf-8") + (package / "__init__.py").write_text("raise AssertionError('must not be imported')") + + spec = importlib.util.spec_from_file_location( + "xformers", package / "__init__.py", submodule_search_locations = [str(package)] + ) + monkeypatch.setattr( + importlib.util, "find_spec", lambda name, *a, **k: spec if name == "xformers" else None + ) + # The __init__ above blows up if executed; getting a result proves we only located it. + assert xc.xformers_build_metadata() == _cpp_lib() + assert xc.xformers_build_summary() == { + "torch": "2.10.0+cu128", + "cuda": "12.8", + "hip": None, + "python": "3.10.11", + } + + +def test_build_metadata_is_none_when_xformers_is_absent(monkeypatch): + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: None) + assert xc.xformers_build_metadata() is None + assert xc.xformers_build_summary() is None + + +def test_build_metadata_survives_a_raising_find_spec(monkeypatch): + def boom(name, *args, **kwargs): + raise ValueError("half-removed distribution") + + monkeypatch.setattr(importlib.util, "find_spec", boom) + # Diagnostics must never be the thing that takes the caller down. + assert xc.xformers_build_metadata() is None + + +def test_build_metadata_ignores_a_malformed_cpp_lib(tmp_path, monkeypatch): + package = tmp_path / "xformers" + package.mkdir() + (package / "cpp_lib.json").write_text("{ not json", encoding = "utf-8") + spec = types.SimpleNamespace( + submodule_search_locations = [str(package)], origin = str(package / "__init__.py") + ) + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: spec) + assert xc.xformers_build_metadata() is None + + +# ---------------------------------------------------------------- mismatch detection + + +def test_the_nvidia_p0_is_detected(): + # NVIDIA QA P0-1 exactly: the managed Windows wheel (xformers 0.0.34/0.0.35, built + # for torch 2.10.0+cu128 on Python 3.10.11) beside a cu130 / Python 3.13.2 runtime. + reason = xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu130", + torch_cuda = "13.0", + xformers_version = "0.0.34", + build_metadata = _cpp_lib(), + python_version = "3.13.2", + ) + assert reason is not None + assert "2.10.0+cu128" in reason + assert "3.10.11" in reason + assert "2.10.0+cu130" in reason + assert "3.13.2" in reason + + +def test_torch_release_mismatch_is_detected(): + # Below the stable-ABI floor a torch-release difference really is a mismatch. + reason = xc.describe_xformers_mismatch( + torch_version = "2.9.1+cu128", + xformers_version = "0.0.33.post1", + build_metadata = _cpp_lib(torch = "2.9.0+cu128"), + ) + assert reason is not None + assert "2.9.0+cu128" in reason and "2.9.1+cu128" in reason + + +def test_a_stable_abi_build_on_a_later_torch_is_not_a_mismatch(): + # xFormers moved to the PyTorch stable API/ABI in 0.0.34, and its release notes say + # 2.10+ binary builds are compatible with any later version. Calling that a mismatch + # blames a working wheel for an unrelated failure (a missing VC++ runtime, say) and + # sends the user hunting for a release that does not exist. + assert ( + xc.describe_xformers_mismatch( + torch_version = "2.12.1+cu128", + torch_cuda = "12.8", + xformers_version = "0.0.35", + build_metadata = _cpp_lib(torch = "2.10.0+cu128"), + ) + is None + ) + + +def test_the_stable_abi_guarantee_only_runs_forwards(): + # A 2.11-built wheel on torch 2.10 is not covered by anything: the guarantee is about + # LATER versions. + reason = xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu128", + torch_cuda = "12.8", + xformers_version = "0.0.35", + build_metadata = _cpp_lib(torch = "2.11.0+cu128"), + ) + assert reason is not None and "2.11.0+cu128" in reason + + +def test_a_stable_abi_build_still_reports_a_cuda_major_mismatch(): + # The ABI guarantee says nothing about CUDA: a cu128 wheel beside a cu130 runtime is + # the NVIDIA case, and it must still be named. + reason = xc.describe_xformers_mismatch( + torch_version = "2.12.1+cu130", + torch_cuda = "13.0", + xformers_version = "0.0.35", + build_metadata = _cpp_lib(torch = "2.10.0+cu128", cuda = 1208), + ) + assert reason is not None and "CUDA 13" in reason + + +def test_matching_build_reports_nothing(): + assert ( + xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu128", + torch_cuda = "12.8", + xformers_version = "0.0.34", + build_metadata = _cpp_lib(), + python_version = "3.13.2", + ) + is None + ) + + +def test_python_difference_alone_is_not_a_mismatch(): + # The wheels are abi3/none-tagged and _C is loaded via torch.ops.load_library, not + # the CPython ABI, so 3.10-built kernels run fine on 3.13. Flagging this would send + # every Studio user chasing the wrong thing. + assert ( + xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu128", + build_metadata = _cpp_lib(python = "3.10.11"), + python_version = "3.13.2", + ) + is None + ) + + +def test_cuda_minor_difference_alone_is_not_a_mismatch(): + # CUDA minor version compatibility: a cu126-built wheel loads against a cu128 torch. + assert ( + xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu128", + torch_cuda = "12.8", + build_metadata = _cpp_lib(cuda = 1206, torch = "2.10.0+cu126"), + ) + is None + ) + + +def test_cuda_major_mismatch_is_detected_without_a_local_tag(): + # Conda/source torch has no "+cuXXX", so torch.version.cuda is the only signal. + reason = xc.describe_xformers_mismatch( + torch_version = "2.10.0", + torch_cuda = "13.0", + build_metadata = _cpp_lib(torch = "2.10.0", cuda = 1208), + ) + assert reason is not None + assert "(CUDA 13.x)" in reason + + +def test_unknown_runtime_reports_nothing(monkeypatch): + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: None) + assert xc.describe_xformers_mismatch(torch_version = None) is None + assert xc.describe_xformers_mismatch(torch_version = "garbage") is None + + +def test_pin_fallback_when_the_wheel_ships_no_build_metadata(monkeypatch): + # Source/editable installs have no cpp_lib.json; the declared pin still catches a + # wholesale torch-release mismatch. + monkeypatch.setattr(xc, "xformers_build_metadata", lambda: None) + monkeypatch.setattr(xc, "declared_torch_pin", lambda version = None: "2.9.1") + reason = xc.describe_xformers_mismatch( + torch_version = "2.10.0+cu128", xformers_version = "0.0.33.post2" + ) + assert reason is not None + assert "2.9.1" in reason + + +def test_pin_fallback_stays_quiet_when_the_pin_matches(monkeypatch): + monkeypatch.setattr(xc, "xformers_build_metadata", lambda: None) + monkeypatch.setattr(xc, "declared_torch_pin", lambda version = None: "2.10.0") + assert ( + xc.describe_xformers_mismatch(torch_version = "2.10.0+cu128", xformers_version = "0.0.34") + is None + ) + + +def test_the_pin_fallback_honours_the_stable_abi_too(monkeypatch): + """Same exemption the build-metadata branch applies, on the path that has no metadata. + + A source or editable 0.0.34+ build carries the 2.10 pin it was cut against, so without + this every later torch is reported as a mismatch -- turning an unrelated extension + failure into a torch-version diagnosis with reinstall instructions that fix nothing.""" + monkeypatch.setattr(xc, "xformers_build_metadata", lambda: None) + monkeypatch.setattr(xc, "declared_torch_pin", lambda version = None: "2.10.0") + assert ( + xc.describe_xformers_mismatch(torch_version = "2.11.0+cu130", xformers_version = "0.0.34") + is None + ) + # One-directional, as ever: a 2.11-built pin on 2.10 is still a mismatch. + monkeypatch.setattr(xc, "declared_torch_pin", lambda version = None: "2.11.0") + assert ( + xc.describe_xformers_mismatch(torch_version = "2.10.0+cu130", xformers_version = "0.0.35") + is not None + ) + # And below the floor there is no guarantee to lean on. + monkeypatch.setattr(xc, "declared_torch_pin", lambda version = None: "2.9.1") + assert ( + xc.describe_xformers_mismatch(torch_version = "2.10.0+cu128", xformers_version = "0.0.33") + is not None + ) + + +def test_declared_pin_falls_back_to_the_table_for_a_different_version(): + # Asking about a version other than the resident one must use the table, not the + # resident METADATA, which describes a different wheel entirely. + assert xc.declared_torch_pin("0.0.30") == "2.7.0" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7d67356d89..b22cddff47 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -22,6 +22,8 @@ __all__ = [ "xformers", "xformers_attention", "xformers_version", + "XFORMERS_BUILD_METADATA", + "XFORMERS_BROKEN_REASON", "__version__", "importlib_version", "HAS_FLASH_ATTENTION", @@ -96,7 +98,7 @@ __all__ = [ import torch from typing import Union, Optional, List, Any, Callable, Tuple, Iterator -from platform import system as platform_system +from platform import python_version as _python_version, system as platform_system platform_system = platform_system() import numpy as np @@ -2205,13 +2207,184 @@ elif DEVICE_TYPE == "xpu": # ============================================= # Get Xformers -# Silence xformers CUDA mismatch warnings before import +from ..xformers_compat import ( + describe_xformers_mismatch, + xformers_build_metadata, + xformers_for_torch, +) + +# The installed wheel's own cpp_lib.json (torch / CUDA / Python it was compiled against), +# and a one-line reason when its kernels cannot load here. Both are exported so the Studio +# hardware report and any diagnostic can say WHAT broke instead of just going quiet. +XFORMERS_BUILD_METADATA = xformers_build_metadata() +XFORMERS_BROKEN_REASON = None + +# Resolved from the wheel on disk BEFORE importing xformers: reading cpp_lib.json costs no +# import and fires no warning, so the mismatch is known in time to decide whether to +# silence xformers' own diagnostic below. try: - _xformers_logger = logging.getLogger("xformers") - _xformers_logger.setLevel(logging.ERROR) - del _xformers_logger -except: - pass + _xformers_dist_version = importlib_version("xformers") +except Exception: + _xformers_dist_version = None +try: + _xformers_predicted_break = describe_xformers_mismatch( + torch_version = torch.__version__, + torch_cuda = getattr(torch.version, "cuda", None), + xformers_version = _xformers_dist_version, + build_metadata = XFORMERS_BUILD_METADATA, + python_version = _python_version(), + ) +except Exception: + # Diagnostics must never be the thing that stops unsloth importing. + _xformers_predicted_break = None + +# Silence xformers' CUDA mismatch chatter before the import -- EXCEPT when the wheel on +# disk already tells us it does not match this runtime. That warning is the only +# first-party diagnostic for a cu128-built wheel sitting beside a cu130 torch, so +# suppressing it in exactly that case is suppressing the bug report. +if _xformers_predicted_break is None: + try: + _xformers_logger = logging.getLogger("xformers") + _xformers_logger.setLevel(logging.ERROR) + del _xformers_logger + except: + pass + +_XFORMERS_BREAKAGE_ANNOUNCED = False + + +def _announce_xformers_breakage( + reason, + error = None, + build_mismatch = False, +): + """Print the xformers breakage once per process, on the default path. + + Deliberately not behind UNSLOTH_ENABLE_LOGGING: silently dropping to SDPA is how a + wheel built for the wrong torch shipped to users unnoticed. + + Two shapes, because this arm also catches failures that are not ABI mismatches at all + -- the sm_100/110/120 FA3 guard and the old-torch guards above both raise here with + their own multi-line, fenced, already-actionable messages. Reflowing one of those into + "its optimized kernels cannot load ... install the matching build" states the wrong + cause and mangles the text. Only ``build_mismatch`` gets the version-pin treatment; + everything else is passed through verbatim. + """ + global _XFORMERS_BREAKAGE_ANNOUNCED + if _XFORMERS_BREAKAGE_ANNOUNCED: + return + _XFORMERS_BREAKAGE_ANNOUNCED = True + # What actually picks up the work. select_attention_backend puts FlashAttention ahead of + # xformers, so on a dual install a broken xformers costs nothing and "falling back to SDPA, + # it uses more memory" is simply false -- a performance warning about a run that is on the + # faster kernel of the two. The probe runs on that host anyway (it has to, to report the + # install honestly), so this is the common case, not a corner. + fallback = ( + "FlashAttention is handling attention instead, so this costs no memory" + if HAS_FLASH_ATTENTION + else "Falling back to PyTorch SDPA attention - training still works, but it uses " + "more memory" + ) + if build_mismatch: + print( + "Unsloth: Xformers is installed but its optimized kernels cannot load.\n" + f"{str(reason).strip().rstrip('.')}.\n" + f"{fallback}.\n" + f"{_xformers_fix_hint()}" + ) + else: + print(f"Unsloth: Xformers is installed but could not be used. {fallback}.\n" f"{reason}") + if UNSLOTH_ENABLE_LOGGING and error is not None: + print(str(error)) + + +def _xformers_torch_index_url(): + """The download.pytorch.org index serving wheels for the resident CUDA build, or None. + + ``torch.version.cuda`` ('13.0') is the authority for the family, not the local tag, + which a source or nightly build may not carry. None on a CPU / ROCm / XPU torch, where + there is no CUDA-matched xformers index to point at. + """ + try: + cuda = getattr(getattr(torch, "version", None), "cuda", None) + if not cuda: + return None + major, _, minor = str(cuda).partition(".") + family = f"cu{int(major)}{int(minor or 0)}" + except Exception: + return None + import os + + base = os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" + redacted = _redact_index_url(base.rstrip("/")) + return f"{redacted}/{family}" if redacted else None + + +def _redact_index_url(url): + """Strip userinfo, query and fragment from a wheel index before it is PRINTED. + + UNSLOTH_PYTORCH_MIRROR can carry credentials (``https://user:token@host/simple``) or a + signed-URL query, and this hint goes to stdout on the default import path -- into logs, + CI output and shared notebooks. pip redacts the same thing in its own output; the URL is + still copy-pasteable, since the user's own mirror config supplies the secret back. + + Returns None when the URL cannot be parsed: the caller then omits ``--index-url`` + entirely rather than printing an authority we could not take apart. + """ + try: + from urllib.parse import urlsplit, urlunsplit + + parts = urlsplit(str(url)) + if not parts.netloc and not parts.query and not parts.fragment: + return str(url) + host = parts.hostname or "" + if parts.port: + host = f"{host}:{parts.port}" + if parts.username: + host = f"***@{host}" + return urlunsplit((parts.scheme, host, parts.path, "", "")) + except Exception: # noqa: BLE001 -- an unparseable mirror must not break the hint + # And must not be echoed back either. A malformed authority (a bad IPv6 bracket, a + # non-numeric port) is exactly the input whose credentials or signed query this + # function exists to keep out of stdout, so the fallback drops the URL rather than + # printing the one thing it was asked to hide. The caller omits --index-url on None. + return None + + +def _xformers_fix_hint(): + """How to repair the install, pinned when we know the matching release. + + On a torch newer than any xformers release there IS no version to pin -- naming one + would send the user straight back into the same mismatch -- so say so instead. + """ + try: + matching = xformers_for_torch(torch.__version__) + except Exception: + matching = None + if matching is not None: + # WHERE from, not just which version. cu126 / cu128 / cu130 publish the same + # xformers version string, and PyPI carries only the CUDA-12.8 flavour -- so a + # bare version pin on the reported case (torch cu130 beside a cu128-built 0.0.34) + # force-reinstalls the very wheel that is already broken and changes nothing. + # Naming the matching CUDA index is what actually repairs it. + index = _xformers_torch_index_url() + index_flag = f" --index-url {index}" if index else "" + return ( + "To fix, install the xformers build that matches your torch:\n" + f'\npip install --no-deps --force-reinstall{index_flag} "xformers=={matching}"\n' + "\nRun `python -m xformers.info` to see xformers' own report." + ) + return ( + f"No xformers release is built for torch {torch.__version__.split('+')[0]} yet. " + "Either downgrade torch to a version xformers ships wheels for, or build xformers " + "from source:\n" + "\npip install ninja\n" + "pip install -v --no-build-isolation -U " + "git+https://github.com/facebookresearch/xformers.git@main#egg=xformers\n" + "\nRun `python -m xformers.info` to see xformers' own report." + ) + + try: from xformers import __version__ as xformers_version @@ -2284,14 +2457,36 @@ try: import xformers.ops.fmha as xformers xformers_attention = xformers.memory_efficient_attention + # The extensions loaded, so whatever the tables predicted did not happen. Never cry + # wolf: a working xformers must not be reported as broken. + XFORMERS_BROKEN_REASON = None + # And put the logger back where the unpredicted case leaves it. The silencing above is + # skipped on a PREDICTED break so xformers' own diagnostic can get through, but the + # prediction is read off cpp_lib.json and is wrong in the healthy direction whenever a + # wheel records a torch patch it still loads against. Leaving the logger open then is + # permanent and process-wide, so unrelated xformers warnings -- flash3's "package can't + # be used", which fires on exactly this hardware -- start reaching users who have + # nothing wrong with their install. + if _xformers_predicted_break is not None: + try: + logging.getLogger("xformers").setLevel(logging.ERROR) + except Exception: + pass except ModuleNotFoundError: + # Not installed at all. Nothing to warn about - SDPA is the expected path here. xformers = None xformers_attention = None xformers_version = None except Exception as e: - if UNSLOTH_ENABLE_LOGGING: - print("========\nSwitching to PyTorch attention since your Xformers is broken.\n========\n") - print(str(e)) + # Installed but dead. Prefer the reason read off cpp_lib.json (it names the torch and + # CUDA the wheel was built for) over the raised message, which for the ABI case is + # xformers' own text with the CUDA version printed as a raw integer like 1208. + XFORMERS_BROKEN_REASON = _xformers_predicted_break or str(e) + _announce_xformers_breakage( + XFORMERS_BROKEN_REASON, + e, + build_mismatch = _xformers_predicted_break is not None, + ) xformers = None xformers_attention = None xformers_version = None diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c80f274ecc..d73942c27f 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -40,6 +40,7 @@ from ..utils.packing import ( from ..utils.attention_dispatch import ( AttentionConfig, AttentionContext, + HAS_XFORMERS, run_attention, SDPA, select_attention_backend, @@ -127,7 +128,11 @@ except: from huggingface_hub.utils._token import get_token from triton import __version__ as triton_version -HAS_XFORMERS = xformers is not None +# Not `xformers is not None`: attention_dispatch probes the install and turns HAS_XFORMERS +# off when the library imports but has no kernel that runs here. Recomputing it from the +# import alone left the model code on the xFormers path the dispatcher had already left -- +# and Mistral answers "xFormers" by skipping the 4D sliding-window mask, so every sequence +# longer than config.sliding_window attended to the whole causal history on SDPA instead. BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 4350565fe2..5462d354c2 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -109,7 +109,14 @@ def MistralAttention_fast_forward( # Attention module sw_cfg = getattr(self.config, "sliding_window", None) - sw = kv_seq_len if (sw_cfg is None or sw_cfg == "null") else sw_cfg + # A non-positive window means "no local attention", the same as absent -- which is how + # the mask builders above read it. Passing 0 through made window_size (0, 0) for flash + # and, once the dispatcher started honouring the window, an all-false SDPA mask: the + # lower bound q_pos - (0 - 1) sits above the causal upper bound, so nothing is visible. + if sw_cfg is None or sw_cfg == "null" or (isinstance(sw_cfg, int) and sw_cfg <= 0): + sw = kv_seq_len + else: + sw = sw_cfg window_size = (-1, -1) if (kv_seq_len <= sw) else (sw, sw) use_varlen = seq_info is not None and past_key_value is None and window_size == (-1, -1) @@ -138,6 +145,12 @@ def MistralAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + # The window the flash path already gets through window_size. SDPA needs it too, and + # not only in the branch above: training takes `elif self.training: pass`, so no 4D + # mask is synthesized, and with xformers off and flash absent the local window was the + # one thing nothing carried -- every sequence past config.sliding_window silently + # attended its whole causal history. + sliding_window = None if window_size == (-1, -1) else sw, prefix_seg_info = _pg_seg, ) diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 54f8100ca1..89bfd71540 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -26,6 +26,10 @@ from torch import Tensor from torch.nn.functional import scaled_dot_product_attention from ..models._utils import * +from ..models._utils import _announce_xformers_breakage # not in __all__, needed by the probe gate +from ..models._utils import ( + UNSLOTH_ENABLE_LOGGING, +) # ditto: the inconclusive branch reads it at import time from ..utils.packing import ( build_sdpa_packed_attention_mask, build_xformers_block_causal_mask, @@ -38,43 +42,201 @@ if HAS_FLASH_ATTENTION: HAS_XFORMERS = xformers is not None +# Why the on-device probe last failed, or None when it passed or never ran. Cached +# alongside the boolean so callers can report WHY xformers went away, not just that it did. +XFORMERS_PROBE_REASON: Optional[str] = None + +# True when the probe failed for a reason that says nothing about the build: the GPU was +# busy, out of memory, or otherwise unavailable to us right now. +XFORMERS_PROBE_INCONCLUSIVE = False + +# Failures that mean "ask again later", not "this xformers is broken". Disabling +# memory-efficient attention for the whole process because device 0 happened to be full, +# or was claimed by another rank in EXCLUSIVE_PROCESS mode, is a silent 2x memory +# regression caused by the diagnostic itself. +_INCONCLUSIVE_PROBE_ERRORS = ( + "out of memory", + "busy or unavailable", + "all cuda-capable devices are busy", + "no cuda-capable device", + "cuda_error_not_permitted", + "insufficient driver", + "initialization error", + # Belt and braces for the device index. It is clamped below, so this should be + # unreachable -- but if it ever is reached, "we aimed at a device that is not there" + # must not be recorded as "your xformers is broken" and disable it process-wide. + "invalid device ordinal", + "invalid device id", + # EXCLUSIVE_PROCESS wording that the phrases above do not cover. The driver says this + # when another process holds the device; the wheel was never tested, so turning + # xformers off process-wide on the strength of it is the same silent regression. + "currently in use", + "in use by another process", + "exclusive", +) + + +# Which device to probe. Under torchrun each rank owns a different GPU, and on a mixed box +# device 0 is often the small display card, so probing 0 for everyone lets a wheel with no +# kernel for the weakest GPU disable xformers on the good ones. +# +# LOCAL_RANK is NOT an index into the devices this process can see. Slurm with +# --gpus-per-task=1, and anything that narrows CUDA_VISIBLE_DEVICES per rank, gives every +# rank one visible device while still exporting its global rank -- so rank 3 sees exactly +# one GPU and LOCAL_RANK says 3. accelerate and transformers also use -1 as their "not +# distributed" sentinel. Both are out of range, torch.cuda.get_device_capability raises on +# an invalid ordinal, and that call is at module scope, so an unclamped index turns +# `import unsloth` into a crash on an ordinary launch. Fall back to 0, which is the only +# device such a rank has. +# +# With no usable LOCAL_RANK, the device the CALLER already selected: a single-process +# application that runs torch.cuda.set_device(1) before importing us is telling us where its +# work goes, and probing 0 anyway can disable xformers over a card nothing will touch -- and +# creates a context on it while doing so. +def _resolve_probe_device_index() -> int: + count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + if count <= 0: + return 0 + try: + rank = int(os.environ.get("LOCAL_RANK", "") or -1) + except ValueError: + rank = -1 + if 0 <= rank < count: + return rank + try: + current = int(torch.cuda.current_device()) + except Exception: + return 0 + return current if 0 <= current < count else 0 + + +_PROBE_DEVICE_INDEX = _resolve_probe_device_index() + + def _xformers_runs_on_device() -> bool: - """One tiny attention forward; True iff the xformers kernel actually runs here.""" + """One tiny attention forward; True iff the xformers kernel actually runs here. + + Never raises. Every failure becomes False plus a one-line XFORMERS_PROBE_REASON, + because this runs at import and a diagnostic must not be what breaks the import. + """ + global XFORMERS_PROBE_REASON, XFORMERS_PROBE_INCONCLUSIVE try: # Pre-Ampere GPUs (sm < 80: Turing/Volta) have no bfloat16 attention kernel # but run xformers fine in float16, so pick the dtype the device supports. - dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16 - q = torch.zeros((1, 8, 1, 64), device = "cuda", dtype = dtype) - attn_bias = xformers.attn_bias.BlockDiagonalCausalMask.from_seqlens([8]) - xformers_attention(q, q, q, attn_bias = attn_bias) - # Launches are async; synchronize so a deferred kernel failure fails the probe here. - torch.cuda.synchronize() + # + # Read off the device being PROBED, not the module-level SUPPORTS_BFLOAT16, which + # describes device 0. On a mixed box where 0 is Ampere-or-newer and this rank owns + # a Turing card, the global says bf16, the kernel rejects it, and a healthy + # xformers is recorded as broken for the whole process. + dtype = ( + torch.bfloat16 + if torch.cuda.get_device_capability(_PROBE_DEVICE_INDEX)[0] >= 8 + else torch.float16 + ) + device = f"cuda:{_PROBE_DEVICE_INDEX}" + # Under the device context, not just device= on the tensor. BlockDiagonalCausalMask + # builds its seqstart tensors on the CURRENT device, and at import time that is + # still cuda:0 on every rank -- launchers set LOCAL_RANK in the environment but + # torch.cuda.set_device happens later, inside the trainer. So q lands on cuda:N and + # the bias on cuda:0, xformers rejects the pair, and the probe fails on every rank + # but zero. That is the silent drop to SDPA this whole gate exists to prevent, on a + # healthy install, caused by the diagnostic itself -- and it allocates on cuda:0 + # from every rank as well, pinning a second context per rank. + with torch.cuda.device(_PROBE_DEVICE_INDEX): + q = torch.zeros((1, 8, 1, 64), device = device, dtype = dtype) + attn_bias = xformers.attn_bias.BlockDiagonalCausalMask.from_seqlens([8]) + xformers_attention(q, q, q, attn_bias = attn_bias) + # Launches are async; synchronize so a deferred kernel failure fails the probe here. + torch.cuda.synchronize(device) + XFORMERS_PROBE_REASON = None + XFORMERS_PROBE_INCONCLUSIVE = False return True - except Exception: + except Exception as error: + XFORMERS_PROBE_REASON = f"{type(error).__name__}: {error}".strip() + text = str(error).lower() + XFORMERS_PROBE_INCONCLUSIVE = any(marker in text for marker in _INCONCLUSIVE_PROBE_ERRORS) return False -def _xformers_disabled_for_capability(capability, probe = _xformers_runs_on_device) -> bool: - # At sm_120 (RTX 50-series) xformers' cutlass op is capability-rejected (caps at - # sm_90) and its flash-2 op runs only if the build ships an sm_120 kernel, so run - # one real forward to decide. Below sm_120 xformers always works; skip the probe. - if capability[0] < 12: - return False +def _xformers_disabled(probe = _xformers_runs_on_device) -> bool: + # Probe on EVERY capability, not just sm_120+. The old gate returned early below + # sm_120 on the assumption that xformers always works there, which only holds when + # the wheel matches the runtime: a cu128-built xformers on a cu130 torch is just as + # dead on an sm_90 Hopper, and never probing there is what let a mismatched managed + # Windows package ship (NVIDIA P0-1). + # + # At sm_120 (RTX 50-series) there is a second, unrelated reason to probe: xformers' + # cutlass op is capability-rejected (it caps at sm_90) and its flash-2 op runs only + # if the build ships an sm_120 kernel (unslothai/unsloth#4631). + # + # No capability argument: the answer is always the real op now, and reading the + # capability at the call site put an UNGUARDED torch.cuda.get_device_capability() at + # module scope. CUDA refuses that query when the device is busy, in exclusive mode or + # temporarily unavailable, and there it raised before the probe could classify the + # failure as inconclusive -- turning `import unsloth` into a crash over a diagnostic + # whose worst answer is "keep xformers on and let the forward decide". return not probe() -# FlashAttention always wins in select_attention_backend and nothing downgrades -# flash -> xformers, so when it's installed xformers is never selected: skip the probe. -if HAS_XFORMERS and not HAS_FLASH_ATTENTION and torch.cuda.is_available(): - if _xformers_disabled_for_capability(torch.cuda.get_device_capability()): - HAS_XFORMERS = False +# Probe whenever xformers imported, including when flash-attn is installed and will win +# select_attention_backend anyway: a dead xformers is worth knowing about either way, and +# reporting it is the point. Cost is one 1x8x1x64 forward on a CUDA context torch has +# already initialised (_XFORMERS_FP32_UNSUPPORTED below forces the same lazy init). +XFORMERS_DISABLED_REASON = XFORMERS_BROKEN_REASON +if HAS_XFORMERS and torch.cuda.is_available(): + if _xformers_disabled(): + if XFORMERS_PROBE_INCONCLUSIVE: + # The GPU was busy or full, which says nothing about the build. Keep xformers + # and let the real forward pass decide; turning it off here would be a silent + # memory regression caused by the diagnostic. + if UNSLOTH_ENABLE_LOGGING: + print( + f"Unsloth: Could not probe xformers ({XFORMERS_PROBE_REASON}); keeping it on." + ) + else: + HAS_XFORMERS = False + XFORMERS_DISABLED_REASON = XFORMERS_PROBE_REASON + # Say so. A probe that turns off memory-efficient attention and prints nothing + # is the same silent downgrade this whole change exists to remove. + # + # First line by default, the rest behind UNSLOTH_ENABLE_LOGGING. This reason is + # a captured exception, and xformers answers a capability rejection with a dump + # of every operator it considered and why -- a dozen lines. Announcing that + # verbatim on the default path would put a wall of text in front of every user + # of an affected card and bury the one sentence that matters. Truncating HERE + # rather than in the announcer, because the announcer's other callers pass + # deliberately multi-line, fenced, copy-pasteable instructions that have to + # arrive intact. + _probe_head, _, _probe_rest = str(XFORMERS_PROBE_REASON).strip().partition("\n") + _announce_xformers_breakage( + _probe_head, + _probe_rest.strip() or None, + ) + # On sm_100+ (B200, sm_120) xformers' fp32-capable cutlass op is capability-rejected and # only its fp16/bf16 flash-2 op runs, so fp32 Q/K/V (DoRA, #1013) must be downcast there; -# below sm_100 cutlass handles fp32 natively. Read once from device 0, like the probe gate. -_XFORMERS_FP32_UNSUPPORTED = ( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 -) +# below sm_100 cutlass handles fp32 natively. Read once from the same device the probe gate +# above used, so the two answers describe the same GPU: on a mixed box, reading the fp32 +# capability off device 0 while probing this rank's device is how a display card ends up +# deciding downcast policy for a compute card. +def _probe_device_major() -> Optional[int]: + """Major compute capability of the probed device, or None if CUDA will not say. + + Guarded for the same reason the probe is: a busy or exclusive-mode device makes this + query raise, and at module scope that is an import crash rather than a missing answer. + """ + if not torch.cuda.is_available(): + return None + try: + return torch.cuda.get_device_capability(_PROBE_DEVICE_INDEX)[0] + except Exception: + return None + + +# None (unknown) is treated as supported: downcasting fp32 that did not need it is a +# quality regression, and the fp16/bf16 paths are unaffected either way. +_XFORMERS_FP32_UNSUPPORTED = (_probe_device_major() or 0) >= 10 SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") # PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load @@ -234,6 +396,11 @@ def run_attention( kv_seq_len = context.kv_seq_len requires_grad = context.requires_grad sliding_window = context.sliding_window + # A non-positive window means "no local attention", not "a window of nothing": a config + # spelling it 0 would otherwise put the mask's lower bound above its causal upper bound + # and hide every position from every other. + if sliding_window is not None and sliding_window <= 0: + sliding_window = None # DoRA promotes q/k/v_proj outputs to fp32, which FlashAttention rejects (and so does # the xformers flash-2 op on sm_100+, see _XFORMERS_FP32_UNSUPPORTED), so downcast any @@ -393,6 +560,18 @@ def run_attention( no_allowed = ~local_mask.any(dim = -1, keepdim = True) # (bsz,1,q_len,1) local_mask = local_mask | no_allowed + if local_mask is None and sliding_window is not None and k_len_local > sliding_window: + # SDPA's is_causal is FULL causal; it has no window. With no padding mask to + # hang the window off, a model whose config declares one attended its whole + # history the moment neither the xformers bias nor flash's window_size was the + # thing running -- which is exactly the SDPA fallback this probe can now cause. + q_pos = torch.arange(k_len_local - q_len_local, k_len_local, device = Q.device) + k_pos = torch.arange(k_len_local, device = Q.device) + local_mask = ( + (k_pos[None, :] <= q_pos[:, None]) + & (k_pos[None, :] >= (q_pos[:, None] - (sliding_window - 1))) + )[None, None, :, :] + is_causal_local = local_mask is None and q_len_local == k_len_local kwargs = dict(sdpa_kwargs) diff --git a/unsloth/xformers_compat.py b/unsloth/xformers_compat.py new file mode 100644 index 0000000000..49571c451c --- /dev/null +++ b/unsloth/xformers_compat.py @@ -0,0 +1,470 @@ +# 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 General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +"""torch <-> xformers ABI compatibility, resolved without importing xformers or torch. + +Every xformers wheel is compiled against ONE torch release AND one CUDA major. Getting +either wrong does not fail the install: the wheels are ``cp39-abi3`` (or ``py39-none`` +from 0.0.35), so a wheel built on CPython 3.10 drops cleanly into 3.13, and +``pip install --no-deps`` never reads the wheel's ``Requires-Dist: torch==X`` at all. +What fails is ``torch.ops.load_library("xformers/_C.pyd")`` at import time, which +``xformers/_cpp_lib.py`` catches and downgrades to a ``logging`` warning -- so the package +imports, reports a version, and silently has no memory-efficient attention. That is +unslothai's NVIDIA P0: a ``cu128``-built wheel shipped next to a ``cu130`` runtime. + +Two independent signals live here, both cheap and both offline: + +1. ``XFORMERS_TORCH_PINS`` / ``XFORMERS_BUILT_FOR_TORCH`` -- what each xformers release + declares, and what it was actually compiled against. Usable BEFORE anything is + installed, which is what picks the wheel. +2. ``xformers_build_metadata()`` -- the installed wheel's own ``xformers/cpp_lib.json``, + which records the torch version, CUDA version and Python it was built with, e.g. + ``{"version": {"cuda": 1208, "torch": "2.10.0+cu128", "python": "3.10.11"}}``. This is + the authority once a wheel is on disk. xformers itself does not *compare* these fields + -- ``_register_extensions()`` just tries the load and, on ``OSError``, quotes them back + in ``xFormersInvalidLibException`` -- so reading the file is the only way to know + *before* the load, and it costs no import and fires no warning. + +Pure module: stdlib only, no torch and no xformers import at any point. It lives at the +top of ``unsloth/`` (next to device_type.py / import_fixes.py) rather than in +``unsloth/utils/``, because ``unsloth/utils/__init__.py`` imports attention_dispatch, +which imports ``unsloth.models._utils`` -- the very module that needs this table. + +``studio/backend/utils/hardware/hardware.py`` deliberately re-implements the +``cpp_lib.json`` read instead of importing this module: the studio backend runs with +``studio/backend`` on ``sys.path`` and importing ``unsloth.xformers_compat`` would execute +``unsloth/__init__.py``, which drags in torch. Only the ~15-line file read is duplicated, +not the tables. +""" + +from __future__ import annotations + +import json +import os +import re +from typing import Any, Dict, Optional + +__all__ = [ + "XFORMERS_TORCH_PINS", + "XFORMERS_BUILT_FOR_TORCH", + "TORCH_TO_XFORMERS", + "normalize_release", + "normalize_release_with_post", + "cuda_major_from_torch_version", + "format_build_cuda", + "expected_torch_for_xformers", + "xformers_for_torch", + "declared_torch_pin", + "xformers_build_metadata", + "xformers_build_summary", + "describe_xformers_mismatch", + "stable_abi_covers", +] + + +# xformers release -> the exact torch release its wheels declare in Requires-Dist. +# Read off the published wheel METADATA (pypi.org/pypi/xformers//json -> +# info.requires_dist), not guessed; do not extend this by pattern-matching version +# numbers, read the wheel. +# +# 0.0.35 is deliberately absent: it is the first release to declare a RANGE +# (``torch>=2.10``) rather than ``torch==X``, so its declared pin no longer says what it +# was built against and pip can no longer keep the pair consistent on its own. Its build +# torch lives in XFORMERS_BUILT_FOR_TORCH below. +XFORMERS_TORCH_PINS: Dict[str, str] = { + "0.0.28": "2.4.1", + "0.0.28.post1": "2.4.1", + "0.0.28.post2": "2.5.0", + "0.0.28.post3": "2.5.1", + "0.0.29": "2.5.1", + "0.0.29.post1": "2.5.1", + "0.0.29.post2": "2.6.0", + "0.0.29.post3": "2.6.0", + "0.0.30": "2.7.0", + "0.0.31": "2.7.1", + "0.0.31.post1": "2.7.1", + "0.0.32.post1": "2.8.0", + "0.0.32.post2": "2.8.0", + "0.0.33": "2.9.0", + "0.0.33.post1": "2.9.0", + "0.0.33.post2": "2.9.1", + "0.0.34": "2.10.0", +} + +# xformers release -> the torch release its ``_C`` extension was actually COMPILED +# against, i.e. ``cpp_lib.json``'s ``version.torch`` with the ``+cuXXX`` tag dropped +# (the tag varies per wheel variant, the release does not). +# +# For everything up to 0.0.34 this equals the declared pin, which is the point: the pin +# was trustworthy. 0.0.35 is where they diverge -- it declares ``torch>=2.10`` but its +# ``_C`` is still a single build against 2.10.0, so pip will happily pair it with torch +# 2.11/2.12/2.13 and the extension will not load. There is no xformers release built for +# torch 2.11 or later, so ``xformers_for_torch`` returns None there rather than guessing. +XFORMERS_BUILT_FOR_TORCH: Dict[str, str] = { + "0.0.33": "2.9.0", + "0.0.33.post1": "2.9.0", + "0.0.33.post2": "2.9.1", + "0.0.34": "2.10.0", + "0.0.35": "2.10.0", +} + +# torch release -> the NEWEST xformers built for it. The inverse of the tables above with +# the post-releases winning, which is what an installer wants: 2.9.0 has both 0.0.33 and +# 0.0.33.post1, and the post is the one to ship. 0.0.35 never wins 2.10.0 over 0.0.34: +# both are built against 2.10.0, and 0.0.34's ``torch==2.10.0`` pin is the one pip can +# still enforce. +TORCH_TO_XFORMERS: Dict[str, str] = { + "2.4.1": "0.0.28.post1", + "2.5.0": "0.0.28.post2", + "2.5.1": "0.0.29.post1", + "2.6.0": "0.0.29.post3", + "2.7.0": "0.0.30", + "2.7.1": "0.0.31.post1", + "2.8.0": "0.0.32.post2", + "2.9.0": "0.0.33.post1", + "2.9.1": "0.0.33.post2", + "2.10.0": "0.0.34", +} + + +def normalize_release(version: Any) -> Optional[str]: + """``'2.10.0+cu130'`` / ``'2.11.0.dev20260101'`` -> ``'2.10.0'`` / ``'2.11.0'``. + + Drops the local tag and any pre-release/dev suffix, matching the release part that + ``Requires-Dist: torch==X`` compares against. None when unparseable. + """ + if not version: + return None + release = str(version).strip().split("+", 1)[0] + match = re.match(r"^(\d+(?:\.\d+)*)", release) + if match is None: + return None + return match.group(1) + + +# xFormers moved to the PyTorch stable API/ABI in 0.0.34: the v0.0.34 release notes state +# that "binary builds targeting PyTorch 2.10+ will be compatible with any later version". +# So a build recorded against 2.10 running on 2.11 or 2.12 is the DESIGN, not a mismatch. +_STABLE_ABI_TORCH_FLOOR = (2, 10) + + +def _release_tuple(release: Optional[str]) -> Optional[tuple]: + if not release: + return None + try: + return tuple(int(part) for part in release.split(".")) + except ValueError: + return None + + +def stable_abi_covers(built_release: Optional[str], running_release: Optional[str]) -> bool: + """Whether a wheel built for ``built_release`` is expected to run on ``running_release``. + + True only in the direction the guarantee actually runs: a 2.10-built wheel on 2.11 is + fine, a 2.11-built wheel on 2.10 is not. Below the floor there is no stable ABI, so a + difference there is still a mismatch. Saying otherwise turns an unrelated failure (a + missing VC++ runtime, say) into a confident wrong diagnosis, and tells the user to go + hunting for a release that does not exist. + """ + built = _release_tuple(built_release) + running = _release_tuple(running_release) + if built is None or running is None: + return False + return built[:2] >= _STABLE_ABI_TORCH_FLOOR and running >= built + + +def normalize_release_with_post(version: Any) -> Optional[str]: + """Like normalize_release but KEEPS ``.postN``. None for a pre-release. + + xformers keys on the post: 0.0.33 and 0.0.33.post2 are built for different torch + releases (2.9.0 and 2.9.1), so dropping it would answer the wrong question. + + A ``.dev`` / ``a`` / ``b`` / ``rc`` build is reported as unknown rather than folded + onto its release: the fourteen ``0.0.35.devNNNN`` wheels on PyPI are built against + torch nightlies, so answering "0.0.35, therefore torch 2.10.0" for one of them is + confidently wrong, and unknown is the honest answer. + """ + if not version: + return None + release = str(version).strip().split("+", 1)[0] + match = re.match(r"^(\d+(?:\.\d+)*(?:\.post\d+)?)$", release) + if match is None: + return None + return match.group(1) + + +def cuda_major_from_torch_version(torch_version: Any) -> Optional[int]: + """``'2.10.0+cu130'`` -> 13, ``'2.10.0+cu128'`` -> 12. None for rocm/cpu/tagless. + + Mirrors ``_cuda_major_from_torch_version`` in studio/install_python_stack.py. + """ + if not torch_version: + return None + parts = str(torch_version).split("+", 1) + if len(parts) < 2 or not parts[1].startswith("cu"): + return None + digits = re.sub(r"[^0-9].*", "", parts[1][2:]) # 'cu130' -> '130' + if not digits: + return None + return int(digits) // 10 # '130' -> 13, '128' -> 12, '118' -> 11 + + +def format_build_cuda(build_cuda: Any) -> Optional[str]: + """cpp_lib.json's integer CUDA version -> ``'12.8'``. None when absent (ROCm/CPU). + + xformers' setup.py stores ``major * 100 + minor``: 1208 is CUDA 12.8, 1300 is 13.0. + Its own exception message prints the raw integer, which is unreadable, so format it. + """ + if not isinstance(build_cuda, int) or isinstance(build_cuda, bool): + return None + return f"{build_cuda // 100}.{build_cuda % 100}" + + +def expected_torch_for_xformers(xformers_version: Any) -> Optional[str]: + """The torch release ``xformers_version`` was built for, or None if unknown to us. + + Prefers the recorded build torch over the declared pin: from 0.0.35 the pin is a + range and no longer names a single release. + """ + release = normalize_release_with_post(xformers_version) + if release is None: + return None + built_for = XFORMERS_BUILT_FOR_TORCH.get(release) + if built_for is not None: + return built_for + return XFORMERS_TORCH_PINS.get(release) + + +# The newest release whose ``_C`` is a stable-ABI build, and the torch release it was +# compiled against. Above the floor these two are what covers a torch nobody has shipped a +# row for yet. +_STABLE_ABI_XFORMERS = "0.0.35" +_STABLE_ABI_BUILT_FOR = XFORMERS_BUILT_FOR_TORCH[_STABLE_ABI_XFORMERS] + + +def xformers_for_torch(torch_version: Any) -> Optional[str]: + """The newest xformers release that runs on ``torch_version``, or None if unknown. + + Takes a full torch version (local tag and all) so callers can pass + ``torch.__version__`` straight in. + + Beyond the exact table, the stable ABI answers the rest: 0.0.34+ binaries target + PyTorch 2.10+ and the release notes say such builds are "compatible with any later + version". Returning None for torch 2.11 contradicted this module's own + ``describe_xformers_mismatch``, which accepts exactly that pairing -- so the diagnosis + said "your CUDA family is wrong" and the fix hint said "no release exists, downgrade + torch or build from source". Only for a plain release: a dev/rc torch is not something + to make a compatibility promise about. + """ + release = normalize_release(torch_version) + if release is None: + return None + exact = TORCH_TO_XFORMERS.get(release) + if exact is not None: + return exact + if normalize_release_with_post(torch_version) is None: + return None # pre-release: unknown is the honest answer + if stable_abi_covers(_STABLE_ABI_BUILT_FOR, release): + return _STABLE_ABI_XFORMERS + return None + + +def declared_torch_pin(xformers_version: Any = None) -> Optional[str]: + """The installed xformers distribution's own ``Requires-Dist: torch==X`` pin. + + Prefers the resident metadata over the static table, so a release we have never seen + still answers correctly -- but only when it describes the same version the caller + asked about, otherwise a stale table lookup is the honest answer. + + When the resident pin is a range rather than ``==`` (0.0.35 and later) there is no + declared pin to report, so this falls back to what the wheel was actually built + against. Callers must phrase that as "is built for", never "declares": the range + release deliberately does not declare a single torch. + """ + wanted = normalize_release_with_post(xformers_version) + try: + from importlib.metadata import requires as _requires, version as _version + resident = normalize_release_with_post(_version("xformers")) + requirements = _requires("xformers") or () + except Exception: + resident, requirements = None, () + if wanted is None or resident is None or wanted == resident: + for requirement in requirements: + match = re.match(r"^\s*torch\s*==\s*([0-9][0-9A-Za-z.\-+]*)", str(requirement)) + if match is not None: + return normalize_release(match.group(1)) + return expected_torch_for_xformers(xformers_version) + + +def xformers_build_metadata() -> Optional[Dict[str, Any]]: + """The installed xformers wheel's ``cpp_lib.json``, WITHOUT importing xformers. + + ``importlib.util.find_spec`` only locates the package, it does not execute + ``xformers/__init__.py`` -- which matters because importing xformers is what emits the + warning we are trying to explain, and because it drags in torch. Returns None when + xformers is absent, is an editable/source checkout with no built extension, or ships + no cpp_lib.json. + """ + try: + import importlib.util + spec = importlib.util.find_spec("xformers") + except Exception: + # find_spec raises (not returns None) on a half-removed dist, and ImportError + # here must never take the caller down: this is diagnostics. + return None + if spec is None: + return None + locations = list(getattr(spec, "submodule_search_locations", None) or ()) + origin = getattr(spec, "origin", None) + if origin: + locations.append(os.path.dirname(origin)) + for location in locations: + path = os.path.join(location, "cpp_lib.json") + try: + with open(path, "r", encoding = "utf-8") as handle: + metadata = json.load(handle) + except Exception: + continue + if isinstance(metadata, dict) and isinstance(metadata.get("version"), dict): + return metadata + return None + + +def xformers_build_summary( + build_metadata: Optional[Dict[str, Any]] = None, +) -> Optional[Dict[str, Optional[str]]]: + """``cpp_lib.json`` -> ``{"torch": ..., "cuda": "12.8", "python": ...}`` for reporting. + + None when there is no build metadata to summarise. Every value is a string or None, + so this drops straight into a JSON API response. + """ + if build_metadata is None: + build_metadata = xformers_build_metadata() + version_block = (build_metadata or {}).get("version") + if not isinstance(version_block, dict): + return None + built_torch = version_block.get("torch") + built_python = version_block.get("python") + built_hip = version_block.get("hip") + return { + "torch": str(built_torch) if built_torch else None, + "cuda": format_build_cuda(version_block.get("cuda")), + "hip": str(built_hip) if built_hip else None, + "python": str(built_python) if built_python else None, + } + + +def _running_detail(torch_version: Any, python_version: Any = None) -> str: + """``'torch 2.10.0+cu130'``, plus Python when the caller supplied it.""" + detail = f"torch {torch_version}" + if python_version: + detail += f" / Python {python_version}" + return detail + + +def _built_detail(summary: Dict[str, Optional[str]]) -> str: + """``'torch 2.10.0+cu128 / Python 3.10.11'`` from a build summary.""" + parts = [] + if summary.get("torch"): + parts.append(f"torch {summary['torch']}") + elif summary.get("cuda"): + parts.append(f"CUDA {summary['cuda']}") + if summary.get("python"): + parts.append(f"Python {summary['python']}") + return " / ".join(parts) if parts else "an unknown build" + + +def describe_xformers_mismatch( + torch_version: Any, + torch_cuda: Any = None, + xformers_version: Any = None, + build_metadata: Optional[Dict[str, Any]] = None, + python_version: Any = None, +) -> Optional[str]: + """One sentence naming why this xformers cannot load its kernels here, else None. + + ``torch_version`` is ``torch.__version__`` (local tag included -- it carries the CUDA + family) and ``torch_cuda`` is ``torch.version.cuda``. Both the recorded build metadata + and the declared pin are consulted; the build metadata wins because it describes the + binary that is actually on disk. Returns None when the two agree, when xformers is + absent, or when there is not enough information to be sure -- this must never cry + wolf, the caller prints it as a warning on the default path. + + A Python-version difference alone is NOT a mismatch and never triggers this: the + wheels are abi3/none-tagged and ``_C`` is loaded through ``torch.ops.load_library``, + not the CPython ABI, so 3.10-built kernels run fine on 3.13. It is still reported as + context, because xformers' own message leads with it and users chase it first. + """ + running_release = normalize_release(torch_version) + if running_release is None: + return None + + if build_metadata is None: + build_metadata = xformers_build_metadata() + summary = xformers_build_summary(build_metadata) or {} + built_torch = summary.get("torch") + built_release = normalize_release(built_torch) + running = _running_detail(torch_version, python_version) + + if ( + built_release is not None + and built_release != running_release + and not stable_abi_covers(built_release, running_release) + ): + return ( + f"xformers was built for {_built_detail(summary)} but you are running " + f"{running}; its C++/CUDA extensions cannot load, so memory-efficient " + f"attention is unavailable" + ) + + # Same torch release, different CUDA major: the case NVIDIA hit (a cu128 wheel beside + # a cu130 runtime). Majors only -- CUDA minor version compatibility means a cu126 + # wheel loads fine against a cu128 torch, and flagging that would be crying wolf. + built_cuda = summary.get("cuda") + running_cuda_major = cuda_major_from_torch_version(torch_version) + if running_cuda_major is None and torch_cuda: + try: + running_cuda_major = int(str(torch_cuda).split(".", 1)[0]) + except (TypeError, ValueError): + running_cuda_major = None + if built_cuda is not None and running_cuda_major is not None: + built_cuda_major = int(built_cuda.split(".", 1)[0]) + if built_cuda_major != running_cuda_major: + return ( + f"xformers was built for {_built_detail(summary)} but you are running " + f"{running} (CUDA {running_cuda_major}.x); its C++/CUDA extensions cannot " + f"load, so memory-efficient attention is unavailable" + ) + + # No build metadata (source install, or a wheel that ships none): fall back to the + # declared torch pin, which at least catches a wholesale torch-release mismatch. + if built_release is None: + pinned = declared_torch_pin(xformers_version) + # Same stable-ABI exemption the build-metadata branch above applies. Without it a + # source or editable 0.0.34+ build, whose pin is the 2.10 it was cut against, reports + # every later torch as a mismatch -- so an unrelated extension failure (a missing + # runtime DLL) is diagnosed as a torch-version problem with reinstall instructions + # that fix nothing. + if ( + pinned is not None + and pinned != running_release + and not stable_abi_covers(pinned, running_release) + ): + return ( + f"xformers {xformers_version or 'installed'} is built for torch {pinned} " + f"but you are running {running}; its C++/CUDA extensions cannot load, so " + f"memory-efficient attention is unavailable" + ) + return None