diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 14fc0933d..503bae8da 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -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", diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 2d3dc5fbf..e4775a10a 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -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"), diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd09..7e7fc1af4 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -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) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83..1b5926fd4 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -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( diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e033a56a0..19c492dea 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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: diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 26ff03505..bf3ed5778 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -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: