Add CUDA index to torch2110 auto-install and a tomli test fallback

The cu126/cu128 torch2110 extras pin torch==2.11.0+cuNNN, which only resolves
from the matching PyTorch CUDA index (torch 2.11's default PyPI wheel is CUDA
13.0). _auto_install.py now appends --extra-index-url download.pytorch.org/whl/
cuNNN for exactly the CUDA-12 torch 2.11 selections so the printed install
command resolves; cu130 and non-torch2110 selections are unchanged.

tests/test_torch2110_cuda_extras.py falls back to the tomli backport when
tomllib is unavailable (Python 3.9 / 3.10).
This commit is contained in:
Daniel Han 2026-07-08 08:45:42 +00:00
commit c216d27ac4
2 changed files with 10 additions and 2 deletions

View file

@ -11,12 +11,16 @@ Hermetic: only parses pyproject.toml, no network or install.
from __future__ import annotations
import tomllib
from pathlib import Path
import pytest
from packaging.requirements import Requirement
try: # tomllib is stdlib on Python 3.11+; older interpreters need the tomli backport.
import tomllib
except ModuleNotFoundError: # pragma: no cover - Python 3.9 / 3.10
tomllib = pytest.importorskip("tomli")
PYPROJECT = Path(__file__).resolve().parents[1] / "pyproject.toml"
_TORCH_TRIO = ("torch", "torchvision", "torchaudio")

View file

@ -41,4 +41,8 @@ else: raise RuntimeError(f"Torch = {v} too new!")
if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8", "13.0"): raise RuntimeError(f"CUDA = {cuda} not supported!")
if v >= V('2.10.0') and cuda not in ("12.6", "12.8", "13.0"): raise RuntimeError(f"Torch 2.10 requires CUDA 12.6, 12.8, or 13.0! Got CUDA = {cuda}")
x = x.format(cuda.replace(".", ""), "-ampere" if False else "") # is_ampere is broken due to flash-attn
print(f'pip install --upgrade pip setuptools wheel && pip install --no-deps git+https://github.com/unslothai/unsloth-zoo.git && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git" --no-build-isolation')
# The cu126/cu128 torch2110 extras pin torch==2.11.0+cuNNN, a local build that only
# resolves from the matching PyTorch CUDA index (torch 2.11's default PyPI wheel is
# CUDA 13.0), so add that index for exactly those CUDA-12 torch 2.11 environments.
extra_index = f' --extra-index-url https://download.pytorch.org/whl/cu{cuda.replace(".", "")}' if (x.endswith('-torch2110') and cuda in ("12.6", "12.8")) else ''
print(f'pip install --upgrade pip setuptools wheel && pip install --no-deps git+https://github.com/unslothai/unsloth-zoo.git && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git" --no-build-isolation{extra_index}')