unsloth/tests/python
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
..
__init__.py Consolidate dual venvs and separate install from update (#4530) 2026-03-25 05:24:21 -07:00
conftest.py fix: add tokenizers to no-torch deps and TORCH_CONSTRAINT for arm64 macOS py313+ (#4748) 2026-04-01 06:12:17 -07:00
test_cross_platform_parity.py Add configurable PyTorch mirror via UNSLOTH_PYTORCH_MIRROR env var (#5024) 2026-04-15 11:39:11 +04:00
test_dpo_vision_processor_passthrough.py Fix DPO trainer multi process hang (#5199) 2026-04-29 04:15:34 -07:00
test_e2e_no_torch_sandbox.py tests: add no-torch / Intel Mac test suite (#4646) 2026-03-27 02:33:45 -07:00
test_fast_sentence_transformer_redirect_lifecycle.py Add Studio PR-time CI: pin enforcement, frontend, backend, wheel smoke (#5298) 2026-05-06 04:41:57 -07:00
test_flash_attn_install_python_stack.py studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420) 2026-05-14 18:13:50 +04:00
test_gpu_init_ldconfig_guard.py feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265) 2026-05-05 23:54:58 -07:00
test_install_python_stack.py Consolidate dual venvs and separate install from update (#4530) 2026-03-25 05:24:21 -07:00
test_no_torch_filtering.py [Studio] Install flash attn at setup time for linux (#4979) 2026-04-14 16:40:17 +04:00
test_patch_trl_rl_trainers_defensive.py fix: unblock 4 tests deselected/skipped in #5312 (real bugs) (#5359) 2026-05-11 02:39:17 -07:00
test_studio_import_no_torch.py tests: add no-torch / Intel Mac test suite (#4646) 2026-03-27 02:33:45 -07:00
test_tokenizers_and_torch_constraint.py fix: add tokenizers to no-torch deps and TORCH_CONSTRAINT for arm64 macOS py313+ (#4748) 2026-04-01 06:12:17 -07:00
test_unsloth_run_tool_policy_resolver.py unsloth run: add --enable-tools/--disable-tools server-side tool policy (#5277) 2026-05-05 12:45:15 +04:00