diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 8cd1d1f9c8..6b3b3b6609 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -30,6 +30,7 @@ from utils.hardware import apply_gpu_ids from utils.wheel_utils import ( direct_wheel_url, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -313,6 +314,12 @@ def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: 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 + if has_blackwell_gpu(): + _send_status( + event_queue, + "Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + ) + return installed = _install_package_wheel_first( event_queue = event_queue, diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 5900af4e3d..98c7bdaa55 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,6 +37,7 @@ 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", diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 41a7c87df1..0737bdc82f 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -37,6 +37,7 @@ 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, @@ -65,6 +66,7 @@ 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, @@ -112,6 +114,29 @@ 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) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 3ed9bda827..5c42e890d1 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations +import functools import json import logging import platform @@ -22,6 +23,49 @@ FLASH_ATTN_RELEASE_BASE_URL = ( ) +@functools.lru_cache(maxsize = 1) +def has_blackwell_gpu() -> bool: + """Return True if any visible NVIDIA GPU has compute capability >= 10.0 + (Blackwell: sm_100, sm_120, sm_121, ...). + + Dao-AILab does not publish prebuilt flash-attention wheels for these + architectures, and the older-arch wheels fail to load on Blackwell, so + callers use this gate to skip the flash-attn install/upgrade path. + + Result is cached for the process lifetime since GPU hardware does not + change. Tests that mock subprocess/nvidia-smi must call + ``has_blackwell_gpu.cache_clear()`` before each invocation. + """ + exe = shutil.which("nvidia-smi") + if not exe: + return False + try: + result = subprocess.run( + [exe, "--query-gpu=compute_cap", "--format=csv,noheader"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + env = child_env_without_native_path_secret(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + if result.returncode != 0: + return False + for line in result.stdout.splitlines(): + cap = line.strip() + if not cap: + continue + major_part = cap.split(".", 1)[0] + try: + major = int(major_part) + except ValueError: + continue + if major >= 10: + return True + return False + + def linux_wheel_platform_tag() -> str | None: machine = platform.machine().lower() if sys.platform.startswith("linux"): diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3fd1e6af66..ab234ad566 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -28,6 +28,7 @@ if str(_BACKEND_DIR) not in sys.path: from backend.utils.wheel_utils import ( flash_attn_package_version, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -628,10 +629,19 @@ def _flash_attn_install_disabled() -> bool: def _ensure_flash_attn() -> None: - if NO_TORCH or IS_WINDOWS or IS_MACOS: - return if _flash_attn_install_disabled(): return + if NO_TORCH: + return + if has_blackwell_gpu(): + _step( + "warning", + "Skipping flash-attn: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + _cyan, + ) + return + if IS_WINDOWS or IS_MACOS: + return if ( subprocess.run( [sys.executable, "-c", "import flash_attn"], diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 9881f2258e..49f4350a7b 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -10,8 +10,133 @@ from unittest import mock 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 +from backend.utils import wheel_utils + + +def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess: + return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "") + + +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 class TestFlashAttnWheelSelection: @@ -234,6 +359,76 @@ 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)) + + 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: def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int: