ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/install_python_stack.py
Daniel Han d671d8fb2f install: apply an explicit custom torch-index pin on the first update 
Round 6: an explicitly-set custom (unknown-family) UNSLOTH_TORCH_INDEX_URL
was silently ignored on the first `studio update` of a venv that predates
the marker feature, on both platforms, because the no-marker case was
treated as "do nothing" and the version-tag heuristics cannot judge an
unknown leaf.

1. install_python_stack.py _ensure_verbatim_torch_index now reinstalls
   verbatim when the marker is ABSENT (None), not only when it differs, and
   short-circuits only when the marker already records this exact pin. It
   then writes the marker, so every later update is a no-op. A user who did
   not set the override gets pin=None and is untouched, so an out-of-band
   torch install is never clobbered.

2. setup.ps1: for an unknown-family pin on a marker-less venv the stale-venv
   check now sets PinChangedForceReinstall so the torch block reinstalls in
   place from the pin. It deliberately does NOT set shouldRebuild, which
   would wipe the venv and strand a direct `studio update`.

3. setup.sh (the Linux `studio update` entry point) skipped
   install_python_stack.py entirely when unsloth was already current, so the
   marker-driven reinstall (both the verbatim custom pin and the cu/rocm
   flavor and family-change repair, e.g. gfx1151 to gfx120X-all) never ran.
   It now forces the dependency pass when a torch-index pin env var is set;
   the pass is idempotent and no-ops when the marker already matches. This
   mirrors setup.ps1's stale-venv pre-check.

Tests: 3 new parity assertions.
2026-07-12 12:39:29 +00:00

