vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-10 08:18:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29
unsloth/studio/install_python_stack.py
Daniel Han 55c392ff7c
studio: fix sentence-transformers RAG embedder on Windows ROCm (torchao) (#6608) 
torchao has no working Windows ROCm build. transformers.quantizers imports it,
and it loads torch's c10d distributed backend at module level, which the AMD
Windows wheels omit (no RCCL). The import aborts, transformers can no longer
expose PreTrainedModel, and the sentence-transformers embedder silently falls
back to the llama-server GGUF embedder. Linux ROCm and NVIDIA are unaffected
(the c10d ops are present / torchao is real there).

The training and export workers already install the shared torchao stub before
importing transformers, but the RAG embedder runs in the main backend process,
which never did. Two fixes, both no-ops off Windows ROCm:

- embeddings.py: install_torchao_windows_rocm_stub() before the first
  sentence-transformers import, so an already-installed torchao is neutralized
  (fixes existing venvs).
- install_python_stack.py: stop installing torchao on Windows ROCm; it can only
  crash on import there, so new venvs never ship it.

Add tests covering the embedder stub call and the install skip.
2026-06-23 05:39:02 -07:00

2375 lines
92 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,
)
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",
}
# 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("/")
# 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_PKG_SPEC: tuple[str, str, str] = (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
)
# torchao's C++ extensions are built against ONE exact torch release; a newer
# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions
# due to incompatible torch version ...") and fall back to slow Python. Because
# the torch pin above is a range (and every CUDA index now tops out at torch
# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So
# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919.
# torch 2.9.x -> torchao 0.14.0 (today's pin; built for torch 2.9.0)
# torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0)
# torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2)
# Unknown/older torch keeps the conservative default (no regression vs today).
_TORCHAO_DEFAULT_SPEC = "torchao==0.14.0"
_TORCHAO_BY_TORCH_MINOR: dict[int, str] = {
10: "torchao==0.16.0",
11: "torchao==0.17.0",
}
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_BY_TORCH_MINOR[11] # newest known build; covers 2.11+
return _TORCHAO_BY_TORCH_MINOR.get(minor, _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
# 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.
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).
"""
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 _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
# 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
# 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.
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if _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.
if 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; "
"hip = getattr(torch.version, 'hip', '') or ''; "
"cuda = getattr(torch.version, 'cuda', '') or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
"print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
),
],
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 or _marker_lines[-1] != "hip":
return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
index_url = _detect_cuda_torch_index_url()
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
print(
f" torch is a ROCm build on an NVIDIA host -- reinstalling "
f"CUDA torch from {index_url}\n"
f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch "
f"on a mixed AMD+NVIDIA host)"
)
pip_install(
"CUDA torch repair",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
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
# 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
# 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
# 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:
print(" ROCm detected but version unreadable -- skipping torch reinstall")
return
# Probe whether torch already links against HIP (ROCm already working).
# Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts
# (the NVIDIA check above already handled mixed AMD+NVIDIA setups).
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=getattr(torch,'__version__','').lower(); "
# Print the HIP version when present (back-compat), else a
# "rocm" sentinel when only torch.__version__ flags ROCm
# (AMD SDK / Radeon wheels). Empty string = CPU/CUDA.
"print(hip if hip else ('rocm' if 'rocm' in ver else ''))"
),
],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
timeout = 90,
)
except (OSError, subprocess.TimeoutExpired):
probe = None
has_hip_torch = (
probe is not None and probe.returncode == 0 and probe.stdout.decode().strip() != ""
)
rocm_torch_ready = has_hip_torch
# 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
if ver < (7, 2):
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,
)
rocm_torch_ready = True
elif not has_hip_torch:
# Select best matching wheel tag (newest ROCm version <= installed)
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]} -- " f"skipping torch reinstall")
else:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
_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,
)
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,
)
def _uv_safe_path(path: object) -> str:
# uv 0.11.x: `-c <path with space>` truncates at the space; use 8.3 short form.
s = str(path)
if not IS_WINDOWS or " " not in s:
return s
try:
import ctypes
from ctypes import wintypes
get_short = ctypes.windll.kernel32.GetShortPathNameW
get_short.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD]
get_short.restype = wintypes.DWORD
buf = ctypes.create_unicode_buffer(32768)
rc = get_short(s, buf, 32768)
if 0 < rc < 32768 and " " not in buf.value:
return buf.value
except Exception:
pass
return s
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()
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"
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
if IS_MAC_ARM and _MLX_OVERRIDES.is_file():
os.environ.setdefault("UV_OVERRIDE", str(_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,
**_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 = {"open_spiel", "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",
"transformers-cfg",
"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.
_tb = os.environ.get("UV_TORCH_BACKEND", "")
if _tb:
cmd.append(f"--torch-backend={_tb}")
return cmd
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,
)
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,
**_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 = 10 if IS_WINDOWS else 11
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).
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",
"mlx-lm",
"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()
# 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,
# which has no working torchao build (see below).
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
elif _rocm_windows_torch_installed:
# torchao has no working Windows ROCm build: it loads torch's c10d
# distributed backend at import, which the AMD Windows wheels omit (no
# RCCL), so `import torchao` raises and takes transformers.quantizers
# down with it. Studio stubs torchao at runtime (core/_torchao_stub.py),
# 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",
)
# 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()
# 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