f5f9e217c1* studio: select torchao version from the installed torch
The Studio installer pins CUDA torch to torch>=2.4,<2.11 and its driver
ladder selects the cu130 wheel index on recent NVIDIA drivers, so pip
resolves torch 2.10.0. overrides.txt hard-pinned torchao==0.14.0, whose
C++ extensions are built against torch 2.9.0, so torchao skipped its cpp
kernels ("Skipping import of cpp extensions due to incompatible torch
version 2.10.0+cu130 for torchao version 0.14.0") and fell back to the
slow Python path. Every CUDA index now tops out at torch 2.10.0, so this
hit most modern installs, not just cu130.
Pick the torchao version matching the torch actually installed in the
venv (table: pytorch/ao#2919): torch 2.10.x -> torchao 0.16.0, 2.11.x ->
torchao 0.17.0, otherwise the previous 0.14.0 (so torch <=2.9 is
unchanged). The installer reads torch.__version__ from the venv via a
cross-platform sys.executable probe (probe_torch_wheel_env is Linux-only)
and passes the computed spec positionally to the existing force-reinstall
override step; overrides.txt becomes a pointer to that logic. torchao's
Python API (Float8Tensor, used by unsloth/kernels/utils.py) imports
cleanly on 0.16.0/0.17.0, verified against torch 2.9.1.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: address review on torchao selection
- Clean the torch minor of pre-release/dev suffixes before parsing
(e.g. '2.10rc1' -> minor 10), matching wheel_utils.probe_torch_wheel_env.
- Pass _windows_hidden_subprocess_kwargs() to the torch-version probe so
it does not flash a console window on Windows (no-op elsewhere).
- Use _safe_print for the selection log line, consistent with the file's
other status output (safe on non-UTF-8 consoles).
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for _select_torchao_spec in install_python_stack.py.
|
|
|
|
torchao's C++ extensions are built against one exact torch release, so the
|
|
installer must pick the torchao version matching the torch installed in the
|
|
venv (otherwise the cpp kernels are skipped). This pins that mapping.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# install_python_stack.py lives at repo_root/studio/install_python_stack.py
|
|
_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py"
|
|
|
|
|
|
def _load_module(monkeypatch):
|
|
"""(Re-)import install_python_stack and return it (mirrors test_pytorch_mirror)."""
|
|
sys.modules.pop("install_python_stack", None)
|
|
monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent))
|
|
import install_python_stack
|
|
|
|
return install_python_stack
|
|
|
|
|
|
@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"),
|
|
("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.
|
|
("2.10.0rc1", "torchao==0.16.0"),
|
|
("2.10.0.dev20250804+cu130", "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"),
|
|
("2.11.0", "torchao==0.17.0"),
|
|
("2.12.0", "torchao==0.17.0"),
|
|
# torch <=2.9 keeps today's pin (already a correct match for 2.9.0).
|
|
("2.9.0+cu128", "torchao==0.14.0"),
|
|
("2.9.1", "torchao==0.14.0"),
|
|
("2.8.0", "torchao==0.14.0"),
|
|
("2.4.0", "torchao==0.14.0"),
|
|
# Unparseable / missing / non-2.x major -> conservative default.
|
|
(None, "torchao==0.14.0"),
|
|
("", "torchao==0.14.0"),
|
|
("garbage", "torchao==0.14.0"),
|
|
("2", "torchao==0.14.0"),
|
|
("3.0.0", "torchao==0.14.0"),
|
|
],
|
|
)
|
|
def test_select_torchao_spec(monkeypatch, torch_version, expected):
|
|
mod = _load_module(monkeypatch)
|
|
assert mod._select_torchao_spec(torch_version) == expected
|
|
|
|
|
|
def test_default_spec_matches_table(monkeypatch):
|
|
"""The default/floor stays the historical pin so older torch is unchanged."""
|
|
mod = _load_module(monkeypatch)
|
|
assert mod._TORCHAO_DEFAULT_SPEC == "torchao==0.14.0"
|
|
assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC
|