6 commits

Author	SHA1	Message	Date
Daniel Han	947b4bbd7c	Studio: verify the flash-attn import after installing it (#8465) ... * Studio: verify the flash-attn import after installing it A prebuilt wheel can install with exit code 0 and then fail to load, so the exit code on its own is not proof the install is usable. Both Blackwell incidents were that shape: in #5420 the older-arch wheels installed and raised on import, and the arch gate added to work around it became the bug in #6961 once Dao-AILab started shipping sm_100 wheels. Verify the import after a zero exit code in the two paths that did not, and treat a failure as not installed. ssm_runtime._install_kernel already did exactly this, for the same reason, so this is that check applied to the setup installer and the long-context training worker. worker._is_importable() catches any exception rather than only ImportError: an arch/ABI mismatch surfaces as OSError or RuntimeError ("undefined symbol"), which the old pre-check would have let escape mid-training. Also drop has_blackwell_gpu() and its two call sites. It has returned False unconditionally since 03cbe21, so both blocks were already unreachable and removing them changes no behaviour. An arch gate encodes a snapshot of what upstream publishes and goes stale silently in both directions, whereas the import check catches a wheel that will not load whatever the cause. Checked on an 8x B200 host (compute_cap 10.0): the wheel this resolver builds for cu13/torch2.10/cp313 carries sm_100 and sm_120 cubins, installs through install_wheel(), and runs, matching SDPA to 0.0078 max abs difference on a 2x4096x16x128 bf16 causal forward. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten the comments added by this PR (no behaviour change) * Stub utils.native_tls in the MLX worker config test worker.py imports six utils.* submodules and calls activate_native_tls() at import time; _load_worker_module stubbed five of them. The file therefore loaded only when another test had already imported the real utils.native_tls, so it passed in a full run and failed on its own with "No module named 'utils.native_tls'". Reproduced on upstream main at aaf99488, so this predates this branch: running the file alone, or in any narrow selection, errors during collection. Cross-platform CI runs the PR's own test files as a narrow selection, which is how it surfaced. * Replace a rejected wheel on the fallback, and verify that install too Rejecting a wheel that installed but would not import left the broken distribution in site-packages, and the fallback installed over it rather than replacing it. pip reports it as already satisfied and uv audits it as no change, both exiting 0, and that exit code reached an unconditional return True. Long-context training then continued with an unusable extension, which is the failure the wheel check was meant to stop. Measured against a real flash_attn 2.8.1 install: pip install --no-build-isolation --no-deps flash-attn -> "Requirement already satisfied: flash-attn ... (2.8.1)", exit 0 uv pip install --no-build-isolation --no-deps flash-attn -> "Audited 1 package", "Would make no changes", exit 0 So pass --force-reinstall (pip) / --reinstall (uv) on the fallback when a wheel was rejected, and verify the import afterwards instead of trusting rc=0. The flag is gated on the rejection so the ordinary path keeps installing over nothing and does not rebuild. test_runtime_flash_attn_falls_back_to_pypi mocked an install that exits 0 without ever making the module importable, which under the post-install check is a failed install, not a successful one. It now flips the import stub when the install runs. * Uninstall a rejected wheel instead of reinstalling over it, and probe out of process Two problems with the previous commit. --force-reinstall was the wrong tool. pip documents it as "Reinstall all packages even if they are already up-to-date" and uv's --reinstall as "Reinstall all packages": both scope to the whole resolved transaction, not the named one. flash-attn depends on torch, so on the plain fallback path (which carries no --no-deps) that could reinstall or downgrade the torch the worker is currently running on. Measured with a stand-in whose metadata resolves, since flash-attn's sdist cannot build metadata here: pip install --dry-run requests -> 0 packages pip install --force-reinstall --dry-run requests -> 5 packages (requests + its whole closure) It also fails outright on that path: --force-reinstall makes pip rebuild flash-attn from sdist, and the plain branch passes no --no-build-isolation, so the build dies on "No module named 'torch'". This file already pairs --force-reinstall with --no-deps elsewhere for the same reason. So remove the rejected distribution and leave the fallback command untouched. Nothing is reported as already satisfied afterwards, the transaction stays exactly as wide as it was before, and no rebuild flags change. Second, the post-install probe now runs in a child. A wheel built for the wrong arch can abort or segfault inside the extension's initialiser rather than raising, and except Exception cannot catch that: it would kill the worker and take the fallback with it. The child turns it into a return code, negative for a fatal signal. install_python_stack already probes this way; the worker now matches it. The cheap in-process check stays for "is it already installed", which costs no spawn on every training start and imports the module the worker wants imported anyway. The module name is passed as argv rather than formatted into the -c body. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Bound the installer's flash-attn import probe The worker bounds this same untrusted import at 300s and handles TimeoutExpired; the installer probe had no timeout, so a native extension that hangs in its initialiser rather than failing would leave setup waiting forever and never reach the warning. The asymmetry arrived with the post-install probe added earlier in this branch, which is exactly the call that hands an unvalidated wheel to the import machinery. Bound it the same way and treat a hang, or a probe that cannot be spawned, as an unusable wheel. The timeout is a named constant on both sides so the two paths cannot drift apart again silently. * Tighten the comments added in review (no behaviour change) * Remove a rejected flash-attn wheel in the installer too The setup path warned and continued while leaving the unusable distribution in site-packages, so "Continuing without flash-attn" was not true. unsloth/models/_utils.py gates on _package_available("flash_attn"), which reads metadata rather than importing, and then does an in-process "from flash_attn.flash_attn_interface import ...". A wheel that aborts during native initialisation therefore takes the training process down, and that is exactly the wheel this branch has just rejected. Confirmed with a package that calls os.abort() at import: importing it in process kills the interpreter with SIGABRT, while the out-of-process probe returns rc -6 and survives. Uninstalling a properly installed distribution removes it cleanly. The worker path already uninstalls before its fallback; the installer now matches. * Uninstall with the mode uv was installed with, and stop claiming a failed removal Two problems in the cleanup added in the previous commit. It hard-coded "uv pip uninstall --python". _bootstrap_uv sets UV_NEEDS_SYSTEM exactly when the --python probe FAILED and --system succeeded, so on those hosts the cleanup used the one mode already known not to work there: the uninstall fails, the unusable wheel stays in site-packages, and unsloth/models/_utils.py still finds it by metadata and imports it in process. Mirror the install mode instead. uv documents --system as "Use the system Python to uninstall packages". It also printed "removed it" unconditionally, so a failed removal produced two contradicting warnings in the same run: "Could not remove the unusable flash-attn install" followed by "...; removed it". The helper now returns whether the package is actually gone, and the caller says which happened. A wheel still installed is not the same state as never having installed one, so the failure case says so and points at the manual uninstall. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route every rejected install through one discard path The PyPI fallback returned False without removing the distribution it had just rejected, so the unusable extension stayed where unsloth/models/_utils.py finds it: that gate reads package METADATA and then imports the native module in process, which is exactly the import the isolated probe could not survive. Two ways in: a failed uninstall leaves the wheel "already satisfied" so the fallback no-ops, or the source build itself produces an incompatible extension. This is the third round on the same defect, in a third place, so fix the shape rather than the instance. _reject_install is now the single discard path: it uninstalls, and it reports which state we actually ended in. _uninstall_package returns whether the distribution is gone rather than only logging. The pre-fallback uninstall stays a plain call: a failure there is not fatal, because the fallback then no-ops on "already satisfied", the probe rejects it, and _reject_install reports the real state. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enforce the discard invariant in one place instead of at each return The timeout and install-failure exits returned without discarding the rejected distribution, so a wheel that had already failed its import check stayed where unsloth/models/_utils.py finds it: that gate reads METADATA and only then imports the native module in process. The comment added last round claimed the fallback path would clean up on its own, and that is only true when the fallback exits 0. This is the fourth defect of the same shape, each one an exit somebody did not think to clean up, so enforce the invariant structurally rather than adding a fourth call. _install_package_wheel_first now keeps the two "touch nothing" guards (already importable, offline) and delegates the rest to _attempt_package_install, discarding whatever is left in a finally. Any unsuccessful exit, including ones added later, is covered. The discard is state-based so it is safe to run everywhere: _distribution_present reads metadata via importlib.metadata without importing, so it never loads an extension that would abort, and _reject_install no-ops when there is nothing installed. Verified on the B200 that a working flash-attn 2.8.1 is still untouched: three consecutive runs return True with zero install or uninstall subprocesses, and it stays importable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * State the installed-but-broken state in the rejection tests Two rejection tests asserted on the discard status message without pinning _distribution_present, so after the discard became state-based they only passed where flash-attn happened to be installed. That was true in the venv I had been running them in and false on a clean runner, where the discard correctly no-ops and no status is sent. Cross-platform CI caught it on ubuntu-latest; the same two tests fail in a fresh venv here and pass with the state pinned. No production change: the invariant holds either way. The tests were reading the machine rather than the code. * Uninstall from the interpreter install_wheel actually installed into install_wheel always targets sys.executable: its uv command passes --python in addition to --system, and its pip fallback runs that interpreter directly. The cleanup passed --system INSTEAD of --python, so on a UV_NEEDS_SYSTEM host it uninstalled from the system Python while the wheel sat in the venv, and setup then reported the wheel removed when it was still there for the metadata gate to import. Pass both, so the removal mirrors the install rather than half of it. The earlier test asserting the --system-only command is replaced, since it pinned the broken command. The worker path already targets --python sys.executable on both sides and needs no change. * Tighten the review comments (no behaviour change) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>	2026-08-12 05:43:51 -07:00
Daniel Han	5e6b4aa1ee	Reuse the torch2.10 prebuilt accelerator wheels on torch 2.12 (#7495)	2026-08-07 01:19:14 -03:00
Thomas Eric 🇧🇷	03cbe211a3	Studio: fix flash-attn and torchao install on Blackwell (sm_100+) GPUs (Closes #6961) (#6970) ... * fix: Remove moot has_blackwell_gpu() function Fixes unslothai/unsloth#6961. This function skipped flash-attn on Blackwell GPUs because no prebuilt wheel existed; Dao-AILab now ships one and url_exists() already gates resolution. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix: use torchao 0.17.0 for Blackwell Fixes #6961. Torchao 0.16.0's cpp extensions are built against CUDA 12, so on a CUDA-13 torch (cu130 / Blackwell) they fail to load with "libcudart.so.12: cannot open shared object file". Select 0.17.0 there instead: its cpp targets torch 2.11, so it is skipped cleanly rather than crashing. CUDA-12 / ROCm / CPU torch 2.10 keeps 0.16.0 and its working kernels. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Condense torchao version-selection comments (no behavior change) * Support torch 2.11 in the Studio installer via the torch2.10 prebuilt wheels Map torch 2.11 to the torch2.10 prebuilt wheels for flash-attn, causal-conv1d, and mamba through wheel_utils.prebuilt_wheel_torch_mm, applied in direct_wheel_url (filename) and flash_attn_wheel_url (version). Those torch2.10 CUDA wheels load and pass each project's own test suite on torch 2.11 (verified on B200), so a torch 2.11 environment gets the prebuilt accelerators instead of skipping or building from source. Raise _CUDA_TORCH_PKG_SPEC to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so the CUDA torch repair path can install torch 2.11, where torchao 0.17's cpp kernels load cleanly. Add tests for the mapping. * Keep has_blackwell_gpu as a False stub for future arch gating * Restore has_blackwell_gpu as a return-False probe kept for future arch gating Keep the nvidia-smi compute_cap detection and its two call sites, but short-circuit with return False at the top so flash-attn is no longer skipped on Blackwell (sm_100+ now has prebuilt wheels and url_exists gates resolution). Drop the early return to re-enable arch-based detection later. --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>	2026-07-08 06:38:10 -07:00
Daniel Han	3ce187da02	Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079) ... Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.	2026-06-08 04:24:13 -07:00
Roland Tannous	79adfd9c71	studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420) ... * studio: skip flash-attn install on Blackwell GPUs (sm_100+) Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120, or sm_121, and the older-arch wheels fail to load on Blackwell. Add a shared has_blackwell_gpu() helper and gate both the install-time (install_python_stack._ensure_flash_attn) and runtime (worker._ensure_flash_attn_for_long_context) paths on it. Detection uses nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows. * test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests prefers_prebuilt_wheel and falls_back_to_pypi exercise the install paths that the Blackwell guard now short-circuits. Make them explicit about non-Blackwell so they pass on real Blackwell hosts. * studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH - Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a single process avoid redundant nvidia-smi spawns. Tests clear the cache via setup_method/teardown_method. - In _ensure_flash_attn, run the NO_TORCH short-circuit before the Blackwell check so GGUF-only users (who never install torch anyway) do not see a Blackwell warning. Blackwell check still runs above the IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see the explicit reason rather than a silent OS skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: add has_blackwell_gpu to mlx worker test wheel_utils stub test_mlx_training_worker_config loads worker.py against a hand-rolled utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol list so worker's import line resolves. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>	2026-05-14 18:13:50 +04:00
Datta Nimmaturi	da78c6be71	[Studio] Install flash attn at setup time for linux (#4979) ... * [Studio] Install flash attn at setup time for linux * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup changes Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Test cases * wheel_utils: narrow url_exists exceptions and log at debug level --------- Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>	2026-04-14 16:40:17 +04:00