mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-22 06:20:56 +00:00
16e32f3820
Stacked on #14336. Merge only after the Linux torch 2.10 to 2.12 causal-conv1d / mamba-ssm cells from #14336 are published on `prebuilt-wheels-cu13`; before that these URLs would 404 in front of the upstream fallback. ## Why Linux users on torch 2.10, 2.11 or 2.12 with a CUDA 13 torch download upstream's causal-conv1d (205 MB) and mamba-ssm (351 MB) wheels, built for torch 2.10 with `-lineinfo` and every arch. Our builds of the same packages for torch 2.13 / 2.14 are about 29 / 54 MB on x86_64 and 42 / 79 MB on aarch64. ## Changes - `_UNSLOTH_PREBUILT_EXTRA_CELLS` gains linux_x86_64 and linux_aarch64 torch 2.10 / 2.11 / 2.12 on cp311 to cp313. - `_UNSLOTH_PREBUILT_EXTRA_CELL_PREFIXES` limits those Linux cells to causal_conv1d and mamba_ssm, so flash-attn keeps upstream 2.8.1 there (125 MB, smaller than ours). - Unchanged: a CUDA 12 torch, cp310 and cp314 still resolve upstream, and torch 2.13 / 2.14 and Windows resolve as before. ## Tests Updated the tests that pinned upstream URLs for cu13 Linux torch 2.10 to 2.12. They now assert our wheel for cu13 cp311 to cp313 and upstream for cu12 / cp310. Added tests that flash-attn stays upstream and that every Linux extra cell the resolver asks for is one the workflow builds. Affected suites: 1125 passed.
899 lines
35 KiB
Python
899 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",
|
|
# cu12: cu13 cp311-cp313 resolves to our own build instead.
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "causal_conv1d-1.6.1+cu12torch2.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",
|
|
# cu12: cu13 cp311-cp313 resolves to our own build instead.
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
assert url is not None
|
|
assert "mamba_ssm-2.3.1+cu12torch2.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}
|