mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-31 19:20:08 +00:00
7729c37fd0
Follow up to #13861 and #13808. An audit of the merged on demand SSM
kernel change plus a T4 benchmark found a few gaps, fixed here.
## Why
- On a Kaggle T4 (sm75, torch 2.11 cu130, fp16, Qwen3.5-0.8B), Studio
chat decode was 4.4% slower without causal-conv1d (19.10 vs 19.95 tok/s,
6 alternating repeats, slower in all 9 pairs). Training speed, losses
and peak memory were unchanged. The zoo Triton conv covers
`causal_conv1d_fn` only, so single token decode fell back to the torch
update path. Training with trainable conv weights also allocates an fp32
dW scratch in the Triton backward, and packed `seq_idx` is not accepted
by the zoo dispatcher under transformers 5.20.
- Mamba-family models that have a working transformers torch path
(mamba, mamba2, falcon_mamba, jamba, zamba, zamba2, bamba) were made
fatal when mamba-ssm could not be installed.
- Export, local LoRA continuation, and inference model switching did not
provision or pick up kernels correctly once setup stopped installing
them.
- Setup on pre-Ampere GPUs (T4) downloaded flash-attn even though
Unsloth only enables FA2 at sm80+.
## Changes
1. Qwen3.5 / Qwen3.6 / Qwen3-Next request causal-conv1d on demand again
in training, inference and Clef (same as before #13861, minus the setup
time install). The zoo Triton platform probe is removed.
2. A missing mamba-ssm is fatal only for `nemotron_h`, `falcon_h1`,
`granitemoehybrid`. Other Mamba families install best effort and
continue on the torch path with a status message.
3. Local LoRA folders (`adapter_config.json`, no `config.json`) resolve
kernel needs from `base_model_name_or_path`; the training
`adapter_base_pin` wins when set.
4. The export worker provisions SSM kernels before importing Unsloth /
the export backend, with the same required vs optional split.
5. Inference re-enables the modeling source scan by reading
`transformers/models/<type>/modeling_*.py` as text, without importing
transformers.
6. If a kernel is newly installed after the model's
`transformers.models.<type>` module was already imported, the inference
worker asks for a restart and the orchestrator respawns it once.
7. Torch repair removes causal-conv1d / mamba-ssm whose compiled
extension no longer imports (never installs anything new). A hung import
check keeps the package.
8. Setup and the long context runtime install skip flash-attn below sm80
(and remove a stale copy at setup). No visible GPU and ROCm keep current
behavior.
## Tests
New and updated tests fail on main and pass here (`test_ssm_runtime.py`,
`test_ssm_restart_respawn.py`, `test_export_ssm_kernels.py`,
`test_clef_fast_paths.py`, `test_kernel_install.py`,
`test_training_worker_flash_attn.py`,
`test_ssm_kernels_install_python_stack.py`,
`test_flash_attn_install_python_stack.py`,
`test_windows_torch_flavor_invariant.py`). With
`test_install_python_stack.py` and `test_cuda_repair.py`: 966 passed
after merging current main.
## Hardening (second commit)
- Stale SSM kernel removal after torch repair needs proof: torch must
import first, and only an ImportError / OSError from the extension
itself removes the package. Signals, hangs, launch errors and a broken
torch keep it.
- `unsloth install-kernels --without-gpu` (or
`UNSLOTH_INSTALL_KERNELS_WITHOUT_GPU=1`) installs GPU gated kernels on a
GPU-less image builder. An already working kernel counts as success
before any skip, skipped kernels are listed, and the exit code is
unchanged (0) unless `--strict` is passed. The CLI respects `UV_OFFLINE`
/ HF offline variables and never falls back from uv to pip offline.
- Inference respawn refuses to spawn while the old worker survives
shutdown. SSM restart and the Xet fallback each get their own one-use
retry (3 spawns max).
- Restart detection only counts native `modeling_*` modules and the
selected checkpoint's remote code modules from its `auto_map`, matched
on the exact `transformers_modules` namespace.
- Export provisions SSM kernels inside the existing transformers
activation window (no extra Hub reachability probe), and skips
provisioning on macOS and for the MLX / CPU export backends.
- `UV_OFFLINE` counts as offline for on-demand installs (no URL probes,
source builds, or uv to pip fallback).
- Long context flash-attn install returns early when flash_attn already
imports, reuses the worker's own torch for the capability check, and
never applies the sm80 floor on ROCm.
## Install robustness (platform matrix review)
- A wheel is downloaded whenever its HEAD check is not a definite 404
(proxies that refuse HEAD no longer push required hybrids into a
failure), then import-verified; fallback only after a real failure.
- A failed install is remembered for the worker process (keyed by
kernel, torch, toolchain), so preflight and load share one attempt.
Optional Mamba families get mamba-ssm only as a prebuilt wheel, never a
source build. ROCm source builds get an 1800 s deadline and CUDA source
builds 7200 s (`UNSLOTH_SOURCE_BUILD_TIMEOUT` overrides); a timeout
kills the whole build tree.
- A runtime install whose extension does not import is uninstalled, so
transformers never finds a broken package.
- One cross-process lock per kernel and environment covers check,
install, verify and cleanup (training and runtime), with a recheck after
acquiring it.
- Only a CUDA or ROCm torch provisions anything; CPU, XPU, MPS and
unknown torch builds skip (export included).
- Optional mamba-ssm follows setup's sm80 floor (sm75 with Triton 3.4+);
the required hybrids keep the previous behavior.
- The resolver only returns URLs for cells that are actually published
(Unsloth release: torch 2.13/2.14 cp313, win_amd64 torch 2.10 to 2.12
cp311 to cp313; upstream causal-conv1d / mamba-ssm release cells), and
there is no source build of the pinned 1.6.1 / 2.3.1 sources on torch
2.13+, which need C++20.
- Training stays best effort for mamba-ssm, as before #13861.
- Export runs the malware and remote code consent gates before any
kernel provisioning.
- The extension is probed again after the install lock is acquired, so a
concurrent replacement is not treated as broken.
- The automatic HF offline window (Hub unreachable) no longer blocks
package installs; only user set offline variables do.
- `unsloth install-kernels --dry-run` prints the resolved URL while
offline.
- A Mamba model whose code needs mamba-ssm fails with a clear error on a
torch build without CUDA or ROCm.
## GPU checks (merge base vs head 4a6429a65)
Kaggle T4, torch 2.11.0+cu130, Qwen3.5-0.8B fp16:
| | base | head |
|---|---|---|
| causal-conv1d on demand | not installed | installed,
`causal_conv1d_update_kernel` engaged |
| decode tok/s (6 alternating repeats) | 18.70 | 19.41 (+3.8%, 6/6
pairs) |
| decode peak alloc | 1.748 GB | 1.748 GB |
| LoRA step time (3 repeats) | 1975.5 ms | 1970.7 ms (inside spread) |
| LoRA peak alloc | 2.338 GB | 2.338 GB |
| flash-attn at setup on sm75 | downloaded | skipped |
AMD gfx1151 (torch 2.11.0+rocm7.13.0): Qwen3.5 LoRA losses within 1.2e-4
relative (head vs merge commit bit identical), peak 2.389 GB both,
backend suites 262/262 and 275/275 passing.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
897 lines
35 KiB
Python
897 lines
35 KiB
Python
"""Tests for the optional FlashAttention installer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
|
|
def _shared_setup_1(fake_run, step_messages):
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))],
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"_step",
|
|
side_effect = lambda label, value, color_fn = None: step_messages.append((label, value)),
|
|
),
|
|
mock.patch("subprocess.run", side_effect = fake_run),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
|
|
STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
|
|
sys.path.insert(0, str(STUDIO_DIR))
|
|
sys.path.insert(0, str(STUDIO_DIR / "backend"))
|
|
|
|
import install_python_stack as ips
|
|
import pytest
|
|
from utils import wheel_utils
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _an_sm80_gpu(monkeypatch):
|
|
# The host's real GPU must not decide these tests; the capability cases below override it.
|
|
monkeypatch.setattr(ips, "_flash_attn_gpu_below_floor", lambda: False)
|
|
|
|
|
|
class TestPrebuiltWheelTorchMapping:
|
|
def test_torch_211_maps_to_torch210(self):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10"
|
|
|
|
def test_torch_212_maps_to_torch210(self):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm("2.12") == "2.10"
|
|
|
|
def test_other_versions_pass_through(self):
|
|
# 2.13 stays unmapped until its wheels have been measured.
|
|
for torch_mm in ("2.9", "2.10", "2.13"):
|
|
assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm
|
|
|
|
def test_reuse_never_targets_a_pre_210_wheel(self):
|
|
# torch broke extension ABI between 2.9 and 2.10.
|
|
assert set(wheel_utils._PREBUILT_WHEEL_TORCH_MM.values()) == {"2.10"}
|
|
|
|
def test_direct_wheel_url_reuses_torch210_on_211(self):
|
|
url = wheel_utils.direct_wheel_url(
|
|
filename_prefix = "causal_conv1d",
|
|
package_version = "1.6.1",
|
|
release_tag = "v1.6.1.post4",
|
|
release_base_url = "https://example.test/download",
|
|
env = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.11",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_direct_wheel_url_reuses_torch210_on_212(self):
|
|
url = wheel_utils.direct_wheel_url(
|
|
filename_prefix = "mamba_ssm",
|
|
package_version = "2.3.1",
|
|
release_tag = "v2.3.1",
|
|
release_base_url = "https://example.test/download",
|
|
env = {
|
|
"python_tag": "cp312",
|
|
"torch_mm": "2.12",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "mamba_ssm-2.3.1+cu13torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" in url
|
|
|
|
|
|
class TestFlashAttnWheelSelection:
|
|
def test_torch_210_maps_to_v281(self):
|
|
assert ips._select_flash_attn_version("2.10") == "2.8.1"
|
|
|
|
def test_selected_version_is_never_a_post_release(self):
|
|
# v2.8.3.post1 dropped torch2.10 assets; a naive newest-release bump must fail here.
|
|
for torch_mm in ("2.4", "2.7", "2.9", "2.10"):
|
|
version = ips._select_flash_attn_version(torch_mm)
|
|
assert version is not None
|
|
assert ".post" not in version
|
|
|
|
def test_torch_29_maps_to_v283(self):
|
|
assert ips._select_flash_attn_version("2.9") == "2.8.3"
|
|
|
|
def test_torch_211_has_no_native_version_entry(self):
|
|
assert ips._select_flash_attn_version("2.11") is None
|
|
|
|
def test_torch_211_reuses_torch210_wheel(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.11",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_torch_212_reuses_torch210_wheel(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.12",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_exact_wheel_url_uses_full_env_tuple(self):
|
|
url = ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
assert url is not None
|
|
assert "v2.8.1" in url
|
|
assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
|
|
|
|
def test_missing_cuda_major_disables_wheel_lookup(self):
|
|
assert (
|
|
ips._build_flash_attn_wheel_url(
|
|
{
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
class TestFlashAttnImportProbe:
|
|
"""The probe is bounded: a native extension can hang in its initialiser, not just fail."""
|
|
|
|
def test_clean_exit_is_importable(self):
|
|
with mock.patch(
|
|
"subprocess.run",
|
|
return_value = subprocess.CompletedProcess(["python"], 0),
|
|
):
|
|
assert ips._flash_attn_importable() is True
|
|
|
|
def test_non_zero_exit_is_not_importable(self):
|
|
with mock.patch(
|
|
"subprocess.run",
|
|
return_value = subprocess.CompletedProcess(["python"], 1),
|
|
):
|
|
assert ips._flash_attn_importable() is False
|
|
|
|
def test_a_hung_import_is_not_importable(self):
|
|
with mock.patch(
|
|
"subprocess.run",
|
|
side_effect = subprocess.TimeoutExpired(cmd = "python", timeout = 300),
|
|
):
|
|
assert ips._flash_attn_importable() is False
|
|
|
|
def test_the_probe_is_bounded(self):
|
|
with mock.patch(
|
|
"subprocess.run",
|
|
return_value = subprocess.CompletedProcess(["python"], 0),
|
|
) as run:
|
|
ips._flash_attn_importable()
|
|
|
|
assert run.call_args.kwargs["timeout"] == ips._FLASH_ATTN_IMPORT_PROBE_TIMEOUT
|
|
|
|
|
|
class TestEnsureFlashAttn:
|
|
def _import_check(self, code: int = 1):
|
|
return subprocess.CompletedProcess(["python", "-c", "import flash_attn"], code)
|
|
|
|
def _import_fails_removal_works(self, cmd, **_kwargs):
|
|
"""flash_attn never imports, but uninstalling it succeeds."""
|
|
if "uninstall" in cmd:
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
return self._import_check()
|
|
|
|
def test_prefers_exact_match_wheel(self):
|
|
install_calls = []
|
|
|
|
def fake_install_wheel(*args, **kwargs):
|
|
install_calls.append((args, kwargs))
|
|
return [("uv", subprocess.CompletedProcess(["uv"], 0, ""))]
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(ips, "install_wheel", side_effect = fake_install_wheel),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
assert len(install_calls) == 1
|
|
args, kwargs = install_calls[0]
|
|
assert args == (
|
|
"https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl",
|
|
)
|
|
assert kwargs["python_executable"] == sys.executable
|
|
assert kwargs["use_uv"] is True
|
|
assert kwargs["uv_needs_system"] is False
|
|
|
|
def test_uv_install_respects_system_flag(self):
|
|
install_calls = []
|
|
|
|
def fake_install_wheel(*args, **kwargs):
|
|
install_calls.append((args, kwargs))
|
|
return [("uv", subprocess.CompletedProcess(["uv"], 0, ""))]
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", True),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(ips, "install_wheel", side_effect = fake_install_wheel),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
assert len(install_calls) == 1
|
|
_, kwargs = install_calls[0]
|
|
assert kwargs["uv_needs_system"] is True
|
|
|
|
def test_wheel_that_does_not_import_is_not_trusted(self):
|
|
"""pip exits 0 on a wrong-arch or ABI wheel, so only a successful import shows flash-attn is
|
|
ready."""
|
|
step_messages: list[tuple[str, str]] = []
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))],
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"_step",
|
|
side_effect = lambda label, value, color_fn = None: step_messages.append(
|
|
(label, value)
|
|
),
|
|
),
|
|
mock.patch("subprocess.run", side_effect = self._import_fails_removal_works),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
assert (
|
|
"warning",
|
|
"flash-attn wheel installed but is not importable on this GPU; removed it",
|
|
) in step_messages
|
|
assert ("warning", "Continuing without flash-attn") in step_messages
|
|
|
|
def test_rejected_wheel_is_uninstalled(self):
|
|
"""Rejected wheels must be uninstalled: _package_available trusts metadata and imports natively."""
|
|
step_messages: list[tuple[str, str]] = []
|
|
removals: list[list[str]] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
if "uninstall" in cmd:
|
|
removals.append(list(cmd))
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
return self._import_check()
|
|
|
|
_shared_setup_1(fake_run, step_messages)
|
|
|
|
assert removals, "the rejected wheel must be uninstalled, not left in site-packages"
|
|
assert any("flash-attn" in cmd for cmd in removals), removals
|
|
assert ("warning", "Continuing without flash-attn") in step_messages
|
|
|
|
def test_uninstall_targets_the_interpreter_install_wheel_used(self):
|
|
"""install_wheel passes --python as well as --system, and its pip fallback runs
|
|
sys.executable directly, so --system alone would uninstall from the wrong Python."""
|
|
commands: list[list[str]] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
commands.append(list(cmd))
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
|
|
with (
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", True),
|
|
mock.patch.object(ips.shutil, "which", return_value = "/usr/bin/uv"),
|
|
mock.patch("subprocess.run", side_effect = fake_run),
|
|
):
|
|
assert ips._remove_rejected_flash_attn() is True
|
|
|
|
assert commands == [
|
|
["uv", "pip", "uninstall", "--system", "--python", sys.executable, "flash-attn"]
|
|
]
|
|
|
|
def test_uninstall_targets_the_interpreter_without_system_mode(self):
|
|
commands: list[list[str]] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
commands.append(list(cmd))
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
|
|
with (
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(ips.shutil, "which", return_value = "/usr/bin/uv"),
|
|
mock.patch("subprocess.run", side_effect = fake_run),
|
|
):
|
|
assert ips._remove_rejected_flash_attn() is True
|
|
|
|
assert commands == [["uv", "pip", "uninstall", "--python", sys.executable, "flash-attn"]]
|
|
|
|
def test_a_failed_removal_is_not_reported_as_removed(self):
|
|
"""Still importable in process, so it must not read like a clean skip."""
|
|
step_messages: list[tuple[str, str]] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
if "uninstall" in cmd:
|
|
return subprocess.CompletedProcess(cmd, 1)
|
|
return self._import_check()
|
|
|
|
_shared_setup_1(fake_run, step_messages)
|
|
|
|
warnings = [value for _, value in step_messages]
|
|
assert any("could not be removed" in value for value in warnings), warnings
|
|
assert not any("removed it" in value for value in warnings), warnings
|
|
|
|
def test_working_wheel_reports_no_warning(self):
|
|
"""The happy path stays silent: install exits 0 and the module imports."""
|
|
step_messages: list[tuple[str, str]] = []
|
|
import_calls: list[int] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
import_calls.append(1)
|
|
return self._import_check(1 if len(import_calls) == 1 else 0)
|
|
|
|
_shared_setup_1(fake_run, step_messages)
|
|
|
|
assert step_messages == []
|
|
assert len(import_calls) == 2, "expected a verification import after the install"
|
|
|
|
def test_wheel_failure_warns_and_continues(self):
|
|
step_messages: list[tuple[str, str]] = []
|
|
printed_failures: list[str] = []
|
|
|
|
def fake_step(
|
|
label: str,
|
|
value: str,
|
|
color_fn = None,
|
|
):
|
|
step_messages.append((label, value))
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [
|
|
("uv", subprocess.CompletedProcess(["uv"], 1, "uv wheel failed")),
|
|
(
|
|
"pip",
|
|
subprocess.CompletedProcess(["pip"], 1, "pip wheel failed"),
|
|
),
|
|
],
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"_print_optional_install_failure",
|
|
side_effect = lambda label, result: printed_failures.append(label),
|
|
),
|
|
mock.patch.object(ips, "_step", side_effect = fake_step),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
assert printed_failures == [
|
|
"Installing flash-attn prebuilt wheel with uv",
|
|
"Installing flash-attn prebuilt wheel with pip",
|
|
]
|
|
assert ("warning", "Continuing without flash-attn") in step_messages
|
|
|
|
def test_wheel_missing_skips_install_at_setup_time(self):
|
|
step_messages: list[tuple[str, str]] = []
|
|
|
|
def fake_step(
|
|
label: str,
|
|
value: str,
|
|
color_fn = None,
|
|
):
|
|
step_messages.append((label, value))
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = False),
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch.object(ips, "_step", side_effect = fake_step),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_install_wheel.assert_not_called()
|
|
assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
|
|
|
|
def test_an_unchecked_wheel_is_not_reported_as_unpublished(self):
|
|
step_messages = []
|
|
|
|
def fake_step(kind, message, *args, **kwargs):
|
|
step_messages.append((kind, message))
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
),
|
|
mock.patch.object(ips, "url_exists", return_value = None),
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch.object(ips, "_step", side_effect = fake_step),
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_install_wheel.assert_not_called()
|
|
assert ("warning", "No published flash-attn prebuilt wheel found") not in step_messages
|
|
assert any("could not" in message.lower() for _kind, message in step_messages)
|
|
|
|
def _run_with_resident_copy(
|
|
self,
|
|
*,
|
|
resident: str | None = "2.8.1",
|
|
env: dict[str, str] | None = None,
|
|
torch_imports: bool = True,
|
|
url_available: bool | None = False,
|
|
install_returncode: int = 0,
|
|
fake_run = None,
|
|
) -> tuple[list[tuple[str, str]], list[list[str]], mock.Mock]:
|
|
"""#13244 leaves cu13torch2.10 under torch 2.9+cu126 on cp313, unsupported by v2.8.3."""
|
|
step_messages: list[tuple[str, str]] = []
|
|
removals: list[list[str]] = []
|
|
|
|
def default_run(cmd, **kwargs):
|
|
if "uninstall" in cmd:
|
|
removals.append(list(cmd))
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
return self._import_check()
|
|
|
|
if env is None:
|
|
env = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.9",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
}
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips, "probe_torch_wheel_env", return_value = env if torch_imports else None
|
|
),
|
|
mock.patch.object(ips, "_installed_distribution_version", return_value = resident),
|
|
mock.patch.object(ips, "url_exists", return_value = url_available),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [("uv", subprocess.CompletedProcess(["uv"], install_returncode, ""))],
|
|
) as mock_install_wheel,
|
|
mock.patch.object(
|
|
ips,
|
|
"_step",
|
|
side_effect = lambda label, value, color_fn = None: step_messages.append(
|
|
(label, value)
|
|
),
|
|
),
|
|
mock.patch("subprocess.run", side_effect = fake_run or default_run),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
return step_messages, removals, mock_install_wheel
|
|
|
|
def test_stale_copy_is_removed_when_no_wheel_is_published(self):
|
|
"""if kept, xformers imports it and every diffusers pipeline fails (#13244)."""
|
|
step_messages, removals, mock_install_wheel = self._run_with_resident_copy()
|
|
|
|
mock_install_wheel.assert_not_called()
|
|
assert any("flash-attn" in cmd for cmd in removals), removals
|
|
assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
|
|
assert (
|
|
"warning",
|
|
"flash-attn was built for another torch or CUDA and does not import; removed it",
|
|
) in step_messages
|
|
|
|
def test_stale_copy_is_removed_when_no_wheel_matches(self):
|
|
step_messages, removals, _ = self._run_with_resident_copy(
|
|
env = {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.3",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
|
|
assert any("flash-attn" in cmd for cmd in removals), removals
|
|
assert ("warning", "No compatible flash-attn prebuilt wheel found") in step_messages
|
|
|
|
def test_stale_copy_is_removed_when_the_replacement_fails_to_install(self):
|
|
step_messages, removals, _ = self._run_with_resident_copy(
|
|
url_available = True, install_returncode = 1
|
|
)
|
|
|
|
assert any("flash-attn" in cmd for cmd in removals), removals
|
|
assert ("warning", "Continuing without flash-attn") in step_messages
|
|
|
|
def test_stale_copy_is_force_replaced(self):
|
|
"""a same-version build for another CUDA counts as installed unless forced."""
|
|
imports: list[int] = []
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
imports.append(1)
|
|
return self._import_check(1 if len(imports) == 1 else 0)
|
|
|
|
step_messages, _, mock_install_wheel = self._run_with_resident_copy(
|
|
url_available = True, fake_run = fake_run
|
|
)
|
|
|
|
assert mock_install_wheel.call_args.kwargs["reinstall"] is True
|
|
assert step_messages == []
|
|
|
|
def test_a_failed_stale_removal_is_reported(self):
|
|
def fake_run(cmd, **kwargs):
|
|
if "uninstall" in cmd:
|
|
return subprocess.CompletedProcess(cmd, 1)
|
|
return self._import_check()
|
|
|
|
step_messages, _, _ = self._run_with_resident_copy(fake_run = fake_run)
|
|
|
|
warnings = [value for _, value in step_messages]
|
|
assert any("could not be removed" in value for value in warnings), warnings
|
|
assert not any("removed it" in value for value in warnings), warnings
|
|
|
|
def test_resident_copy_is_kept_when_torch_does_not_import(self):
|
|
"""the import probe loads torch first, so a broken torch says nothing about flash-attn."""
|
|
_, removals, mock_install_wheel = self._run_with_resident_copy(torch_imports = False)
|
|
|
|
assert removals == []
|
|
mock_install_wheel.assert_not_called()
|
|
|
|
def test_nothing_is_uninstalled_when_flash_attn_is_absent(self):
|
|
step_messages, removals, _ = self._run_with_resident_copy(resident = None)
|
|
|
|
assert removals == []
|
|
assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
|
|
|
|
def test_skip_env_disables_setup_install(self):
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", False),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.dict(os.environ, {"UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL": "1"}),
|
|
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch("subprocess.run", return_value = self._import_check()),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_probe.assert_not_called()
|
|
mock_install_wheel.assert_not_called()
|
|
|
|
def test_windows_probes_with_windows_opt_in_and_installs_the_wheel(self):
|
|
url = "https://example.invalid/flash_attn-2.8.4+cu13torch2.13cxx11abiTRUE-cp313-cp313-win_amd64.whl"
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", True),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(
|
|
ips, "probe_torch_wheel_env", return_value = {"platform_tag": "win_amd64"}
|
|
) as mock_probe,
|
|
mock.patch.object(ips, "_build_flash_attn_wheel_url", return_value = url),
|
|
mock.patch.object(ips, "url_exists", return_value = True),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_wheel",
|
|
return_value = [("uv", subprocess.CompletedProcess(["uv"], 0, ""))],
|
|
) as mock_install_wheel,
|
|
mock.patch.object(ips, "_flash_attn_importable", side_effect = [False, True]),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_probe.assert_called_once_with(
|
|
timeout = ips._FLASH_ATTN_IMPORT_PROBE_TIMEOUT, include_windows = True
|
|
)
|
|
assert mock_install_wheel.call_args.args[0] == url
|
|
|
|
def test_windows_without_a_wheel_notes_it_quietly(self):
|
|
notes: list[str] = []
|
|
steps: list[tuple[str, str]] = []
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", True),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(
|
|
ips, "probe_torch_wheel_env", return_value = {"platform_tag": "win_amd64"}
|
|
),
|
|
mock.patch.object(ips, "_build_flash_attn_wheel_url", return_value = None),
|
|
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
|
|
mock.patch.object(ips, "_flash_attn_importable", return_value = False),
|
|
mock.patch.object(ips, "_note", side_effect = lambda message, *a: notes.append(message)),
|
|
mock.patch.object(
|
|
ips, "_step", side_effect = lambda kind, message, *a: steps.append((kind, message))
|
|
),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
|
|
mock_install_wheel.assert_not_called()
|
|
assert steps == []
|
|
assert any("flash-attn" in note for note in notes)
|
|
|
|
|
|
class TestInstallPythonStackFlashAttnIntegration:
|
|
def _run_install(
|
|
self,
|
|
*,
|
|
no_torch: bool,
|
|
is_macos: bool,
|
|
is_windows: bool,
|
|
flavor_moves_torch: bool = False,
|
|
) -> int:
|
|
flash_attn_calls = 0
|
|
torch_version = ["2.13.0+cpu"]
|
|
|
|
def fix_flavor(tag):
|
|
if flavor_moves_torch:
|
|
torch_version[0] = "2.13.0+cu130"
|
|
return True
|
|
|
|
def fake_run(cmd, **kw):
|
|
return subprocess.CompletedProcess(cmd, 0, b"", b"")
|
|
|
|
def count_flash_attn():
|
|
nonlocal flash_attn_calls
|
|
flash_attn_calls += 1
|
|
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", no_torch),
|
|
mock.patch.object(ips, "IS_MACOS", is_macos),
|
|
mock.patch.object(ips, "IS_WINDOWS", is_windows),
|
|
mock.patch.object(ips, "USE_UV", True),
|
|
mock.patch.object(ips, "UV_NEEDS_SYSTEM", False),
|
|
mock.patch.object(ips, "VERBOSE", False),
|
|
mock.patch.object(ips, "_bootstrap_uv", return_value = True),
|
|
mock.patch.object(ips, "_ensure_flash_attn", side_effect = count_flash_attn),
|
|
mock.patch("subprocess.run", side_effect = fake_run),
|
|
mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False),
|
|
mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
|
|
mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
|
|
mock.patch("pathlib.Path.is_dir", return_value = True),
|
|
mock.patch("pathlib.Path.is_file", return_value = True),
|
|
mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False),
|
|
mock.patch.object(ips, "_ensure_expected_torch_flavor", side_effect = fix_flavor),
|
|
mock.patch.object(
|
|
ips, "_probe_installed_torch_version", side_effect = lambda: torch_version[0]
|
|
),
|
|
):
|
|
ips.install_python_stack()
|
|
|
|
return flash_attn_calls
|
|
|
|
def test_linux_torch_install_calls_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = False, is_windows = False) == 1
|
|
|
|
def test_no_torch_install_skips_flash_attn_step(self):
|
|
assert self._run_install(no_torch = True, is_macos = False, is_windows = False) == 0
|
|
|
|
def test_macos_install_skips_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = True, is_windows = False) == 0
|
|
|
|
def test_windows_install_calls_flash_attn_step(self):
|
|
assert self._run_install(no_torch = False, is_macos = False, is_windows = True) == 1
|
|
|
|
def test_flash_attn_is_rechecked_after_the_linux_torch_repair(self):
|
|
"""step 13 can move torch from 2.12+cu130 to 2.9+cu126 under cu13torch2.10 (#13244)."""
|
|
source = (STUDIO_DIR / "install_python_stack.py").read_text(encoding = "utf-8")
|
|
step = source.split('_progress(_torch_step_label("final"))', 1)[1]
|
|
step = step.split("# 13w.", 1)[0]
|
|
moved = step.split(
|
|
"if _torch_after_repair and _torch_after_repair != _torch_before_repair:", 1
|
|
)[1]
|
|
moved = moved.split("_evict_xformers_built_for_another_torch(", 1)[0]
|
|
assert "_ensure_flash_attn()" in moved
|
|
|
|
def test_flash_attn_is_rechecked_when_the_windows_flavor_invariant_moves_torch(self):
|
|
assert (
|
|
self._run_install(
|
|
no_torch = False, is_macos = False, is_windows = True, flavor_moves_torch = True
|
|
)
|
|
== 2
|
|
)
|
|
|
|
|
|
class TestWarningColor:
|
|
def test_optional_install_warning_uses_shell_installer_warn_color(self, capsys):
|
|
import re
|
|
|
|
install_sh = (STUDIO_DIR.parent / "install.sh").read_text(encoding = "utf-8")
|
|
code = re.search(r'C_WARN="\$\{_ESC\}\[([0-9;]+)m"', install_sh).group(1)
|
|
result = subprocess.CompletedProcess(["uv"], 1, "", "")
|
|
with mock.patch.object(ips, "_HAS_COLOR", True):
|
|
ips._print_optional_install_failure("flash-attn install", result)
|
|
out = capsys.readouterr().out
|
|
assert f"\033[{code}mflash-attn install failed (exit code 1)\033[0m" in out
|
|
assert "\033[96m" not in out
|
|
|
|
|
|
class TestFlashAttnCapabilityFloor:
|
|
"""Unsloth only enables FlashAttention 2 from sm80: below it the wheel is never downloaded."""
|
|
|
|
def _run(self, monkeypatch, capability):
|
|
from utils import kernel_install
|
|
|
|
monkeypatch.setattr(
|
|
ips,
|
|
"_flash_attn_gpu_below_floor",
|
|
lambda: kernel_install.gpu_below_flash_attn_floor(lambda *a, **k: None),
|
|
)
|
|
monkeypatch.setattr(kernel_install, "_gpu_capability", lambda run: capability)
|
|
calls = {"head": 0, "install": 0}
|
|
notes = []
|
|
with (
|
|
mock.patch.object(ips, "NO_TORCH", False),
|
|
mock.patch.object(ips, "IS_WINDOWS", True),
|
|
mock.patch.object(ips, "IS_MACOS", False),
|
|
mock.patch.object(ips, "_flash_attn_importable", return_value = False),
|
|
mock.patch.object(ips, "_installed_distribution_version", return_value = None),
|
|
mock.patch.object(
|
|
ips,
|
|
"probe_torch_wheel_env",
|
|
return_value = {
|
|
"python_tag": "cp313",
|
|
"abi_tag": "cp313",
|
|
"torch_mm": "2.13",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "win_amd64",
|
|
"hip_version": "",
|
|
},
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"url_exists",
|
|
side_effect = lambda url: calls.__setitem__("head", calls["head"] + 1) or True,
|
|
),
|
|
mock.patch.object(
|
|
ips,
|
|
"install_prebuilt",
|
|
side_effect = lambda *a, **k: calls.__setitem__("install", calls["install"] + 1)
|
|
or "installed",
|
|
),
|
|
mock.patch.object(ips, "_note", side_effect = notes.append),
|
|
mock.patch.object(ips, "_count_install_action"),
|
|
):
|
|
ips._ensure_flash_attn()
|
|
return calls, notes
|
|
|
|
def test_a_t4_does_not_download_flash_attn(self, monkeypatch):
|
|
calls, notes = self._run(monkeypatch, (7, 5))
|
|
assert calls == {"head": 0, "install": 0}
|
|
assert any("sm80" in note for note in notes)
|
|
|
|
@pytest.mark.parametrize("capability", [(8, 0), (12, 0), None])
|
|
def test_sm80_and_newer_or_no_visible_gpu_still_install(self, monkeypatch, capability):
|
|
calls, _ = self._run(monkeypatch, capability)
|
|
assert calls == {"head": 1, "install": 1}
|