mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
* 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>
This commit is contained in:
Thomas Eric 🇧🇷
GitHub
parent
62a6eb2a3d
commit
03cbe211a3
6 changed files with 123 additions and 217 deletions
|
|
@ -37,7 +37,6 @@ def _load_worker_module():
|
|||
for name in (
|
||||
"direct_wheel_url",
|
||||
"flash_attn_wheel_url",
|
||||
"has_blackwell_gpu",
|
||||
"install_wheel",
|
||||
"probe_torch_wheel_env",
|
||||
"url_exists",
|
||||
|
|
|
|||
|
|
@ -32,16 +32,23 @@ def _load_module(monkeypatch):
|
|||
@pytest.mark.parametrize(
|
||||
"torch_version, expected",
|
||||
[
|
||||
# torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0,
|
||||
# independent of the local +cuXXX/+rocm/+cpu suffix or patch level.
|
||||
("2.10.0+cu130", "torchao==0.16.0"),
|
||||
# torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and
|
||||
# loads against the CUDA-12 PyPI wheel). Independent of patch level.
|
||||
("2.10.0+cu128", "torchao==0.16.0"),
|
||||
("2.10.0+cu126", "torchao==0.16.0"),
|
||||
("2.10.0+rocm6.4", "torchao==0.16.0"),
|
||||
("2.10.0+cpu", "torchao==0.16.0"),
|
||||
("2.10.1", "torchao==0.16.0"),
|
||||
("2.10.0", "torchao==0.16.0"),
|
||||
# Pre-release / dev / rc builds: the minor is cleaned of non-digits.
|
||||
# torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't
|
||||
# load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0.
|
||||
("2.10.0+cu130", "torchao==0.17.0"),
|
||||
("2.10.0+cu140", "torchao==0.17.0"),
|
||||
# Pre-release / dev / rc builds: the minor is cleaned of non-digits; the
|
||||
# CUDA tag still decides 0.16.0 vs 0.17.0.
|
||||
("2.10.0rc1", "torchao==0.16.0"),
|
||||
("2.10.0.dev20250804+cu130", "torchao==0.16.0"),
|
||||
("2.10.0.dev20250804+cu130", "torchao==0.17.0"),
|
||||
("2.10.0.dev20250804+cu128", "torchao==0.16.0"),
|
||||
("2.10rc1", "torchao==0.16.0"),
|
||||
# torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0.
|
||||
("2.11.0+cu130", "torchao==0.17.0"),
|
||||
|
|
|
|||
|
|
@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch):
|
|||
statuses: list[str] = []
|
||||
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
|
||||
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
|
|
@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
|
|||
statuses: list[str] = []
|
||||
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False)
|
||||
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
|
|
@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch):
|
|||
worker._sp.run.assert_not_called()
|
||||
|
||||
|
||||
def test_runtime_flash_attn_skips_on_blackwell(monkeypatch):
|
||||
statuses: list[str] = []
|
||||
install_mock = mock.Mock()
|
||||
|
||||
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True)
|
||||
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"_send_status",
|
||||
lambda queue, message: statuses.append(message),
|
||||
)
|
||||
|
||||
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536)
|
||||
|
||||
install_mock.assert_not_called()
|
||||
assert len(statuses) == 1
|
||||
assert "Blackwell" in statuses[0]
|
||||
|
||||
|
||||
def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch):
|
||||
install_mock = mock.Mock(return_value = True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
|
|
|
|||
|
|
@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele
|
|||
def has_blackwell_gpu() -> bool:
|
||||
"""Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell).
|
||||
|
||||
Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels
|
||||
fail to load, so callers use this to skip the flash-attn install path. Cached
|
||||
for the process lifetime; tests mocking nvidia-smi must call
|
||||
Cached for the process lifetime; tests mocking nvidia-smi must call
|
||||
``has_blackwell_gpu.cache_clear()`` first.
|
||||
"""
|
||||
# Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn
|
||||
# wheels and url_exists() already gates resolution, so we no longer skip
|
||||
# flash-attn on Blackwell. The nvidia-smi probe below is kept for possible
|
||||
# future arch-based gating; drop this early return to re-enable it.
|
||||
return False
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe:
|
||||
return False
|
||||
|
|
@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
|
|||
return env
|
||||
|
||||
|
||||
# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba
|
||||
# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites
|
||||
# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective
|
||||
# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11
|
||||
# install still gets these prebuilt accelerators instead of building from source.
|
||||
_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"}
|
||||
|
||||
|
||||
def prebuilt_wheel_torch_mm(torch_mm: str) -> str:
|
||||
"""Map a torch major.minor to the one whose prebuilt accelerator wheels to use."""
|
||||
return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm)
|
||||
|
||||
|
||||
def direct_wheel_url(
|
||||
*,
|
||||
filename_prefix: str,
|
||||
|
|
@ -130,7 +146,7 @@ def direct_wheel_url(
|
|||
|
||||
filename = (
|
||||
f"{filename_prefix}-{package_version}"
|
||||
f"+cu{env['cuda_major']}torch{env['torch_mm']}"
|
||||
f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}"
|
||||
f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}"
|
||||
f"-{env['platform_tag']}.whl"
|
||||
)
|
||||
|
|
@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None:
|
|||
def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None:
|
||||
if env is None:
|
||||
return None
|
||||
package_version = flash_attn_package_version(env["torch_mm"])
|
||||
package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"]))
|
||||
if package_version is None:
|
||||
return None
|
||||
return direct_wheel_url(
|
||||
|
|
|
|||
|
|
@ -103,33 +103,47 @@ _PYTORCH_WHL_BASE = (
|
|||
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
|
||||
).rstrip("/")
|
||||
|
||||
# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are
|
||||
# pinned to the torch<2.11 family rather than left bare: the install uses an
|
||||
# exclusive --index-url (no PyPI fallback), so a bare name could resolve a
|
||||
# torchvision built against a different torch major (e.g. 0.27 for torch 2.12)
|
||||
# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm
|
||||
# spec above, which targets the same torch family.
|
||||
# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its
|
||||
# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn
|
||||
# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on
|
||||
# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned
|
||||
# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so
|
||||
# a bare name could resolve one built against a different torch major (e.g. 0.27 for
|
||||
# torch 2.12) and fail at runtime with an ABI mismatch.
|
||||
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
|
||||
"torch>=2.4,<2.11.0",
|
||||
"torchvision>=0.19,<0.26.0",
|
||||
"torchaudio>=2.4,<2.11.0",
|
||||
"torch>=2.4,<2.12.0",
|
||||
"torchvision>=0.19,<0.27.0",
|
||||
"torchaudio>=2.4,<2.12.0",
|
||||
)
|
||||
|
||||
# torchao's C++ extensions are built against ONE exact torch release; a newer
|
||||
# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions
|
||||
# due to incompatible torch version ...") and fall back to slow Python. Because
|
||||
# the torch pin above is a range (and every CUDA index now tops out at torch
|
||||
# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So
|
||||
# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919.
|
||||
# torch 2.9.x -> torchao 0.14.0 (today's pin; built for torch 2.9.0)
|
||||
# torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0)
|
||||
# torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2)
|
||||
# Unknown/older torch keeps the conservative default (no regression vs today).
|
||||
# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch
|
||||
# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails
|
||||
# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a
|
||||
# range, so match torchao to the installed torch (table: pytorch/ao#2919):
|
||||
# 2.9.x -> 0.14.0
|
||||
# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel)
|
||||
# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0
|
||||
# targets torch 2.11 so its cpp is cleanly skipped, not crashed)
|
||||
# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2)
|
||||
# Unknown/older torch keeps the conservative default.
|
||||
_TORCHAO_DEFAULT_SPEC = "torchao==0.14.0"
|
||||
_TORCHAO_BY_TORCH_MINOR: dict[int, str] = {
|
||||
10: "torchao==0.16.0",
|
||||
11: "torchao==0.17.0",
|
||||
}
|
||||
_TORCHAO_TORCH_210_SPEC = "torchao==0.16.0"
|
||||
_TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0"
|
||||
_TORCHAO_TORCH_211_PLUS_SPEC = "torchao==0.17.0"
|
||||
# torch 2.10 built against CUDA >= this major can't load 0.16.0's CUDA-12 cpp.
|
||||
_TORCHAO_CUDA13_MIN_MAJOR = 13
|
||||
|
||||
|
||||
def _cuda_major_from_torch_version(torch_version: str) -> int | None:
|
||||
"""Extract the CUDA major from a torch local version tag, e.g. '2.10.0+cu130'
|
||||
-> 13, '2.10.0+cu128' -> 12. Returns None for rocm/cpu/tagless builds."""
|
||||
local = str(torch_version).split("+", 1)
|
||||
if len(local) < 2 or not local[1].startswith("cu"):
|
||||
return None
|
||||
digits = re.sub(r"[^0-9].*", "", local[1][2:]) # 'cu130' -> '130'
|
||||
if not digits:
|
||||
return None
|
||||
return int(digits) // 10 # '130' -> 13, '128' -> 12, '118' -> 11
|
||||
|
||||
|
||||
def _select_torchao_spec(torch_version: str | None) -> str:
|
||||
|
|
@ -151,8 +165,14 @@ def _select_torchao_spec(torch_version: str | None) -> str:
|
|||
if major != 2:
|
||||
return _TORCHAO_DEFAULT_SPEC
|
||||
if minor >= 11:
|
||||
return _TORCHAO_BY_TORCH_MINOR[11] # newest known build; covers 2.11+
|
||||
return _TORCHAO_BY_TORCH_MINOR.get(minor, _TORCHAO_DEFAULT_SPEC)
|
||||
return _TORCHAO_TORCH_211_PLUS_SPEC # newest known build; covers 2.11+
|
||||
if minor == 10:
|
||||
# cu130+ can't load 0.16.0's CUDA-12 cpp; use 0.17.0 (cpp skipped, not crashed).
|
||||
cuda_major = _cuda_major_from_torch_version(str(torch_version))
|
||||
if cuda_major is not None and cuda_major >= _TORCHAO_CUDA13_MIN_MAJOR:
|
||||
return _TORCHAO_TORCH_210_CUDA13_SPEC
|
||||
return _TORCHAO_TORCH_210_SPEC
|
||||
return _TORCHAO_DEFAULT_SPEC
|
||||
|
||||
|
||||
def _probe_installed_torch_version() -> str | None:
|
||||
|
|
|
|||
|
|
@ -13,102 +13,35 @@ sys.path.insert(0, str(STUDIO_DIR))
|
|||
sys.path.insert(0, str(STUDIO_DIR / "backend"))
|
||||
|
||||
import install_python_stack as ips
|
||||
from backend.utils import wheel_utils
|
||||
from utils import wheel_utils
|
||||
|
||||
|
||||
def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess:
|
||||
return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "")
|
||||
class TestPrebuiltWheelTorchMapping:
|
||||
def test_torch_211_maps_to_torch210(self):
|
||||
assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10"
|
||||
|
||||
def test_other_versions_pass_through(self):
|
||||
for torch_mm in ("2.9", "2.10", "2.12"):
|
||||
assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm
|
||||
|
||||
class TestHasBlackwellGpu:
|
||||
def setup_method(self):
|
||||
wheel_utils.has_blackwell_gpu.cache_clear()
|
||||
|
||||
def teardown_method(self):
|
||||
wheel_utils.has_blackwell_gpu.cache_clear()
|
||||
|
||||
def test_returns_false_when_nvidia_smi_missing(self):
|
||||
with mock.patch.object(wheel_utils.shutil, "which", return_value = None):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
|
||||
def test_returns_true_for_sm_100(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is True
|
||||
|
||||
def test_returns_true_for_sm_120(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is True
|
||||
|
||||
def test_returns_true_for_sm_121(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is True
|
||||
|
||||
def test_returns_false_for_sm_90(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
|
||||
def test_returns_false_for_sm_89(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
|
||||
def test_mixed_gpus_with_one_blackwell_returns_true(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(
|
||||
wheel_utils.subprocess,
|
||||
"run",
|
||||
return_value = _smi_result("8.0\n10.0\n"),
|
||||
),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is True
|
||||
|
||||
def test_returns_false_when_nvidia_smi_fails(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(
|
||||
wheel_utils.subprocess,
|
||||
"run",
|
||||
return_value = _smi_result("", returncode = 1),
|
||||
),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
|
||||
def test_returns_false_on_subprocess_timeout(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(
|
||||
wheel_utils.subprocess,
|
||||
"run",
|
||||
side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10),
|
||||
),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
|
||||
def test_returns_false_on_malformed_output(self):
|
||||
with (
|
||||
mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
|
||||
mock.patch.object(
|
||||
wheel_utils.subprocess,
|
||||
"run",
|
||||
return_value = _smi_result("not-a-number\n\n"),
|
||||
),
|
||||
):
|
||||
assert wheel_utils.has_blackwell_gpu() is False
|
||||
def test_direct_wheel_url_reuses_torch210_on_211(self):
|
||||
# causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the
|
||||
# torch2.10 wheel filename just like flash-attn does.
|
||||
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
|
||||
|
||||
|
||||
class TestFlashAttnWheelSelection:
|
||||
|
|
@ -118,9 +51,24 @@ class TestFlashAttnWheelSelection:
|
|||
def test_torch_29_maps_to_v283(self):
|
||||
assert ips._select_flash_attn_version("2.9") == "2.8.3"
|
||||
|
||||
def test_unsupported_torch_has_no_wheel_mapping(self):
|
||||
def test_torch_211_has_no_native_version_entry(self):
|
||||
# The raw version table has no torch2.11-tagged wheel; the URL builder
|
||||
# reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel).
|
||||
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_exact_wheel_url_uses_full_env_tuple(self):
|
||||
url = ips._build_flash_attn_wheel_url(
|
||||
{
|
||||
|
|
@ -333,83 +281,22 @@ class TestEnsureFlashAttn:
|
|||
mock_probe.assert_not_called()
|
||||
mock_install_wheel.assert_not_called()
|
||||
|
||||
def test_blackwell_gpu_skips_install_with_warning(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, "has_blackwell_gpu", return_value = True),
|
||||
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
||||
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_probe.assert_not_called()
|
||||
mock_install_wheel.assert_not_called()
|
||||
assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
|
||||
|
||||
def test_blackwell_gpu_on_windows_emits_blackwell_warning(self):
|
||||
step_messages: list[tuple[str, str]] = []
|
||||
|
||||
def fake_step(
|
||||
label: str,
|
||||
value: str,
|
||||
color_fn = None,
|
||||
):
|
||||
step_messages.append((label, value))
|
||||
|
||||
def test_windows_skips_install_without_probing(self):
|
||||
# flash-attn is Linux-only: on Windows the installer returns before
|
||||
# probing the torch env or resolving a wheel (no Windows wheels are
|
||||
# published upstream).
|
||||
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, "has_blackwell_gpu", return_value = True),
|
||||
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
||||
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_probe.assert_not_called()
|
||||
mock_install_wheel.assert_not_called()
|
||||
assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
|
||||
|
||||
def test_non_blackwell_windows_does_not_emit_blackwell_warning(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", True),
|
||||
mock.patch.object(ips, "IS_MACOS", False),
|
||||
mock.patch.object(ips, "has_blackwell_gpu", return_value = False),
|
||||
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
|
||||
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_probe.assert_not_called()
|
||||
mock_install_wheel.assert_not_called()
|
||||
assert not any("Blackwell" in msg for _, msg in step_messages)
|
||||
|
||||
|
||||
class TestInstallPythonStackFlashAttnIntegration:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue