studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420)
* studio: skip flash-attn install on Blackwell GPUs (sm_100+) Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120, or sm_121, and the older-arch wheels fail to load on Blackwell. Add a shared has_blackwell_gpu() helper and gate both the install-time (install_python_stack._ensure_flash_attn) and runtime (worker._ensure_flash_attn_for_long_context) paths on it. Detection uses nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows. * test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests prefers_prebuilt_wheel and falls_back_to_pypi exercise the install paths that the Blackwell guard now short-circuits. Make them explicit about non-Blackwell so they pass on real Blackwell hosts. * studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH - Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a single process avoid redundant nvidia-smi spawns. Tests clear the cache via setup_method/teardown_method. - In _ensure_flash_attn, run the NO_TORCH short-circuit before the Blackwell check so GGUF-only users (who never install torch anyway) do not see a Blackwell warning. Blackwell check still runs above the IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see the explicit reason rather than a silent OS skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: add has_blackwell_gpu to mlx worker test wheel_utils stub test_mlx_training_worker_config loads worker.py against a hand-rolled utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol list so worker's import line resolves. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Roland Tannous
GitHub
parent
000ca89301
commit
79adfd9c71
6 changed files with 284 additions and 2 deletions
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue