mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-25 08:42:25 +00:00
Say when xFormers is installed but its kernels cannot load (#8157)
* Detect and report a mismatched xformers build at import xformers wheels are compiled against one torch release and one CUDA major, but nothing at install time enforces it: the wheels are abi3/none-tagged and --no-deps skips Requires-Dist entirely, so a cu128 wheel drops cleanly next to a cu130 runtime and only fails inside torch.ops.load_library, which xformers catches and downgrades to a logger warning. _utils.py made that worse in two ways: it silenced the xformers logger to ERROR immediately before the import, and it reported the resulting failure only under UNSLOTH_ENABLE_LOGGING. The net effect was a silent fall back to SDPA attention. Its torch<->xformers version table also stopped at torch 2.4. New unsloth/xformers_compat.py (stdlib only, imports neither torch nor xformers) carries the version knowledge up to torch 2.10 / xformers 0.0.35 and reads the installed wheel's own xformers/cpp_lib.json to learn what it was really built against. The tables and the cpp_lib.json shape were read off the published wheels, not inferred. 0.0.35 is the case worth calling out: it is the first release to declare a range (torch>=2.10) while its _C is still a single build against 2.10.0, so pip will pair it with torch 2.11+ and the extension will not load. No xformers release is built for torch 2.11 or later, so xformers_for_torch returns None there rather than recommending the install that produces the bug. Because the check runs off the file before the import, the mismatch is known in time to leave xformers' own diagnostic un-silenced in exactly the case where it is the bug report, and to print a single actionable message naming the torch, CUDA and Python the wheel was built for versus what is running. A Python version difference alone is never reported as a break: _C loads through torch.ops.load_library, not the CPython ABI. XFORMERS_BUILD_METADATA and XFORMERS_BROKEN_REASON are exported so callers can report what broke instead of guessing. * Probe the xformers kernel on every CUDA capability, and keep the reason _xformers_disabled_for_capability returned early below sm_120 with the comment "Below sm_120 xformers always works; skip the probe". That holds only when the wheel matches the runtime. A cu128-built xformers on a cu130 torch is dead on an sm_90 Hopper exactly as it is on an sm_120, and the early return meant the one piece of code that actually runs the kernel never looked. The gate also skipped the probe whenever flash-attn was installed, so those machines could never report a dead xformers at all. Now the real op decides on every capability, including with flash-attn present. The extra cost is one 1x8x1x64 forward on a CUDA context torch has already initialised, since _XFORMERS_FP32_UNSUPPORTED just below forces the same lazy init unconditionally. The probe also caches why it failed in XFORMERS_PROBE_REASON, surfaced as XFORMERS_DISABLED_REASON alongside XFORMERS_BROKEN_REASON from _utils, so a caller can report what is wrong rather than only that xformers is off. It still never raises: it runs at import, so a diagnostic must not be what breaks the import. The #4631 regressions are kept and extended: a working kernel is still kept on sm_120, and now the probe-means-disable failure mode is pinned on every other capability too. * Report the accelerator stack in the Studio hardware detail get_package_versions() covered unsloth, torch, transformers and torch's CUDA/HIP version, and nothing else. An xformers built for a different torch reports its version happily and has no memory-efficient attention, so the report a user sends with a bug looked identical whether the acceleration stack worked or not. There was no Python version either, which is half of what a mismatch report needs. get_package_versions() now also carries python, xformers, flash_attn, torchao and bitsandbytes as plain version strings. Still flat str-or-None, so every existing consumer is untouched. The part that costs a real import lives in the new get_accelerator_report(), returned under /api/system/hardware?include_details=true only -- the default response is polled for training-method selection and must stay cheap. Per package it separates installed / imports / runs, and for xformers reads the wheel's own cpp_lib.json to say what it was compiled against, which is the part that names the fix. The cpp_lib.json read is a deliberate copy of the one in unsloth.xformers_compat rather than a call to it: this module has to keep working with no torch installed, and importing unsloth would execute unsloth/__init__.py. Cached per process, deep-copied out, never raises, and skippable with UNSLOTH_SKIP_ACCELERATOR_PROBE=1 for an install where importing a broken native wheel is worse than not knowing. Deliberately does not gate startup. NVIDIA asked to block release when optimized kernels cannot load; blocking the app instead would turn a slower-but-working install into a dead one. Detection and loud reporting, not refusal to run. Ran it here and it immediately found a real one: torchao 0.18.0+cu130 against torch 2.9.1+cu128 does not import. * Surface dead optimized kernels in Settings The About tab showed a version string per runtime and nothing else, which a mismatched xformers reports exactly as happily as a working one, and no Python version at all -- the other half of every "built for 3.10, running 3.13" report. About gains a Python row and an Optimized kernels section listing xFormers, FlashAttention, torchao and bitsandbytes, each as installed / not installed / working / not loading rather than as a version. A broken one shows what its wheel was built for, with the full reason behind the row's info hint. The banner sits in the Settings dialog rather than on the About tab, above whichever tab is showing: a kernel that is installed and cannot load produces no symptom other than being slower, so a user has no reason to go looking in About. It mounts inside DialogContent, which Radix only renders while the dialog is open, so the detail fetch and the native imports behind it stay off app startup. Status is never colour-only -- each state carries its own word -- so it survives a greyscale screenshot in a bug report, which is where this will mostly be read. Parsing moved to hooks/accelerator-report.ts so it is importable without react or the auth chain, which is what makes it unit testable. A backend that predates the field parses to null, not to an empty healthy report: "cannot tell" must not render as "all fine". Also drops the misleading half of the CLI fix hint. On a torch newer than any xformers release there is no version to pin, and naming one would send the user back into the same mismatch, so it now says so and points at a source build. * Move the accelerator-report tests into the hardware test namespace The suite is routinely run as `pytest studio/backend/tests -k "hardware or system"`, and only 1 of the 15 matched under the old filename. It is the hardware module's report, so name it that way. * Never report a degraded stack on a host these kernels do not run on bitsandbytes is a dependency on every platform but only loads on CUDA/XPU, so on a Mac or a CPU-only host its import failure is the expected outcome, not a broken acceleration stack. Probing there pinned a permanent false banner to those installs, which is exactly the credibility problem that makes people stop reading banners. The probe now runs only when get_device() is CUDA or XPU, and "we did not look" stays distinguishable from "you told us not to look": "not used on this device" versus "not probed". The probe tests pinned get_device to CUDA as well. Without that they passed on a CPU-only runner by never probing at all -- green for the wrong reason. * Keep the About Python row on a host with no GPU It sat inside the Hardware section, which only renders when there is a GPU or an accelerator runtime to show. On a CPU-only host that hid the one field a bug report most needs. * Probe the accelerator stack out of process, and stop guessing at the reason Review of the first pass turned up four things that were worse than the problem being reported, three of them measured on this host. Importing the packages in the web-server process was wrong three ways. A package whose __init__ raises leaves every submodule it already executed in sys.modules, so the next import returns a half-built module -- this repo already documents that hazard and ships purge_partial_import for it (unslothai/unsloth#7580); torchao left 40 stale submodules behind here. `import bitsandbytes` latches a CUDA context, and the backend deliberately never had 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 was about to take several hundred MB off every later VRAM reading for the rest of the session. And a badly broken native wheel can abort the interpreter rather than raise, which in the server means the app dies on exactly the installs this is meant to describe. It now runs in one throwaway child with CUDA_VISIBLE_DEVICES emptied. Measured after the change: no CUDA context, no leaked modules, 3.5s once per process. _describe_xformers_break asserted a version mismatch whenever build metadata merely existed, without ever comparing it to the runtime, and dropped the real exception on the floor. [WinError 126] with a perfectly matching wheel -- a missing VC++ runtime or CUDA DLL, the most common Windows xformers failure that is NOT a mismatch -- was reported as "built for X but runs X". It now compares first and hands back the real error otherwise. The Mac/CPU gate let ROCm and Intel XPU through: ROCm hosts report DeviceType.CUDA internally, and the stock bitsandbytes is CUDA-only on both. Same permanent false banner, just moved to AMD and Intel users. The report also moves off include_details onto its own include_accelerators flag. The detail path is read by Export, Video and onboarding, none of which should pay for this, so the frontend gets a separate hook used only by the Settings surfaces. Also: read xformers' already-recorded _cpp_library_load_exception instead of re-calling the private _register_extensions (whose absence was being reported as "broken"), cap reason strings that land in an aria-label, name the subprocess encoding so Windows paths survive, and carry hip through to the UI so a ROCm wheel shows a build at all. Tests: the probe seam is faked everywhere, so none of them import the host's native wheels; added the matching-build, python-only, cuda-minor, ROCm, XPU and could-not-be-checked cases that let those defects through. * Do not let the probe silently disable xformers, or mislabel what it caught Review of the first pass found the new probe reintroducing the failure mode it was written to remove, plus two ways of stating the wrong cause. The gate turned any probe exception into HAS_XFORMERS = False for the whole process, and printed nothing. Before, capability[0] < 12 returned early, so sm_70-sm_100 could not lose xformers this way at all. A full device 0, a device claimed by another rank under EXCLUSIVE_PROCESS, a MIG slice: each of those now cost memory-efficient attention for the session, silently, because of the diagnostic. Those failures are now inconclusive -- xformers stays on and the real forward decides -- and a conclusive one is announced through the same path as an ABI mismatch, so XFORMERS_DISABLED_REASON is no longer written and never read. The probe also targeted device 0 for everybody. Under torchrun each rank owns a different GPU, and on a mixed box device 0 is often the small display card, so a wheel with no kernel for the weakest GPU disabled xformers on the good ones. It now probes LOCAL_RANK's device. The loud warning was reflowing every exception into "its optimized kernels cannot load ... install the matching build". That arm also catches the sm_100/110/120 FA3 guard and the three old-torch guards, whose messages are multi-line, fenced and already actionable; the template stated the wrong cause and mangled the text. Only a real build mismatch gets the version-pin treatment now. xformers_compat: a .dev/rc build no longer folds 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 was confidently wrong; unknown is the honest answer. And declared_torch_pin's docstring claimed it returns None for a range pin when it falls back to the build table -- the code is right, the docstring was not. Tests: test_warning_is_printed_once was vacuous (a fresh subprocess reaches the announcement once whether or not the guard exists), so it now counts in-process; _mismatched_build synthesizes a CUDA-major difference and is skipped where there is no running CUDA major; the probe tests patch the globals they write through instead of leaking them into the session; and there is now a test that a healthy install prints nothing at all, which is the half of "never cry wolf" nothing covered. * Let the P0 xformers tests run on a host with no GPU Both subprocess tests in tests/utils/test_xformers_broken_warning.py failed on the cross-platform runners for two independent reasons, neither of them the warning they are meant to prove. First, the child ran a bare `python -c`. Everything else under tests/ gets tests/conftest.py, the GPU-free harness that forces DEVICE_TYPE to "cuda" and stubs the torch.cuda probes unsloth fires at import time; the child got none of it, so unsloth/_gpu_init.py's module-scope torch.cuda.get_device_capability() raised "Found no NVIDIA driver" and the import exited 1. That is also why the in-process `import unsloth` in test_xformers_capability_gate.py was green on the same runner. The child now loads that same harness by path before importing unsloth, so the subprocess and the parent agree about the host. On a real accelerator the harness is a no-op. Second, the stub xformers was an importable module with no .dist-info, and `import unsloth` asks importlib.metadata for the xformers version twice. With no xformers on the host that raised PackageNotFoundError straight out of fix_xformers_performance_issue(); with one, it silently answered from the HOST's xformers, which is exactly what this file's docstring promises not to depend on. The stub now ships its own metadata and the child asserts the version it reads back is the stub's. test_a_healthy_install_says_nothing additionally needs a healthy install to exist, so it now skips when importlib.util.find_spec("xformers") is None instead of reaching its "this host's xformers is genuinely broken" skip, which would have been the wrong reason for a host that simply has no xformers. Tests only - no change to the detection in unsloth/models/_utils.py, unsloth/xformers_compat.py or unsloth/utils/attention_dispatch.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the probe on this rank's device, and stop it crashing the import Widening the probe from sm_120-only to every capability put it on the standard multi-GPU path for the first time, and it does not survive there. Both of these reproduce on a real two-GPU torchrun; neither happens on main. The probe placed q on cuda:{LOCAL_RANK} but built the attn_bias with BlockDiagonalCausalMask.from_seqlens, which allocates on the CURRENT device -- still cuda:0 on every rank at import time, because launchers set LOCAL_RANK in the environment while torch.cuda.set_device happens later inside the trainer. So the two tensors disagreed, xformers rejected the pair with "Attention bias and Query/Key/Value should be on the same device", and since that text is not an inconclusive marker the gate hard-disabled xformers. Measured on two healthy identical B200s: rank 0 kept xformers, rank 1 lost it. On an 8-GPU node that is 7 ranks silently dropping to SDPA at double the attention memory -- the exact regression the inconclusive guard was written to prevent, caused by the diagnostic itself, and reported to the user as a broken install. It also allocated on cuda:0 from every rank, pinning a second context per rank, about 700 MB each. Running the whole probe under torch.cuda.device fixes both. Separately, LOCAL_RANK is a rank and not an index into the devices a process can see. 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 both use -1 to mean "not distributed". torch.cuda.get_device_capability raises on an invalid ordinal and that call is at module scope, so `import unsloth` died outright with AssertionError: Invalid device id. Verified for LOCAL_RANK of 3, 1, 07 and -1 against one visible device; all four now import and fall back to device 0, which is the only device such a rank has. Clamping is the fix rather than a try/except around the capability read: with an out-of-range index the probe fails with "invalid device ordinal", which is not an inconclusive marker either, so catching the crash alone would have converted it into the silent disable above. Those two strings are added to the inconclusive list regardless, so that path can never again be read as evidence about the build. Two more things this turned up. The probe's reason is a captured exception, and xformers answers a capability rejection with a dozen lines listing every operator it considered -- announced verbatim, that is what every RTX 50-series user would see on every import. Print the first line and put the rest behind UNSLOTH_ENABLE_LOGGING, truncating at the call site because the announcer's other callers pass deliberately fenced, copy-pasteable text that has to arrive intact. And re-silence the xformers logger once the extensions turn out to have loaded: skipping the silencing on a PREDICTED break is right, 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, and leaving the logger open is permanent and process-wide. _XFORMERS_FP32_UNSUPPORTED now reads the same device the probe used, which is what its comment already claimed. Tests: the rank-targeting test stubs torch.cuda.device so it runs on a host with any GPU count and asserts the context index, not just the tensor's; new tests cover the clamp invariant and a real subprocess import under five hostile LOCAL_RANK values, which is the regression the mocked test could not catch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tell the accelerator probe child to emit UTF-8 The probe decodes its child as UTF-8, but the child's own stdout defaults to locale.getpreferredencoding(), which on Windows is the ANSI code page. That is exactly where this probe's value is: native import errors quote DLL and file paths, and a non-ASCII path would come back mangled. Route the env through utf8_child_env, as the repo's own contract test requires. * Report what the accelerator stack can actually do, not what imported - the xFormers probe picked bf16 from device 0 while probing this rank's GPU, so a mixed box with a Turing card wrote off a healthy install. - the fix hint pinned a version but no index; all three CUDA families publish the same version string and PyPI carries only cu128, so on the reported case it reinstalled the identical broken wheel. - both mismatch diagnoses (unsloth and Studio) called a stable-ABI wheel on a later torch a mismatch, and the Studio one compared full version strings, so 2.10.0+cu126 vs 2.10.0+cu128 hid the real error before the CUDA-major check could give the accurate one. - the probe reported bitsandbytes and torchao as working on an import alone, which is exactly the state where 4-bit dies mid-run and torchao has no kernels; both now check their native side, bitsandbytes through the same leaf module the loader gates on. - xFormers loading its library is not the same as having a kernel for this GPU; a capability rejection now reads as degraded instead of Working. - ROCm hosts probed nothing at all, though Unsloth enables flash-attn there. - and a child that dies without answering no longer erases the whole report: the survivors keep their answers and the package that killed it is named. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Ask the GPU the app can actually use, and let each row say if it was checked Five things the probe got wrong once it started answering "does this run". nvidia-smi lists physical devices and ignores CUDA_VISIBLE_DEVICES, so on a mixed host masked to the sm_120 card the verdict was computed against row 0's sm_90 op table, and xFormers read as working on a GPU that has no kernel for it. The mask is resolved in the parent now, by index or by (abbreviated) UUID, and handed to the child as UNSLOTH_PROBE_DEVICE_CC, because the child is started with the mask cleared and cannot recover it. bitsandbytes picks its native library from torch.cuda.is_available(), so hiding the GPUs from the child made a healthy CUDA install report dead handles and light the banner. It gets its own child now, with the mask the app itself runs under. The server still never holds a context: that child exits immediately. On Intel XPU the applicable set was empty, so a broken XPU bitsandbytes wheel was "not used on this device" and never reached degraded, while the loader gates 4-bit on native_kernels_ready(_bnb_probe, DEVICE_TYPE) with XPU symbols. It is probed. torch.ops.<ns> materialises an empty _OpNamespace on attribute access, and its dir() already lists __name__ and friends, so dir(torch.ops.torchao) was non-empty for the no-extension case the check exists to catch. Counted off the dispatcher instead. And the About table read one report-wide "probed" flag onto every row, so a package the backend deliberately skipped rendered "Not loading" as soon as any other package was probed -- red for something that is fine, and absent from the banner. Each entry carries its own probed flag; the report-wide one is the fallback for an older backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the capability read out of module scope, and let the model code see the verdict Two ways the probe could hurt the thing it was added to protect. The gate read torch.cuda.get_device_capability() at the call site, at module scope, and CUDA refuses that query when the device is busy, in exclusive-compute mode or otherwise unavailable. There it raised before the probe could classify the failure as inconclusive, so a diagnostic whose worst answer is "keep xformers on and let the forward decide" turned `import unsloth` into a crash. The argument is gone (the gate asks the kernel, not the number), and the fp32 read below it now degrades to unknown through a guarded helper. A test walks the module's top-level statements and fails on any capability query that runs at import. And llama.py recomputed HAS_XFORMERS as `xformers is not None`, so turning the dispatcher's flag off never reached the model code. Mistral answers "xFormers is on" by skipping the 4D sliding-window mask and letting the xFormers bias carry the window; with the dispatcher already on SDPA, that mask is the only thing making the window local, so a sequence longer than config.sliding_window attended to the whole causal history. It takes the probed flag from the dispatcher now. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Answer for every visible GPU, and say unknown where the probe cannot tell Five follow-ups, all the same shape: a diagnostic that overstates what it knows. Mistral's training branch takes `elif self.training: pass`, so no 4D mask is synthesized, and its AttentionContext never carried sliding_window at all. With xFormers off and flash-attn absent, nothing was left holding the window and every sequence past config.sliding_window attended its whole causal history. The context carries the window now, and the SDPA path builds a windowed keep mask even with no padding mask to hang it off, since is_causal is full causal and has no window of its own. The capability resolver truncated the mask to its first entry, so with CUDA_VISIBLE_DEVICES=0,1 across a mixed pair a rank on the second card fell back to SDPA while the verdict from the first said Working. Every visible device is resolved and handed to the child now, and one uncovered card makes the whole install degraded. flash-attn recorded runs=true on the strength of importing its interface. Our own installer refuses prebuilt wheels on sm_100+ because none covers those cards, so a wheel that got there another way could load and still fail on its first launch. That is unknown, with the reason, rather than a verdict this child would have to launch a kernel to earn. The About table read runs=null as Working. A probe that ran and could not decide -- an unrecognised xformers layout, a missing bitsandbytes checker, a torch with no dispatcher table, the flash-attn case above -- is unknown, and nothing else would have corrected it since those packages are not in degraded either. And xformers_for_torch returned None above 2.10, which contradicted this module's own describe_xformers_mismatch: it accepts a 2.10-built wheel on 2.11 under the stable ABI, then the fix hint told the user no release exists and to downgrade torch or build from source. Above the floor the answer is 0.0.35, and the hint already names the CUDA index that repairs the family. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Stop the probe calling a supported fallback broken, and normalise a zero window Five more, all cases where the report claimed more than it knew, plus one bug of mine. torchao with no registered operators is the MANAGED stack: install_python_stack pins 0.17 on torch 2.10+cu130 knowing its extension is cleanly skipped, and torchao keeps running through its Python fallbacks with only the optimized kernels gone. Calling that degraded put a destructive banner on the standard configuration. Unknown, with the reason, so the row says it without the banner. flash-attn is unknown on every card now, not just above sm_100. A build with no cubin or PTX image for this architecture imports fine and fails its first launch, which is as true of a source build or an older wheel as of a Blackwell card; flash-attn exposes no list of the architectures it was compiled for, and this child never launches a kernel. A failed import is still broken. "device is currently in use by another process" is the driver saying the GPU is someone else's right now. It matched none of the inconclusive markers, so a transient turned into a process-wide fallback to SDPA on a wheel that was never tested. The probe device honoured LOCAL_RANK and otherwise took device 0, ignoring a caller that had already run torch.cuda.set_device(1). That can disable xformers over a card nothing will touch, and creates a context on it to do so. LOCAL_RANK still wins when usable. And sliding_window = 0 means "no local attention", the same as absent. Passing it through made window_size (0, 0) for flash and, now that the dispatcher honours the window, an all-false SDPA mask: the lower bound q_pos - (0 - 1) sits above the causal upper bound, so the layer returned nothing at all. Normalised in Mistral and again in the dispatcher. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the xFormers kernel verdict unknown until a launch proves it 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, and a source build compiled for other architectures registers the same op and fails its first launch with 'no kernel image is available'. The probe holds no CUDA context by design, so it cannot establish coverage. The negative verdict is sound and stays; every other exit is now unknown, matching probe_flash_attn. The About row renders the reason on an unknown result too. Most unknowns are deliberate and come with an explanation, and the row was showing 'Not checked' and discarding it, so a skipped native extension looked like a probe that never ran. Also imports UNSLOTH_ENABLE_LOGGING explicitly in attention_dispatch: it is not in _utils.__all__, so the newly reachable inconclusive-probe branch raised NameError during import unsloth instead of keeping xformers on. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Stop three xFormers diagnoses from stating more than they know The no-metadata path (a source or editable 0.0.34+ build) fell back to the recorded torch pin and called every later release a mismatch, contradicting the stable-ABI rule the build-metadata branch two blocks up already applies. An unrelated extension failure was therefore diagnosed as a torch-version problem with reinstall instructions that fix nothing. The breakage announcement claimed a fallback to SDPA and more memory even when FlashAttention is installed, which select_attention_backend prefers over xformers anyway. The consequence sentence now follows what actually picks up the work; the breakage is still reported either way. A numeric CUDA_VISIBLE_DEVICES entry is a CUDA ordinal in the current CUDA_DEVICE_ORDER, and nvidia-smi's index is the PCI_BUS_ID ordering. The pin is a setdefault, so a user who exported FASTEST_FIRST keeps it and the ordinal cannot be inverted without initialising CUDA. That case now answers unknown instead of probing the wrong card; UUID masks are unaffected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Four more honest answers from the accelerator report An unresolvable device mask reached the probe child as an empty UNSLOTH_PROBE_DEVICE_CC, which the child reads as no override at all and answers from nvidia-smi over every physical GPU: the whole-box verdict the resolver returned None to avoid. It now travels as an explicit unknown sentinel, and the child refuses to fall back on it. An xformers whose _cpp_lib raises a native load error left runs at None, so the package stayed out of degraded and the About row read 'Not checked' with no banner, on exactly the corrupt install this report is for. That is now broken; a layout with no _cpp_lib at all stays unknown. bitsandbytes is probed on ROCm. device_type.py imports it on every backend and gates 4-bit and prequantized loading on native_kernels_ready, so a dead ROCm wheel silently costs quantized loading while the report called the package 'not used on this device'. And the fix hint redacts the mirror URL. UNSLOTH_PYTORCH_MIRROR can carry credentials or a signed-URL query, and that hint prints to stdout on the default import path, into logs and shared notebooks. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Find the bitsandbytes checker from the checkout too find_spec is the right answer for an installed unsloth, but CI installs a published wheel that predates bnb_availability.py, so the lookup landed on a site-packages copy without the file and every bitsandbytes verdict came back unknown. This file's own location settles it: studio/backend/utils/hardware to the repo root. * Stop three accelerator messages from overstating what is known A transient probe failure is not a broken 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 healthy install -- and that answer is cached for the life of the backend, so the banner sat there until restart. Busy, out-of-memory and unavailable now read as not probed, with the reason shown, using the same classification the in-process xformers probe makes. The mirror redactor no longer returns its input when it cannot parse it. A malformed authority is exactly the URL whose credentials or signed query this was added to keep out of stdout, so it drops the index instead and the hint omits --index-url. And the Settings banner no longer promises a slower fallback. When xFormers is the dead package and FlashAttention loaded, the dispatcher prefers FlashAttention and there is no slower path; the banner now names what cannot load and points at the list, in every locale. * Pin two accelerator-report tests to a CUDA host Both assert on the probe's classification, and on a CPU-only runner the report never probes at all: every row comes back "not used on this device" and degraded is empty, so the two assertions failed on CI while passing on any machine with a GPU. The on_accelerator fixture the rest of the file already uses is what makes the probe applicable. --------- Co-authored-by: Daniel Han <moonshotaisubstack@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
5d461d7458
commit
72aec0fa0a
34 changed files with 5247 additions and 53 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
360
studio/backend/tests/test_accelerator_probe_kernels.py
Normal file
360
studio/backend/tests/test_accelerator_probe_kernels.py
Normal file
|
|
@ -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.<ns> 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"]
|
||||
754
studio/backend/tests/test_hardware_accelerator_report.py
Normal file
754
studio/backend/tests/test_hardware_accelerator_report.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
426
studio/backend/utils/hardware/accelerator_probe.py
Normal file
426
studio/backend/utils/hardware/accelerator_probe.py
Normal file
|
|
@ -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.<ns>)``: 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:]))
|
||||
|
|
@ -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) ==========
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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<string, string> = {
|
||||
xformers: "xFormers",
|
||||
flash_attn: "FlashAttention",
|
||||
torchao: "torchao",
|
||||
bitsandbytes: "bitsandbytes",
|
||||
};
|
||||
|
||||
function packageLabel(name: string): string {
|
||||
return PACKAGE_LABELS[name] ?? name;
|
||||
}
|
||||
|
||||
const HEALTH_LABEL_KEYS: Record<Health, TranslationKey> = {
|
||||
working: "settings.about.accelerator.working",
|
||||
broken: "settings.about.accelerator.notLoading",
|
||||
absent: "settings.about.accelerator.notInstalled",
|
||||
unknown: "settings.about.accelerator.notChecked",
|
||||
};
|
||||
|
||||
const HEALTH_CLASSES: Record<Health, string> = {
|
||||
// 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 (
|
||||
<SettingsRow
|
||||
label={packageLabel(pkg.name)}
|
||||
description={detail}
|
||||
alignTop={detail != null}
|
||||
// The full reason stays reachable without dominating the row.
|
||||
hint={explained && pkg.reason ? pkg.reason : undefined}
|
||||
>
|
||||
<span className="flex items-baseline gap-2">
|
||||
<code className="font-mono text-xs text-muted-foreground">
|
||||
{pkg.version ?? "—"}
|
||||
</code>
|
||||
<span className={cn("text-xs", HEALTH_CLASSES[health])}>
|
||||
{t(HEALTH_LABEL_KEYS[health])}
|
||||
</span>
|
||||
</span>
|
||||
</SettingsRow>
|
||||
);
|
||||
}
|
||||
|
||||
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 (
|
||||
<SettingsSection
|
||||
title={t("settings.about.accelerator.sectionTitle")}
|
||||
description={t("settings.about.accelerator.sectionDescription")}
|
||||
>
|
||||
{report.packages.map((pkg) => (
|
||||
<AcceleratorRow key={pkg.name} pkg={pkg} probed={report.probed} />
|
||||
))}
|
||||
</SettingsSection>
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 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 (
|
||||
<Alert variant="destructive" className="mb-4">
|
||||
<HugeiconsIcon icon={Alert01Icon} />
|
||||
<AlertTitle>{t("settings.about.accelerator.bannerTitle")}</AlertTitle>
|
||||
<AlertDescription>
|
||||
{t("settings.about.accelerator.bannerBody", { packages })}
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
);
|
||||
}
|
||||
|
|
@ -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. */}
|
||||
<AcceleratorBanner />
|
||||
{renderTab(panelTab)}
|
||||
</div>
|
||||
</main>
|
||||
|
|
|
|||
|
|
@ -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() {
|
|||
</SettingsSection>
|
||||
</div>
|
||||
|
||||
{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 ? (
|
||||
<SettingsSection title={t("settings.about.hardware")}>
|
||||
{hw.gpus.map((gpu, i) => (
|
||||
<SettingsRow
|
||||
|
|
@ -203,9 +206,20 @@ export function AboutTab() {
|
|||
</code>
|
||||
</SettingsRow>
|
||||
))}
|
||||
{/* 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 ? (
|
||||
<SettingsRow label={t("settings.about.python")}>
|
||||
<code className="font-mono text-xs text-muted-foreground">
|
||||
{hw.python}
|
||||
</code>
|
||||
</SettingsRow>
|
||||
) : null}
|
||||
</SettingsSection>
|
||||
) : null}
|
||||
|
||||
<AcceleratorSection />
|
||||
|
||||
<SettingsSection title={t("settings.about.help")}>
|
||||
<SettingsRow label={t("settings.about.documentation")}>
|
||||
<a
|
||||
|
|
|
|||
158
studio/frontend/src/hooks/accelerator-report.ts
Normal file
158
studio/frontend/src/hooks/accelerator-report.ts
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
/**
|
||||
* Parsing for the accelerator-health block of `GET /api/system/hardware?include_details=true`.
|
||||
*
|
||||
* Split out of use-hardware-info so it stays importable without react or the auth chain,
|
||||
* which is what lets it be unit tested directly.
|
||||
*/
|
||||
|
||||
/** What an xformers wheel records in its own cpp_lib.json. */
|
||||
export interface AcceleratorBuild {
|
||||
torch: string | null;
|
||||
cuda: string | null;
|
||||
// ROCm wheels record hip and leave cuda null; without this a ROCm xformers renders no
|
||||
// build detail at all.
|
||||
hip: string | null;
|
||||
python: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* One optimized-kernel package. `installed` and `imports` and `runs` are three separate
|
||||
* questions on purpose: an xformers built for a different torch is installed, imports,
|
||||
* reports a version -- and has no working kernels (NVIDIA QA P0-1). `runs` is null where
|
||||
* the question does not apply (absent, unprobed, or no separate kernel-load step).
|
||||
*/
|
||||
export interface AcceleratorPackage {
|
||||
name: string;
|
||||
version: string | null;
|
||||
installed: boolean;
|
||||
/**
|
||||
* Whether the probe child actually ran for THIS package. Null from a backend too old to
|
||||
* say, which is the only case the report-wide flag is allowed to answer for: the probe
|
||||
* set is per device, so a package the backend deliberately skipped is unknown, not dead.
|
||||
*/
|
||||
probed: boolean | null;
|
||||
imports: boolean;
|
||||
runs: boolean | null;
|
||||
reason: string | null;
|
||||
builtFor: AcceleratorBuild | null;
|
||||
}
|
||||
|
||||
export interface AcceleratorReport {
|
||||
pythonVersion: string | null;
|
||||
torchVersion: string | null;
|
||||
torchCuda: string | null;
|
||||
/** False when UNSLOTH_SKIP_ACCELERATOR_PROBE is set: unknown, not healthy. */
|
||||
probed: boolean;
|
||||
packages: AcceleratorPackage[];
|
||||
/** Import names of packages that are installed but cannot load. */
|
||||
degraded: string[];
|
||||
}
|
||||
|
||||
// Fixed order so the About table does not reshuffle between reads, and so the most
|
||||
// consequential package (the one in the P0) is first.
|
||||
const ACCELERATOR_ORDER = ["xformers", "flash_attn", "torchao", "bitsandbytes"];
|
||||
|
||||
function parseAcceleratorBuild(raw: unknown): AcceleratorBuild | null {
|
||||
if (typeof raw !== "object" || raw === null) return null;
|
||||
const build = raw as Record<string, unknown>;
|
||||
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<string, unknown>;
|
||||
const rawPackages = (report.packages ?? {}) as Record<
|
||||
string,
|
||||
Record<string, unknown>
|
||||
>;
|
||||
// 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;
|
||||
}
|
||||
69
studio/frontend/src/hooks/use-accelerator-report.ts
Normal file
69
studio/frontend/src/hooks/use-accelerator-report.ts
Normal file
|
|
@ -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<AcceleratorReport | null> | null = null;
|
||||
|
||||
export function invalidateAcceleratorReport() {
|
||||
cached = null;
|
||||
fetchPromise = null;
|
||||
}
|
||||
|
||||
async function fetchOnce(): Promise<AcceleratorReport | null> {
|
||||
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<AcceleratorReport | null>(cached);
|
||||
|
||||
useEffect(() => {
|
||||
let cancelled = false;
|
||||
fetchOnce().then((next) => {
|
||||
if (!cancelled) setReport(next);
|
||||
});
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, []);
|
||||
|
||||
return report;
|
||||
}
|
||||
|
|
@ -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<HardwareInfo> {
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -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: "الوثائق",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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: "दस्तावेज़",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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: "ドキュメント",
|
||||
|
|
|
|||
|
|
@ -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: "문서",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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: "Документация",
|
||||
|
|
|
|||
|
|
@ -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: "文档",
|
||||
|
|
|
|||
270
studio/frontend/tests/accelerator-report.test.ts
Normal file
270
studio/frontend/tests/accelerator-report.test.ts
Normal file
|
|
@ -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);
|
||||
});
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
351
tests/utils/test_xformers_broken_warning.py
Normal file
351
tests/utils/test_xformers_broken_warning.py
Normal file
|
|
@ -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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""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
|
||||
|
|
@ -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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""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"
|
||||
|
|
|
|||
405
tests/utils/test_xformers_compat.py
Normal file
405
tests/utils/test_xformers_compat.py
Normal file
|
|
@ -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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""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/<version>/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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
470
unsloth/xformers_compat.py
Normal file
470
unsloth/xformers_compat.py
Normal file
|
|
@ -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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""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/<version>/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
|
||||
Loading…
Add table
Add a link
Reference in a new issue