3172 lines
134 KiB
Python
Raw Blame History

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cross-platform Python dependency installer for Unsloth Studio.
Called by setup.sh (Linux/WSL) and setup.ps1 (Windows) after the venv is
activated. Expects `pip` and `python` on PATH to point at the venv.
"""
from __future__ import annotations
import glob
import os
import platform
import re
import shutil
import subprocess
import sys
import sysconfig
import tempfile
import textwrap
import urllib.request
from pathlib import Path
_BACKEND_DIR = Path(__file__).resolve().parent / "backend"
if str(_BACKEND_DIR) not in sys.path:
sys.path.insert(1, str(_BACKEND_DIR))
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,
)
from backend.utils.uv_path_safety import uv_safe_path as _uv_safe_path
IS_WINDOWS = sys.platform == "win32"
IS_MACOS = sys.platform == "darwin"
IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64"
IS_LINUX = sys.platform.startswith("linux")
# DiskPart-prompt suppression: amd-smi auto-elevates on Windows, popping a
# UAC/DiskPart prompt mid-install. This installer only spawns probes and pip/uv
# (none need elevation), so set __COMPAT_LAYER=RunAsInvoker process-wide -- every
# amd-smi subprocess then runs un-elevated, no per-call guard needed. setup.ps1
# keeps per-call guards since it ALSO spawns winget installers that need elevation.
if IS_WINDOWS:
os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker")
# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64,
# and win_amd64. On other hosts the audio extras must be filtered out (the
# extras-no-deps step would otherwise fail), regardless of NO_TORCH.
PLATFORM_LACKS_TORCHCODEC_WHEEL = (
(IS_LINUX and platform.machine() in {"aarch64", "arm64"})
or (IS_WINDOWS and platform.machine().lower() in {"arm64", "aarch64"})
or IS_MAC_INTEL
)
# ── ROCm / AMD GPU support ─────────────────────────────────────────────────────
# Detected ROCm (major, minor) -> best PyTorch wheel tag on
# download.pytorch.org. Checked newest-first (>=).
_ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(7, 2): "rocm7.2", # torch 2.11.0
(7, 1): "rocm7.1", # torch 2.10.0
(7, 0): "rocm7.0",
(6, 4): "rocm6.4",
(6, 3): "rocm6.3",
(6, 2): "rocm6.2",
(6, 1): "rocm6.1",
(6, 0): "rocm6.0",
}
# AMD per-arch index leaves that need the torch 2.11 floor (the torch._C._grouped_mm
# null-ptr bug lives in the <2.11 wheels for these arches). Mirrors the gfx keys in
# _WINDOWS_ROCM_TORCH_PKG_SPECS and the *FloorMap sets in install.ps1 / setup.ps1.
# Other per-arch indexes (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels and must
# stay bare, so an override to one of them must NOT be forced onto the 2.11 line.
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
# The pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (verified against
# download.pytorch.org): rocm7.2 -> torch 2.11.0 is the ONLY stable 2.11 rocm
# index today (rocm7.1 -> 2.10.0, rocm6.4 -> 2.9.1; rocm7.3 / torch 2.12 do NOT
# exist as stable). Do NOT treat an unknown newer rocm (rocm7.3, rocm8.0, ...) as
# 2.11 speculatively -- that is exactly the mismatch bug tracked in the review.
# MUST match the rocm leaf in the KNOWN-2.11 case in install.sh / setup.ps1 /
# install.ps1 (rocm7.2 there too). Bump alongside those when a new stable rocm
# index publishes torch 2.11+. Stored as (major, minor) tuples for exact compares.
_ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)})
# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x).
_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"rocm7.2": (
"torch>=2.11.0,<2.12.0",
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
),
# Default for rocm7.1 and earlier: torch 2.x below 2.11
"_default": (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
),
}
# Windows AMD per-arch companion pins for the repo.amd.com index, mirroring the
# install.ps1 / setup.ps1 floor maps (gfx120X and Strix Halo/Point use the rocm7.2
# torch 2.11 trio). Pinning the companions keeps AMD's per-arch index -- which
# publishes each independently -- from resolving an ABI-mismatched one. Unlisted
# arches have no published floor, so stay bare. Bump with the PS maps at 2.12.x.
_WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"gfx1201": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1151": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1150": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
}
_PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
).rstrip("/")
# ── Torch-index marker ─────────────────────────────────────────────────────────
# After a successful torch install/reinstall, record the exact wheel --index-url
# used at a stable per-venv path. On `studio update`/repair the marker turns the
# "did the pinned index change?" decision into an EXACT string compare instead of
# inferring it from the wheel's +rocm/+cu version tag (which cannot encode the AMD
# per-arch gfx family: two gfx 2.11 indexes both install +rocm7.13.0). This
# dissolves the per-arch-switch (gfx1151 -> gfx120X-all) and custom-URL
# (/simple, /current) cases the version-tag heuristics cannot see.
#
# The path/format MUST match install.sh, install.ps1 and setup.ps1:
# <venv_prefix>/.unsloth-torch-index (single line = the resolved index URL)
# Written atomically (temp file + os.replace). A missing/empty/corrupt marker is
# treated as absent, so old venvs (and torch installed out-of-band) fall back to
# the existing +rocm/version-tag heuristics -- backward compatibility is required.
_TORCH_INDEX_MARKER_NAME = ".unsloth-torch-index"
def _normalize_family_leaf(leaf: str) -> str:
"""Lowercase ONLY a known wheel-family leaf (rocm<digit>* / gfx* / cpu / cuXXX).
The canonical gfx120X-all (capital X) must match AMD's lowercase gfx120x-all, so
known-family leaves are lowercased. A custom mirror leaf (/Current, /simple, ...)
keeps its case: an unknown-family URL pin is applied verbatim, so /Current and
/current must NOT compare equal. The rocm prefix is digit-gated like
_is_pip_rocm_family_leaf: rocm7.2 is a family leaf, but rocm-rel-7.2.1 /
rocm-Current are verbatim pins whose case must survive normalization (URL
paths can be case-sensitive). Mirrors the gate in install.sh / setup.ps1.
Pure function.
"""
low = leaf.lower()
if low.startswith("gfx") or low == "cpu" or re.match(r"^(rocm|cu)[0-9]", low):
return low
return leaf
def _normalize_index_url(url: "str | None") -> "str | None":
"""Canonicalise a wheel index URL for exact marker/pin comparison.
Trims surrounding whitespace, strips ALL trailing slashes, and lowercases the
FINAL path segment ONLY when it is a known wheel-family leaf (cu128 / cpu /
rocm7.2 / gfx1151 / gfx120X-all) -- see _normalize_family_leaf. The host part is
left untouched (case-sensitive on some mirrors), and a custom (unknown-family)
leaf keeps its case so a verbatim URL pin is not falsely matched equal. MUST
match the same normalization in install.sh / setup.ps1. Returns None for an
empty/whitespace-only input. Pure function.
"""
if url is None:
return None
url = url.strip()
if not url:
return None
url = url.rstrip("/")
if not url:
return None
head, sep, leaf = url.rpartition("/")
if sep:
return f"{head}/{_normalize_family_leaf(leaf)}"
return _normalize_family_leaf(url)
def _torch_index_marker_path() -> Path:
"""Path to the per-venv torch-index marker (see _TORCH_INDEX_MARKER_NAME).
Anchored at sys.prefix (the venv the installer targets via --python
sys.executable), so it matches install.sh's $VENV_DIR and the PowerShell
$VenvDir marker location.
"""
return Path(sys.prefix) / _TORCH_INDEX_MARKER_NAME
def _read_torch_index_marker() -> "str | None":
"""Return the recorded torch --index-url from the marker, else None.
None when the marker is missing, empty, or unreadable (corrupt/permission) --
all treated as "no marker" so the caller falls back to the version-tag
heuristics. The stored URL is returned VERBATIM (not normalized); callers
normalize both sides before comparing.
"""
try:
text = _torch_index_marker_path().read_text(encoding = "utf-8")
except (OSError, ValueError):
return None
line = text.strip()
return line or None
def _write_torch_index_marker(index_url: "str | None") -> None:
"""Record the resolved torch wheel --index-url at the per-venv marker path.
Best-effort and atomic (temp file in the same dir + os.replace). Never raises:
a marker write failure must not abort an otherwise-successful install (the
repair path then falls back to the heuristics, same as an old venv). A blank
index_url is ignored (nothing meaningful to record).
"""
if not index_url or not index_url.strip():
return
marker = _torch_index_marker_path()
payload = index_url.strip() + "\n"
try:
marker.parent.mkdir(parents = True, exist_ok = True)
fd, tmp = tempfile.mkstemp(prefix = ".unsloth-torch-index.", dir = str(marker.parent))
try:
with os.fdopen(fd, "w", encoding = "utf-8") as fh:
fh.write(payload)
os.replace(tmp, str(marker))
except OSError:
try:
os.unlink(tmp)
except OSError:
pass
raise
except OSError:
# Non-fatal: fall back to writing directly, then give up silently.
try:
marker.write_text(payload, encoding = "utf-8")
except OSError:
pass
def _marker_pin_mismatch(pin_url: str) -> "bool | None":
"""Compare an explicit torch-index pin against the recorded marker.
Returns:
* True -> the pin resolves to a DIFFERENT index than the marker records
(exact compare after _normalize_index_url) -> reinstall needed.
* False -> the pin matches the marker exactly -> no reinstall (no loop).
* None -> no usable marker (missing/empty/corrupt) -> the caller must fall
back to the +rocm/version-tag heuristics (backward compatibility).
This is the EXACT signal the version-tag heuristics cannot provide: switching
one 2.11 gfx pin to another (gfx1151 -> gfx120X-all) changes the pinned URL but
not the installed +rocm7.13.0 wheel tag. Pure w.r.t. its arg (reads the marker).
"""
marker = _read_torch_index_marker()
if marker is None:
return None
return _normalize_index_url(pin_url) != _normalize_index_url(marker)
# 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.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
)
# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as the
# CUDA spec above: the /cpu index also publishes newer torch, so a bare trio
# from the exclusive --index-url would resolve outside the repo's supported
# range (and can pick a torchvision built against a different torch major).
_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC
# 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_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:
"""Map an installed torch version string (e.g. '2.10.0+cu130') to the torchao
pip spec whose cpp extensions match it. Falls back to _TORCHAO_DEFAULT_SPEC for
torch <=2.9, a non-2.x major, or an unparseable/missing version. Pure function.
"""
if not torch_version:
return _TORCHAO_DEFAULT_SPEC
release = str(torch_version).split("+", 1)[0] # drop +cu130/+rocm6.4/+cpu
parts = release.split(".")
try:
# Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10'),
# matching wheel_utils.probe_torch_wheel_env.
minor_str = re.sub(r"[^0-9].*", "", parts[1]) if len(parts) > 1 else ""
major, minor = int(parts[0]), int(minor_str)
except (IndexError, ValueError):
return _TORCHAO_DEFAULT_SPEC
if major != 2:
return _TORCHAO_DEFAULT_SPEC
if minor >= 11:
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:
"""Return torch.__version__ from the target venv (sys.executable), or None if
torch is absent/unimportable. Cross-platform (unlike probe_torch_wheel_env,
which is Linux-only); mirrors the subprocess probe in _ensure_cuda_torch.
"""
try:
probe = subprocess.run(
[
sys.executable,
"-c",
"import torch, sys; sys.stdout.write(getattr(torch, '__version__', ''))",
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 90,
**_windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired):
return None
if probe.returncode != 0:
return None
lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()]
return lines[-1] if lines else None
def _installed_torch_is_windows_rocm() -> bool:
"""Return True when the target venv currently has a Windows ROCm torch build.
This is a belt-and-suspenders guard for the torchao override step: if the
earlier ROCm install path failed to set _rocm_windows_torch_installed but the
venv already contains a ROCm torch wheel, still skip torchao because it
crashes on import on Windows ROCm.
"""
if not IS_WINDOWS:
return False
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import sys, torch; "
"hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
"sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 90,
**_windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired):
return False
lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()]
return probe.returncode == 0 and bool(lines and lines[-1] == "yes")
# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install
# from before the cap existed can already be stuck at 4.14+, which later
# constrained installs won't touch since it already satisfies mcp/fastmcp.
_ANYIO_BAD_FLOOR = (4, 14)
def _installed_anyio_version() -> tuple[int, int] | None:
try:
from importlib.metadata import version as _pkg_version
raw = _pkg_version("anyio")
except Exception:
return None
parts = raw.split(".")
try:
major = int(parts[0])
minor = int(re.sub(r"[^0-9].*", "", parts[1])) if len(parts) > 1 else 0
except (IndexError, ValueError):
return None
return (major, minor)
def _repair_bad_anyio() -> None:
installed = _installed_anyio_version()
if installed is None or installed < _ANYIO_BAD_FLOOR:
return
_safe_print(_dim(f" anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14..."))
pip_install(
"Repairing anyio version",
"--no-cache-dir",
"--force-reinstall",
"anyio<4.14.0",
constrain = False,
)
# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/).
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs.
_ROCM_WINDOWS_INDEX_BASE = (
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR") or "https://repo.amd.com/rocm/whl"
).rstrip("/")
# gfx arch → AMD index arch-family suffix; each family is a separate
# pip index on repo.amd.com.
_GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx1201": "gfx120X-all",
"gfx1200": "gfx120X-all", # RDNA 4
"gfx1151": "gfx1151",
"gfx1150": "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1103": "gfx110X-all",
"gfx1102": "gfx110X-all", # RDNA 3
"gfx1101": "gfx110X-all",
"gfx1100": "gfx110X-all",
"gfx90a": "gfx90a",
"gfx908": "gfx908", # MI200/MI100
}
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"x86_64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl"
),
"aarch64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl"
),
# Windows ROCm wheel ships libbitsandbytes_rocm{VER}.dll. BNB's HIP
# auto-detect may mismatch the DLL suffix, so we scan the wheel and set
# BNB_ROCM_VERSION in _install_bnb_windows_rocm() and worker.py.
"win_amd64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl"
),
}
_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1"
def _bnb_rocm_prerelease_url() -> str | None:
"""Return the continuous-release_main bnb wheel URL for the current arch,
or None when no pre-release wheel is available.
"""
arch = platform.machine().lower()
arch = {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch)
return _BNB_ROCM_PRERELEASE_URLS.get(arch)
def _amd_smi_env() -> dict[str, str] | None:
"""On Windows, env with __COMPAT_LAYER=RunAsInvoker; None elsewhere.
NB: RunAsInvoker doesn't stop amd-smi's runtime elevation (its manifest is
asInvoker -- it elevates a child via ShellExecute). The real guard is
_amd_smi_allowed() below; this is harmless belt-and-suspenders."""
if platform.system() != "Windows":
return None
return {**os.environ, "__COMPAT_LAYER": "RunAsInvoker"}
def _path_inside_venv(path: str) -> bool:
"""True if ``path`` is inside the active venv (sys.prefix).
The venv hipInfo.exe (AMD wheel, put on PATH by the bnb fix) is NOT a HIP SDK
(_amd_smi_allowed)."""
try:
# realpath (not abspath): resolve symlinks/8.3 names so an aliased venv matches.
_root = os.path.normcase(os.path.realpath(sys.prefix))
# Guard a root-dir prefix (C:\ or /): commonpath would match every path on
# it. A venv is never at root, so treat that as outside.
if os.path.dirname(_root) == _root:
return False
return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root
except (ValueError, OSError):
# Different drive / unresolvable -> treat as outside the venv.
return False
def _external_hipinfo_on_path() -> bool:
"""True if a hipinfo OUTSIDE the venv is on PATH.
shutil.which returns only the first hit, so the venv hipInfo could shadow a
real HIP SDK's; scan every PATH entry and skip the venv copy."""
for _dir in os.environ.get("PATH", "").split(os.pathsep):
_dir = _dir.strip('"') # PATH entries can be quoted on Windows
if not _dir:
continue
_candidate = os.path.join(_dir, "hipinfo.exe")
if os.path.isfile(_candidate) and not _path_inside_venv(_candidate):
return True
return False
def _amd_smi_allowed() -> bool:
"""Whether it is safe to spawn amd-smi here.
On Windows w/o a working HIP runtime, amd-smi elevates a child and pops a
UAC/DiskPart prompt RunAsInvoker can't suppress. Only call it on Windows with
a HIP SDK (hipinfo present) or UNSLOTH_ENABLE_AMD_SMI=1; Linux/macOS always.
"""
if platform.system() != "Windows":
return True
flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower()
if flag in ("1", "true", "yes", "on"):
return True
if flag in ("0", "false", "no", "off"):
return False
# A real HIP SDK lets amd-smi run un-elevated; hipinfo-on-PATH is the proxy.
# Ignore the venv hipInfo.exe (AMD wheel via bnb fix): not a HIP SDK, doesn't
# stop amd-smi's DiskPart UAC.
if _external_hipinfo_on_path():
return True
for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
_root = os.environ.get(_var)
if not _root:
continue
_candidate = os.path.join(_root, "bin", "hipinfo.exe")
if os.path.isfile(_candidate) and not _path_inside_venv(_candidate):
return True
return False
def _detect_rocm_version() -> tuple[int, int] | None:
"""Return (major, minor) of the installed ROCm stack, or None."""
rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm"
for path in (
os.path.join(rocm_root, ".info", "version"),
os.path.join(rocm_root, "lib", "rocm_version"),
):
try:
with open(path) as fh:
parts = fh.read().strip().split("-")[0].split(".")
# Explicit length guard so we don't rely on the broad except
# below to swallow IndexError when the version file has a
# single component (e.g. "6\n" on a partial install).
if len(parts) >= 2:
return int(parts[0]), int(parts[1])
except Exception:
pass
# Try amd-smi version (outputs "... | ROCm version: X.Y.Z").
# Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt);
# hipconfig below covers that case.
amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None
if amd_smi:
try:
result = subprocess.run(
[amd_smi, "version"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 5,
env = _amd_smi_env(),
)
if result.returncode == 0:
m = re.search(r"ROCm version:\s*(\d+)\.(\d+)", result.stdout)
if m:
return int(m.group(1)), int(m.group(2))
except Exception:
pass
# Try hipconfig --version (outputs bare version like "6.3.21234.2")
hipconfig = shutil.which("hipconfig")
if hipconfig:
try:
result = subprocess.run(
[hipconfig, "--version"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 5,
)
if result.returncode == 0:
raw = result.stdout.decode().strip().split("\n")[0]
parts = raw.split(".")
if len(parts) >= 2 and parts[0].isdigit() and parts[1].split("-")[0].isdigit():
return int(parts[0]), int(parts[1].split("-")[0])
except Exception:
pass
# Distro package-manager fallbacks. Package-managed ROCm installs can
# expose GPUs via rocminfo/amd-smi but lack /opt/rocm/.info/version and
# hipconfig, so probe dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE)
# for the rocm-core version. Matches install.sh::get_torch_index_url so
# `unsloth studio update` behaves like a fresh `curl | sh` install.
for cmd in (
["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"],
["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"],
):
exe = shutil.which(cmd[0])
if not exe:
continue
try:
result = subprocess.run(
[exe, *cmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 5,
)
except Exception:
continue
if result.returncode != 0 or not result.stdout.strip():
continue
raw = result.stdout.strip()
# dpkg can prepend an epoch ("1:6.3.0-1"); strip it before parsing.
raw = re.sub(r"^\d+:", "", raw)
m = re.match(r"(\d+)[.-](\d+)", raw)
if m:
return int(m.group(1)), int(m.group(2))
return None
def _pick_visible_index(num_tokens: int) -> int:
"""Resolve HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES to an index into a
list of length num_tokens. Returns 0 (first GPU) for unset, empty, '-1',
UUID-style, or out-of-range values."""
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
_val = os.environ.get(_env)
if _val is None:
continue
_val = _val.strip()
if _val == "" or _val == "-1":
return 0
_first = _val.split(",")[0].strip()
try:
_idx = int(_first)
if 0 <= _idx < num_tokens:
return _idx
except ValueError:
pass
return 0
return 0
def _detect_windows_gfx_arch() -> str | None:
"""Return the gcnArchName on Windows (e.g. 'gfx1200'), or None.
Probe order matches the PowerShell installer: env-var override, then
hipinfo (PATH or HIP_PATH/ROCM_PATH bin), then amd-smi. Without the
amd-smi fallback, runtime-only AMD installs lacking hipinfo on PATH
return early and `studio update` cannot repair a CPU-only venv.
On multi-GPU hosts, detected gfx tokens are deduplicated (preserving
enumeration order) and HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES picks
which to install for. The first GPU is used when no env var is set.
"""
# 1. Explicit override (matches PowerShell installer's env-var path).
_override = os.environ.get("UNSLOTH_ROCM_GFX_ARCH")
if _override and _override.strip():
return _override.strip().lower()
def _dedup_pick(tokens: list[str]) -> "str | None":
if not tokens:
return None
# Index into the full ordered list so HIP_VISIBLE_DEVICES addresses
# GPU N on mixed-arch hosts, then return that arch.
return tokens[_pick_visible_index(len(tokens))]
# 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
hipinfo = shutil.which("hipinfo")
if not hipinfo:
for _env_var in ("HIP_PATH", "ROCM_PATH"):
_root = os.environ.get(_env_var)
if _root:
_candidate = os.path.join(_root, "bin", "hipinfo.exe")
if os.path.isfile(_candidate):
hipinfo = _candidate
break
if not hipinfo:
# 2b. AMD torch wheels ship hipInfo.exe into the venv Scripts dir
# (next to python.exe); resolvable even on driver-only hosts with no
# SDK install at all. Lets `studio update` re-detect the arch on a
# venv that already has the AMD wheel.
_venv_hipinfo = os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")
if os.path.isfile(_venv_hipinfo):
hipinfo = _venv_hipinfo
if hipinfo:
try:
result = subprocess.run(
[hipinfo],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 10,
)
# Accept partial output even when hipinfo crashes (e.g. exit code
# 0xC0000005 / STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): if
# gcnArchName is present in stdout the device was enumerated before
# the crash, so the arch is trustworthy. Ignoring it causes a
# silent CPU PyTorch fallback (issue #6043).
text = result.stdout.decode(errors = "replace")
# findall gets every gcnArchName line so multi-GPU hosts are
# enumerable and HIP_VISIBLE_DEVICES selects correctly.
_tokens = [
t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
]
_pick = _dedup_pick(_tokens)
if _pick:
return _pick
except Exception:
pass
# 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo.
# Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); the arch
# arrives via --rocm-gfx / name inference there, so this is only needed when safe.
amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None
if amd_smi:
for _args in (("static", "--asic"), ("list",)):
try:
result = subprocess.run(
[amd_smi, *_args],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 10,
env = _amd_smi_env(),
)
if result.returncode != 0:
continue
text = result.stdout.decode(errors = "replace")
# Prefer labelled gfx lines; fall back to bare tokens.
_labelled = re.findall(
r"(?im)^\s*(?:target_graphics_version|gfx|arch|asic)\b[^:\r\n]*:\s*(gfx[1-9][0-9a-z]{2,3})\b",
text,
)
_tokens = [t.lower() for t in _labelled]
if not _tokens:
_tokens = re.findall(r"\bgfx[1-9][0-9a-z]{2,3}\b", text.lower())
_pick = _dedup_pick(_tokens)
if _pick:
return _pick
except Exception:
continue
# 4. Last resort: GPU marketing name via WMI → arch table. Driver-only
# hosts (Adrenalin, no HIP SDK) have neither hipinfo nor amd-smi
# (amd-smi does not exist on Windows at all), but the display driver
# always knows the GPU name. Mirrors setup.ps1's $nameArchTable so a
# standalone `studio update` can repair a CPU-only venv on such hosts.
try:
result = subprocess.run(
[
"powershell.exe",
"-NoProfile",
"-NonInteractive",
"-Command",
"(Get-CimInstance Win32_VideoController).Name",
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 30,
creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0),
)
if result.returncode == 0:
_tokens = []
for _name in result.stdout.decode(errors = "replace").splitlines():
_arch = _gfx_arch_from_gpu_name(_name.strip())
if _arch:
_tokens.append(_arch)
_pick = _dedup_pick(_tokens)
if _pick:
print(f" gfx arch inferred from GPU name (WMI): {_pick}")
return _pick
except Exception:
pass
return None
# GPU marketing-name → gfx arch table, mirroring setup.ps1's $nameArchTable.
# Most-specific first; first match wins. Covers only arches the ROCm
# prebuilts / AMD Windows torch indexes support; unknown names return None
# (callers then fall back cleanly to CPU).
_WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [
(r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080)
(r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060)
# RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
(r"8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"),
# RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
(
r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33",
"gfx1150",
),
# RDNA 3 desktop / workstation (Navi 31)
(r"RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700", "gfx1100"),
(r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710", "gfx1102"), # Navi 33
# RDNA 3 iGPU (Phoenix / Hawk Point)
(r"780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme", "gfx1103"),
(r"RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900", "gfx1030"), # Navi 21
(r"RX 6650|RX 6600|PRO W6600|PRO W6650", "gfx1032"), # Navi 23
(r"RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500", "gfx1034"), # Navi 24
]
def _gfx_arch_from_gpu_name(name: str) -> "str | None":
"""Map a GPU marketing name to its gfx arch via _WIN_GPU_NAME_ARCH_TABLE."""
if not name:
return None
for _pat, _arch in _WIN_GPU_NAME_ARCH_TABLE:
if re.search(_pat, name, re.IGNORECASE):
return _arch
return None
def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
"""Return the AMD pip index URL for the given GPU arch, or None if unsupported."""
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
if arch_family is None:
return None
return f"{_ROCM_WINDOWS_INDEX_BASE}/{arch_family}/"
def _detect_bnb_rocm_dll_ver() -> str | None:
"""Scan the installed bitsandbytes package for libbitsandbytes_rocm{VER}.dll.
Returns the version suffix (e.g. ``"72"``, ``"713"``) or ``None`` if
bitsandbytes is not installed or no ROCm DLL is found. Does NOT import
bitsandbytes — uses importlib.util.find_spec, so it is safe to call
before BNB is imported.
"""
import importlib.util
spec = importlib.util.find_spec("bitsandbytes")
if spec is None or not spec.submodule_search_locations:
return None
all_vers: list[str] = []
for pkg_dir in spec.submodule_search_locations:
for dll in glob.glob(os.path.join(pkg_dir, "libbitsandbytes_rocm*.dll")):
m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll))
if m:
all_vers.append(m.group(1))
# Pick the highest numeric suffix so e.g. "713" wins over "72" when both
# variants are present. Glob order is not guaranteed, so always sort
# rather than stopping at the first match.
return max(all_vers, key = lambda v: int(v)) if all_vers else None
_BNB_ROCM_SITECUSTOMIZE_BEGIN = "# BEGIN Unsloth BNB_ROCM_VERSION"
_BNB_ROCM_SITECUSTOMIZE_END = "# END Unsloth BNB_ROCM_VERSION"
_BNB_ROCM_VERSION_SOURCE_ENV = "UNSLOTH_BNB_ROCM_VERSION_SOURCE"
_BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE = "sitecustomize"
_BNB_ROCM_VERSION_SOURCE_DETECTED = "detected"
def _persist_bnb_rocm_version(version: str) -> bool:
"""Persist BNB_ROCM_VERSION for future Python processes in this venv."""
version = str(version).strip()
if not version:
return False
site_packages = sysconfig.get_path("purelib")
if not site_packages:
return False
sitecustomize_path = Path(site_packages) / "sitecustomize.py"
block = (
f"{_BNB_ROCM_SITECUSTOMIZE_BEGIN}\n"
"import os as _unsloth_os\n"
"_unsloth_existing_bnb_rocm = _unsloth_os.environ.get('BNB_ROCM_VERSION')\n"
f"_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', {version!r})\n"
"if _unsloth_existing_bnb_rocm is None and "
f"_unsloth_os.environ.get('BNB_ROCM_VERSION') == {version!r}:\n"
" _unsloth_os.environ.setdefault("
f"{_BNB_ROCM_VERSION_SOURCE_ENV!r}, "
f"{_BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE!r})\n"
"del _unsloth_existing_bnb_rocm\n"
f"{_BNB_ROCM_SITECUSTOMIZE_END}\n"
)
try:
sitecustomize_path.parent.mkdir(parents = True, exist_ok = True)
existing = (
sitecustomize_path.read_text(encoding = "utf-8") if sitecustomize_path.exists() else ""
)
# Strip all managed regions, including one whose END marker was lost to
# an interrupted write, then append exactly one fresh block.
pattern = re.compile(
rf"{re.escape(_BNB_ROCM_SITECUSTOMIZE_BEGIN)}.*?"
rf"(?:{re.escape(_BNB_ROCM_SITECUSTOMIZE_END)}\n?|\Z)",
re.DOTALL,
)
remainder = pattern.sub("", existing)
separator = "" if not remainder or remainder.endswith("\n") else "\n"
updated = f"{remainder}{separator}{block}"
tmp_path = sitecustomize_path.with_name(
f"{sitecustomize_path.name}.unsloth-tmp{os.getpid()}"
)
try:
tmp_path.write_text(updated, encoding = "utf-8")
if sitecustomize_path.exists():
shutil.copymode(sitecustomize_path, tmp_path)
os.replace(tmp_path, sitecustomize_path)
finally:
tmp_path.unlink(missing_ok = True)
except (OSError, UnicodeDecodeError) as exc:
print(
f" Warning: could not persist BNB_ROCM_VERSION={version} "
f"to {sitecustomize_path}: {exc}"
)
return False
return True
def _has_rocm_gpu() -> bool:
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed).
Always returns False when an NVIDIA GPU is present -- NVIDIA takes
priority on mixed hosts and prevents every detection path below
(rocminfo, amd-smi, KFD sysfs) from producing a false positive even
if ROCm tools are installed alongside the NVIDIA driver.
"""
if _has_usable_nvidia_gpu():
return False
for cmd, check_fn in (
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
# gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like
# "gfx11-generic"/"gfx9-4-generic" with only 1-2 digits before the dash,
# which must not be treated as a real GPU.
(
["rocminfo"],
lambda out: bool(re.search(r"gfx[1-9][0-9a-z]{2,3}", out.lower())),
),
# amd-smi list: require "GPU: <number>" data rows, not just a header
(
["amd-smi", "list"],
lambda out: bool(re.search(r"(?im)^gpu\s*[:\[]\s*\d", out)),
),
):
exe = shutil.which(cmd[0])
if not exe:
continue
# Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt);
# rely on rocminfo / the sysfs fallback there.
if cmd[0] == "amd-smi" and not _amd_smi_allowed():
continue
try:
result = subprocess.run(
[exe, *cmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
env = _amd_smi_env() if cmd[0] == "amd-smi" else None,
)
except Exception:
continue
if result.returncode == 0 and result.stdout.strip():
if check_fn(result.stdout):
return True
# sysfs KFD topology fallback (Linux only) -- matches install.sh's
# runtime-only detection. On minimal package-managed installs (no
# rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via
# /sys/class/kfd so `studio update` can still detect and repair.
#
# Guard: reject any KFD node whose properties file reports a non-AMD
# vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs
# can register KFD topology nodes with a non-zero gpu_id; those nodes
# have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002).
# Without this check the fallback returns True on NVIDIA-only systems,
# causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware.
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
if os.path.isdir(kfd_nodes):
for entry in os.listdir(kfd_nodes):
gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id")
try:
with open(gpu_id_path) as fh:
gpu_id = fh.read().strip()
except OSError:
continue
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
# Require AMD vendor_id 4098 (0x1002) in the properties file.
# KFD properties files exist on every kernel that exposes
# /sys/class/kfd, so absence of the file means we cannot
# confirm AMD ownership -- skip the node rather than risk a
# false positive (e.g. NVIDIA open driver KFD nodes that
# lack a properties file on some kernel versions).
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path) as fh:
props = fh.read()
except OSError:
continue # can't confirm vendor -- skip
if not re.search(r"\bvendor_id\s+4098\b", props):
continue
return True
except OSError:
pass
return False
def _has_usable_nvidia_gpu() -> bool:
"""Return True when an NVIDIA GPU is present and usable.
Primary probe: nvidia-smi -L (subprocess).
Fallback: /proc/driver/nvidia/gpus/ sysfs (Linux only) -- handles the
case where nvidia-smi is present but the subprocess fails (PATH gap,
timeout, driver initialisation race). If either probe confirms an
NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked.
CUDA_VISIBLE_DEVICES set to "" or "-1" hides every NVIDIA device (mixed
AMD+NVIDIA hosts steering work to the AMD card); neither probe honours
that env var, so check it first and report the GPU as not usable. Unset
means all devices visible.
"""
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None and cvd.strip() in ("", "-1"):
return False
exe = shutil.which("nvidia-smi")
if exe:
try:
result = subprocess.run(
[exe, "-L"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
if result.returncode == 0 and "GPU " in result.stdout:
return True
except Exception:
pass
# Fallback: the NVIDIA driver exposes one subdirectory per GPU under
# /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state.
if sys.platform != "win32":
try:
gpu_dir = "/proc/driver/nvidia/gpus"
if os.path.isdir(gpu_dir) and os.listdir(gpu_dir):
return True
except OSError:
pass
return False
def _detect_amd_gfx_codes() -> list[str]:
"""Return the AMD gfx ISA strings visible to ROCm (e.g. ['gfx1151']).
Probes rocminfo, then falls back to ``amd-smi list`` and ``amd-smi
static --asic`` for runtime-only Radeon hosts that ship amd-smi but no
rocminfo. Returns an empty list when no probe yields a gfx target.
"""
def _extract(text: str) -> list[str]:
codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower())
return list(dict.fromkeys(f"gfx{c}" for c in codes))
probes: list[list[str]] = []
if shutil.which("rocminfo"):
probes.append(["rocminfo"])
# Gate amd-smi off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt).
if shutil.which("amd-smi") and _amd_smi_allowed():
probes.append(["amd-smi", "list"])
probes.append(["amd-smi", "static", "--asic"])
for cmd in probes:
try:
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 15,
env = _amd_smi_env() if cmd[0] == "amd-smi" else None,
)
except Exception:
continue
if result.returncode != 0 or not result.stdout.strip():
continue
codes = _extract(result.stdout)
if codes:
return codes
return []
# Set by _ensure_rocm_torch() on success; suppresses the post-install AMD warning.
_rocm_windows_torch_installed: bool = False
def _install_bnb_windows_rocm() -> bool:
"""Install the AMD Windows BNB prerelease wheel. Returns True on success.
The continuous-release wheel is intentionally mismatched: the filename
encodes 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the wheel
metadata reports 0.50.0.dev0. uv rejects this filename/metadata mismatch,
and bypassing it with UV_SKIP_WHEEL_FILENAME_CHECK still leaves uv mangling
the bitsandbytes install. Per the AMD install guide
(https://unsloth.ai/docs/get-started/install/amd/amd-hackathon) the wheel
must be installed with plain pip, not uv, so we force pip (force_pip=True);
plain pip performs no wheel filename/metadata check.
"""
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is None:
return False
_ok = pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
force_pip = True,
)
if not _ok:
return False
# After install: detect the actual ROCm DLL suffix shipped in the wheel and
# set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of
# what torch.version.hip reports. The wheel may ship an older suffix (e.g.
# "72") while torch reports a newer HIP version (e.g. 7.13); the env var
# override ensures bitsandbytes does not fail looking for a non-existent DLL.
# The worker subprocess inherits this env var automatically.
# Fall back to "72" if detection fails (e.g. install was a no-op / dry-run).
_env_ver = os.environ.get("BNB_ROCM_VERSION")
_env_is_persisted_default = (
os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE
)
_persist_detected_version = False
if _env_ver and not _env_is_persisted_default:
_ver = _env_ver
else:
_ver = _detect_bnb_rocm_dll_ver() or "72"
os.environ["BNB_ROCM_VERSION"] = _ver
os.environ[_BNB_ROCM_VERSION_SOURCE_ENV] = _BNB_ROCM_VERSION_SOURCE_DETECTED
_persist_detected_version = True
if _persist_detected_version:
_persist_bnb_rocm_version(_ver)
# Make hipInfo.exe (shipped into the venv Scripts dir by the AMD torch
# wheel) resolvable via PATH for this process and every child python the
# installer spawns (import checks, precompile): bitsandbytes runs
# `hipinfo.exe` at import time to detect the GPU arch and logs a scary
# (harmless) ERROR + WARNING on every import when it is missing. The venv
# Scripts dir is on PATH only when the venv is activated, which neither
# Studio nor the installer's child processes ever do.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which(
"hipinfo.exe"
):
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
return True
def _detect_cuda_torch_index_url() -> str:
"""Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver.
Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
repairs to the same wheel family a fresh `curl | sh` install would pick.
Honours the same explicit overrides first (UNSLOTH_TORCH_INDEX_URL /
UNSLOTH_TORCH_INDEX_FAMILY) so a headless / container / CI install never lets
the host GPU decide the wheel family. Otherwise probes nvidia-smi (PATH, then
/usr/bin/nvidia-smi) and parses both the legacy "CUDA Version:" and the newer
"CUDA UMD Version:" spellings. Defaults to cu126 when nvidia-smi is missing or
the version is unreadable (e.g. NVIDIA detected only via the
/proc/driver/nvidia/gpus fallback).
"""
_override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
if _override_url:
return _override_url.rstrip("/")
_override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
if _override_family:
return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}"
exe = shutil.which("nvidia-smi")
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
exe = "/usr/bin/nvidia-smi"
tag = "cu126" # default when the driver CUDA version cannot be read
if exe:
try:
result = subprocess.run(
[exe],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
if result.returncode == 0:
m = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", result.stdout)
if m:
major, minor = int(m.group(1)), int(m.group(2))
if major >= 13:
tag = "cu130"
elif major == 12 and minor >= 8:
tag = "cu128"
elif major == 12 and minor >= 6:
tag = "cu126"
elif major >= 12:
tag = "cu124"
elif major >= 11:
tag = "cu118"
else:
tag = "cpu" # ancient driver: no usable CUDA wheels
except Exception:
pass
return f"{_PYTORCH_WHL_BASE}/{tag}"
def _explicit_torch_index_url() -> "str | None":
"""The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None.
Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead
of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override.
"""
url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
if url:
return url.rstrip("/")
family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
if family:
return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}"
return None
def _is_pip_rocm_family_leaf(leaf: str) -> bool:
"""True when a lowercased index leaf names a pip --index-url ROCm family: an
actual rocm<digit>... leaf (download.pytorch.org/whl/rocm7.2) or a repo.amd.com
per-arch gfx leaf (gfx120x-all). A Radeon find-links directory leaf
(repo.radeon.com/.../rocm-rel-7.2.1, which install.sh records in the marker)
starts with "rocm" but is NOT a pip index -- it must route to the verbatim/marker
path, not a --index-url reinstall that fails against a find-links listing.
Mirrors install.sh's rocm[0-9]* / setup.ps1's ^(rocm[0-9]|gfx) gate. Pure function.
"""
return bool(re.match(r"^rocm\d", leaf)) or leaf.startswith("gfx")
def _explicit_rocm_torch_index_url() -> "str | None":
"""The pinned wheel index URL when it names a pip ROCm family (rocm<d>/gfx*), else None."""
url = _explicit_torch_index_url()
if url is None:
return None
leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
return url if _is_pip_rocm_family_leaf(leaf) else None
def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
"""True when an explicit ROCm pin names a different ROCm family than the
already-installed ROCm torch, so the pin needs a reinstall to be applied.
Mirrors setup.ps1's stale-venv ROCm comparison. The pin leaf classifies into
three cases, matching the install-spec path in _ensure_rocm_torch:
* rocmX.Y leaf -> compare the pinned rocm version to the installed one
exactly when both are readable; else fall back to the torch 2.11 line
(rocm>=7.2 serves 2.11, older rocm does not).
* gfx leaf in _ROCM_GFX_TORCH211_LEAVES (gfx120x-all/gfx1151/gfx1150) ->
the install path pulls AMD's per-arch wheel (tagged with a THREE-part
+rocmA.B.C local version, e.g. 2.11.0+rocm7.13.0). A generic pytorch.org
rocm wheel (two-part +rocmA.B, e.g. +rocm7.2) or any pre-2.11 build is a
mismatch even when both are torch 2.11 -- the user asked for the per-arch
index. An already-installed per-arch wheel (three-part tag) is NOT a
mismatch, so a satisfied gfx pin does not reinstall-loop.
* gfx leaf NOT in the 2.11 allowlist (gfx110X-all/gfx90a/gfx908) -> the
install path uses the default <2.11 specs, so a correct 2.10+rocm wheel
must NOT be flagged. Mismatch only when the installed torch is 2.11+.
A pin that resolves to the same family as what is installed is NOT a mismatch,
so a correct ROCm venv is never needlessly reinstalled. Pure function.
"""
leaf = pin_url.rstrip("/").rsplit("/", 1)[-1].lower()
# Pinned ROCm version (from a rocmX.Y leaf).
_pin_rocm = re.match(r"^rocm(\d+)\.(\d+)", leaf)
_pin_ver = (int(_pin_rocm.group(1)), int(_pin_rocm.group(2))) if _pin_rocm else None
# Installed ROCm version (+rocmX.Y) and whether the installed wheel carries a
# THREE-part local version (+rocmA.B.C) -- the AMD per-arch signature that
# distinguishes a repo.amd.com/gfx* wheel from a two-part pytorch.org one.
_inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver)
_inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None
_inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None
# A ROCm build MUST carry a +rocm local tag. An untagged CPU/CUDA wheel (no
# +rocm, e.g. "2.10.0" / "2.11.0") never satisfies a ROCm pin -- always a
# mismatch -- mirroring setup.ps1's Get-RocmPinStaleTags. (In practice
# _ensure_rocm_torch only calls this when has_hip_torch is True, but keep the
# pure function correct for any input so it stays in lockstep with the PS side.)
_inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None
# Whether the installed torch RELEASE (before "+") is 2.11+.
_inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver)
_inst_is_211 = (
(int(_inst_rel.group(1)), int(_inst_rel.group(2))) >= (2, 11) if _inst_rel else False
)
if leaf.startswith("gfx"):
# gfx per-arch pin: only the _grouped_mm-bug arches (the 2.11 allowlist)
# pull the AMD per-arch wheel; other gfx leaves stay on the default
# <2.11 specs (see _ROCM_TORCH_PKG_SPECS selection below).
if leaf in _ROCM_GFX_TORCH211_LEAVES:
# Expect the AMD per-arch wheel (three-part +rocmA.B.C, torch 2.11+).
# A satisfied per-arch install is NOT a mismatch (no reinstall loop);
# a generic rocm wheel or any pre-2.11 build IS a mismatch even at 2.11.
return not (_inst_is_211 and _inst_is_perarch)
# Non-2.11 gfx leaf: install path uses default <2.11 specs, so a correct
# <2.11 wheel must stay. An untagged (no +rocm) wheel never satisfies the
# pin -> mismatch; otherwise mismatch only when the installed torch is 2.11+.
return (not _inst_has_rocm) or _inst_is_211
# rocmX.Y pin. Only the KNOWN-2.11 rocm indexes are treated as the 2.11 line;
# an unknown newer rocm (rocm7.3, rocm8.0, ...) is NOT floored speculatively
# (it does not exist yet). Aligns with install.sh / setup.ps1 / install.ps1.
_pin_is_211 = _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False
if _pin_ver is not None and _inst_ver is not None:
# Both ROCm versions readable: exact (major, minor) comparison. A generic
# rocm7.2 pin over the AMD per-arch (+rocm7.13.x) wheel compares (7, 2) vs
# (7, 13) -> mismatch, which correctly reinstalls the generic wheel the
# user pinned instead of leaving the per-arch one in place.
return _pin_ver != _inst_ver
# rocm pin with an unreadable installed version: compare on the torch 2.11 line,
# but an untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> mismatch.
if not _inst_has_rocm:
return True
return _pin_is_211 != _inst_is_211
def _explicit_cpu_torch_index_url() -> "str | None":
"""The pinned wheel index URL when it names the CPU family (leaf == cpu), else None.
An explicit CPU pin (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu)
is authoritative -- see _ensure_cpu_torch.
"""
url = _explicit_torch_index_url()
if url is None:
return None
leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
return url if leaf == "cpu" else None
def _is_cuda_family_leaf(leaf: str) -> bool:
"""True only for a real CUDA wheel-family leaf: "cu" followed by digits
(cu118, cu126, cu128, cu130, ...).
A bare startswith("cu") wrongly matches arbitrary mirror leaves like "custom"
or "current", which would let _ensure_cuda_torch treat a generic mirror pin as
CUDA authority and force a CUDA reinstall over a CPU/ROCm venv on a non-NVIDIA
host -- exactly what _explicit_cuda_torch_index_url's contract forbids.
"""
return re.match(r"^cu[0-9]", leaf) is not None
def _explicit_cuda_torch_index_url() -> "str | None":
"""The pinned wheel index URL when it names a CUDA family (leaf cuXXX), else None.
Mirrors _explicit_rocm/cpu_torch_index_url so _ensure_cuda_torch only treats a
*CUDA* pin as authority to override the NVIDIA-presence gate. An arbitrary
mirror URL (or a ROCm/CPU pin) must not force a CUDA reinstall over a working
ROCm/CPU venv on a non-NVIDIA host, so match cuXXX (cu + digits) narrowly
rather than any leaf starting with "cu" (which would catch custom/current).
"""
url = _explicit_torch_index_url()
if url is None:
return None
leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
return url if _is_cuda_family_leaf(leaf) else None
def _explicit_unknown_family_torch_index_url() -> "str | None":
"""The pinned index URL when its leaf names NO known torch family, else None.
A "known" family leaf is rocm* / gfx* / cpu / cuXXX -- the ones the dedicated
_explicit_{rocm,cpu,cuda}_torch_index_url helpers already classify. Anything
else (a private PEP 503 mirror ending in /simple, /current, /custom, ...) is an
UNKNOWN family: the version-tag heuristics cannot infer whether it is stale, so
the marker drives the decision and, when it differs (or is absent), the URL is
reinstalled VERBATIM -- "URL wins verbatim". Matches the unknown-leaf branch in
install.sh / setup.ps1 / install.ps1.
"""
url = _explicit_torch_index_url()
if url is None:
return None
leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
if _is_pip_rocm_family_leaf(leaf) or leaf == "cpu" or _is_cuda_family_leaf(leaf):
return None
return url
def _ensure_verbatim_torch_index() -> None:
"""Reinstall torch/vision/audio VERBATIM from an explicit custom index pin.
Handles the "URL wins verbatim" case for an explicit UNSLOTH_TORCH_INDEX_URL
(or _FAMILY) whose leaf names no known family (e.g. a private mirror ending in
/simple or /current). The other _ensure_* helpers all return None for such a
pin (it is neither rocm/gfx nor cpu nor cuXXX), so without this the pin would be
silently ignored and the GPU-probed default index used instead.
Fires when the marker differs from the pin (True) OR is ABSENT (None): a venv
predating the marker feature has no record, and the version-tag heuristics
cannot judge an unknown-family pin, so an explicitly-set URL would otherwise be
silently ignored on the first `studio update` -- the user asked for this index,
so apply it verbatim ONCE and record it. The write below makes every later
update a no-op (marker == pin -> False). Skips only when the marker already
records this exact pin (False). A user who did NOT set the override gets
pin=None and is never touched, so an out-of-band torch install is safe. macOS/
no-torch: skipped (no torch to repair). The install uses the pinned URL
exclusively (--index-url) with bare specs so it "wins verbatim" -- an incomplete
mirror that cannot serve the trio fails loudly here, same as the marker-present
path (that is the cost of honouring an explicit pin).
"""
if NO_TORCH or IS_MACOS:
return
pin = _explicit_unknown_family_torch_index_url()
if pin is None:
return
_mismatch = _marker_pin_mismatch(pin)
if _mismatch is False:
# Marker already records this exact pin -> no reinstall (no per-update loop).
# True (marker differs) or None (no marker yet) both fall through to apply
# the explicit pin verbatim once, then _write_torch_index_marker below makes
# the next update a no-op.
return
_why = "differs from the recorded index" if _mismatch is True else "has no recorded index yet"
print(f" explicit torch index pin ({pin}) {_why} -- reinstalling torch verbatim from it")
pip_install(
"torch (pinned custom index)",
"--force-reinstall",
"--no-cache-dir",
"torch",
"torchvision",
"torchaudio",
"--index-url",
pin,
constrain = False,
)
_write_torch_index_marker(pin)
def _ensure_cuda_torch() -> None:
"""Repair a venv whose torch is a ROCm build on an NVIDIA host.
Counterpart to _ensure_rocm_torch. A venv poisoned by the pre-fix KFD
gpu_id false positive (ROCm torch installed on an NVIDIA-only machine)
keeps that broken torch on `studio update`, because a torch+rocm wheel
satisfies the version constraint and nothing force-reinstalls it. This
detects that exact case and reinstalls CUDA torch.
Only repairs when torch actually links against HIP/ROCm. Healthy CUDA
torch and deliberate CPU-only torch are left untouched.
"""
# Respect an explicit backend choice from install.sh: only "" (standalone
# `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu"
# (or any unrecognised value) are deliberate and must not be overridden.
if _TORCH_BACKEND not in ("", "cuda"):
return
# An explicit custom-index pin whose leaf names no known torch family wins
# VERBATIM (_ensure_verbatim_torch_index applies it); do not override it with
# auto-detected CUDA wheels here.
if _explicit_unknown_family_torch_index_url() is not None:
return
# No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
# install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
if IS_MACOS or IS_WINDOWS or NO_TORCH:
return
# Never undo a deliberate ROCm install (setup.ps1 sets this marker).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
return
# An explicit CUDA wheel-index pin (headless / container / CI cross-install)
# commits to CUDA wheels and, like install.sh's get_torch_index_url override,
# skips ALL host-GPU probing -- so it must clear BOTH the CUDA_VISIBLE_DEVICES
# hide gate and the NVIDIA-presence gate below, not just the latter. Otherwise
# `CUDA_VISIBLE_DEVICES=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update`
# (the exact GPU-less CI case this override targets) would still bail here.
_cuda_pinned = _explicit_cuda_torch_index_url() is not None
# CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for
# example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card);
# never force CUDA wheels over that choice unless a CUDA index is pinned.
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"):
return
# Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
# covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
# The explicit CUDA pin overrides the GPU-presence gate too.
if not _cuda_pinned and not _has_usable_nvidia_gpu():
return
# Classify the installed torch: "hip" (ROCm build -- the poisoning
# signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A
# non-zero exit means torch is missing or un-importable; the base install
# step handles that, so leave it alone.
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch, re; "
"hip = getattr(torch.version, 'hip', '') or ''; "
"cuda = getattr(torch.version, 'cuda', '') or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
"m = re.search(r'\\+(cu\\d+)', ver); "
"marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); "
"print(marker + '|' + (m.group(1) if m else ''))"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
return
# Take the last non-empty stdout line: stray output from sitecustomize or
# an import hook must not mask the marker (fail-closed either way).
_marker_lines = [
line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
]
if not _marker_lines:
return
_marker, _, _installed_cu = _marker_lines[-1].partition("|")
# Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host
# (the poisoning signature), or when an explicit CUDA index is pinned but the
# venv has the wrong family -- a CPU wheel, or a different cuXXX than pinned.
# This covers the headless cross-install (`studio update` with
# UNSLOTH_TORCH_INDEX_FAMILY=cu128): the update path preserves torch rather
# than preinstalling it from install.sh, so without this an explicit CUDA pin
# stays ineffective. A healthy CUDA torch matching the pin, or a CPU wheel
# with no CUDA pin, is deliberate and left alone.
_pin = _explicit_torch_index_url()
_pin_leaf = _pin.rstrip("/").rsplit("/", 1)[-1].lower() if _pin else ""
_pinned_cuda = _is_cuda_family_leaf(_pin_leaf)
if _marker == "hip":
_why = "torch is a ROCm build on an NVIDIA host"
elif _marker == "cpu" and _pinned_cuda:
_why = "torch is a CPU build but an explicit CUDA index is pinned"
elif _marker == "cuda" and _pinned_cuda and _installed_cu != _pin_leaf:
# Mismatch when the installed cuXXX differs from the pin. An UNTAGGED cuda
# build (empty _installed_cu -- e.g. torch re-resolved from default PyPI
# into a CUDA wheel with no +cuXXX local tag) also counts: the family
# cannot be confirmed to match the pin, so reinstall to enforce it. The
# reinstall targets the pinned family, so an already-matching untagged
# build simply re-lands on the same family (idempotent).
_installed_desc = _installed_cu if _installed_cu else "an untagged CUDA build"
_why = f"torch is {_installed_desc} but the pinned CUDA index is {_pin_leaf}"
elif _marker == "cuda" and _pinned_cuda and _marker_pin_mismatch(_pin) is True:
# Same cuXXX leaf but the marker records a DIFFERENT full index URL (e.g. the
# official cu128 index vs an internal mirror's cu128). The +cuXXX tag cannot
# see the host change, so consult the exact-URL marker and reinstall from the
# pinned URL (via _detect_cuda_torch_index_url, which honours the override) so
# an explicit mirror pin is applied and re-recorded, not skipped.
_why = f"the pinned CUDA index URL differs from the recorded marker (leaf {_pin_leaf})"
else:
return # healthy CUDA torch matching the pin, or a deliberate CPU wheel
index_url = _detect_cuda_torch_index_url()
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
print(
f" {_why} -- reinstalling CUDA torch from {index_url}\n"
f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate "
f"non-CUDA torch)"
)
pip_install(
"CUDA torch repair",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
_write_torch_index_marker(index_url)
def _ensure_cpu_torch() -> None:
"""Reinstall CPU torch when an explicit CPU pin is set but the venv has a GPU build.
Counterpart to _ensure_cuda_torch / _ensure_rocm_torch for the explicit-CPU
case (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu). Those helpers
both treat a CPU backend as a skip signal, so on a standalone `unsloth studio
update` -- which does not run install.sh's post-install flavor enforcement --
an existing CUDA/ROCm torch satisfies the version constraint and is never
replaced, ignoring the authoritative CPU pin. Only fires for an EXPLICIT pin:
a CPU backend that came from auto-detection (genuine CPU host via install.sh)
already installed CPU wheels, so there is nothing to repair.
"""
if NO_TORCH:
return
pin = _explicit_cpu_torch_index_url()
if pin is None:
return
# Classify the installed torch family. A non-zero exit means torch is missing
# or un-importable; the base install step handles that, so leave it alone.
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch, re; "
"hip = getattr(torch.version, 'hip', '') or ''; "
"cuda = getattr(torch.version, 'cuda', '') or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
"gpu = bool(hip) or 'rocm' in ver or bool(cuda) or bool(re.search(r'\\+cu\\d+', ver)); "
"print('gpu' if gpu else 'cpu')"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
return
_lines = [
line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
]
if not _lines:
return # unreadable -- the base install step handles a missing torch
if _lines[-1] != "gpu":
# torch is already a CPU build. Normally nothing to repair, BUT a standalone
# update may change the CPU index URL itself (official /cpu -> a private
# UNSLOTH_PYTORCH_MIRROR /cpu) with the same +cpu tag. The tag cannot see the
# host change, so consult the exact-URL marker and reinstall from the new pin
# only when it records a DIFFERENT index -- mirroring the CUDA/ROCm
# same-family-URL handling. No marker (or a matching one) -> leave it alone.
if _marker_pin_mismatch(pin) is not True:
return
_why = "the pinned CPU index URL differs from the recorded marker"
else:
_why = "torch is a GPU build but an explicit CPU index is pinned"
print(f" {_why} -- reinstalling CPU torch from {pin}")
# Pin to the supported torch<2.11 family (same bounds as the CUDA/ROCm repair
# specs). The /cpu index now also serves torch 2.11+, so a bare trio off the
# exclusive --index-url could resolve outside the supported range or drag in
# an ABI-mismatched torchvision/torchaudio.
_torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
pip_install(
"CPU torch repair",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
pin,
constrain = False,
)
_write_torch_index_marker(pin)
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
On Linux x86_64: uses pytorch.org ROCm wheel index tags.
On Windows: uses AMD's repo.amd.com arch-specific pip index.
No-op on macOS, non-x86_64 Linux, NVIDIA-primary hosts, or when torch
already links against HIP.
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family
# ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh
# already selected a non-ROCm backend -- this is the authoritative signal
# and avoids re-running GPU detection in a subprocess that may see a
# different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
if _TORCH_BACKEND in ("cuda", "cpu"):
return
# An explicit custom-index pin whose leaf names no known torch family (a private
# PEP 503 mirror, /simple, /current, ...) wins VERBATIM -- _ensure_verbatim_torch_index()
# applies it. Never override such a pin with the auto-detected ROCm index here: the
# user chose that index (a valid cross-install / custom-mirror case), and installing
# ROCm wheels over it before the verbatim pass runs would leave the URL override
# unhonored (the marker then matches the pin, so verbatim skips the restore).
if _explicit_unknown_family_torch_index_url() is not None:
return
# setup.ps1 sets this after installing AMD wheels; skip the probe only when
# torch is actually importable as ROCm. If the venv was wiped between runs,
# the stale env-var would suppress a needed reinstall.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_torch_ok = False
try:
_probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"import sys; "
"sys.exit(0 if (hip or 'rocm' in torch.__version__.lower()) else 1)"
),
],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
timeout = 90,
)
_torch_ok = _probe.returncode == 0
except (OSError, subprocess.TimeoutExpired):
pass
if _torch_ok:
_rocm_windows_torch_installed = True
# setup.ps1 already installed ROCm torch, but we still need the AMD
# Windows BNB wheel here -- the PyPI bitsandbytes wheel ships only
# CUDA DLLs and fails to load on ROCm.
_install_bnb_windows_rocm()
return
# torch was wiped between runs; fall through to the full install path
if IS_MACOS:
return
if IS_WINDOWS:
if _has_usable_nvidia_gpu():
return
gfx_arch = _detect_windows_gfx_arch()
if not gfx_arch:
return # no AMD GPU visible via hipinfo
# Probe whether torch already links against HIP.
_torch_already_rocm = False
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=torch.__version__; "
"print('yes' if hip or 'rocm' in ver.lower() else '')"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
if probe.returncode == 0 and probe.stdout.decode().strip() == "yes":
_torch_already_rocm = True
except (OSError, subprocess.TimeoutExpired):
pass
if not _torch_already_rocm:
index_url = _windows_rocm_index_url(gfx_arch)
if index_url is None:
print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping")
return
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
# Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X /
# Strix) so the per-arch index resolves an ABI-consistent trio; other
# arches stay bare (no published floor), matching the PowerShell side.
_torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get(
gfx_arch, ("torch", "torchvision", "torchaudio")
)
# Nonfatal: a transient AMD-index failure must not abort the whole
# install once the PowerShell side has fallen back to CPU torch.
# --force-reinstall resolves before uninstalling, so a failed index
# leaves the existing build intact; keep it and let the user retry.
if not pip_install_try(
f"ROCm torch (Windows, {gfx_arch})",
"--force-reinstall",
"--index-url",
index_url,
_torch_pkg,
_vision_pkg,
_audio_pkg,
constrain = False,
):
print(
f" Warning: AMD Windows ROCm torch install failed for {gfx_arch}; "
"keeping the existing torch build. Re-run 'unsloth studio update' "
"later to retry ROCm."
)
return
_write_torch_index_marker(index_url)
# ROCm torch is installed (or already was); flag it so later phases
# do not overwrite it with the generic CPU torch wheel. BNB is a
# separate dependency -- a BNB install failure must NOT roll back the
# torch ROCm install.
_rocm_windows_torch_installed = True
# Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
# CUDA DLLs and fails on ROCm. Install even when torch was already a
# ROCm build so `studio update` repairs a broken bnb.
if not _install_bnb_windows_rocm():
print(
" Warning: AMD Windows bitsandbytes install failed; "
"ROCm torch is installed but bitsandbytes may need manual install"
)
return
# ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──
if platform.machine().lower() not in {"x86_64", "amd64"}:
return
# An explicit ROCm wheel-index pin (UNSLOTH_TORCH_INDEX_URL/_FAMILY naming a
# rocm*/gfx* leaf) commits to ROCm wheels regardless of which GPU is visible
# here -- the headless / container / CI cross-install case. Mirror
# _ensure_cuda_torch's explicit-pin bypass: skip the NVIDIA-present / no-AMD /
# unreadable-ROCm gates so the pinned index is honoured. Without this, a
# standalone `studio update` with an explicit ROCm pin on an NVIDIA-only or
# GPU-less box returned here and left the CPU/CUDA torch in place.
_rocm_pin = _explicit_rocm_torch_index_url()
if _rocm_pin is None:
# NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable.
if _has_usable_nvidia_gpu():
return
# Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the
# authoritative "is this an AMD ROCm host?" signal. The old gate required
# /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs
# (minimal package-managed installs, Radeon software) that ship
# amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio
# update` unable to repair a CPU-only venv on those systems.
if not _has_rocm_gpu():
return # no AMD GPU visible
ver = _detect_rocm_version()
if ver is None:
if _rocm_pin is None:
print(" ROCm detected but version unreadable -- skipping torch reinstall")
return
# Explicit pin: the pinned index leaf (not the host ROCm version) drives
# the install below, so a missing/unreadable host ROCm version is fine.
# Use a sentinel so the version-gated Strix reroute (skipped for explicit
# pins anyway) and any ver comparisons stay well-defined.
ver = (0, 0)
# Probe whether torch already links against HIP (ROCm already working), and
# capture the installed ROCm build tag so a pin mismatch can be detected.
# Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts
# (the NVIDIA check above already handled mixed AMD+NVIDIA setups).
# Emit ONE "<hip_marker>|<version>" line (mirrors _ensure_cuda_torch) so the
# parse is positional and robust: the HIP marker is the field before "|" (HIP
# version, "rocm" sentinel, or empty for CPU/CUDA torch), the installed wheel
# version (e.g. "2.10.0+rocm6.4") is the field after it. Do NOT filter empty
# LINES and take slot 0 -- for CPU/CUDA torch the marker field IS empty, and
# dropping it would shift the version into slot 0 and wrongly flag has_hip_torch.
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=getattr(torch,'__version__','').lower(); "
# HIP version when present (back-compat), else a "rocm"
# sentinel when only torch.__version__ flags ROCm (AMD SDK /
# Radeon wheels). Empty marker before "|" = CPU/CUDA torch.
"marker=hip if hip else ('rocm' if 'rocm' in ver else ''); "
"print(marker + '|' + ver)"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
probe = None
# Take the last non-empty stdout line so stray sitecustomize / import-hook
# output cannot mask the marker; then split positionally on the FIRST "|" --
# the empty HIP-marker field for CPU/CUDA torch is preserved (has_hip_torch
# is driven by that field, not by "first non-empty line").
_marker_lines = (
[ln.strip() for ln in probe.stdout.decode(errors = "replace").splitlines() if ln.strip()]
if (probe is not None and probe.returncode == 0)
else []
)
_hip_marker, _sep, _installed_torch_ver = (
_marker_lines[-1].partition("|") if _marker_lines else ("", "", "")
)
# A "|"-delimited marker line is required: without the separator the probe
# output is unrecognised (old torch, injected noise), so treat HIP as absent
# and fall through to a reinstall rather than trusting an ambiguous string.
has_hip_torch = bool(_sep) and _hip_marker != ""
# An explicit ROCm pin whose family differs from the already-installed ROCm
# torch must reinstall, mirroring _ensure_cuda_torch (installed cuXXX != pin).
# Without this, `studio update` with UNSLOTH_TORCH_INDEX_FAMILY=rocm7.2 (or a
# gfx* URL) on a venv that already carries an OLDER ROCm build (+rocm6.4 /
# +rocm7.1) short-circuits on has_hip_torch and never applies the override.
#
# Prefer the torch-index MARKER when present: an EXACT compare of the pinned
# index against the index the last install recorded. This is the ONLY signal
# that catches a per-arch switch between two 2.11 gfx indexes (gfx1151 ->
# gfx120X-all) -- both install a +rocm7.13.0 wheel, so the version-tag
# heuristic below sees no difference and would leave the old arch in place. A
# matching marker also guarantees a correctly-pinned venv does NOT reinstall
# (no loop). When there is NO marker (old venv, or torch installed out-of-band)
# fall back to the +rocm/version-tag heuristic -- backward compatibility.
_rocm_pin_mismatch = False
if has_hip_torch and _rocm_pin is not None:
_marker_verdict = _marker_pin_mismatch(_rocm_pin)
if _marker_verdict is True:
# Marker records a DIFFERENT index than the pin -> reinstall. This is the
# only signal that catches a per-arch gfx switch (gfx1151 -> gfx120X-all,
# both +rocm7.13.0) the version-tag heuristic below cannot see.
_rocm_pin_mismatch = True
else:
# Marker matches OR is absent: the marker is an ADDITIONAL rebuild signal,
# not a substitute for validating the installed wheel. Still run the
# family/version check so a stale wheel (torch swapped after the marker was
# written -- e.g. marker records gfx1151 but the venv now carries generic
# +rocm7.2 or an older +rocm6.4) is caught. Mirrors setup.ps1, which keeps
# the flavor check alongside the marker for this reason.
_rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
# Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1
# in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0
# with the real fix, so route those hosts there instead of the generic
# pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override.
# On mixed hosts (Strix iGPU + non-Strix dGPU), route to the AMD per-gfx
# index only when HIP's runtime GPU is the Strix one -- else the dGPU gets
# an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
# An explicit ROCm wheel-index pin is authoritative: never auto-reroute it to
# the AMD per-gfx index (the caller already chose the family/URL).
if ver < (7, 2) and _explicit_rocm_torch_index_url() is None:
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
if _detected_strix:
# Pick the runtime-visible GPU: use the HIP_VISIBLE_DEVICES index
# into gfx_codes, else default to the first GPU. Skip the override
# unless the resolved GPU is Strix.
_runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
if _runtime_gfx in _strix_gfx:
_selected_gfx = _runtime_gfx
_amd_mirror = (
os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") or "https://repo.amd.com/rocm/whl"
).rstrip("/")
_strix_override_url = f"{_amd_mirror}/{_selected_gfx}/"
_strix_override_pkgs = (
"torch>=2.11.0,<2.12.0",
# Pin torchvision/torchaudio to the 2.11.x-compatible range.
# The install uses --index-url (exclusive, no PyPI fallback),
# so bare unversioned names risk resolving an AMD-index build
# targeting a different torch major (e.g. 0.27 built against
# torch 2.12), which fails at runtime with an ABI/version
# mismatch. Matches _ROCM_TORCH_CONSTRAINT["rocm7.2"].
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
)
print(
f"\n {_selected_gfx} (AMD Strix) is the runtime target with ROCm "
f"{ver[0]}.{ver[1]}.\n"
f" ROCm 7.1 has a known _grouped_mm segfault on this GPU;\n"
f" routing torch install to AMD's arch-specific index\n"
f" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n"
f" with the upstream fix.\n"
)
else:
_gfx_str = ", ".join(sorted(_detected_strix))
print(
f"\n Strix GPU ({_gfx_str}) present but HIP_VISIBLE_DEVICES "
f"selects a non-Strix runtime target ({_runtime_gfx});\n"
f" skipping AMD per-gfx index override.\n"
)
# Strix override on ROCm 7.1 must fire even when has_hip_torch is True --
# an existing torch with `torch.version.hip == "7.1"` is exactly the broken
# combo the override repairs, so skipping it leaves users on the known
# _grouped_mm segfault.
if _strix_override_url is not None and _strix_override_pkgs is not None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
print(f" Strix ROCm 7.1 override -- installing torch from {index_url}")
pip_install(
"ROCm torch (Strix arch-specific)",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
_write_torch_index_marker(index_url)
rocm_torch_ready = True
elif not has_hip_torch or _rocm_pin_mismatch:
# Reinstall when torch is not ROCm yet, OR when a ROCm build is present but
# its family differs from an explicit pin (_rocm_pin_mismatch -- the ROCm
# analogue of _ensure_cuda_torch's installed-cuXXX != pin reinstall).
# Honour an explicit ROCm wheel-index pin verbatim instead of re-detecting
# the host ROCm version; otherwise select the best wheel tag (newest ROCm
# version <= installed). gfx*/rocm7.2 indexes serve torch 2.11+, so match
# the constraints to the index leaf when overridden.
_override_idx = _explicit_rocm_torch_index_url()
if _override_idx is not None:
index_url = _override_idx
tag = index_url.rstrip("/").rsplit("/", 1)[-1].lower()
else:
tag = next(
(
t
for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
if ver >= (maj, mn)
),
None,
)
if tag is None:
print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall")
else:
if _override_idx is None:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm torch -- installing from {index_url}")
# Only the gfx arches with the _grouped_mm bug (gfx120X-all, gfx1151,
# gfx1150) need the torch 2.11 spec; other gfx per-arch indexes
# (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels, so a pinned
# override to one of those stays on the default range. Matches the
# gfx floor gating in install.ps1 / setup.ps1.
if tag in _ROCM_GFX_TORCH211_LEAVES:
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"]
elif tag.startswith("gfx"):
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["_default"]
else:
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
tag, _ROCM_TORCH_PKG_SPECS["_default"]
)
pip_install(
f"ROCm torch ({tag})",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
_write_torch_index_marker(index_url)
rocm_torch_ready = True
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix), falling back
# to PyPI when the pre-release wheel won't install. Use pip for the
# pre-release wheel because uv rejects its filename/metadata version mismatch.
if rocm_torch_ready:
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
if _bnb_url is not None:
_bnb_installed = pip_install_try(
"bitsandbytes (AMD, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_url,
constrain = False,
force_pip = True,
)
if not _bnb_installed:
print(
_red(
" bnb pre-release install failed; falling back to PyPI "
"(4-bit decode will be broken on ROCm)"
)
)
if not _bnb_installed:
pip_install(
"bitsandbytes (AMD)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_BNB_ROCM_PYPI_FALLBACK,
constrain = False,
)
# _uv_safe_path is imported from backend.utils.uv_path_safety (shared with mlx_repair).
def _windows_hidden_subprocess_kwargs() -> dict[str, object]:
"""Return Windows-only subprocess kwargs that suppress console windows."""
if not IS_WINDOWS:
return {}
kwargs: dict[str, object] = {}
create_no_window = getattr(subprocess, "CREATE_NO_WINDOW", 0)
if create_no_window:
kwargs["creationflags"] = create_no_window
startupinfo_factory = getattr(subprocess, "STARTUPINFO", None)
startf_use_showwindow = getattr(subprocess, "STARTF_USESHOWWINDOW", 0)
sw_hide = getattr(subprocess, "SW_HIDE", 0)
if startupinfo_factory is not None and startf_use_showwindow:
startupinfo = startupinfo_factory()
startupinfo.dwFlags |= startf_use_showwindow
startupinfo.wShowWindow = sw_hide
kwargs["startupinfo"] = startupinfo
return kwargs
def _infer_no_torch() -> bool:
"""Determine whether to run in no-torch (GGUF-only) mode.
Checks UNSLOTH_NO_TORCH first. When unset, falls back to platform
detection so Intel Macs use GGUF-only mode even when invoked from
``unsloth studio update`` (which does not inject the env var).
"""
env = os.environ.get("UNSLOTH_NO_TORCH")
if env is not None:
return env.strip().lower() in ("1", "true")
return IS_MAC_INTEL
NO_TORCH = _infer_no_torch()
# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so
# that this script knows which torch variant was selected without re-running
# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
# (standalone `unsloth studio update` runs, where we re-detect normally).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
# When install.sh did not run (standalone `unsloth studio update`) but the caller
# pinned the wheel index explicitly, derive the backend from that override so the
# CUDA/ROCm repair helpers honour it instead of re-probing the GPU and possibly
# reinstalling a different family. Classify on the final URL/family segment,
# mirroring install.sh's UNSLOTH_TORCH_BACKEND case.
if not _TORCH_BACKEND:
_idx_override = (
os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
)
_idx_leaf = _idx_override.rstrip("/").rsplit("/", 1)[-1].lower()
if _idx_leaf.startswith(("rocm", "gfx")):
_TORCH_BACKEND = "rocm"
elif _idx_leaf == "cpu":
_TORCH_BACKEND = "cpu"
elif _is_cuda_family_leaf(_idx_leaf):
# Require a digit after "cu" (^cu[0-9]) so a full-override URL ending in
# /current or /custom is NOT branded CUDA. A wrong "cuda" backend makes
# _ensure_rocm_torch() return early on AMD hosts and leaves a CPU/wrong
# torch unrepaired; falling through here keeps _TORCH_BACKEND="" so the
# helpers probe the GPU instead.
_TORCH_BACKEND = "cuda"
def _torch_step_label(suffix: str) -> str:
"""Return a progress label like 'torch check (cuda)' using the known backend.
Falls back to GPU detection when UNSLOTH_TORCH_BACKEND is not set (e.g.
standalone `unsloth studio update` runs that bypass install.sh).
"""
backend = _TORCH_BACKEND
if not backend:
if _has_usable_nvidia_gpu():
backend = "cuda"
elif _has_rocm_gpu():
backend = "rocm"
else:
backend = "cpu"
return f"torch {suffix} ({backend})"
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal in-place one-line progress bar.
# Set UNSLOTH_VERBOSE=1 to restore full per-step output:
# CLI: unsloth studio setup --verbose
# Linux/Mac: UNSLOTH_VERBOSE=1 ./studio/setup.sh
# Windows: $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1
VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"
# Progress bar state -- updated by _progress() per install step.
# Update _TOTAL if you add/remove steps in install_python_stack().
_STEP: int = 0
_TOTAL: int = 0 # set at runtime in install_python_stack() based on platform
_PROGRESS_LINE_ACTIVE: bool = False
# -- Paths --------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
REQ_ROOT = SCRIPT_DIR / "backend" / "requirements"
SINGLE_ENV = REQ_ROOT / "single-env"
CONSTRAINTS = SINGLE_ENV / "constraints.txt"
LOCAL_DD_UNSTRUCTURED_PLUGIN = (
SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed"
)
LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the
# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242.
MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3"
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides).
# _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
if IS_MAC_ARM and _MLX_OVERRIDES.is_file() and "UV_OVERRIDE" not in os.environ:
os.environ["UV_OVERRIDE"] = _uv_safe_path(_MLX_OVERRIDES)
# -- Unicode-safe printing ---------------------------------------------
# On Windows the console encoding may be a legacy code page (e.g. CP1252)
# that cannot represent glyphs like ✅ or ❌. _safe_print() degrades to ASCII
# equivalents so the installer never crashes over a status glyph.
_UNICODE_TO_ASCII: dict[str, str] = {
"\u2705": "[OK]", # ✅
"\u274c": "[FAIL]", # ❌
"\u26a0\ufe0f": "[!]", # ⚠️ (warning + variation selector)
"\u26a0": "[!]", # ⚠ (warning without variation selector)
}
def _safe_print(*args: object, **kwargs: object) -> None:
"""Drop-in print() replacement that survives non-UTF-8 consoles and detached stdout."""
try:
print(*args, **kwargs)
except OSError:
return
except UnicodeEncodeError:
# Stringify, then swap emoji for ASCII equivalents.
text = " ".join(str(a) for a in args)
for uni, ascii_alt in _UNICODE_TO_ASCII.items():
text = text.replace(uni, ascii_alt)
# Final fallback: replace any remaining unencodable chars.
print(
text.encode(sys.stdout.encoding or "ascii", errors = "replace").decode(
sys.stdout.encoding or "ascii", errors = "replace"
),
**kwargs,
)
# ── Color support ──────────────────────────────────────────────────────
# Same logic as startup_banner: NO_COLOR disables, FORCE_COLOR or TTY enables.
def _stdout_supports_color() -> bool:
"""True if we should emit ANSI colors (matches startup_banner)."""
if os.environ.get("NO_COLOR", "").strip():
return False
if os.environ.get("FORCE_COLOR", "").strip():
return True
try:
if not sys.stdout.isatty():
return False
except (AttributeError, OSError, ValueError):
return False
if IS_WINDOWS:
try:
import ctypes
kernel32 = ctypes.windll.kernel32
handle = kernel32.GetStdHandle(-11)
mode = ctypes.c_ulong()
kernel32.GetConsoleMode(handle, ctypes.byref(mode))
kernel32.SetConsoleMode(handle, mode.value | 0x0004)
except (ImportError, AttributeError, OSError):
return False
return True
_HAS_COLOR = _stdout_supports_color()
# Column layout — matches setup.sh step() helper:
# 2-space indent, 15-char label (dim), then value.
_LABEL = "deps"
_COL = 15
_INDENT = 2
def _green(msg: str) -> str:
return f"\033[38;5;108m{msg}\033[0m" if _HAS_COLOR else msg
def _cyan(msg: str) -> str:
return f"\033[96m{msg}\033[0m" if _HAS_COLOR else msg
def _red(msg: str) -> str:
return f"\033[91m{msg}\033[0m" if _HAS_COLOR else msg
def _dim(msg: str) -> str:
return f"\033[38;5;245m{msg}\033[0m" if _HAS_COLOR else msg
def _title(msg: str) -> str:
return f"\033[38;5;150m{msg}\033[0m" if _HAS_COLOR else msg
_RULE = "\u2500" * 52
def _step(
label: str,
value: str,
color_fn = None,
) -> None:
"""Print a single step line in the column format."""
global _PROGRESS_LINE_ACTIVE
if color_fn is None:
color_fn = _green
padded = label[:_COL]
plain_prefix_width = _INDENT + _COL
prefix = f"{' ' * _INDENT}{_dim(padded)}{' ' * (_COL - len(padded))}"
wrap_width = max(
24,
shutil.get_terminal_size((100, 20)).columns - plain_prefix_width,
)
lines = textwrap.wrap(
value,
width = wrap_width,
break_long_words = False,
break_on_hyphens = False,
) or [""]
if _PROGRESS_LINE_ACTIVE and not VERBOSE:
try:
sys.stdout.write("\n")
sys.stdout.flush()
except OSError:
pass
_PROGRESS_LINE_ACTIVE = False
_safe_print(f"{prefix}{color_fn(lines[0])}")
continuation_prefix = " " * plain_prefix_width
for line in lines[1:]:
_safe_print(f"{continuation_prefix}{color_fn(line)}")
def _progress(label: str) -> None:
"""Print an in-place progress bar aligned to the step column layout."""
global _STEP, _PROGRESS_LINE_ACTIVE
_STEP += 1
if VERBOSE:
return
width = 20
filled = int(width * _STEP / _TOTAL)
bar = "=" * filled + "-" * (width - filled)
pad = " " * (_COL - len(_LABEL))
end = "\n" if _STEP >= _TOTAL else ""
try:
sys.stdout.write(f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}")
sys.stdout.flush()
_PROGRESS_LINE_ACTIVE = end == ""
except OSError:
pass
def run(
label: str,
cmd: list[str],
*,
quiet: bool = True,
) -> subprocess.CompletedProcess[bytes]:
"""Run a command; on failure print output and exit."""
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE if quiet else None,
stderr = subprocess.STDOUT if quiet else None,
env = _install_env_for_cmd(cmd),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
_step("error", f"{label} failed (exit code {result.returncode})", _red)
if result.stdout:
print(result.stdout.decode(errors = "replace"))
sys.exit(result.returncode)
return result
# Packages to skip on Windows (require special build steps)
WINDOWS_SKIP_PACKAGES = {"triton_kernels"}
# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode).
# These either *are* torch extensions or have unconditional
# ``Requires-Dist: torch``, so installing them would pull torch back in.
# ``librosa`` is here too despite not requiring torch: upstream ``llvmlite``
# dropped its macOS x86_64 wheel between 0.42.0 and 0.46.0+ (see
# https://pypi.org/project/llvmlite/0.47.0/#files -- only
# macosx_arm64 / manylinux / win_amd64 remain), so on Intel Mac the
# librosa -> numba -> llvmlite chain triggers a from-source build that fails
# in CI and on hosts without LLVM 14/15 headers. Tracked in unslothai/unsloth#5046.
NO_TORCH_SKIP_PACKAGES = {
"torch-stoi",
"timm",
"torchcodec",
"torch-c-dlpack-ext",
"openai-whisper",
"librosa",
}
def _select_flash_attn_version(torch_mm: str) -> str | None:
return flash_attn_package_version(torch_mm)
def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None:
return flash_attn_wheel_url(env)
def _print_optional_install_failure(label: str, result: subprocess.CompletedProcess[str]) -> None:
_step("warning", f"{label} failed (exit code {result.returncode})", _cyan)
if result.stdout:
print(result.stdout.strip())
def _flash_attn_install_disabled() -> bool:
return os.getenv("UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL") == "1"
def _ensure_flash_attn() -> None:
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"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
).returncode
== 0
):
return
env = probe_torch_wheel_env()
wheel_url = _build_flash_attn_wheel_url(env) if env else None
if wheel_url and url_exists(wheel_url):
for installer, wheel_result in install_wheel(
wheel_url,
python_executable = sys.executable,
use_uv = USE_UV,
uv_needs_system = UV_NEEDS_SYSTEM,
):
if wheel_result.returncode == 0:
return
_print_optional_install_failure(
f"Installing flash-attn prebuilt wheel with {installer}",
wheel_result,
)
_step("warning", "Continuing without flash-attn", _cyan)
return
if wheel_url is None:
_step("warning", "No compatible flash-attn prebuilt wheel found", _cyan)
else:
_step("warning", "No published flash-attn prebuilt wheel found", _cyan)
# -- uv bootstrap ------------------------------------------------------
USE_UV = False # Set by _bootstrap_uv() at the start of install_python_stack()
UV_NEEDS_SYSTEM = False # Set by _bootstrap_uv() via probe
def _bootstrap_uv() -> bool:
"""Check if uv is available and probe whether --system is needed."""
global UV_NEEDS_SYSTEM
if not shutil.which("uv"):
return False
# Probe: try a dry-run install targeting the current Python explicitly.
# Without --python, uv can ignore the activated venv on some platforms.
probe = subprocess.run(
["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
**_windows_hidden_subprocess_kwargs(),
)
if probe.returncode != 0:
# Retry with --system (some envs need it when uv can't find a venv)
probe_sys = subprocess.run(
["uv", "pip", "install", "--dry-run", "--system", "pip"],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
**_windows_hidden_subprocess_kwargs(),
)
if probe_sys.returncode != 0:
return False # uv is broken, fall back to pip
UV_NEEDS_SYSTEM = True
return True
def _filter_requirements(req: Path, skip: set[str]) -> Path:
"""Return a temp copy of a requirements file with certain packages removed."""
lines = req.read_text(encoding = "utf-8").splitlines(keepends = True)
filtered = [
line for line in lines if not any(line.strip().lower().startswith(pkg) for pkg in skip)
]
tmp = tempfile.NamedTemporaryFile(
mode = "w",
suffix = ".txt",
delete = False,
encoding = "utf-8",
)
tmp.writelines(filtered)
tmp.close()
return Path(tmp.name)
def _translate_pip_args_for_uv(args: tuple[str, ...]) -> list[str]:
"""Translate pip flags to their uv equivalents."""
translated: list[str] = []
for arg in args:
if arg == "--no-cache-dir":
continue # uv cache is fast; drop this flag
elif arg == "--force-reinstall":
translated.append("--reinstall")
else:
translated.append(arg)
return translated
def _build_pip_cmd(args: tuple[str, ...]) -> list[str]:
"""Build a standard pip install command.
Strips uv-only flags like --upgrade-package that pip doesn't understand.
"""
cmd = [sys.executable, "-m", "pip", "install"]
skip_next = False
for arg in args:
if skip_next:
skip_next = False
continue
if arg == "--upgrade-package":
skip_next = True # skip the flag and its value
continue
cmd.append(arg)
return cmd
def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
"""Build a uv pip install command with translated flags."""
cmd = ["uv", "pip", "install"]
if UV_NEEDS_SYSTEM:
cmd.append("--system")
# Always pass --python so uv targets the right environment. Without it, uv
# can ignore an activated venv and install into the system Python (seen on
# Colab and similar).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
# Torch is pre-installed by install.sh/setup.ps1. Do not add
# --torch-backend by default -- it can cause solver dead-ends on CPU-only
# machines. Callers that need it can set UV_TORCH_BACKEND. Never add it to a
# pinned-index command: uv's torch backend redirects torch resolution to its
# own per-backend index (verified: --index-url .../cu128 with UV_TORCH_BACKEND
# =cpu resolves torch+cpu), which would defeat the explicit pin exactly like
# the index env vars _install_env_for_cmd() strips.
_tb = os.environ.get("UV_TORCH_BACKEND", "")
if _tb and not _is_pinned_index_cmd(cmd):
cmd.append(f"--torch-backend={_tb}")
return cmd
# uv resolves the default index (passed as --index-url / --default-index) at the
# LOWEST priority: an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate or CPU
# mirror) is searched FIRST and, under uv's default first-index strategy, wins for
# any package it also serves. So a pinned torch repair (--index-url <cuXXX/rocm/cpu>)
# could silently resolve torch from that mirror instead of the pinned wheel index,
# and then _write_torch_index_marker() would record the pinned URL that was never
# actually used. install.sh (run_install_cmd, #6898), install.ps1 and setup.ps1
# already neutralise these vars for their pinned installs; do the same here so the
# pin wins on every platform. Non-pinned installs (no --index-url) keep the mirror.
# UV_TORCH_BACKEND is stripped too: uv's torch backend redirects torch resolution
# to its own per-backend index even when a --index-url is given, so it defeats a
# pin the same way ( _build_uv_cmd also refuses to turn it into a --torch-backend
# flag for pinned commands; this covers uv reading the env var directly).
# PIP_EXTRA_INDEX_URL matters for the pip FALLBACK (uv missing/failed): pip's
# --extra-index-url env form adds indexes "in addition to --index-url", so an
# inherited mirror could still satisfy torch off the pin. (PIP_INDEX_URL needs no
# strip: the command's explicit --index-url flag overrides it.)
_UV_INDEX_ENV_VARS = (
"UV_DEFAULT_INDEX",
"UV_INDEX_URL",
"UV_INDEX",
"UV_EXTRA_INDEX_URL",
"UV_TORCH_BACKEND",
"UV_FIND_LINKS",
"PIP_EXTRA_INDEX_URL",
"PIP_FIND_LINKS",
)
def _is_pinned_index_cmd(cmd: "list[str] | tuple[str, ...]") -> bool:
"""True when the command pins an index via --index-url / --default-index."""
return any(arg in ("--index-url", "--default-index") for arg in cmd)
def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None":
"""Return an env with the uv index vars stripped for a pinned-index install.
Returns None (inherit the caller's environment unchanged) when the command
does NOT pin an index, so ordinary installs still honour a user's UV_INDEX /
UV_EXTRA_INDEX_URL mirror. When the command passes --index-url / --default-index
(the torch repair paths), the uv index env vars (and UV_TORCH_BACKEND) are
removed so the pinned index is not overridden by an inherited mirror or torch
backend (uv treats the default index as lowest priority). Mirrors install.sh's
run_install_cmd gate (#6898).
"""
if not _is_pinned_index_cmd(cmd):
return None
env = os.environ.copy()
for name in _UV_INDEX_ENV_VARS:
env.pop(name, None)
return env
def pip_install_try(
label: str,
*args: str,
constrain: bool = True,
force_pip: bool = False,
) -> bool:
"""Like pip_install but returns False on failure instead of exiting.
For optional installs that have a follow-up fallback.
"""
constraint_args_pip: list[str] = []
constraint_args_uv: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args_pip = ["-c", str(CONSTRAINTS)]
constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)]
if USE_UV and not force_pip:
cmd = _build_uv_cmd(args) + constraint_args_uv
else:
cmd = _build_pip_cmd(args) + constraint_args_pip
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
env = _install_env_for_cmd(cmd),
)
if result.returncode == 0:
return True
if VERBOSE and result.stdout:
print(result.stdout.decode(errors = "replace"))
return False
def pip_install(
label: str,
*args: str,
req: Path | None = None,
constrain: bool = True,
) -> None:
"""Build and run a pip install command (uses uv when available, falls back to pip)."""
constraint_args_pip: list[str] = []
constraint_args_uv: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args_pip = ["-c", str(CONSTRAINTS)]
constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)]
actual_req = req
temp_reqs: list[Path] = []
if req is not None and IS_WINDOWS and WINDOWS_SKIP_PACKAGES:
actual_req = _filter_requirements(req, WINDOWS_SKIP_PACKAGES)
temp_reqs.append(actual_req)
if actual_req is not None and NO_TORCH and NO_TORCH_SKIP_PACKAGES:
actual_req = _filter_requirements(actual_req, NO_TORCH_SKIP_PACKAGES)
temp_reqs.append(actual_req)
if actual_req is not None and PLATFORM_LACKS_TORCHCODEC_WHEEL:
# Linux aarch64 / Windows ARM64 / Intel Mac have no torchcodec
# wheel. `unsloth studio update --local` does not pass
# --no-torch, so the NO_TORCH filter above does not fire; do
# the targeted skip independently so the audio extras step
# does not take down the whole update.
actual_req = _filter_requirements(actual_req, {"torchcodec"})
temp_reqs.append(actual_req)
req_args_pip: list[str] = []
req_args_uv: list[str] = []
if actual_req is not None:
req_args_pip = ["-r", str(actual_req)]
req_args_uv = ["-r", _uv_safe_path(actual_req)]
try:
if USE_UV:
uv_cmd = _build_uv_cmd(args) + constraint_args_uv + req_args_uv
if VERBOSE:
print(f" {label}...")
result = subprocess.run(
uv_cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
env = _install_env_for_cmd(uv_cmd),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode == 0:
return
print(_red(f" uv failed, falling back to pip..."))
if result.stdout:
print(result.stdout.decode(errors = "replace"))
pip_cmd = _build_pip_cmd(args) + constraint_args_pip + req_args_pip
run(f"{label} (pip)" if USE_UV else label, pip_cmd)
finally:
for temp_req in temp_reqs:
temp_req.unlink(missing_ok = True)
def download_file(url: str, dest: Path) -> None:
"""Download a file using urllib (no curl dependency)."""
urllib.request.urlretrieve(url, dest)
def patch_package_file(package_name: str, relative_path: str, url: str) -> None:
"""Download a file from url and overwrite a file inside an installed package."""
result = subprocess.run(
[sys.executable, "-m", "pip", "show", package_name],
capture_output = True,
text = True,
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
_step(_LABEL, f"package {package_name} not found, skipping patch", _red)
return
location = None
for line in result.stdout.splitlines():
if line.lower().startswith("location:"):
location = line.split(":", 1)[1].strip()
break
if not location:
_step(_LABEL, f"could not locate {package_name}", _red)
return
dest = Path(location) / relative_path
_step(_LABEL, f"patching {dest.name} in {package_name}...", _dim)
download_file(url, dest)
# -- Main install sequence ---------------------------------------------
def install_python_stack() -> int:
global USE_UV, _STEP, _TOTAL
_STEP = 0
# install.sh (which already installed unsloth) sets SKIP_STUDIO_BASE=1 to
# avoid reinstalling base packages. "unsloth studio update" does NOT set it,
# so base packages (unsloth + unsloth-zoo) are reinstalled to pick up new
# versions.
skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1"
# --package installs a different package name (for testing).
package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
# --local overlays a local repo checkout after updating deps.
local_repo = os.environ.get("STUDIO_LOCAL_REPO", "")
base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b)
if IS_MACOS:
base_total -= 1 # triton step is skipped on macOS
if not IS_MACOS and not NO_TORCH:
base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms
if not IS_WINDOWS:
base_total += 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try uv for faster installs (before pip upgrade -- uv venvs don't
# include pip by default).
USE_UV = _bootstrap_uv()
# 2. Ensure pip is available (uv venvs from install.sh omit pip).
_progress("pip bootstrap")
if USE_UV:
run(
"Bootstrapping pip via uv",
[
"uv",
"pip",
"install",
"--python",
sys.executable,
"pip",
],
)
else:
# pip may not exist yet (uv-created venvs omit it). Try ensurepip,
# then upgrade. Direct upgrade only when pip is already present.
_has_pip = (
subprocess.run(
[sys.executable, "-m", "pip", "--version"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
**_windows_hidden_subprocess_kwargs(),
).returncode
== 0
)
if not _has_pip:
run(
"Bootstrapping pip via ensurepip",
[sys.executable, "-m", "ensurepip", "--upgrade"],
)
else:
run(
"Upgrading pip",
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the
# mlx-vlm / mlx-lm transformers pin -- set at module load).
# Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke
# gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242.
if IS_MAC_ARM and not skip_base:
_progress("MLX stack (Apple Silicon)")
pip_install(
"Installing MLX stack (mlx + mlx-lm + mlx-vlm)",
"--no-cache-dir",
"--upgrade",
"mlx",
"mlx-metal",
f"mlx-lm{MLX_LM_BAD_VERSION_EXCLUSION}",
"mlx-vlm",
)
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
pass
elif NO_TORCH:
# No-torch update path: install unsloth + unsloth-zoo with --no-deps
# (PyPI metadata still declares torch as a hard dep), then runtime deps
# with --no-deps (avoids transitive torch).
_progress("base packages (no torch)")
pip_install(
f"Updating {package_name} + unsloth-zoo (no-torch mode)",
"--no-cache-dir",
"--no-deps",
"--upgrade-package",
package_name,
"--upgrade-package",
"unsloth-zoo",
package_name,
"unsloth-zoo",
)
# Resolve pydantic WITH deps so pip pins pydantic-core to the exact
# version pydantic's metadata declares. Under --no-deps pip picks the
# latest of each and trips pydantic's _ensure_pydantic_core_version
# check. Transitive deps are torch-free.
pip_install(
"Installing pydantic (with deps for compatible core)",
"--no-cache-dir",
"pydantic",
)
pip_install(
"Installing no-torch runtime deps",
"--no-cache-dir",
"--no-deps",
req = REQ_ROOT / "no-torch-runtime.txt",
)
if local_repo:
_step(_LABEL, f"overlaying local repo (editable): {local_repo}")
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
_step(_LABEL, "overlaying unsloth-zoo from git main")
pip_install(
"Overlaying unsloth-zoo from git main",
"--no-cache-dir",
"--no-deps",
"--force-reinstall",
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
constrain = False,
)
elif local_repo:
# Local dev install: update deps from base.txt, then overlay the local
# checkout as an editable install (--no-deps so torch is not re-resolved).
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
_step(_LABEL, f"overlaying local repo (editable): {local_repo}")
pip_install(
"Overlaying local repo (editable)",
"--no-cache-dir",
"--no-deps",
"-e",
local_repo,
constrain = False,
)
_step(_LABEL, "overlaying unsloth-zoo from git main")
pip_install(
"Overlaying unsloth-zoo from git main",
"--no-cache-dir",
"--no-deps",
"--force-reinstall",
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
constrain = False,
)
elif package_name != "unsloth":
# Custom package name (for testing): install directly.
_progress("base packages")
pip_install(
f"Installing {package_name}",
"--no-cache-dir",
package_name,
)
else:
# Update path: upgrade only unsloth + unsloth-zoo, preserving existing
# torch/CUDA installs. Torch is pre-installed by install.sh/setup.ps1;
# --upgrade-package targets only base pkgs.
_progress("base packages")
pip_install(
"Updating base packages",
"--no-cache-dir",
"--upgrade-package",
"unsloth",
"--upgrade-package",
"unsloth-zoo",
req = REQ_ROOT / "base.txt",
)
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv got CPU-only torch (common when pip resolves torch from PyPI).
# Must follow base packages so torch is present for inspection.
if not IS_MACOS and not NO_TORCH:
_progress(_torch_step_label("check"))
_ensure_cuda_torch()
_ensure_rocm_torch()
_ensure_cpu_torch()
_ensure_verbatim_torch_index()
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
# version or unknown ROCm version).
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
# Validate actual AMD GPU presence (not just tool existence).
import re as _re_win
def _win_amd_smi_has_gpu(stdout: str) -> bool:
return bool(_re_win.search(r"(?im)^gpu\s*[:\[]\s*\d", stdout))
_win_amd_gpu = False
for _wcmd, _check_fn in (
(["hipinfo"], lambda out: "gcnarchname" in out.lower()),
(["amd-smi", "list"], _win_amd_smi_has_gpu),
):
_wexe = shutil.which(_wcmd[0])
if not _wexe:
continue
# Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart
# prompt), as _has_rocm_gpu()/_detect_amd_gfx_codes do. The only loss
# is the best-effort "AMD GPU detected" note; ROCm-torch state below
# comes from the install itself.
if _wcmd[0] == "amd-smi" and not _amd_smi_allowed():
continue
try:
_wr = subprocess.run(
[_wexe, *_wcmd[1:]],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
env = _amd_smi_env() if _wcmd[0] == "amd-smi" else None,
)
except Exception:
continue
if _wr.returncode == 0 and _check_fn(_wr.stdout):
_win_amd_gpu = True
break
if _win_amd_gpu and not _rocm_windows_torch_installed:
_safe_print(
_dim(" Note:"),
"AMD GPU detected but ROCm PyTorch could not be auto-installed.",
)
_safe_print(
" " * 8,
"Manual install may be required. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
)
# 3. Extra dependencies
_progress("unsloth extras")
pip_install(
"Installing additional unsloth dependencies",
"--no-cache-dir",
req = REQ_ROOT / "extras.txt",
)
# 3b. Extra dependencies (no-deps) -- audio model support etc.
_progress("extra codecs")
pip_install(
"Installing extras (no-deps)",
"--no-deps",
"--no-cache-dir",
req = REQ_ROOT / "extras-no-deps.txt",
)
# 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to
# match the torch installed in the venv so its C++ extensions load (see
# _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac
# GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm
# (no working build; see below).
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm():
# No working Windows ROCm torchao build: it imports an absent c10d backend
# and crashes transformers.quantizers. Studio stubs it at runtime, so
# installing it only ships a package that crashes on import -- skip it.
_progress("dependency overrides (skipped, Windows ROCm)")
_safe_print(" Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
else:
_progress("dependency overrides")
_torch_ver = _probe_installed_torch_version()
_torchao_spec = _select_torchao_spec(_torch_ver)
_safe_print(f" torch {_torch_ver or 'unknown'} detected -- installing {_torchao_spec}")
pip_install(
"Installing dependency overrides",
"--force-reinstall",
"--no-cache-dir",
_torchao_spec,
)
# 5. Triton kernels (no-deps, from source). Skip on Windows and macOS
# (no support).
if not IS_WINDOWS and not IS_MACOS:
_progress("triton kernels")
pip_install(
"Installing triton kernels",
"--no-deps",
"--no-cache-dir",
req = REQ_ROOT / "triton-kernels.txt",
constrain = False,
)
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress("flash-attn")
_ensure_flash_attn()
# # 6. Patch: override llama_cpp.py with fix from unsloth-zoo feature/llama-cpp-windows-support branch
# patch_package_file(
# "unsloth-zoo",
# os.path.join("unsloth_zoo", "llama_cpp.py"),
# "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py",
# )
# # 7a. Patch: override vision.py with fix from unsloth PR #4091
# patch_package_file(
# "unsloth",
# os.path.join("unsloth", "models", "vision.py"),
# "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py",
# )
# # 7b. Patch : override save.py with fix from feature/llama-cpp-windows-support
# patch_package_file(
# "unsloth",
# os.path.join("unsloth", "save.py"),
# "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py",
# )
# 8. Studio dependencies
_progress("studio deps")
pip_install(
"Installing studio dependencies",
"--no-cache-dir",
req = REQ_ROOT / "studio.txt",
)
# 8b. anyio repair (#6483)
_progress("anyio check")
_repair_bad_anyio()
# 9. Data-designer dependencies
_progress("data designer deps")
pip_install(
"Installing data-designer base dependencies",
"--no-cache-dir",
req = SINGLE_ENV / "data-designer-deps.txt",
)
# 10. Data-designer packages (no-deps to avoid conflicts)
_progress("data designer")
pip_install(
"Installing data-designer",
"--no-cache-dir",
"--no-deps",
req = SINGLE_ENV / "data-designer.txt",
)
# 11. Local Data Designer seed plugins
local_dd_plugins = [
("unstructured", LOCAL_DD_UNSTRUCTURED_PLUGIN),
("github", LOCAL_DD_GITHUB_PLUGIN),
]
for _plugin_name, plugin_dir in local_dd_plugins:
if not plugin_dir.is_dir():
_safe_print(
_red(
f"❌ Missing local plugin directory: {plugin_dir}",
),
)
return 1
_progress("local plugin")
for plugin_name, plugin_dir in local_dd_plugins:
pip_install(
f"Installing local data-designer {plugin_name} plugin",
"--no-cache-dir",
"--no-deps",
str(plugin_dir),
constrain = False,
)
# 12. Patch metadata for single-env compatibility
_progress("finalizing")
run(
"Patching single-env metadata",
[sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
)
# 13. AMD ROCm: final torch repair. Several steps above can pull in CUDA
# torch from PyPI (base packages, extras, overrides, studio deps, etc.).
# Running the repair last ensures ROCm torch is in place at runtime,
# whichever intermediate step clobbered it.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress(_torch_step_label("final"))
_ensure_cuda_torch()
_ensure_rocm_torch()
_ensure_cpu_torch()
_ensure_verbatim_torch_index()
# 14. Final check (silent; third-party conflicts are expected)
subprocess.run(
[sys.executable, "-m", "pip", "check"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
**_windows_hidden_subprocess_kwargs(),
)
_step(_LABEL, "installed")
return 0
if __name__ == "__main__":
sys.exit(install_python_stack())
Reference in a new issue View git blame Copy permalink