diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index c3d3daf07e..8ab2b5b2be 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -16,26 +16,31 @@ from __future__ import annotations import structlog from loggers import get_logger import os -import platform import shutil import sys import time import traceback -import json import subprocess as _sp from pathlib import Path -from typing import Any -import urllib.error -import urllib.request +from typing import Any, Callable logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +from utils.wheel_utils import ( + direct_wheel_url, + flash_attn_wheel_url, + install_wheel, + probe_torch_wheel_env, + url_exists, +) _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" _CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1" _MAMBA_SSM_RELEASE_TAG = "v2.3.1" _MAMBA_SSM_PACKAGE_VERSION = "2.3.1" +_FLASH_ATTN_RUNTIME_MIN_SEQ_LEN = 32768 +_FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL" def _model_wants_causal_conv1d(model_name: str) -> bool: @@ -59,184 +64,60 @@ def _model_wants_causal_conv1d(model_name: str) -> bool: ) -def _causal_conv1d_platform_tag() -> str | None: - machine = platform.machine().lower() - if sys.platform.startswith("linux"): - if machine in {"x86_64", "amd64"}: - return "linux_x86_64" - if machine in {"aarch64", "arm64"}: - return "linux_aarch64" - return None - # No prebuilt wheels published for macOS or Windows - return None - - -def _probe_causal_conv1d_env() -> dict[str, str] | None: - try: - probe = _sp.run( - [ - sys.executable, - "-c", - ( - "import json, sys, re, torch; " - "parts = torch.__version__.split('+', 1)[0].split('.')[:2]; " - "minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; " - "torch_mm = parts[0] + '.' + minor; " - "print(json.dumps({" - "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " - "'torch_mm': torch_mm, " - "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " - "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', " - "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" - "}))" - ), - ], - stdout = _sp.PIPE, - stderr = _sp.PIPE, - text = True, - timeout = 30, - ) - except _sp.TimeoutExpired: - logger.warning("Torch environment probe timed out after 30s") - return None - if probe.returncode != 0: - logger.warning( - "Failed to probe torch environment for causal-conv1d wheel:\n%s", - probe.stdout, - ) - return None - - try: - return json.loads(probe.stdout.strip()) - except json.JSONDecodeError: - logger.warning( - "Failed to parse torch environment probe output: %s", probe.stdout - ) - return None - - -def _direct_wheel_url( - *, - filename_prefix: str, - package_version: str, - release_tag: str, - release_base_url: str, - env: dict[str, str] | None = None, -) -> str | None: - env = env or _probe_causal_conv1d_env() - platform_tag = _causal_conv1d_platform_tag() - if env is None or platform_tag is None or not env.get("cuda_major"): - return None - - filename = ( - f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" - f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}-{platform_tag}.whl" - ) - return f"{release_base_url}/{release_tag}/{filename}" - - -def _url_exists(url: str) -> bool: - try: - request = urllib.request.Request(url, method = "HEAD") - with urllib.request.urlopen(request, timeout = 10): - return True - except urllib.error.HTTPError as exc: - if exc.code == 404: - return False - logger.warning("Unexpected HTTP error while probing %s: %s", url, exc) - return False - except Exception as exc: - logger.warning("Failed to probe %s: %s", url, exc) - return False - - def _install_package_wheel_first( *, event_queue: Any, import_name: str, display_name: str, pypi_name: str, - pypi_version: str, - filename_prefix: str, - release_tag: str, - release_base_url: str, -) -> None: + pypi_version: str | None = None, + filename_prefix: str | None = None, + release_tag: str | None = None, + release_base_url: str | None = None, + wheel_url_builder: Callable[[dict[str, str] | None], str | None] | None = None, + pypi_spec: str | None = None, + pypi_status_message: str | None = None, +) -> bool: try: __import__(import_name) logger.info("%s already installed", display_name) - return + return True except ImportError: pass - env = _probe_causal_conv1d_env() - wheel_url = _direct_wheel_url( - filename_prefix = filename_prefix, - package_version = pypi_version, - release_tag = release_tag, - release_base_url = release_base_url, - env = env, - ) + env = probe_torch_wheel_env(timeout = 30) + if wheel_url_builder is not None: + wheel_url = wheel_url_builder(env) + else: + wheel_url = direct_wheel_url( + filename_prefix = filename_prefix, + package_version = pypi_version, + release_tag = release_tag, + release_base_url = release_base_url, + env = env, + ) if wheel_url is None: logger.info("No compatible %s wheel candidate", display_name) - else: - if _url_exists(wheel_url): - _send_status(event_queue, f"Installing prebuilt {display_name} wheel...") - installed = False - # Try uv first if available, then fall back to pip - if shutil.which("uv"): - uv_cmd = [ - "uv", - "pip", - "install", - "--python", - sys.executable, - "--no-deps", - wheel_url, - ] - result = _sp.run( - uv_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) - if result.returncode == 0: - installed = True - else: - logger.warning( - "uv failed to install %s wheel:\n%s", - display_name, - result.stdout, - ) - if not installed: - pip_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-deps", - wheel_url, - ] - result = _sp.run( - pip_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) - if result.returncode == 0: - installed = True - else: - logger.warning( - "pip failed to install %s wheel:\n%s", - display_name, - result.stdout, - ) - if installed: + elif url_exists(wheel_url): + _send_status(event_queue, f"Installing prebuilt {display_name} wheel...") + for installer, result in install_wheel( + wheel_url, + python_executable = sys.executable, + use_uv = bool(shutil.which("uv")), + run = _sp.run, + ): + if result.returncode == 0: logger.info("Installed prebuilt %s wheel successfully", display_name) - return - else: - logger.info("No published %s wheel found: %s", display_name, wheel_url) + return True + logger.warning( + "%s failed to install %s wheel:\n%s", + installer, + display_name, + result.stdout, + ) + else: + logger.info("No published %s wheel found: %s", display_name, wheel_url) is_hip = env and env.get("hip_version") if is_hip and not shutil.which("hipcc"): @@ -249,43 +130,62 @@ def _install_package_wheel_first( event_queue, f"{display_name}: hipcc not found (ROCm HIP SDK required)", ) - return + return False - if is_hip: - _send_status( - event_queue, - f"Compiling {display_name} from source for ROCm " - "(this may take several minutes)...", - ) - else: - _send_status(event_queue, f"Installing {display_name} from PyPI...") + if pypi_spec is None: + pypi_spec = f"{pypi_name}=={pypi_version}" + + if pypi_status_message is None: + if is_hip: + pypi_status_message = ( + f"Compiling {display_name} from source for ROCm " + "(this may take several minutes)..." + ) + else: + pypi_status_message = f"Installing {display_name} from PyPI..." + + _send_status(event_queue, pypi_status_message) # Prefer uv for faster dependency resolution when available - if shutil.which("uv"): - pypi_cmd = [ - "uv", - "pip", - "install", - "--python", - sys.executable, - "--no-build-isolation", - "--no-deps", - ] - # Avoid stale cache artifacts from partial HIP source builds - if is_hip: - pypi_cmd.append("--no-cache") - pypi_cmd.append(f"{pypi_name}=={pypi_version}") + plain_pypi_install = pypi_version is None + if plain_pypi_install: + if shutil.which("uv"): + pypi_cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + pypi_spec, + ] + else: + pypi_cmd = [sys.executable, "-m", "pip", "install", pypi_spec] else: - pypi_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-build-isolation", - "--no-deps", - "--no-cache-dir", - f"{pypi_name}=={pypi_version}", - ] + if shutil.which("uv"): + pypi_cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "--no-build-isolation", + "--no-deps", + ] + # Avoid stale cache artifacts from partial HIP source builds + if is_hip: + pypi_cmd.append("--no-cache") + pypi_cmd.append(pypi_spec) + else: + pypi_cmd = [ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-deps", + "--no-cache-dir", + pypi_spec, + ] # Source compilation on ROCm can take 10-30 minutes; use a generous # timeout. Non-HIP installs preserve the pre-existing "no timeout" @@ -313,7 +213,7 @@ def _install_package_wheel_first( f"{display_name} installation timed out after " f"{_run_kwargs.get('timeout')}s", ) - return + return False if result.returncode != 0: if is_hip: @@ -337,12 +237,13 @@ def _install_package_wheel_first( display_name, result.stdout, ) - return + return False if is_hip: logger.info("Compiled and installed %s from source for ROCm", display_name) else: logger.info("Installed %s from PyPI", display_name) + return True def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: @@ -389,6 +290,31 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None: ) +def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: + if os.getenv(_FLASH_ATTN_SKIP_ENV) == "1": + return False + if max_seq_length < _FLASH_ATTN_RUNTIME_MIN_SEQ_LEN: + return False + return sys.platform.startswith("linux") + + +def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None: + if not _should_try_runtime_flash_attn_install(max_seq_length): + return + + installed = _install_package_wheel_first( + event_queue = event_queue, + import_name = "flash_attn", + display_name = "flash-attn", + pypi_name = "flash-attn", + wheel_url_builder = flash_attn_wheel_url, + pypi_spec = "flash-attn", + pypi_status_message = "Installing flash-attn from PyPI for long-context training...", + ) + if not installed: + _send_status(event_queue, "Continuing without flash-attn") + + def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports @@ -471,6 +397,10 @@ def run_training_process( try: _ensure_causal_conv1d_fast_path(event_queue, model_name) _ensure_mamba_ssm(event_queue, model_name) + _ensure_flash_attn_for_long_context( + event_queue, + int(config.get("max_seq_length", 2048)), + ) except Exception as exc: event_queue.put( { diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py new file mode 100644 index 0000000000..986958408e --- /dev/null +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import builtins +import subprocess +import sys +from unittest import mock + +from core.training import worker + + +def _missing_flash_attn_import(): + real_import = builtins.__import__ + + def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + if name == "flash_attn": + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + return fake_import + + +def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + assert worker._should_try_runtime_flash_attn_install(32767) is False + assert worker._should_try_runtime_flash_attn_install( + 32768 + ) is sys.platform.startswith("linux") + + monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") + assert worker._should_try_runtime_flash_attn_install(32768) is False + + +def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): + statuses: list[str] = [] + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) + monkeypatch.setattr( + worker, + "flash_attn_wheel_url", + lambda env: "https://example.com/fa.whl", + ) + monkeypatch.setattr(worker, "url_exists", lambda url: True) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + monkeypatch.setattr( + worker, + "install_wheel", + lambda *args, **kwargs: [("pip", subprocess.CompletedProcess(["pip"], 0, ""))], + ) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + assert statuses == ["Installing prebuilt flash-attn wheel..."] + + +def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): + calls: list[list[str]] = [] + statuses: list[str] = [] + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "python_tag": "cp313", + "torch_mm": "2.10", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr( + worker, + "flash_attn_wheel_url", + lambda env: "https://example.com/fa.whl", + ) + monkeypatch.setattr(worker, "url_exists", lambda url: False) + monkeypatch.setattr(worker.shutil, "which", lambda name: None) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + monkeypatch.setattr(worker, "install_wheel", mock.Mock()) + + def fake_run(cmd, stdout = None, stderr = None, text = None): + calls.append(list(cmd)) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + assert statuses == ["Installing flash-attn from PyPI for long-context training..."] + assert calls == [[sys.executable, "-m", "pip", "install", "flash-attn"]] + + +def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): + monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") + monkeypatch.setattr(worker._sp, "run", mock.Mock()) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + worker._sp.run.assert_not_called() + + +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) + + worker._ensure_causal_conv1d_fast_path( + event_queue = [], + model_name = "tiiuae/Falcon-H1-0.5B-Instruct", + ) + + install_mock.assert_called_once_with( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = worker._CAUSAL_CONV1D_PACKAGE_VERSION, + filename_prefix = "causal_conv1d", + release_tag = worker._CAUSAL_CONV1D_RELEASE_TAG, + release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download", + ) + + +def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): + install_mock = mock.Mock(return_value = True) + monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) + + worker._ensure_mamba_ssm( + event_queue = [], + model_name = "tiiuae/Falcon-H1-0.5B-Instruct", + ) + + install_mock.assert_called_once_with( + event_queue = [], + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + pypi_version = worker._MAMBA_SSM_PACKAGE_VERSION, + filename_prefix = "mamba_ssm", + release_tag = worker._MAMBA_SSM_RELEASE_TAG, + release_base_url = "https://github.com/state-spaces/mamba/releases/download", + ) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py new file mode 100644 index 0000000000..00240f1e69 --- /dev/null +++ b/studio/backend/utils/wheel_utils.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import logging +import platform +import shutil +import subprocess +import sys +import urllib.error +import urllib.request +from typing import Callable + +_logger = logging.getLogger(__name__) + +FLASH_ATTN_RELEASE_BASE_URL = ( + "https://github.com/Dao-AILab/flash-attention/releases/download" +) + + +def linux_wheel_platform_tag() -> str | None: + machine = platform.machine().lower() + if sys.platform.startswith("linux"): + if machine in {"x86_64", "amd64"}: + return "linux_x86_64" + if machine in {"aarch64", "arm64"}: + return "linux_aarch64" + # No prebuilt wheels published for macOS or Windows + return None + + +def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | None: + platform_tag = linux_wheel_platform_tag() + if platform_tag is None: + return None + + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import json, sys, re, torch; " + "parts = torch.__version__.split('+', 1)[0].split('.')[:2]; " + "minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; " + "torch_mm = parts[0] + '.' + minor; " + "print(json.dumps({" + "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " + "'torch_mm': torch_mm, " + "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " + "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', " + "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" + "}))" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + text = True, + timeout = timeout, + ) + except subprocess.TimeoutExpired: + return None + + if probe.returncode != 0: + return None + + try: + env = json.loads(probe.stdout.strip()) + except json.JSONDecodeError: + return None + env["platform_tag"] = platform_tag + return env + + +def direct_wheel_url( + *, + filename_prefix: str, + package_version: str, + release_tag: str, + release_base_url: str, + env: dict[str, str] | None, +) -> str | None: + if env is None or not env.get("cuda_major"): + return None + + filename = ( + f"{filename_prefix}-{package_version}" + f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" + f"-{env['platform_tag']}.whl" + ) + return f"{release_base_url}/{release_tag}/{filename}" + + +def flash_attn_package_version(torch_mm: str) -> str | None: + if torch_mm == "2.10": + return "2.8.1" + try: + major, minor = (int(part) for part in torch_mm.split(".", 1)) + except ValueError: + return None + if major == 2 and 4 <= minor <= 9: + return "2.8.3" + return 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"]) + if package_version is None: + return None + return direct_wheel_url( + filename_prefix = "flash_attn", + package_version = package_version, + release_tag = f"v{package_version}", + release_base_url = FLASH_ATTN_RELEASE_BASE_URL, + env = env, + ) + + +def install_wheel( + wheel_url: str, + *, + python_executable: str, + use_uv: bool, + uv_needs_system: bool = False, + run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run, +) -> list[tuple[str, subprocess.CompletedProcess[str]]]: + attempts: list[tuple[str, subprocess.CompletedProcess[str]]] = [] + + # Try uv first if available, then fall back to pip + if use_uv and shutil.which("uv"): + uv_cmd = ["uv", "pip", "install"] + if uv_needs_system: + uv_cmd.append("--system") + uv_cmd.extend(["--python", python_executable, "--no-deps", wheel_url]) + result = run( + uv_cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + ) + attempts.append(("uv", result)) + if result.returncode == 0: + return attempts + + pip_cmd = [python_executable, "-m", "pip", "install", "--no-deps", wheel_url] + result = run( + pip_cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + ) + attempts.append(("pip", result)) + return attempts + + +def url_exists(url: str) -> bool: + try: + request = urllib.request.Request(url, method = "HEAD") + with urllib.request.urlopen(request, timeout = 10): + return True + except urllib.error.HTTPError as exc: + _logger.debug("url_exists(%s): HTTP %s", url, exc.code) + except (urllib.error.URLError, TimeoutError) as exc: + _logger.debug("url_exists(%s): %s", url, exc) + return False diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a046f8f892..71a2a5bf83 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -21,6 +21,14 @@ import tempfile import urllib.request from pathlib import Path +from backend.utils.wheel_utils import ( + flash_attn_package_version, + flash_attn_wheel_url, + install_wheel, + probe_torch_wheel_env, + url_exists, +) + IS_WINDOWS = sys.platform == "win32" IS_MACOS = sys.platform == "darwin" IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" @@ -368,7 +376,6 @@ NO_TORCH = _infer_no_torch() VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1" # Progress bar state -- updated by _progress() as each install step runs. -# _TOTAL counts: pip-upgrade + 7 shared steps + triton (non-Windows) + local-plugin + finalize # Update _TOTAL here if you add or remove install steps in install_python_stack(). _STEP: int = 0 _TOTAL: int = 0 # set at runtime in install_python_stack() based on platform @@ -535,6 +542,66 @@ NO_TORCH_SKIP_PACKAGES = { "transformers-cfg", } + +def _select_flash_attn_version(torch_mm: str) -> str | None: + return flash_attn_package_version(torch_mm) + + +def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None: + return flash_attn_wheel_url(env) + + +def _print_optional_install_failure( + label: str, result: subprocess.CompletedProcess[str] +) -> None: + _step("warning", f"{label} failed (exit code {result.returncode})", _cyan) + if result.stdout: + print(result.stdout.strip()) + + +def _flash_attn_install_disabled() -> bool: + return os.getenv("UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL") == "1" + + +def _ensure_flash_attn() -> None: + if NO_TORCH or IS_WINDOWS or IS_MACOS: + return + if _flash_attn_install_disabled(): + return + if ( + subprocess.run( + [sys.executable, "-c", "import flash_attn"], + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + ).returncode + == 0 + ): + return + + env = probe_torch_wheel_env() + wheel_url = _build_flash_attn_wheel_url(env) if env else None + if wheel_url and url_exists(wheel_url): + for installer, wheel_result in install_wheel( + wheel_url, + python_executable = sys.executable, + use_uv = USE_UV, + uv_needs_system = UV_NEEDS_SYSTEM, + ): + if wheel_result.returncode == 0: + return + _print_optional_install_failure( + f"Installing flash-attn prebuilt wheel with {installer}", + wheel_result, + ) + _step("warning", "Continuing without flash-attn", _cyan) + return + + if wheel_url is None: + _step("warning", "No compatible flash-attn prebuilt wheel found", _cyan) + else: + _step("warning", "No published flash-attn prebuilt wheel found", _cyan) + + # -- uv bootstrap ------------------------------------------------------ USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack() @@ -762,10 +829,8 @@ def install_python_stack() -> int: base_total = 10 if IS_WINDOWS else 11 if IS_MACOS: base_total -= 1 # triton step is skipped on macOS - # ROCm torch check steps (Linux only, non-macOS, non-no-torch): - # one early check (step 2b) and one final repair (step 13). if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: - base_total += 2 + base_total += 3 _TOTAL = (base_total - 1) if skip_base else base_total # 1. Try to use uv for faster installs (must happen before pip upgrade @@ -979,6 +1044,10 @@ def install_python_stack() -> int: constrain = False, ) + if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + _progress("flash-attn") + _ensure_flash_attn() + # # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch # patch_package_file( # "unsloth-zoo", diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py new file mode 100644 index 0000000000..9881f2258e --- /dev/null +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -0,0 +1,282 @@ +"""Tests for the optional FlashAttention installer.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path +from unittest import mock + +STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio" +sys.path.insert(0, str(STUDIO_DIR)) + +import install_python_stack as ips + + +class TestFlashAttnWheelSelection: + def test_torch_210_maps_to_v281(self): + assert ips._select_flash_attn_version("2.10") == "2.8.1" + + 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): + assert ips._select_flash_attn_version("2.11") is None + + 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 TestEnsureFlashAttn: + def _import_check(self, code: int = 1): + return subprocess.CompletedProcess(["python", "-c", "import flash_attn"], code) + + 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_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_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() + + +class TestInstallPythonStackFlashAttnIntegration: + def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int: + flash_attn_calls = 0 + + 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), + ): + 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_skips_flash_attn_step(self): + assert self._run_install(no_torch = False, is_macos = False, is_windows = True) == 0 diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py index 5c2926a1f1..29cadf87ae 100644 --- a/tests/python/test_no_torch_filtering.py +++ b/tests/python/test_no_torch_filtering.py @@ -419,6 +419,9 @@ class TestInstallPythonStackSubprocessMock: 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, "_ensure_flash_attn", return_value = None), + mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False), + mock.patch.object(ips, "_has_rocm_gpu", return_value = False), mock.patch("subprocess.run", side_effect = mock_run), mock.patch.object(ips, "_bootstrap_uv", return_value = True), mock.patch.object( @@ -637,7 +640,7 @@ class TestInstallShNoTorchFlag: def test_cpu_hint_message_exists(self): """CPU hint message must exist in install.sh.""" assert ( - "No NVIDIA GPU detected" in self.source + "No GPU detected" in self.source ), "CPU hint message not found in install.sh" assert ( "--no-torch" in self.source