((set, get) => ({
loadedKvCacheDtype: null,
speculativeType: "auto",
loadedSpeculativeType: null,
+ specFallbackReason: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
@@ -977,6 +983,7 @@ export const useChatRuntimeStore = create((set, get) => ({
loadedKvCacheDtype: null,
speculativeType: "auto",
loadedSpeculativeType: null,
+ specFallbackReason: null,
specDraftNMax: null,
loadedSpecDraftNMax: null,
loadedIsMultimodal: false,
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 49524c17b4..a0e5d5355e 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -174,6 +174,12 @@ export interface InferenceStatusResponse {
/** Canonical UI-facing mode currently active. See LoadModelRequest. */
speculative_type?: string | null;
spec_draft_n_max?: number | null;
+ /**
+ * Why MTP was disabled on the loaded model despite being requested.
+ * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable
+ * it; "runtime_error" -> the current build could not run it. Null otherwise.
+ */
+ spec_fallback_reason?: string | null;
}
export interface AudioGenerationResponse {
diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts
index f1966f03fb..97fbd32d57 100644
--- a/studio/frontend/src/features/training/stores/training-runtime-store.ts
+++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts
@@ -274,7 +274,10 @@ export const useTrainingRuntimeStore = create()((set) => (
jobId: payload.job_id || state.jobId,
currentStep: step,
totalSteps: Math.max(payload.total_steps, state.totalSteps),
- currentLoss: currentLoss ?? state.currentLoss,
+ // A null loss at a new step means the backend reported a non-finite
+ // loss; clear the display instead of keeping the stale value.
+ currentLoss:
+ currentLoss ?? (step > state.currentStep ? null : state.currentLoss),
currentLearningRate: currentLearningRate ?? state.currentLearningRate,
progressPercent: payload.progress_percent,
currentEpoch: payload.epoch ?? state.currentEpoch,
diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts
index d30c3f75c2..5fd8286d22 100644
--- a/studio/frontend/src/features/training/types/runtime.ts
+++ b/studio/frontend/src/features/training/types/runtime.ts
@@ -90,7 +90,8 @@ export interface TrainingRuntimeState {
currentStep: number;
totalSteps: number;
currentEpoch: number;
- currentLoss: number;
+ // null while the latest reported loss is non-finite
+ currentLoss: number | null;
currentLearningRate: number;
progressPercent: number;
elapsedSeconds: number | null;
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 797e17420e..0d3b1698e1 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -897,8 +897,7 @@
/* Chat search box: borderless, soft Gemini-style elevation. */
.chat-search-surface {
border: none;
- /* Pin to the dark --radius so the corners are the same (less round) in
- both themes; rounded-4xl here and on the inner Command follow this. */
+ /* Pin to the dark --radius so rounded-3xl corners stay consistent. */
--radius: 0.625rem;
/* Match the chat box: soft elevation. */
box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
@@ -1751,10 +1750,11 @@
margin: 0 !important;
}
-/* Boost shadow on dark surfaces; mirrors .shadow-border / .menu-soft-surface pattern. */
+/* Composer shadow on the dark background color so toasts
+ do not merge into card-colored surfaces behind them. */
.dark [data-sonner-toast][data-styled='true'] {
- background-color: var(--card) !important;
- box-shadow: none !important;
+ background-color: var(--background) !important;
+ box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16) !important;
}
/* Selectable toast text; non-selectable toast buttons. */
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 822a32ddb8..2b42a7a058 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -177,6 +177,7 @@ VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf"
INSTALL_LOCK_TIMEOUT_SECONDS = 300
INSTALL_STAGING_ROOT_NAME = ".staging"
GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"}
+HF_AUTH_HOSTS = {"huggingface.co", "www.huggingface.co"}
RETRYABLE_HTTP_STATUS = {408, 429, 500, 502, 503, 504}
HTTP_FETCH_ATTEMPTS = 4
HTTP_FETCH_BASE_DELAY_SECONDS = 0.75
@@ -484,6 +485,10 @@ def should_send_github_auth(url: str | None) -> bool:
return parsed_hostname(url) in GITHUB_AUTH_HOSTS
+def should_send_hf_auth(url: str | None) -> bool:
+ return parsed_hostname(url) in HF_AUTH_HOSTS
+
+
def auth_headers(url: str | None = None) -> dict[str, str]:
headers = {
"User-Agent": "unsloth-studio-llama-prebuilt",
@@ -491,9 +496,35 @@ def auth_headers(url: str | None = None) -> dict[str, str]:
token = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
if token and should_send_github_auth(url):
headers["Authorization"] = f"Bearer {token}"
+ return headers
+ # Anonymous huggingface.co fetches share a per-IP rate limit that CI
+ # fleets exhaust (HTTP 429), sinking the prebuilt path into a source
+ # build. Authenticate when a token is available.
+ hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+ if hf_token and should_send_hf_auth(url):
+ headers["Authorization"] = f"Bearer {hf_token}"
return headers
+class _CrossHostAuthStrippingRedirectHandler(urllib.request.HTTPRedirectHandler):
+ """Drop Authorization when a redirect leaves the original host.
+
+ huggingface.co redirects file downloads to CDN hosts whose signed URLs
+ can reject foreign Authorization headers; urllib forwards headers to
+ redirect targets by default (requests/huggingface_hub strip them).
+ """
+
+ def redirect_request(self, req, fp, code, msg, headers, newurl):
+ new_request = super().redirect_request(req, fp, code, msg, headers, newurl)
+ if new_request is not None and parsed_hostname(newurl) != parsed_hostname(req.full_url):
+ new_request.headers.pop("Authorization", None)
+ new_request.unredirected_hdrs.pop("Authorization", None)
+ return new_request
+
+
+_URL_OPENER = urllib.request.build_opener(_CrossHostAuthStrippingRedirectHandler())
+
+
def github_api_headers(url: str | None = None) -> dict[str, str]:
return {
"Accept": "application/vnd.github+json",
@@ -936,7 +967,7 @@ def download_bytes(
for attempt in range(1, attempts + 1):
try:
request = urllib.request.Request(url, headers = headers or auth_headers(url))
- with urllib.request.urlopen(request, timeout = timeout) as response:
+ with _URL_OPENER.open(request, timeout = timeout) as response:
total_bytes: int | None = None
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
@@ -1015,7 +1046,7 @@ def download_file(url: str, destination: Path) -> None:
delete = False,
) as handle:
tmp_path = Path(handle.name)
- with urllib.request.urlopen(request, timeout = 120) as response:
+ with _URL_OPENER.open(request, timeout = 120) as response:
total_bytes: int | None = None
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
@@ -1095,6 +1126,20 @@ def commit_source_archive_urls(repo: str, source_commit: str) -> list[str]:
]
+def release_asset_download_url(
+ repo: str | None, release_tag: str | None, asset_name: str | None
+) -> str | None:
+ """Direct download URL for a release asset, or None if any part is missing.
+ A mix build's merged commit is never pushed, so its source tree is only
+ reachable as this asset (codeload would 404 on the merge commit)."""
+ if not repo or not release_tag or not asset_name:
+ return None
+ return (
+ f"https://github.com/{repo}/releases/download/"
+ f"{urllib.parse.quote(release_tag, safe = '')}/{urllib.parse.quote(asset_name, safe = '')}"
+ )
+
+
def github_release_assets(repo: str, tag: str) -> dict[str, str]:
payload = fetch_json(
f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}"
@@ -2890,7 +2935,27 @@ def detect_host() -> HostInfo:
except Exception:
pass
+ # Linux /proc/driver/nvidia/gpus fallback: the NVIDIA driver exposes one
+ # subdir per GPU here regardless of nvidia-smi state, so a host whose
+ # nvidia-smi is absent from PATH, wedged, or failing is still recognised as
+ # NVIDIA. Mirrors the fallback added to install.sh / install_python_stack.py
+ # in PR 6174 so the prebuilt installer does not misroute such hosts to ROCm
+ # or CPU. driver_cuda_version / compute_caps stay unset here; downstream
+ # CUDA asset selection treats unknown SMs as "prefer portable" and an
+ # unknown driver runtime line as "no published CUDA match" (returns None,
+ # no crash), so planning falls back to a source build with GGML_CUDA=ON.
+ if is_linux and not has_physical_nvidia:
+ try:
+ proc_gpu_dir = "/proc/driver/nvidia/gpus"
+ if os.path.isdir(proc_gpu_dir) and os.listdir(proc_gpu_dir):
+ has_physical_nvidia = True
+ has_usable_nvidia = visible_device_tokens != []
+ except OSError:
+ pass
+
# Detect AMD ROCm (HIP) -- require actual GPU, not just tools installed
+ # NVIDIA takes precedence: when an NVIDIA GPU is usable, skip ROCm probing
+ # entirely so co-installed ROCm tools cannot misroute the host (PR 6174).
def _amd_smi_has_gpu(stdout: str) -> bool:
"""Check for 'GPU: ' data rows, not just a table header."""
@@ -2898,7 +2963,7 @@ def detect_host() -> HostInfo:
has_rocm = False
rocm_gfx_target: str | None = None
- if is_linux:
+ if is_linux and not has_usable_nvidia:
# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg
# only when HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and
# rocminfo can live only under /opt/rocm/bin (the profile.d PATH
@@ -2937,7 +3002,7 @@ def detect_host() -> HostInfo:
has_rocm = True
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
- elif is_windows:
+ elif is_windows and not has_usable_nvidia:
# Windows: prefer active probes that validate GPU presence.
# hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer
# sets HIP_PATH / ROCM_PATH but does not always add the bin dir to
@@ -3045,6 +3110,21 @@ def _apply_host_overrides(
return host
+def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str:
+ """The release repo setup.sh / setup.ps1 pick for this host: macOS always the
+ fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows
+ -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA
+ or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux
+ hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork
+ even when the probe cannot confirm an active GPU. Mirrors the shell routing."""
+ if host.is_macos:
+ return DEFAULT_PUBLISHED_REPO
+ has_gpu = (
+ host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present)
+ )
+ return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO
+
+
def pick_windows_cuda_runtime(host: HostInfo) -> str | None:
if not host.driver_cuda_version:
return None
@@ -4460,13 +4540,17 @@ def hydrate_source_tree(
expected_sha256: str | None,
source_label: str | None = None,
exact_source: bool = False,
+ asset_url: str | None = None,
) -> None:
archive_path = work_dir / f"llama.cpp-source-{source_ref}.tar.gz"
- source_urls = (
+ repo_urls = (
commit_source_archive_urls(source_repo, source_ref)
if exact_source
else upstream_source_archive_urls(source_ref)
)
+ # Prefer the published release asset (the only copy of a mix build's merged
+ # tree); fall back to codeload/archive for vanilla builds whose commit is real.
+ source_urls = ([asset_url] if asset_url else []) + repo_urls
label = source_label or f"llama.cpp source tree for {source_ref}"
extract_dir = Path(tempfile.mkdtemp(prefix = "source-extract-", dir = work_dir))
@@ -4817,6 +4901,35 @@ def confirm_install_tree(install_dir: Path, host: HostInfo) -> None:
raise RuntimeError("activated install was missing expected files: " + ", ".join(missing))
+def activate_staged_dir(staging_dir: Path, dst: Path) -> None:
+ """Move a freshly extracted ``staging_dir`` onto ``dst``.
+
+ ``os.replace`` is attempted first as the fast path. On Windows ARM64 the
+ antivirus scanner can transiently hold a freshly extracted DLL open at the
+ moment ``MoveFileEx`` runs, surfacing as ``[WinError 5] Access is denied``;
+ a file-by-file copy bypasses the rename entirely.
+
+ This fallback is intentionally limited to staging trees we just extracted.
+ It must not be used to move an existing/active install aside: there an
+ ``os.replace`` failure means the directory is genuinely in use, and a
+ silent copy + ``rmtree`` could partially delete a live install.
+
+ Only busy/lock errors (``is_busy_lock_error``) trigger the copy; anything
+ else (disk full, cross-device, missing path) re-raises so it cannot leave
+ a partially copied install behind. A copy is preferred over retrying the
+ rename because antivirus scans of large DLLs can outlast any reasonable
+ retry window.
+ """
+ try:
+ os.replace(staging_dir, dst)
+ except OSError as exc:
+ if not is_busy_lock_error(exc):
+ raise
+ log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree")
+ shutil.copytree(staging_dir, dst, dirs_exist_ok = True)
+ remove_tree(staging_dir)
+
+
def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None:
rollback_dir: Path | None = None
failed_dir: Path | None = None
@@ -4828,7 +4941,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo)
log(f"moved existing install to rollback path {rollback_dir.name}")
log(f"activating staged install {staging_dir} -> {install_dir}")
- os.replace(staging_dir, install_dir)
+ activate_staged_dir(staging_dir, install_dir)
log(f"activated staged install at {install_dir}")
log(f"confirming activated install tree at {install_dir}")
confirm_install_tree(install_dir, host)
@@ -6369,6 +6482,15 @@ def validate_prebuilt_choice(
source_repo, source_ref, source_archive, exact_source = preferred_source_archive(
approved_checksums, llama_tag
)
+ # For an exact (mix) source the merge commit lives only in the release asset,
+ # not in any repo, so fetch the asset directly; codeload stays the fallback.
+ asset_url = (
+ release_asset_download_url(
+ approved_checksums.repo, approved_checksums.release_tag, source_archive.asset_name
+ )
+ if exact_source and source_archive is not None
+ else None
+ )
if exact_source:
log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}")
else:
@@ -6385,6 +6507,7 @@ def validate_prebuilt_choice(
else f"llama.cpp source tree for {llama_tag}"
),
exact_source = exact_source,
+ asset_url = asset_url,
)
log(f"overlaying prebuilt bundle {choice.name} into {install_dir}")
server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir)
@@ -6690,6 +6813,16 @@ def parse_args() -> argparse.Namespace:
const = "latest",
help = ("Resolve the source-build fallback plan."),
)
+ resolve_group.add_argument(
+ "--resolve-prebuilt",
+ nargs = "?",
+ const = "latest",
+ help = (
+ "Report whether an official prebuilt exists for this host without "
+ "downloading. Picks the host's published repo when --published-repo "
+ "is left at the default. Use --output-format json."
+ ),
+ )
parser.add_argument(
"--output-format",
choices = ("plain", "json"),
@@ -6774,6 +6907,46 @@ def main() -> int:
)
return EXIT_SUCCESS
+ if args.resolve_prebuilt is not None:
+ # Host-aware "is a prebuilt available" probe, no download. A default repo
+ # means "pick the repo for this host"; PrebuiltFallback == source build.
+ host = _apply_host_overrides(
+ detect_host(),
+ override_has_rocm = args.has_rocm,
+ override_rocm_gfx = args.rocm_gfx,
+ force_cpu = args.cpu_fallback,
+ )
+ # setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU
+ # is probed; mirror that so a HIP source build is not offered a CPU prebuilt.
+ amd_tooling = host.is_linux and any(
+ shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo")
+ )
+ repo = (
+ published_repo_for_host(host, linux_amd_tooling_present = amd_tooling)
+ if args.published_repo == DEFAULT_PUBLISHED_REPO
+ else args.published_repo
+ )
+ try:
+ _requested, plans = resolve_simple_install_release_plans(
+ args.resolve_prebuilt, host, repo, args.published_release_tag or ""
+ )
+ choice = plans[0].attempts[0] if plans and plans[0].attempts else None
+ if choice is None:
+ payload = {"prebuilt_available": False, "repo": repo}
+ else:
+ payload = {
+ "prebuilt_available": True,
+ "repo": repo,
+ "release_tag": plans[0].release_tag,
+ "llama_tag": plans[0].llama_tag,
+ "asset": choice.name,
+ "install_kind": choice.install_kind,
+ }
+ except PrebuiltFallback:
+ payload = {"prebuilt_available": False, "repo": repo}
+ emit_resolver_output(payload, output_format = args.output_format)
+ return EXIT_SUCCESS
+
if not args.install_dir:
raise SystemExit(
"install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag, --resolve-install-tag, or --resolve-source-build is used"
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 0166bfd505..e540aac305 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -90,6 +90,18 @@ _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",
+)
+
# 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 = (
@@ -557,7 +569,15 @@ def _persist_bnb_rocm_version(version: str) -> bool:
def _has_rocm_gpu() -> bool:
- """Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
+ """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
@@ -598,6 +618,13 @@ def _has_rocm_gpu() -> bool:
# 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"
@@ -609,29 +636,69 @@ def _has_rocm_gpu() -> bool:
gpu_id = fh.read().strip()
except OSError:
continue
- if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node
- return True
+ 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 only when nvidia-smi exists AND reports at least one GPU."""
+ """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 not exe:
- return False
- try:
- result = subprocess.run(
- [exe, "-L"],
- stdout = subprocess.PIPE,
- stderr = subprocess.DEVNULL,
- text = True,
- timeout = 10,
- )
- except Exception:
- return False
- return result.returncode == 0 and "GPU " in result.stdout
+ 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]:
@@ -739,6 +806,139 @@ def _install_bnb_windows_rocm() -> bool:
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.
@@ -749,6 +949,13 @@ def _ensure_rocm_torch() -> None:
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.
@@ -1088,6 +1295,29 @@ def _infer_no_torch() -> bool:
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.
@@ -1770,7 +2000,8 @@ def install_python_stack() -> int:
# 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("ROCm torch check")
+ _progress(_torch_step_label("check"))
+ _ensure_cuda_torch()
_ensure_rocm_torch()
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
@@ -1955,7 +2186,8 @@ def install_python_stack() -> int:
# 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("ROCm torch (final)")
+ _progress(_torch_step_label("final"))
+ _ensure_cuda_torch()
_ensure_rocm_torch()
# 14. Final check (silent; third-party conflicts are expected)
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index e707cfda7a..94817b561a 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -299,7 +299,10 @@ function Get-CudaComputeCapability {
if (-not $smiExe) { return $null }
try {
- $raw = & $smiExe --query-gpu=compute_cap --format=csv,noheader 2>$null
+ # Bounded: a wedged nvidia-smi must not hang setup after the initial
+ # -L probe succeeded (the helper merges stderr after stdout, so the
+ # first line is still the compute_cap value).
+ $raw = Invoke-NvidiaSmiBounded $smiExe @('--query-gpu=compute_cap', '--format=csv,noheader')
if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null }
# nvidia-smi may return multiple GPUs; take the first one
@@ -363,10 +366,10 @@ function Get-PytorchCudaTag {
if (-not $smiExe) { return "cu126" }
try {
- # 2>&1 | Out-String merges stderr into stdout then converts to a single
- # string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 --
- # ErrorRecord objects leak into $output and break the -match.
- $output = & $smiExe 2>&1 | Out-String
+ # Bounded: a wedged nvidia-smi must not hang setup. The helper merges
+ # stderr into the returned string, matching the old 2>&1 | Out-String
+ # shape (plain 2>$null leaks ErrorRecord objects in PS 5.1).
+ $output = Invoke-NvidiaSmiBounded $smiExe
# Newer NVIDIA drivers (e.g. 610.x on Windows) print
# "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y".
# Accept both spellings so we don't fall through to the cu126 default.
@@ -667,16 +670,58 @@ try {
# ============================================
# 1a. GPU detection
# ============================================
+# ── Helper: run nvidia-smi under a timeout ──
+# A wedged NVIDIA driver can make nvidia-smi block during init or after a reset;
+# WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate below) so detection
+# cannot hang setup. No RunAsInvoker compat layer: nvidia-smi does not
+# auto-elevate. Returns combined stdout+stderr; "" on timeout/failure.
+function Invoke-NvidiaSmiBounded {
+ param(
+ [Parameter(Mandatory = $true, Position = 0)][string]$Exe,
+ [Parameter(Position = 1)][string[]]$SmiArgs = @(),
+ [int]$TimeoutSec = 10
+ )
+ try {
+ $psi = New-Object System.Diagnostics.ProcessStartInfo
+ $psi.FileName = $Exe
+ $psi.Arguments = ($SmiArgs -join ' ')
+ $psi.UseShellExecute = $false
+ $psi.RedirectStandardOutput = $true
+ $psi.RedirectStandardError = $true
+ $psi.CreateNoWindow = $true
+ $proc = [System.Diagnostics.Process]::Start($psi)
+ $outTask = $proc.StandardOutput.ReadToEndAsync()
+ $errTask = $proc.StandardError.ReadToEndAsync()
+ if (-not $proc.WaitForExit($TimeoutSec * 1000)) {
+ try { $proc.Kill() } catch {}
+ $global:LASTEXITCODE = 124
+ return ""
+ }
+ $global:LASTEXITCODE = $proc.ExitCode
+ return ($outTask.Result + "`n" + $errTask.Result)
+ } catch {
+ $global:LASTEXITCODE = 1
+ return ""
+ }
+}
+
+# ── Helper: nvidia-smi -L lists at least one real GPU ──
+# Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0
+# while listing no GPU, which would mark an AMD host NVIDIA and suppress ROCm
+# detection. Require a "GPU :" data row.
+function Test-NvidiaSmiHasGpu {
+ param([Parameter(Mandatory = $true)][string]$Exe)
+ $out = Invoke-NvidiaSmiBounded $Exe @('-L')
+ return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:')
+}
+
$HasNvidiaSmi = $false
$NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment
try {
$nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
- if ($nvSmiCmd) {
- & $nvSmiCmd.Source *> $null
- if ($LASTEXITCODE -eq 0) {
- $HasNvidiaSmi = $true
- $NvidiaSmiExe = $nvSmiCmd.Source
- }
+ if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) {
+ $HasNvidiaSmi = $true
+ $NvidiaSmiExe = $nvSmiCmd.Source
}
} catch {}
# Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist.
@@ -689,8 +734,7 @@ if (-not $HasNvidiaSmi) {
foreach ($p in $nvSmiDefaults) {
if (Test-Path $p) {
try {
- & $p *> $null
- if ($LASTEXITCODE -eq 0) {
+ if (Test-NvidiaSmiHasGpu $p) {
$HasNvidiaSmi = $true
$NvidiaSmiExe = $p
Write-Host " Found nvidia-smi at $(Split-Path $p -Parent)" -ForegroundColor Gray
@@ -1151,7 +1195,16 @@ function Resolve-CudaToolkit {
$DriverMaxCuda = $null
try {
- $smiOut = & $NvidiaSmiExe 2>&1 | Out-String
+ # Bounded: source-build toolkit resolution must not hang on a wedged smi.
+ # test_resolve_cuda_toolkit.ps1 extracts this function alone into a child
+ # pwsh (no Invoke-NvidiaSmiBounded in scope) and stubs nvidia-smi with a
+ # .ps1 script, so fall back to direct invocation when the bounded runner
+ # is unavailable; production setup.ps1 always has it defined.
+ $smiOut = if (Get-Command Invoke-NvidiaSmiBounded -ErrorAction SilentlyContinue) {
+ Invoke-NvidiaSmiBounded $NvidiaSmiExe
+ } else {
+ & $NvidiaSmiExe 2>&1 | Out-String
+ }
# Newer drivers report "CUDA UMD Version: X.Y" instead of "CUDA Version: X.Y"; accept both.
if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") {
$DriverMaxCuda = "$($Matches[1]).$($Matches[2])"
@@ -1499,10 +1552,48 @@ if ($IsPipInstall) {
}
}
-# 1g. Python (>= 3.11 and < 3.14). Prefer the Studio venv that install.ps1
-# just created, then py.exe so a 3.14 ahead of 3.13 on PATH does not trip the gate.
+# Conda CPython ships modified DLL search paths that break torch's c10.dll
+# loading on Windows; a venv made from conda Python inherits its base_prefix,
+# so check the executable path AND sys.base_prefix.
+$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)'
+function Test-IsConda {
+ param([string]$Exe)
+ if ($Exe -match $CondaSkipPattern) { return $true }
+ try {
+ $basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim()
+ if ($basePrefix -match $CondaSkipPattern) { return $true }
+ } catch { }
+ return $false
+}
+
+# 1g. Python (>= 3.11 and < 3.14). Prefer the interpreter install.ps1 already
+# resolved and built the venv with (UNSLOTH_SETUP_PYTHON), or the existing
+# venv python, before re-probing a system where a 3.14 or a WindowsApps stub
+# ahead on PATH would trip the gate. setup.ps1 only updates packages in that
+# venv, so the handoff is safe to reuse once validated.
+function Resolve-ReusedSetupPython {
+ if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_SETUP_PYTHON) -and
+ (Test-Path -LiteralPath $env:UNSLOTH_SETUP_PYTHON)) {
+ return $env:UNSLOTH_SETUP_PYTHON
+ }
+ # Standalone `unsloth studio setup/update` (install.ps1 did not run): derive
+ # the venv python from the studio root, mirroring the resolver below.
+ $root = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() }
+ elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() }
+ else { Join-Path $env:USERPROFILE ".unsloth\studio" }
+ if ($root -eq "~") {
+ # Join-Path with an empty child throws on Windows PowerShell 5.1.
+ $root = $env:USERPROFILE
+ } elseif ($root -like "~/*" -or $root -like "~\*") {
+ $root = Join-Path $env:USERPROFILE $root.Substring(1).TrimStart('/', '\')
+ }
+ $venvPy = Join-Path $root "unsloth_studio\Scripts\python.exe"
+ if (Test-Path -LiteralPath $venvPy) { return $venvPy }
+ return $null
+}
+$ReusedSetupPython = Resolve-ReusedSetupPython
+
$HasPython = $null -ne (Get-Command python -ErrorAction SilentlyContinue)
-$PyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
$PythonOk = $false
$DetectedPyVer = $null
@@ -1531,28 +1622,24 @@ function Add-PythonDirToProcessPath {
} catch { }
}
-$_prereqStudioHome = $null
-if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) {
- $_prereqStudioHome = $env:UNSLOTH_STUDIO_HOME.Trim()
-} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) {
- $_prereqStudioHome = $env:STUDIO_HOME.Trim()
-} else {
- $_prereqStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
-}
-if ($_prereqStudioHome -eq "~" -or $_prereqStudioHome -like "~/*" -or $_prereqStudioHome -like "~\*") {
- $_prereqStudioHome = (Join-Path $env:USERPROFILE $_prereqStudioHome.Substring(1).TrimStart('/','\'))
-}
-$_prereqVenvPython = Join-Path $_prereqStudioHome "unsloth_studio\Scripts\python.exe"
-if (Test-Path -LiteralPath $_prereqVenvPython) {
- $_venvPyVer = Get-CompatiblePythonVersion $_prereqVenvPython
- if ($_venvPyVer) {
- $DetectedPyVer = $_venvPyVer
- Add-PythonDirToProcessPath $_prereqVenvPython
+# Reuse the install.ps1 / venv interpreter before any system probe.
+if ($ReusedSetupPython) {
+ $_reusedVer = Get-CompatiblePythonVersion $ReusedSetupPython
+ if ($_reusedVer -and -not (Test-IsConda $ReusedSetupPython)) {
+ $DetectedPyVer = $_reusedVer
+ Add-PythonDirToProcessPath $ReusedSetupPython
$PythonOk = $true
}
}
-if (-not $PythonOk -and $PyLauncher) {
+# Fall back to every py.exe on PATH (all-users and per-user launchers can both
+# register). -All is required: Windows PowerShell 5.1 returns only the first
+# launcher without it, and the PowerShell 7 multi-match array breaks the call
+# operator if used directly.
+$PyLaunchers = if ($PythonOk) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) }
+
+foreach ($PyLauncher in $PyLaunchers) {
+ if ($PyLauncher.Source -match $CondaSkipPattern) { continue }
foreach ($minor in @("3.13", "3.12", "3.11")) {
try {
$out = & $PyLauncher.Source "-$minor" --version 2>&1 | Out-String
@@ -1572,6 +1659,7 @@ if (-not $PythonOk -and $PyLauncher) {
}
} catch { }
}
+ if ($PythonOk) { break }
}
if (-not $PythonOk -and $HasPython) {
@@ -1804,36 +1892,33 @@ if (Test-Path $OxcValidatorDir) {
Write-Host ""
substep "setting up Python environment..."
-# Find Python -- skip Anaconda/Miniconda distributions.
-# Conda-bundled CPython ships modified DLL search paths that break
-# torch's c10.dll loading on Windows. Standalone CPython (python.org,
-# winget, uv) does not have this issue.
-# Uses Get-Command -All to look past conda entries that shadow a valid
-# standalone Python further down PATH, and probes py.exe (the Python
-# Launcher) which reliably finds python.org installs.
-#
-# NOTE: A venv created from conda Python inherits conda's base_prefix
-# even though the venv path itself does not contain "conda". We check
-# both the executable path AND sys.base_prefix to catch this case.
-$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)'
+# Find Python -- skip Anaconda/Miniconda distributions ($CondaSkipPattern and
+# Test-IsConda are defined above the 1g gate). Standalone CPython (python.org,
+# winget, uv) does not have conda's torch c10.dll loading issue.
$PythonCmd = $null
-# Helper: check if a Python executable is conda-based by inspecting
-# both the path and sys.base_prefix (catches venvs created from conda).
-function Test-IsConda {
- param([string]$Exe)
- if ($Exe -match $CondaSkipPattern) { return $true }
+# 0. Reuse the interpreter install.ps1 already resolved and built the venv with
+# (UNSLOTH_SETUP_PYTHON, or the existing venv python) before probing the
+# system -- it is already validated as supported and non-conda.
+if ($ReusedSetupPython) {
try {
- $basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim()
- if ($basePrefix -match $CondaSkipPattern) { return $true }
+ $out = & $ReusedSetupPython --version 2>&1 | Out-String
+ if ($out -match 'Python 3\.(\d+)') {
+ $pyMinor = [int]$Matches[1]
+ if ($pyMinor -ge 11 -and $pyMinor -le 13 -and -not (Test-IsConda $ReusedSetupPython)) {
+ $PythonCmd = $ReusedSetupPython
+ }
+ }
} catch { }
- return $false
}
# 1. Try the Python Launcher (py.exe) first -- most reliable on Windows.
-# py.exe is installed by python.org and resolves to standalone CPython.
-$pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
-if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) {
+# Enumerate every launcher with -All (Windows PowerShell 5.1 returns only
+# the first match without it) and search each for a supported, non-conda
+# interpreter.
+$PyLaunchersResolve = if ($PythonCmd) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) }
+foreach ($pyLauncher in $PyLaunchersResolve) {
+ if ($pyLauncher.Source -match $CondaSkipPattern) { continue }
foreach ($minor in @("3.13", "3.12", "3.11")) {
try {
$out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String
@@ -1851,6 +1936,7 @@ if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) {
}
} catch { }
}
+ if ($PythonCmd) { break }
}
# 2. Fall back to scanning python3.x / python3 / python on PATH.
diff --git a/studio/setup.sh b/studio/setup.sh
index 66c03da391..a8603bd0da 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -155,9 +155,60 @@ _nvcc_meets_llama_minimum() {
echo "$_raw"
}
+# Run a GPU probe under a 10s timeout when `timeout` is available so a wedged
+# NVIDIA driver cannot hang setup; fall back to a bare call where it is not.
+_setup_run_smi() {
+ if command -v timeout >/dev/null 2>&1; then
+ timeout 10 "$@"
+ else
+ "$@"
+ fi
+}
+
+# Returns 0 when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every NVIDIA
+# device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to the
+# AMD card). Unset means all devices visible. nvidia-smi ignores this env var,
+# so the probes below cannot see the distinction on their own.
+_setup_cvd_hides_nvidia() {
+ [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
+ _setup_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
+ [ -z "$_setup_cvd_trim" ] || [ "$_setup_cvd_trim" = "-1" ]
+}
+
+# Returns 0 when an NVIDIA GPU is present and usable. Primary probe is
+# `nvidia-smi -L` (timeout-bounded). Fallback is /proc/driver/nvidia/gpus,
+# which the driver populates per GPU regardless of nvidia-smi state -- handles
+# PATH gaps and driver init races. Mirrors install.sh _has_usable_nvidia_gpu
+# (PR 6174) so setup routes the same way as the torch installer. A GPU hidden
+# via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
+# install_llama_prebuilt.py has_usable_nvidia), so the AMD probes still run
+# and a mixed host steered to its AMD card keeps the ROCm route.
+_setup_has_usable_nvidia_gpu() {
+ if _setup_cvd_hides_nvidia; then
+ return 1
+ fi
+ _setup_nvsmi=""
+ if command -v nvidia-smi >/dev/null 2>&1; then
+ _setup_nvsmi="nvidia-smi"
+ elif [ -x "/usr/bin/nvidia-smi" ]; then
+ _setup_nvsmi="/usr/bin/nvidia-smi"
+ fi
+ if [ -n "$_setup_nvsmi" ]; then
+ if _setup_run_smi "$_setup_nvsmi" -L 2>/dev/null \
+ | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
+ return 0
+ fi
+ fi
+ if [ -d /proc/driver/nvidia/gpus ] && \
+ [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
+ return 0
+ fi
+ return 1
+}
+
_cuda_driver_max_version() {
command -v nvidia-smi >/dev/null 2>&1 || return 0
- nvidia-smi 2>/dev/null \
+ _setup_run_smi nvidia-smi 2>/dev/null \
| sed -nE 's/.*CUDA( UMD)? Version:[[:space:]]*([0-9]+)\.([0-9]+).*/\2.\3/p' \
| head -1 || true
}
@@ -815,25 +866,42 @@ _setup_amd_detected=false
_setup_nvidia_usable=false
_setup_gfx_all=""
_setup_mkt=""
-if command -v rocminfo >/dev/null 2>&1 && \
- rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
- _setup_amd_detected=true
- _setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
- _setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \
- '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
-elif command -v amd-smi >/dev/null 2>&1 && \
- amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
- _setup_amd_detected=true
- _setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
- [ -z "$_setup_gfx_all" ] && \
- _setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
- _setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
- '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
+# NVIDIA priority: classify NVIDIA first and skip the AMD probes entirely on
+# a usable-NVIDIA host (mirrors _has_rocm_gpu in install_python_stack.py).
+# This also keeps a wedged rocminfo/amd-smi from hanging setup before the
+# host is classified; the AMD probes themselves run under _setup_run_smi.
+if _setup_has_usable_nvidia_gpu; then
+ _setup_nvidia_usable=true
+fi
+if [ "$_setup_nvidia_usable" != true ]; then
+ if command -v rocminfo >/dev/null 2>&1 && \
+ _setup_run_smi rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
+ _setup_amd_detected=true
+ _setup_gfx_all=$(_setup_run_smi rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ _setup_mkt=$(_setup_run_smi rocminfo 2>/dev/null | awk -F': ' \
+ '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
+ elif command -v amd-smi >/dev/null 2>&1 && \
+ _setup_run_smi amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
+ _setup_amd_detected=true
+ _setup_gfx_all=$(_setup_run_smi amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ [ -z "$_setup_gfx_all" ] && \
+ _setup_gfx_all=$(_setup_run_smi amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ _setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
+ '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
+ elif [ -e /dev/kfd ] && \
+ awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
+ gpu && amd { found=1 } END{ exit !found }' \
+ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
+ # KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh
+ # _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are
+ # missing but the kernel exposes the GPU, so the source-build gate
+ # below does not drop them to a CPU llama.cpp build. No gfx arch is
+ # available from this path; name-based inference handles it.
+ _setup_amd_detected=true
+ fi
fi
-if command -v nvidia-smi >/dev/null 2>&1 && \
- nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
- _setup_nvidia_usable=true
+if [ "$_setup_nvidia_usable" = true ]; then
step "gpu" "NVIDIA GPU detected"
elif [ "$_setup_amd_detected" = true ]; then
_setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
@@ -918,15 +986,15 @@ _HOST_MACHINE="$(uname -m 2>/dev/null || true)"
# use unslothai.
_LINUX_HAS_GPU=false
# Route to the fork only for a usable GPU. NVIDIA counts only when a device is
-# actually enumerated (_setup_nvidia_usable, from the nvidia-smi -L probe above)
-# AND not hidden via CUDA_VISIBLE_DEVICES=-1 -- mirroring install_llama_prebuilt.py's
-# has_usable_nvidia. Mere nvidia-smi presence (CPU-only CUDA-toolkit containers,
-# broken drivers) or a hidden GPU therefore takes the ggml-org CPU prebuilt
-# instead of a slow source build. AMD is deliberately left on tooling presence,
-# not usability: an unusable NVIDIA host has a good CPU prebuilt to fall back to,
-# whereas tightening AMD would regress ROCm hosts exposing only hipconfig/hipinfo
-# into an unnecessary CPU build.
-if [ "$_setup_nvidia_usable" = true ] && [ "${CUDA_VISIBLE_DEVICES:-}" != "-1" ]; then
+# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1
+# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring
+# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence
+# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore
+# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is
+# deliberately left on tooling presence, not usability: an unusable NVIDIA host
+# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress
+# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build.
+if [ "$_setup_nvidia_usable" = true ]; then
_LINUX_HAS_GPU=true
else
for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do
@@ -1271,23 +1339,35 @@ else
GPU_BACKEND=""
NVCC_PATH=""
- if command -v nvcc &>/dev/null; then
- NVCC_PATH="$(command -v nvcc)"
- GPU_BACKEND="cuda"
- elif [ -x /usr/local/cuda/bin/nvcc ]; then
- NVCC_PATH="/usr/local/cuda/bin/nvcc"
- export PATH="/usr/local/cuda/bin:$PATH"
- GPU_BACKEND="cuda"
- elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
- # Pick the newest cuda-XX.X directory
- NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
- export PATH="$(dirname "$NVCC_PATH"):$PATH"
- GPU_BACKEND="cuda"
+ # Gate the CUDA toolkit search on an actually-usable NVIDIA GPU
+ # (_setup_nvidia_usable, computed in the GPU summary block above;
+ # already false when hidden via CUDA_VISIBLE_DEVICES=""/-1).
+ # A CUDA toolkit alone (CPU-only build container, leftover packages)
+ # is not proof of a GPU: building with -DGGML_CUDA=ON there yields a
+ # binary that fails at runtime, so fall through to the CPU build.
+ if [ "$_setup_nvidia_usable" = true ]; then
+ if command -v nvcc &>/dev/null; then
+ NVCC_PATH="$(command -v nvcc)"
+ GPU_BACKEND="cuda"
+ elif [ -x /usr/local/cuda/bin/nvcc ]; then
+ NVCC_PATH="/usr/local/cuda/bin/nvcc"
+ export PATH="/usr/local/cuda/bin:$PATH"
+ GPU_BACKEND="cuda"
+ elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
+ # Pick the newest cuda-XX.X directory
+ NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
+ export PATH="$(dirname "$NVCC_PATH"):$PATH"
+ GPU_BACKEND="cuda"
+ fi
fi
- # Check for ROCm (AMD) only if CUDA was not already selected
+ # Check for ROCm (AMD) only if CUDA was not already selected, and
+ # only when an AMD GPU was actually detected (_setup_amd_detected).
+ # hipcc presence alone (HIP SDK, no GPU) must not select a HIP build.
+ # NVIDIA-usable hosts never build HIP (defense in depth: the AMD
+ # probes above are already skipped when NVIDIA is usable).
ROCM_HIPCC=""
- if [ -z "$GPU_BACKEND" ]; then
+ if [ -z "$GPU_BACKEND" ] && [ "$_setup_nvidia_usable" != true ] && [ "$_setup_amd_detected" = true ]; then
if command -v hipcc &>/dev/null; then
ROCM_HIPCC="$(command -v hipcc)"
GPU_BACKEND="rocm"
@@ -1349,7 +1429,7 @@ else
CUDA_ARCHS=""
if command -v nvidia-smi &>/dev/null; then
- _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
+ _raw_caps=$(_setup_run_smi nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
while IFS= read -r _cap; do
_cap=$(echo "$_cap" | tr -d '[:space:]')
if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
@@ -1455,7 +1535,7 @@ else
CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}"
_BUILD_DESC="building (ROCm, ${GPU_TARGETS//;/+})"
fi
- elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
+ elif [ -d /usr/local/cuda ] || _setup_run_smi nvidia-smi &>/dev/null; then
_BUILD_DESC="building (CPU, CUDA driver found but nvcc missing)"
elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then
_BUILD_DESC="building (CPU, ROCm driver found but hipcc missing)"
diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh
index e20fd1ca86..3dece1f69a 100755
--- a/tests/sh/test_get_torch_index_url.sh
+++ b/tests/sh/test_get_torch_index_url.sh
@@ -13,6 +13,10 @@ FAIL=0
_FUNC_FILE=$(mktemp)
_FAKE_SMI_DIR=$(mktemp -d)
{
+ sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH"
+ echo ""
+ sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH"
+ echo ""
sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
@@ -107,7 +111,7 @@ MOCK
# Build a minimal tools directory with symlinks to essential commands
# (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi.
_TOOLS_DIR=$(mktemp -d)
-for _cmd in uname grep sed head sh bash cat awk printf; do
+for _cmd in uname grep sed head sh bash cat awk printf tr; do
_real=$(command -v "$_cmd" 2>/dev/null || true)
[ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
done
@@ -116,12 +120,19 @@ done
# $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
run_func() {
_mock_dir="$1"
+ # Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak
+ # in; a second argument sets it explicitly (hidden-GPU scenarios).
+ if [ "$#" -ge 2 ]; then
+ _cvd_setup="export CUDA_VISIBLE_DEVICES='$2'"
+ else
+ _cvd_setup="unset CUDA_VISIBLE_DEVICES"
+ fi
if [ "$_mock_dir" = "none" ]; then
# Minimal PATH with only basic tools, no nvidia-smi anywhere
- PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+ PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
else
# Put mock nvidia-smi dir first, then basic tools
- PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+ PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
fi
}
@@ -332,6 +343,40 @@ _result=$(run_func "$_dir")
assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
+# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present)
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "")
+assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
+# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present)
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "-1")
+assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
+# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored
+_cuda_dir=$(make_mock_smi "12.6")
+_amd_dir=$(make_mock_amd_smi "6.4")
+_combined_dir=$(mktemp -d)
+ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi"
+ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi"
+_result=$(run_func "$_combined_dir" "-1")
+assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result"
+rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir"
+
+# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "0")
+assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 38) Whitespace-padded "-1" still hides the GPU
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" " -1 ")
+assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"
diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py
new file mode 100644
index 0000000000..83c2d962e8
--- /dev/null
+++ b/tests/studio/install/test_cuda_repair.py
@@ -0,0 +1,248 @@
+"""Tests for CUDA torch repair on poisoned NVIDIA venvs.
+
+Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA
+torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD
+gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels,
+ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+# ── Load module under test (mirrors test_rocm_support.py) ────────────────────
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
+_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH)
+assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
+stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
+sys.modules[_STACK_SPEC.name] = stack_mod
+_STACK_SPEC.loader.exec_module(stack_mod)
+
+_ensure_cuda_torch = stack_mod._ensure_cuda_torch
+_detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url
+
+
+# ── Helpers ──────────────────────────────────────────────────────────────────
+
+
+def _make_run(
+ torch_state = "hip",
+ cuda_version = "12.8",
+ torch_rc = 0,
+ smi_rc = 0,
+):
+ """Build a subprocess.run side_effect.
+
+ The torch-classify probe runs sys.executable and reads bytes stdout; the
+ nvidia-smi version probe runs the smi path with text=True. Distinguish by
+ the executable.
+ """
+
+ def _run(cmd, *args, **kwargs):
+ result = MagicMock()
+ exe = str(cmd[0]) if cmd else ""
+ if exe == sys.executable:
+ result.returncode = torch_rc
+ result.stdout = (torch_state + "\n").encode()
+ return result
+ # nvidia-smi version probe (text = True)
+ result.returncode = smi_rc
+ out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n"
+ result.stdout = out if kwargs.get("text") else out.encode()
+ return result
+
+ return _run
+
+
+def _run_cuda_repair(
+ *,
+ backend = "",
+ nvidia = True,
+ torch_state = "hip",
+ cuda_version = "12.8",
+ torch_rc = 0,
+ smi_rc = 0,
+ is_macos = False,
+ is_windows = False,
+ no_torch = False,
+ rocm_marker = False,
+ smi_path = "/usr/bin/nvidia-smi",
+ cvd = None,
+):
+ """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
+
+ cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment
+ (the host machine may export one), any string sets it explicitly.
+ """
+ env = {}
+ if rocm_marker:
+ env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
+ if cvd is not None:
+ env["CUDA_VISIBLE_DEVICES"] = cvd
+
+ def _which(name, *a, **k):
+ if name == "nvidia-smi":
+ return smi_path
+ return None
+
+ with (
+ patch.object(stack_mod, "_TORCH_BACKEND", backend),
+ patch.object(stack_mod, "IS_MACOS", is_macos),
+ patch.object(stack_mod, "IS_WINDOWS", is_windows),
+ patch.object(stack_mod, "NO_TORCH", no_torch),
+ patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia),
+ patch.object(stack_mod.shutil, "which", side_effect = _which),
+ patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)),
+ patch.object(stack_mod, "pip_install") as mock_pip,
+ patch.object(
+ stack_mod.subprocess,
+ "run",
+ side_effect = _make_run(torch_state, cuda_version, torch_rc, smi_rc),
+ ),
+ patch.dict(stack_mod.os.environ, env, clear = False),
+ ):
+ if not rocm_marker:
+ stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
+ if cvd is None:
+ stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ _ensure_cuda_torch()
+ return mock_pip
+
+
+def _index_url(mock_pip) -> str:
+ """Return the --index-url value from the recorded pip_install call."""
+ args = [str(a) for a in mock_pip.call_args.args]
+ return args[args.index("--index-url") + 1]
+
+
+# ── Repair fires only on the poisoning signature ─────────────────────────────
+
+
+class TestCudaRepairFires:
+ def test_hip_build_on_nvidia_triggers_repair(self):
+ mock_pip = _run_cuda_repair(torch_state = "hip", cuda_version = "12.8")
+ assert mock_pip.call_count == 1
+ call_args = [str(a) for a in mock_pip.call_args.args]
+ assert "--force-reinstall" in call_args
+ assert "--no-cache-dir" in call_args
+ assert "cu128" in _index_url(mock_pip)
+ assert mock_pip.call_args.kwargs["constrain"] is False
+
+ def test_rocm_in_version_string_triggers_repair(self):
+ # AMD SDK / Radeon wheels may not set torch.version.hip but encode
+ # rocm in __version__; the probe prints "hip" for both.
+ mock_pip = _run_cuda_repair(torch_state = "hip")
+ assert mock_pip.call_count == 1
+
+
+# ── No-op cases ──────────────────────────────────────────────────────────────
+
+
+class TestCudaRepairSkips:
+ def test_healthy_cuda_torch_no_repair(self):
+ mock_pip = _run_cuda_repair(torch_state = "cuda")
+ mock_pip.assert_not_called()
+
+ def test_deliberate_cpu_wheel_no_repair(self):
+ mock_pip = _run_cuda_repair(torch_state = "cpu")
+ mock_pip.assert_not_called()
+
+ def test_backend_rocm_skips(self):
+ mock_pip = _run_cuda_repair(backend = "rocm", torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_backend_cpu_skips(self):
+ mock_pip = _run_cuda_repair(backend = "cpu", torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_unknown_backend_skips(self):
+ mock_pip = _run_cuda_repair(backend = "auto", torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_no_nvidia_gpu_skips(self):
+ mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_torch_missing_skips(self):
+ # Non-zero probe exit = torch missing / un-importable.
+ mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1)
+ mock_pip.assert_not_called()
+
+ def test_macos_skips(self):
+ mock_pip = _run_cuda_repair(is_macos = True, torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_windows_skips(self):
+ mock_pip = _run_cuda_repair(is_windows = True, torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_no_torch_mode_skips(self):
+ mock_pip = _run_cuda_repair(no_torch = True, torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_rocm_install_marker_skips(self):
+ mock_pip = _run_cuda_repair(rocm_marker = True, torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_cvd_minus_one_skips(self):
+ # CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed
+ # AMD+NVIDIA host running ROCm torch on the AMD card).
+ mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_cvd_empty_skips(self):
+ mock_pip = _run_cuda_repair(cvd = "", torch_state = "hip")
+ mock_pip.assert_not_called()
+
+ def test_cvd_explicit_device_still_repairs(self):
+ mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip")
+ assert mock_pip.call_count == 1
+
+
+# ── CUDA index ladder ────────────────────────────────────────────────────────
+
+
+class TestCudaIndexResolution:
+ def test_cuda_128_selects_cu128(self):
+ assert "cu128" in _index_url(_run_cuda_repair(cuda_version = "12.8"))
+
+ def test_cuda_130_selects_cu130(self):
+ assert "cu130" in _index_url(_run_cuda_repair(cuda_version = "13.0"))
+
+ def test_cuda_126_selects_cu126(self):
+ assert "cu126" in _index_url(_run_cuda_repair(cuda_version = "12.6"))
+
+ def test_cuda_124_selects_cu124(self):
+ assert "cu124" in _index_url(_run_cuda_repair(cuda_version = "12.4"))
+
+ def test_cuda_118_selects_cu118(self):
+ assert "cu118" in _index_url(_run_cuda_repair(cuda_version = "11.8"))
+
+ def test_unreadable_version_defaults_cu126(self):
+ # nvidia-smi runs but prints no CUDA version line (or fails).
+ mock_pip = _run_cuda_repair(cuda_version = "", smi_rc = 1)
+ assert "cu126" in _index_url(mock_pip)
+
+ def test_proc_fallback_no_smi_defaults_cu126(self):
+ # NVIDIA usable via /proc fallback, nvidia-smi absent entirely.
+ mock_pip = _run_cuda_repair(smi_path = None)
+ assert "cu126" in _index_url(mock_pip)
+
+ def test_detect_index_url_uses_pytorch_base(self):
+ with (
+ patch.object(stack_mod.shutil, "which", return_value = None),
+ patch.object(stack_mod.os.path, "isfile", return_value = False),
+ ):
+ url = _detect_cuda_torch_index_url()
+ assert url == f"{stack_mod._PYTORCH_WHL_BASE}/cu126"
+
+
+if __name__ == "__main__":
+ sys.exit(pytest.main([__file__, "-q"]))
diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py
new file mode 100644
index 0000000000..983969a4d3
--- /dev/null
+++ b/tests/studio/install/test_gpu_detection_followups.py
@@ -0,0 +1,480 @@
+"""Tests for the GPU-detection follow-ups to PR 6174.
+
+PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus
+fallback in install.sh and studio/install_python_stack.py. These tests cover the
+same hardening ported to the llama.cpp prebuilt installer
+(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh):
+
+ * detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi
+ is unavailable, and skips ROCm probing when NVIDIA is usable.
+ * setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback
+ and only selects a CUDA/ROCm source build when the matching GPU is detected.
+
+All tests use mocks or source-level assertions -- no GPU, network, or real
+nvidia-smi/rocminfo invocation.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+# Load studio/install_llama_prebuilt.py the same way the sibling suite does.
+_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+_SPEC = importlib.util.spec_from_file_location(
+ "studio_install_llama_prebuilt_followups", _MODULE_PATH
+)
+assert _SPEC is not None and _SPEC.loader is not None
+prebuilt_mod = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = prebuilt_mod
+_SPEC.loader.exec_module(prebuilt_mod)
+
+detect_host = prebuilt_mod.detect_host
+_apply_host_overrides = prebuilt_mod._apply_host_overrides
+
+SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
+
+
+def _make_run_capture(rocminfo_stdout: str = ""):
+ """Return a fake run_capture: rocminfo reports rocminfo_stdout, everything
+ else (nvidia-smi, amd-smi) returns empty so only the patched probes matter."""
+
+ def _run_capture(cmd, *args, **kwargs):
+ exe = str(cmd[0]) if cmd else ""
+ result = MagicMock()
+ if exe.endswith("rocminfo"):
+ result.returncode = 0
+ result.stdout = rocminfo_stdout
+ else:
+ result.returncode = 1
+ result.stdout = ""
+ result.stderr = ""
+ return result
+
+ return _run_capture
+
+
+def _run_detect_host(
+ *,
+ machine: str = "x86_64",
+ system: str = "Linux",
+ which_map: dict | None = None,
+ proc_dir_entries: list | None = None,
+ rocminfo_stdout: str = "",
+ env: dict | None = None,
+):
+ """Drive detect_host() against a fully synthetic host."""
+ which_map = which_map or {}
+ proc_dir_entries = proc_dir_entries if proc_dir_entries is not None else []
+
+ real_isdir = prebuilt_mod.os.path.isdir
+ real_listdir = prebuilt_mod.os.listdir
+ proc_path = "/proc/driver/nvidia/gpus"
+
+ def fake_isdir(p):
+ if str(p) == proc_path:
+ return bool(proc_dir_entries)
+ return real_isdir(p)
+
+ def fake_listdir(p):
+ if str(p) == proc_path:
+ if not proc_dir_entries:
+ raise OSError("no such dir")
+ return list(proc_dir_entries)
+ return real_listdir(p)
+
+ patches = [
+ patch.object(prebuilt_mod.platform, "system", return_value = system),
+ patch.object(prebuilt_mod.platform, "machine", return_value = machine),
+ patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")),
+ patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)),
+ patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)),
+ patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir),
+ patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir),
+ patch.object(prebuilt_mod.os, "access", return_value = False),
+ patch.dict(prebuilt_mod.os.environ, env or {}, clear = False),
+ ]
+ for p in patches:
+ p.start()
+ try:
+ # Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless
+ # the scenario sets it explicitly.
+ if env is None or "CUDA_VISIBLE_DEVICES" not in env:
+ prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ return detect_host()
+ finally:
+ for p in patches:
+ p.stop()
+
+
+# ── install_llama_prebuilt.detect_host(): /proc NVIDIA fallback ──────────────
+
+
+class TestDetectHostProcFallback:
+ def test_proc_fallback_marks_physical_nvidia_when_smi_absent(self):
+ """No nvidia-smi, but /proc/driver/nvidia/gpus is populated -> NVIDIA."""
+ host = _run_detect_host(
+ which_map = {}, # nvidia-smi resolves to None
+ proc_dir_entries = ["0000:01:00.0"],
+ )
+ assert host.has_physical_nvidia is True
+
+ def test_proc_fallback_has_usable_nvidia_when_devices_visible(self):
+ """Default CUDA_VISIBLE_DEVICES (unset) -> visible tokens non-empty -> usable."""
+ host = _run_detect_host(
+ which_map = {},
+ proc_dir_entries = ["0000:01:00.0"],
+ )
+ assert host.has_usable_nvidia is True
+
+ def test_proc_fallback_not_usable_when_devices_hidden(self):
+ """CUDA_VISIBLE_DEVICES='' hides all GPUs -> physical yes, usable no."""
+ host = _run_detect_host(
+ which_map = {},
+ proc_dir_entries = ["0000:01:00.0"],
+ env = {"CUDA_VISIBLE_DEVICES": ""},
+ )
+ assert host.has_physical_nvidia is True
+ assert host.has_usable_nvidia is False
+
+ def test_empty_proc_dir_does_not_mark_nvidia(self):
+ """A driver dir that exists but is empty must not assert a GPU."""
+ host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+ assert host.has_physical_nvidia is False
+
+ def test_proc_fallback_is_linux_only(self):
+ """The /proc fallback must not run on Windows (path is Linux-only)."""
+ host = _run_detect_host(
+ system = "Windows",
+ machine = "amd64",
+ which_map = {},
+ proc_dir_entries = ["0000:01:00.0"],
+ )
+ assert host.has_physical_nvidia is False
+
+
+# ── install_llama_prebuilt.detect_host(): NVIDIA precedence over ROCm ────────
+
+
+class TestDetectHostNvidiaPrecedence:
+ def test_rocm_probe_skipped_when_proc_nvidia_present(self):
+ """rocminfo reports gfx1100, but a proc-detected NVIDIA GPU wins."""
+ host = _run_detect_host(
+ which_map = {"rocminfo": "/usr/bin/rocminfo"},
+ proc_dir_entries = ["0000:01:00.0"],
+ rocminfo_stdout = " Name: gfx1100\n",
+ )
+ assert host.has_usable_nvidia is True
+ assert host.has_rocm is False
+
+ def test_rocm_detected_when_no_nvidia(self):
+ """With no NVIDIA signal at all, rocminfo gfx1100 -> has_rocm True."""
+ host = _run_detect_host(
+ which_map = {"rocminfo": "/usr/bin/rocminfo"},
+ proc_dir_entries = [],
+ rocminfo_stdout = " Name: gfx1100\n",
+ )
+ assert host.has_usable_nvidia is False
+ assert host.has_rocm is True
+
+
+# ── _apply_host_overrides: forwarded --rocm-gfx / --has-rocm still win ───────
+
+
+class TestOverridesStillWin:
+ def test_forwarded_gfx_forces_rocm_on_non_nvidia_host(self):
+ host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+ assert host.has_rocm is False
+ overridden = _apply_host_overrides(host, override_rocm_gfx = "gfx1100")
+ assert overridden.has_rocm is True
+ assert overridden.rocm_gfx_target == "gfx1100"
+
+ def test_override_has_rocm_forces_rocm(self):
+ host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+ overridden = _apply_host_overrides(host, override_has_rocm = True)
+ assert overridden.has_rocm is True
+
+ def test_force_cpu_drops_nvidia_attributes(self):
+ host = _run_detect_host(which_map = {}, proc_dir_entries = ["0000:01:00.0"])
+ assert host.has_usable_nvidia is True
+ overridden = _apply_host_overrides(host, force_cpu = True)
+ assert overridden.has_usable_nvidia is False
+ assert overridden.has_physical_nvidia is False
+ assert overridden.has_rocm is False
+
+
+# ── setup.sh source-level guarantees ────────────────────────────────────────
+
+
+class TestSetupShHardening:
+ @pytest.fixture(scope = "class")
+ def setup_src(self) -> str:
+ return SETUP_SH.read_text(encoding = "utf-8")
+
+ def test_has_usable_nvidia_helper_exists(self, setup_src):
+ assert "_setup_has_usable_nvidia_gpu()" in setup_src
+
+ def test_helper_uses_proc_fallback(self, setup_src):
+ start = setup_src.find("_setup_has_usable_nvidia_gpu()")
+ end = setup_src.find("\n}", start)
+ body = setup_src[start:end]
+ assert (
+ "/proc/driver/nvidia/gpus" in body
+ ), "_setup_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus"
+
+ def test_gpu_summary_uses_helper(self, setup_src):
+ assert "if _setup_has_usable_nvidia_gpu; then" in setup_src
+
+ def test_timeout_wrapper_exists(self, setup_src):
+ start = setup_src.find("_setup_run_smi()")
+ assert start >= 0, "_setup_run_smi timeout wrapper must exist"
+ end = setup_src.find("\n}", start)
+ body = setup_src[start:end]
+ assert "timeout 10" in body
+ assert "command -v timeout" in body
+
+ def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src):
+ """The nvcc source-build search must be gated on _setup_nvidia_usable.
+
+ The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside
+ _setup_has_usable_nvidia_gpu, so the gate itself only needs the flag.
+ """
+ anchor = setup_src.find('NVCC_PATH=""\n')
+ assert anchor >= 0
+ window = setup_src[anchor : anchor + 700]
+ assert (
+ 'if [ "$_setup_nvidia_usable" = true ]' in window
+ ), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc"
+
+ def test_nvidia_helper_honours_hidden_cvd(self, setup_src):
+ """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so
+ CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are
+ gated (mixed hosts steered to the AMD card keep the ROCm route)."""
+ assert "_setup_cvd_hides_nvidia()" in setup_src
+ start = setup_src.find("_setup_has_usable_nvidia_gpu() {")
+ end = setup_src.find("\n}", start)
+ body = setup_src[start:end]
+ assert "_setup_cvd_hides_nvidia" in body
+
+ def test_rocm_source_build_gated_on_amd_detected(self, setup_src):
+ """The hipcc source-build search must be gated on _setup_amd_detected."""
+ anchor = setup_src.find('ROCM_HIPCC=""')
+ assert anchor >= 0
+ window = setup_src[anchor : anchor + 400]
+ assert (
+ '[ "$_setup_amd_detected" = true ]' in window
+ ), "ROCm toolkit search must require a detected AMD GPU, not just hipcc"
+
+ def test_compute_cap_probe_timeout_wrapped(self, setup_src):
+ assert "_setup_run_smi nvidia-smi --query-gpu=compute_cap" in setup_src
+
+ def test_driver_version_probe_timeout_wrapped(self, setup_src):
+ start = setup_src.find("_cuda_driver_max_version()")
+ end = setup_src.find("\n}", start)
+ body = setup_src[start:end]
+ assert "_setup_run_smi nvidia-smi" in body
+
+
+# TEST: install.sh -- UNSLOTH_TORCH_BACKEND classified on the final path segment
+
+
+class TestBackendExportLeafClassification:
+ """A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or
+ "gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the
+ final path segment of TORCH_INDEX_URL only."""
+
+ @pytest.fixture(scope = "class")
+ def install_src(self) -> str:
+ return (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+
+ def test_export_block_uses_leaf(self, install_src):
+ anchor = install_src.find("_torch_index_leaf=")
+ assert anchor >= 0, "backend export must classify on the final path segment"
+ window = install_src[anchor : anchor + 500]
+ assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window
+ assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window
+ assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window
+
+ def test_leaf_classification_behaviour(self, tmp_path):
+ import subprocess as sp
+
+ script = tmp_path / "leaf.sh"
+ src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+ anchor = src.find("_torch_index_leaf=")
+ block = src[anchor : src.find("esac", anchor) + 4]
+ # Drive the extracted block with adversarial mirror URLs.
+ script.write_text(
+ "#!/bin/sh\n"
+ 'TORCH_INDEX_URL="$1"\n' + block + "\n"
+ 'printf "%s" "$UNSLOTH_TORCH_BACKEND"\n'
+ )
+ cases = {
+ "https://download.pytorch.org/whl/cu128": "cuda",
+ "https://download.pytorch.org/whl/cpu": "cpu",
+ "https://download.pytorch.org/whl/rocm6.4": "rocm",
+ "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/": "rocm",
+ "https://repo.amd.com/rocm/whl/gfx1151/": "rocm",
+ "https://mirror.local/rocm-cache/cu128": "cuda",
+ "https://mirror.local/gfx-cache/cpu": "cpu",
+ }
+ for url, expected in cases.items():
+ out = sp.run(
+ ["sh", str(script), url], capture_output = True, text = True, timeout = 30
+ ).stdout.strip()
+ assert out == expected, f"{url} classified as {out!r}, expected {expected!r}"
+
+
+# TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper
+
+
+_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
+_STACK_SPEC = importlib.util.spec_from_file_location(
+ "studio_install_python_stack_followups", _STACK_PATH
+)
+assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
+stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
+sys.modules[_STACK_SPEC.name] = stack_mod
+_STACK_SPEC.loader.exec_module(stack_mod)
+
+
+def _stack_nvidia_usable(cvd):
+ """Drive install_python_stack._has_usable_nvidia_gpu with a mocked
+ nvidia-smi that always reports a GPU; cvd = None removes the env var."""
+
+ def fake_run(cmd, *args, **kwargs):
+ result = MagicMock()
+ result.returncode = 0
+ result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
+ return result
+
+ env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd}
+ with (
+ patch.object(
+ stack_mod.shutil,
+ "which",
+ side_effect = lambda n: "/usr/bin/nvidia-smi" if n == "nvidia-smi" else None,
+ ),
+ patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
+ patch.dict(stack_mod.os.environ, env, clear = False),
+ ):
+ if cvd is None:
+ stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ return stack_mod._has_usable_nvidia_gpu()
+
+
+class TestHiddenCvdNotUsable:
+ """CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA
+ device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three
+ _has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh,
+ setup.sh) must report the GPU as not usable so the AMD/CPU routes run,
+ matching install_llama_prebuilt.py's has_usable_nvidia."""
+
+ def test_python_unset_cvd_is_usable(self):
+ assert _stack_nvidia_usable(None) is True
+
+ def test_python_empty_cvd_not_usable(self):
+ assert _stack_nvidia_usable("") is False
+
+ def test_python_minus_one_not_usable(self):
+ assert _stack_nvidia_usable("-1") is False
+
+ def test_python_padded_minus_one_not_usable(self):
+ assert _stack_nvidia_usable(" -1 ") is False
+
+ def test_python_explicit_device_is_usable(self):
+ assert _stack_nvidia_usable("0") is True
+
+ def test_python_device_list_is_usable(self):
+ assert _stack_nvidia_usable("0,1") is True
+
+ def test_hidden_nvidia_restores_rocm_detection(self):
+ """Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100:
+ _has_rocm_gpu must proceed past the NVIDIA guard and return True
+ (before this fix the guard ignored CVD and blocked ROCm)."""
+
+ def fake_run(cmd, *args, **kwargs):
+ result = MagicMock()
+ result.returncode = 0
+ exe = str(cmd[0])
+ if exe.endswith("rocminfo"):
+ result.stdout = " Name: gfx1100\n"
+ else:
+ result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
+ return result
+
+ which_map = {
+ "rocminfo": "/usr/bin/rocminfo",
+ "nvidia-smi": "/usr/bin/nvidia-smi",
+ }
+ with (
+ patch.object(stack_mod.shutil, "which", side_effect = which_map.get),
+ patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
+ patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False),
+ ):
+ assert stack_mod._has_rocm_gpu() is True
+
+ @staticmethod
+ def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd):
+ """Extract shell functions, run the usable-GPU one against a fake
+ nvidia-smi, and return "usable"/"not_usable"."""
+ import os as _os
+ import subprocess as sp
+
+ blocks = []
+ for name in fn_names:
+ start = src.find(f"{name}() {{")
+ assert start >= 0, f"{name} missing"
+ end = src.find("\n}", start) + 2
+ blocks.append(src[start:end])
+ fake_bin = tmp_path / "bin"
+ fake_bin.mkdir(exist_ok = True)
+ smi = fake_bin / "nvidia-smi"
+ smi.write_text("#!/bin/sh\necho 'GPU 0: NVIDIA Fake (UUID: GPU-x)'\n")
+ smi.chmod(0o755)
+ script = tmp_path / "probe.sh"
+ script.write_text(
+ "#!/bin/sh\n" + "\n".join(blocks) + "\n"
+ f"if {fn_names[-1]}; then echo usable; else echo not_usable; fi\n"
+ )
+ env = dict(_os.environ)
+ env["PATH"] = f"{fake_bin}:{env['PATH']}"
+ if cvd is None:
+ env.pop("CUDA_VISIBLE_DEVICES", None)
+ else:
+ env["CUDA_VISIBLE_DEVICES"] = cvd
+ return sp.run(
+ ["sh", str(script)], capture_output = True, text = True, timeout = 30, env = env
+ ).stdout.strip()
+
+ @pytest.mark.parametrize(
+ "cvd, expected",
+ [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
+ )
+ def test_install_sh_helper_cvd(self, tmp_path, cvd, expected):
+ src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+ out = self._run_sh_helper(
+ tmp_path,
+ src,
+ ["_run_bounded", "_cvd_hides_nvidia", "_has_usable_nvidia_gpu"],
+ cvd,
+ )
+ assert out == expected
+
+ @pytest.mark.parametrize(
+ "cvd, expected",
+ [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
+ )
+ def test_setup_sh_helper_cvd(self, tmp_path, cvd, expected):
+ src = SETUP_SH.read_text(encoding = "utf-8")
+ out = self._run_sh_helper(
+ tmp_path,
+ src,
+ ["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"],
+ cvd,
+ )
+ assert out == expected
diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py
new file mode 100644
index 0000000000..7c3296d4fd
--- /dev/null
+++ b/tests/studio/install/test_hf_auth.py
@@ -0,0 +1,114 @@
+"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches.
+
+Anonymous huggingface.co downloads (tiny GGUF validation model) share a
+per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt
+path into a source build. auth_headers now sends HF_TOKEN to huggingface.co
+hosts, and a redirect handler strips Authorization when the download is
+redirected to a different host (CDN signed URLs). All tests are offline.
+"""
+
+import importlib.util
+import sys
+import urllib.request
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+_SPEC = importlib.util.spec_from_file_location(
+ "studio_install_llama_prebuilt_hf_auth", _MODULE_PATH
+)
+assert _SPEC is not None and _SPEC.loader is not None
+mod = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = mod
+_SPEC.loader.exec_module(mod)
+
+_TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN")
+HF_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
+GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
+
+
+def _headers(url, env):
+ """auth_headers under a fully controlled token environment."""
+ with patch.dict(mod.os.environ, env, clear = False):
+ for var in _TOKEN_VARS:
+ if var not in env:
+ mod.os.environ.pop(var, None)
+ return mod.auth_headers(url)
+
+
+class TestAuthHeaderRouting:
+ def test_hf_token_sent_to_huggingface(self):
+ headers = _headers(HF_URL, {"HF_TOKEN": "hf_x"})
+ assert headers.get("Authorization") == "Bearer hf_x"
+
+ def test_hub_token_fallback(self):
+ headers = _headers(HF_URL, {"HUGGING_FACE_HUB_TOKEN": "hf_y"})
+ assert headers.get("Authorization") == "Bearer hf_y"
+
+ def test_hf_token_not_sent_to_github(self):
+ headers = _headers(GH_URL, {"HF_TOKEN": "hf_x"})
+ assert "Authorization" not in headers
+
+ def test_hf_token_not_sent_to_other_hosts(self):
+ headers = _headers("https://cdn-lfs.huggingface.co/x", {"HF_TOKEN": "hf_x"})
+ assert "Authorization" not in headers
+
+ def test_gh_token_not_sent_to_huggingface(self):
+ headers = _headers(HF_URL, {"GH_TOKEN": "gh_x"})
+ assert "Authorization" not in headers
+
+ def test_gh_token_still_wins_on_github(self):
+ headers = _headers(GH_URL, {"GH_TOKEN": "gh_x", "HF_TOKEN": "hf_x"})
+ assert headers.get("Authorization") == "Bearer gh_x"
+
+ def test_no_tokens_no_auth(self):
+ assert "Authorization" not in _headers(HF_URL, {})
+
+ def test_validation_model_url_is_hf(self):
+ assert mod.should_send_hf_auth(mod.TEST_MODEL_URL) is True
+
+
+class TestCrossHostRedirectStripsAuth:
+ def _redirect(self, newurl):
+ req = urllib.request.Request(HF_URL, headers = {"Authorization": "Bearer hf_x"})
+ handler = mod._CrossHostAuthStrippingRedirectHandler()
+ return handler.redirect_request(req, None, 302, "Found", {}, newurl)
+
+ def test_cross_host_redirect_drops_authorization(self):
+ new_request = self._redirect("https://cdn-lfs.huggingface.co/signed/blob")
+ assert new_request is not None
+ assert "Authorization" not in new_request.headers
+ assert "Authorization" not in new_request.unredirected_hdrs
+
+ def test_same_host_redirect_keeps_authorization(self):
+ new_request = self._redirect("https://huggingface.co/elsewhere/blob")
+ assert new_request is not None
+ assert new_request.headers.get("Authorization") == "Bearer hf_x"
+
+
+class TestDownloadBytesWiring:
+ def test_download_bytes_sends_hf_auth(self):
+ response = MagicMock()
+ response.__enter__ = lambda s: s
+ response.__exit__ = lambda s, *a: False
+ response.headers.get.return_value = None
+ response.read.side_effect = [b"data", b""]
+ with (
+ patch.object(mod._URL_OPENER, "open", return_value = response) as opened,
+ patch.dict(mod.os.environ, {"HF_TOKEN": "hf_x"}, clear = False),
+ ):
+ for var in ("GH_TOKEN", "GITHUB_TOKEN"):
+ mod.os.environ.pop(var, None)
+ data = mod.download_bytes(HF_URL)
+ assert data == b"data"
+ request = opened.call_args.args[0]
+ assert request.headers.get("Authorization") == "Bearer hf_x"
+
+
+if __name__ == "__main__":
+ sys.exit(pytest.main([__file__, "-q"]))
diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py
index 057612d27b..000c5a7ace 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -1,3 +1,4 @@
+import errno
import importlib.util
import io
import json
@@ -28,6 +29,7 @@ ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums
hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree
validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice
activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree
+activate_staged_dir = INSTALL_LLAMA_PREBUILT.activate_staged_dir
create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir
sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file
source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name
@@ -206,6 +208,110 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents(
assert not (install_dir / f"llama.cpp-{upstream_tag}").exists()
+def test_release_asset_download_url():
+ fn = INSTALL_LLAMA_PREBUILT.release_asset_download_url
+ assert fn(
+ "unslothai/llama.cpp", "b9000-mix-abc1234", "llama.cpp-source-commit-deadbeef.tar.gz"
+ ) == (
+ "https://github.com/unslothai/llama.cpp/releases/download/"
+ "b9000-mix-abc1234/llama.cpp-source-commit-deadbeef.tar.gz"
+ )
+ # Any missing component -> None (no asset url, caller falls back to codeload).
+ assert fn(None, "b9000", "x.tar.gz") is None
+ assert fn("unslothai/llama.cpp", None, "x.tar.gz") is None
+ assert fn("unslothai/llama.cpp", "b9000", None) is None
+
+
+def _mk_source_tarball(path: Path, tag: str) -> None:
+ with tarfile.open(path, "w:gz") as archive:
+ add_bytes_to_tar(
+ archive, f"llama.cpp-{tag}/CMakeLists.txt", b"cmake_minimum_required(VERSION 3.14)\n"
+ )
+ add_bytes_to_tar(
+ archive,
+ f"llama.cpp-{tag}/convert_hf_to_gguf.py",
+ b"#!/usr/bin/env python3\nimport gguf\n",
+ )
+ add_bytes_to_tar(archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n")
+
+
+def test_hydrate_source_tree_prefers_release_asset_for_mix(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+ # A mix build's merge commit is in no repo, so the codeload/archive URLs 404.
+ # hydrate must fetch the release asset and never touch codeload.
+ commit = "a" * 40
+ archive_path = tmp_path / "merged-source.tar.gz"
+ _mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}")
+ asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
+ "unslothai/llama.cpp", "b9000-mix-abc1234", f"llama.cpp-source-commit-{commit}.tar.gz"
+ )
+ codeload_urls = set(
+ INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
+ )
+ seen = []
+
+ def fake_download_file(url: str, destination: Path) -> None:
+ seen.append(url)
+ if url in codeload_urls:
+ raise AssertionError("codeload was hit even though the release asset was available")
+ assert url == asset_url
+ destination.write_bytes(archive_path.read_bytes())
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+
+ install_dir = tmp_path / "install"
+ work_dir = tmp_path / "work"
+ work_dir.mkdir()
+ hydrate_source_tree(
+ commit,
+ install_dir,
+ work_dir,
+ source_repo = "unslothai/llama.cpp",
+ expected_sha256 = sha256_file(archive_path),
+ exact_source = True,
+ asset_url = asset_url,
+ )
+ assert seen == [asset_url]
+ assert (install_dir / "CMakeLists.txt").exists()
+ assert (install_dir / "convert_hf_to_gguf.py").exists()
+
+
+def test_hydrate_source_tree_falls_back_to_codeload_when_asset_missing(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+ # If the release asset 404s, fall back to codeload/archive (vanilla path).
+ commit = "b" * 40
+ archive_path = tmp_path / "vanilla-source.tar.gz"
+ _mk_source_tarball(archive_path, f"commit-{commit[:7]}")
+ asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
+ "unslothai/llama.cpp", "b9000", f"llama.cpp-source-commit-{commit}.tar.gz"
+ )
+ codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
+
+ def fake_download_file(url: str, destination: Path) -> None:
+ if url == asset_url:
+ raise RuntimeError("404 Not Found")
+ assert url in codeload_urls
+ destination.write_bytes(archive_path.read_bytes())
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+
+ install_dir = tmp_path / "install"
+ work_dir = tmp_path / "work"
+ work_dir.mkdir()
+ hydrate_source_tree(
+ commit,
+ install_dir,
+ work_dir,
+ source_repo = "unslothai/llama.cpp",
+ expected_sha256 = sha256_file(archive_path),
+ exact_source = True,
+ asset_url = asset_url,
+ )
+ assert (install_dir / "CMakeLists.txt").exists()
+
+
def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
@@ -562,6 +668,48 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
assert "removing rollback path" in output
+def test_activate_staged_dir_copies_when_replace_hits_busy_lock(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+):
+ staging_dir = tmp_path / "llama.cpp.staging-test"
+ (staging_dir / "bin").mkdir(parents = True)
+ (staging_dir / "bin" / "ggml-base.dll").write_bytes(b"fake dll")
+ dst = tmp_path / "llama.cpp"
+
+ def denied_replace(src, dst_arg):
+ raise PermissionError(errno.EACCES, "Access is denied", str(src))
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", denied_replace)
+
+ activate_staged_dir(staging_dir, dst)
+
+ assert (dst / "bin" / "ggml-base.dll").read_bytes() == b"fake dll"
+ assert not staging_dir.exists()
+
+ captured = capsys.readouterr()
+ assert "falling back to file-by-file copy" in captured.out + captured.err
+
+
+def test_activate_staged_dir_reraises_non_busy_errors(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+ staging_dir = tmp_path / "llama.cpp.staging-test"
+ staging_dir.mkdir()
+ (staging_dir / "new.txt").write_text("new install\n")
+ dst = tmp_path / "llama.cpp"
+
+ def out_of_space_replace(src, dst_arg):
+ raise OSError(errno.ENOSPC, "No space left on device", str(src))
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", out_of_space_replace)
+
+ with pytest.raises(OSError, match = "No space left on device"):
+ activate_staged_dir(staging_dir, dst)
+
+ assert not dst.exists()
+ assert (staging_dir / "new.txt").read_text() == "new install\n"
+
+
def test_binary_env_linux_includes_binary_parent_in_ld_library_path(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py
new file mode 100644
index 0000000000..acea0ed34d
--- /dev/null
+++ b/tests/studio/install/test_probe_timeouts.py
@@ -0,0 +1,202 @@
+"""Tests that NVIDIA probes in the installers are bounded by a timeout.
+
+Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer,
+and the Windows probe must require a real GPU listing (not just exit code 0).
+
+Source-level assertions verify the guards are present in install.sh / install.ps1
+/ setup.ps1; one behavioral shell test confirms the bash helper actually returns
+within the timeout when nvidia-smi hangs.
+"""
+
+import os
+import shutil
+import stat
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+INSTALL_SH = PACKAGE_ROOT / "install.sh"
+INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
+SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
+
+
+def _extract_sh_function_body(source: str, name: str) -> str:
+ """Return a shell function body from `source` by brace matching."""
+ needle = f"{name}() {{"
+ start = source.find(needle)
+ if start < 0:
+ return ""
+ depth = 0
+ i = start + len(needle) - 1
+ n = len(source)
+ while i < n:
+ ch = source[i]
+ if ch == "{":
+ depth += 1
+ elif ch == "}":
+ depth -= 1
+ if depth == 0:
+ return source[start : i + 1]
+ i += 1
+ return source[start:]
+
+
+# ── install.sh: _run_bounded helper and its use at every nvidia-smi call ──
+
+
+class TestInstallShBoundedProbe:
+ def _src(self) -> str:
+ return INSTALL_SH.read_text(encoding = "utf-8")
+
+ def test_run_bounded_helper_defined(self):
+ body = _extract_sh_function_body(self._src(), "_run_bounded")
+ assert body, "install.sh must define a _run_bounded helper"
+ assert (
+ "command -v timeout" in body
+ ), "_run_bounded must check for the `timeout` binary before using it"
+ assert "timeout 10" in body, "_run_bounded must apply a 10s timeout"
+ # Must fall back to running unbounded when `timeout` is unavailable
+ # (e.g. macOS) so semantics are unchanged there.
+ assert (
+ "else" in body and '"$@"' in body
+ ), "_run_bounded must run the command unbounded when `timeout` is absent"
+
+ def test_nvidia_smi_dash_l_probe_is_bounded(self):
+ body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
+ assert body, "install.sh must define _has_usable_nvidia_gpu"
+ # The -L probe must go through the bounded runner, not call nvidia-smi raw.
+ assert (
+ '_run_bounded "$_nvsmi" -L' in body
+ ), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded"
+ # The /proc fallback from PR 6174 must still be present.
+ assert "/proc/driver/nvidia" in body
+
+ def test_cuda_version_parse_is_bounded(self):
+ body = _extract_sh_function_body(self._src(), "get_torch_index_url")
+ assert body, "install.sh must define get_torch_index_url"
+ assert (
+ "_run_bounded" in body
+ ), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded"
+ # The locale must be forced without depending on `env` being on PATH.
+ assert "LC_ALL=C" in body
+ # _nvidia_detected gating from PR 6174 must remain.
+ assert "_nvidia_detected" in body
+
+ def test_no_unbounded_nvidia_smi_invocation_remains(self):
+ """Every nvidia-smi *execution* in install.sh goes through _run_bounded.
+
+ `command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution
+ checks, not executions, and are allowed. An execution looks like
+ `"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`.
+ """
+ body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
+ body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url")
+ # In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded.
+ assert '"$_nvsmi" -L' not in body_nvidia.replace(
+ '_run_bounded "$_nvsmi" -L', ""
+ ), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu"
+ # In get_torch_index_url the $_smi execution must be bounded.
+ assert (
+ "LC_ALL=C $_smi" not in body_torch
+ ), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url"
+
+
+# ── install.ps1 / setup.ps1: bounded, GPU-row-validated Windows probe ──
+
+
+class TestPowerShellBoundedProbe:
+ @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+ def test_bounded_helper_present(self, path):
+ src = path.read_text(encoding = "utf-8")
+ assert (
+ "function Invoke-NvidiaSmiBounded" in src
+ ), f"{path.name} must define Invoke-NvidiaSmiBounded"
+ assert (
+ "WaitForExit($TimeoutSec * 1000)" in src
+ ), f"{path.name} bounded probe must use WaitForExit with a timeout"
+ # Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate.
+ assert (
+ "$proc.Kill()" in src and "124" in src
+ ), f"{path.name} must kill nvidia-smi and signal a timeout exit code"
+
+ @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+ def test_probe_requires_gpu_row(self, path):
+ src = path.read_text(encoding = "utf-8")
+ assert (
+ "function Test-NvidiaSmiHasGpu" in src
+ ), f"{path.name} must define Test-NvidiaSmiHasGpu"
+ assert "@('-L')" in src, f"{path.name} must probe nvidia-smi with -L"
+ assert (
+ "^GPU\\s+\\d+:" in src
+ ), f"{path.name} must require a 'GPU :' data row, not just exit code 0"
+
+ @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+ def test_detection_uses_validated_probe(self, path):
+ src = path.read_text(encoding = "utf-8")
+ # The exit-code-only pattern must be gone from the detection block.
+ assert (
+ "& $nvSmiCmd.Source *> $null" not in src
+ ), f"{path.name} must not use the exit-code-only nvidia-smi probe"
+ assert (
+ "Test-NvidiaSmiHasGpu $nvSmiCmd.Source" in src
+ ), f"{path.name} PATH probe must use Test-NvidiaSmiHasGpu"
+ assert (
+ "Test-NvidiaSmiHasGpu $p" in src
+ ), f"{path.name} hardcoded-path fallback must use Test-NvidiaSmiHasGpu"
+
+
+# ── Behavioral: a hanging nvidia-smi must not hang _has_usable_nvidia_gpu ──
+
+
+def _have_timeout() -> bool:
+ return shutil.which("timeout") is not None
+
+
+@pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available")
+def test_has_usable_nvidia_gpu_returns_under_timeout():
+ """Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake
+ nvidia-smi that sleeps 30s, and assert the probe returns well under that.
+ """
+ src = INSTALL_SH.read_text(encoding = "utf-8")
+ helper = _extract_sh_function_body(src, "_run_bounded")
+ fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu")
+ assert helper and fn
+
+ workdir = tempfile.mkdtemp(prefix = "pr6174_timeout_", dir = str(PACKAGE_ROOT.parent))
+ try:
+ fake_dir = Path(workdir, "bin")
+ fake_dir.mkdir()
+ fake_smi = fake_dir / "nvidia-smi"
+ fake_smi.write_text("#!/bin/sh\nsleep 30\n")
+ fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
+
+ # Build a minimal PATH that includes the fake nvidia-smi plus the real
+ # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins.
+ real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")}
+ path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins])
+
+ # Force the /proc fallback off so the result depends only on the probe,
+ # and so a host with real NVIDIA does not mask the timeout behaviour.
+ script = (
+ f"{helper}\n{fn}\n"
+ "if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n"
+ )
+ proc = subprocess.run(
+ ["sh", "-c", script],
+ env = {"PATH": path_env},
+ stdout = subprocess.PIPE,
+ stderr = subprocess.DEVNULL,
+ text = True,
+ timeout = 20, # generous: the internal timeout is 10s, sleep is 30s
+ )
+ # The probe must have returned (not hung). On this CI host /proc/driver/
+ # nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA
+ # host the /proc fallback yields DETECTED. Either way it must not hang.
+ assert proc.stdout.strip() in {"NONE", "DETECTED"}
+ finally:
+ shutil.rmtree(workdir, ignore_errors = True)
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 973c0a1e81..9aa06b6ff7 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -719,6 +719,143 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
+ def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
+ """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe."""
+ with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}):
+ # Reload _TORCH_BACKEND from the patched environment.
+ with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
+ _ensure_rocm_torch()
+ mock_pip.assert_not_called()
+
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
+ def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
+ """UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe."""
+ with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}):
+ with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"):
+ _ensure_rocm_torch()
+ mock_pip.assert_not_called()
+
+
+# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
+
+
+class TestHasRocmGpuKfdVendorGuard:
+ """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes.
+
+ These tests are source-level: they verify the regex and logic present in
+ the _has_rocm_gpu implementation rather than running the sysfs traversal
+ (which requires Linux path conventions).
+ """
+
+ def _src(self) -> str:
+ """Return the source of _has_rocm_gpu from install_python_stack.py."""
+ import inspect
+ return inspect.getsource(stack_mod._has_rocm_gpu)
+
+ def test_vendor_id_check_present(self):
+ """_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002)."""
+ src = self._src()
+ assert "vendor_id" in src, (
+ "_has_rocm_gpu KFD sysfs fallback must read the properties file "
+ "to check vendor_id and exclude NVIDIA KFD nodes"
+ )
+ assert "4098" in src, (
+ "_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the "
+ "KFD node properties to avoid false positives on NVIDIA systems"
+ )
+
+ def test_vendor_regex_pattern_anchored(self):
+ """The vendor_id regex must use a word boundary to avoid partial matches."""
+ import re as _re
+
+ src = self._src()
+ # The pattern should have a word boundary before and after the number
+ # so "vendor_id 41098" doesn't match "vendor_id 4098".
+ assert (
+ _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
+ ), "_has_rocm_gpu vendor_id check should use word boundary anchors"
+
+ def test_sysfs_fallback_guarded_by_non_win32(self):
+ """KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
+ src = self._src()
+ assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check"
+
+ def test_cpu_node_excluded(self):
+ """gpu_id == '0' must be excluded (CPU topology nodes)."""
+ src = self._src()
+ assert (
+ '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
+ ), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
+
+ def test_install_sh_has_vendor_check(self):
+ """_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("_has_amd_rocm_gpu()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id"
+ assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)"
+
+ def test_has_rocm_gpu_returns_false_when_nvidia_present(self):
+ """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True.
+
+ This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs
+ produce a false positive, an NVIDIA GPU always wins.
+ """
+ with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True):
+ with patch("shutil.which", return_value = "/usr/bin/rocminfo"):
+ # Simulate rocminfo claiming an AMD GPU is present
+ mock_result = MagicMock()
+ mock_result.returncode = 0
+ mock_result.stdout = "Name: gfx1100\n"
+ with patch("subprocess.run", return_value = mock_result):
+ assert not stack_mod._has_rocm_gpu(), (
+ "_has_rocm_gpu must return False when NVIDIA GPU is detected, "
+ "regardless of what rocminfo reports"
+ )
+
+ def test_install_sh_has_rocm_gpu_nvidia_guard(self):
+ """_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true."""
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("_has_amd_rocm_gpu()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert (
+ "_has_usable_nvidia_gpu" in func_body
+ ), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
+ assert (
+ "return 1" in func_body
+ ), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
+
+ def test_has_usable_nvidia_gpu_proc_fallback_present(self):
+ """`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
+ import inspect
+
+ src = inspect.getsource(stack_mod._has_usable_nvidia_gpu)
+ assert "/proc/driver/nvidia" in src, (
+ "_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when "
+ "nvidia-smi subprocess fails, to handle PATH gaps and driver init races"
+ )
+
+ def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self):
+ """_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback."""
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("_has_usable_nvidia_gpu()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert "/proc/driver/nvidia" in func_body, (
+ "_has_usable_nvidia_gpu in install.sh must fall back to "
+ "/proc/driver/nvidia/gpus when nvidia-smi fails"
+ )
+
# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping
@@ -927,16 +1064,21 @@ class TestInstallShStructure:
source = sh_path.read_text(encoding = "utf-8")
body = _extract_sh_function_body(source, "get_torch_index_url")
nvidia_call = body.find("_has_usable_nvidia_gpu")
- no_nvidia_branch = body.find('if [ -z "$_smi" ]')
+ # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to
+ # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu
+ # returns true via /proc/driver/nvidia/gpus.
+ no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]')
+ if no_nvidia_branch < 0:
+ no_nvidia_branch = body.find('if [ -z "$_smi" ]')
rocm_call = body.find("_has_amd_rocm_gpu")
assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
- assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia-smi"
+ assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
assert (
rocm_call > no_nvidia_branch
- ), "ROCm detection should sit inside the 'no nvidia-smi' branch"
+ ), "ROCm detection should sit inside the 'no NVIDIA' branch"
assert (
nvidia_call < no_nvidia_branch
- ), "NVIDIA detection should run before the no-nvidia-smi branch"
+ ), "NVIDIA detection should run before the no-NVIDIA branch"
def test_bitsandbytes_amd_install(self):
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
@@ -1018,6 +1160,89 @@ class TestInstallShStructure:
rocm_pos = func_body.find("amd-smi")
assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection"
+ def test_unsloth_torch_backend_exported_after_get_torch_index_url(self):
+ """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set.
+
+ This lets install_python_stack.py skip ROCm torch operations on CUDA
+ and CPU hosts without re-running GPU detection in a subprocess.
+ """
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
+ backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
+ assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
+ assert (
+ backend_pos > torch_url_pos
+ ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
+ # Verify all three cases are covered
+ assert '"cuda"' in source[backend_pos : backend_pos + 500]
+ assert '"rocm"' in source[backend_pos : backend_pos + 500]
+ assert '"cpu"' in source[backend_pos : backend_pos + 500]
+ # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it
+ assert "export UNSLOTH_TORCH_BACKEND" in source
+
+ def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self):
+ """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098.
+
+ NVIDIA open kernel module (560+) registers KFD nodes with vendor_id
+ 4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0
+ (true) on NVIDIA-only hosts that have the nvidia-open driver, causing
+ get_torch_index_url to select a ROCm wheel index.
+ """
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("_has_amd_rocm_gpu()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert (
+ "vendor_id" in func_body
+ ), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
+ assert (
+ "4098" in func_body
+ ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
+
+ def test_kfd_awk_resets_state_per_file(self):
+ """KFD sysfs awk must reset gpu/amd state per file (FNR==1).
+
+ Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU
+ agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU
+ (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0
+ sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1
+ before vendor_id 4318 is seen on node 1.
+ """
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("_has_amd_rocm_gpu()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert "FNR==1" in func_body, (
+ "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 "
+ "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes"
+ )
+
+ def test_get_torch_index_url_uses_nvidia_detected_flag(self):
+ """get_torch_index_url must track NVIDIA detection independently of _smi.
+
+ When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback
+ but nvidia-smi is not on PATH, _smi stays empty. Without a separate
+ _nvidia_detected flag, the function falls into the AMD/CPU branch even
+ though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA.
+ """
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ func_start = source.find("get_torch_index_url()")
+ func_end = source.find("\n}", func_start)
+ func_body = source[func_start:func_end]
+ assert "_nvidia_detected" in func_body, (
+ "get_torch_index_url must use a _nvidia_detected flag (separate from "
+ "_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
+ )
+ # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on
+ # _smi being empty.
+ assert (
+ '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
+ ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
+
# TEST: Live regression on current host (NVIDIA B200 expected)
@@ -1389,14 +1614,18 @@ class TestApplyGpuIdsRocmFallback:
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'getattr(_torch.version, "hip", None)' in func_body
- def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
- """apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
+ def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self):
+ """apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited.
+
+ ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices.
+ Overwriting it breaks multi-GPU ROCm systems (see issue #6118).
+ """
hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body
- assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body
+ assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' not in func_body
def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
"""torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""
@@ -2016,7 +2245,22 @@ class TestRuntimeBnbRocmSourceGuards:
"""A failed redetect must not downgrade a persisted suffix to '72'."""
for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
source = path.read_text(encoding = "utf-8")
- assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name
+ assert (
+ '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source
+ ), path.name
+
+ def test_main_requires_found_rocm_dll(self):
+ """HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force
+ a ROCm backend onto a non-ROCm bitsandbytes."""
+ source = self._MAIN_PATH.read_text(encoding = "utf-8")
+ assert "if _found_rocm_bnb:" in source
+ assert "_hip_env" not in source
+
+ def test_worker_requires_found_rocm_dll(self):
+ """No DLL found: the worker must not write any override or touch the
+ seeded marker (later import fixes must still see sitecustomize)."""
+ source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8")
+ assert "if _found_rocm_bnb:" in source
class TestDetectBnbRocmDllVer:
@@ -2218,8 +2462,9 @@ class TestWorkerWindowsRocmPatches:
assert "BNB_ROCM_VERSION" in source
# Detection helper must be used
assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source
- # "72" must appear as the safe fallback
- assert '"72"' in source or "'72'" in source
+ # Falls back to the seeded value, never a blind "72" (which would
+ # force a ROCm backend onto a non-ROCm bitsandbytes wheel)
+ assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source
def test_bnb_rocm_version_set_before_ml_imports(self):
"""BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports."""
diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py
index 20c7bda87c..26c3c244ca 100644
--- a/tests/studio/playwright_extra_ui.py
+++ b/tests/studio/playwright_extra_ui.py
@@ -332,6 +332,21 @@ with sync_playwright() as p:
plus_btn.click(force = True)
page.wait_for_timeout(400)
compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
+ if compare_item.count() == 0:
+ # Compare chat moved into the "More" submenu; hover, then click fallback.
+ more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
+ if more_trigger.count() > 0:
+ more_trigger.hover()
+ page.wait_for_timeout(400)
+ compare_item = page.get_by_role(
+ "menuitem", name = re.compile(r"Compare chat", re.I)
+ ).first
+ if compare_item.count() == 0:
+ more_trigger.click(force = True)
+ page.wait_for_timeout(400)
+ compare_item = page.get_by_role(
+ "menuitem", name = re.compile(r"Compare chat", re.I)
+ ).first
if compare_item.count() > 0:
compare_item.click(force = True)
compare_opened = True
diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py
index 373830ae4a..6c75a4b886 100644
--- a/tests/test_import_fixes_drift.py
+++ b/tests/test_import_fixes_drift.py
@@ -584,3 +584,54 @@ def test_accelerate_utils_imports_module_present():
"accelerate.utils.imports.is_wandb_available is gone; "
"disable_broken_wandb cannot patch the source module."
)
+
+
+# ===========================================================================
+# bitsandbytes -- ROCm arch / warp-size detection shape
+# ===========================================================================
+
+
+def test_bitsandbytes_rocm_detection_helpers_recognizable():
+ """``fix_bitsandbytes_rocm_arch_detection`` swaps bnb's ROCm helpers
+ only when they shell out via subprocess and never consult torch device
+ props; a third shape is declined by design, silently restoring Windows
+ ROCm noise. Fail so the sniff gets updated. Reads source, no import."""
+ spec = importlib.util.find_spec("bitsandbytes")
+ if spec is None:
+ pytest.skip("bitsandbytes not installed -- nothing to drift-check.")
+ cuda_specs_path = None
+ for location in spec.submodule_search_locations or []:
+ candidate = os.path.join(location, "cuda_specs.py")
+ if os.path.isfile(candidate):
+ cuda_specs_path = candidate
+ break
+ if cuda_specs_path is None:
+ pytest.skip("bitsandbytes has no cuda_specs.py (pre-ROCm version).")
+
+ import ast
+
+ with open(cuda_specs_path, "r", encoding = "utf-8") as f:
+ source = f.read()
+ helpers = [
+ node
+ for node in ast.walk(ast.parse(source))
+ if isinstance(node, ast.FunctionDef)
+ and node.name in ("get_rocm_gpu_arch", "get_rocm_warpsize")
+ ]
+ if not helpers:
+ pytest.skip("bitsandbytes cuda_specs has no ROCm detection helpers.")
+ for node in helpers:
+ segment = ast.get_source_segment(source, node) or ""
+ recognized = (
+ "subprocess" in segment
+ or "get_device_properties" in segment
+ or "gcnArchName" in segment
+ )
+ if not recognized:
+ pytest.fail(
+ f"DRIFT DETECTED: bitsandbytes.cuda_specs.{node.name} uses "
+ "neither subprocess nor torch device properties; "
+ "fix_bitsandbytes_rocm_arch_detection's shape sniff will "
+ "decline to patch it and Windows ROCm import-time noise / "
+ "wrong ROCM_GPU_ARCH may return."
+ )
diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py
new file mode 100644
index 0000000000..fae5a7d8a8
--- /dev/null
+++ b/tests/utils/test_rope_scaling_drift.py
@@ -0,0 +1,312 @@
+"""Guard for config.rope_scaling being silently dropped (issue #2405).
+
+Unsloth's replacement rotary classes ignored rope_scaling when constructed
+from a config (the modern-transformers path), so Llama-3.1 ran with unscaled
+RoPE and collapsed into gibberish past ~32K tokens.
+
+Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper
+_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS;
+(3) CUDA checks instantiating the real class (skipped without a real device,
+probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate).
+Layers 2 and 3 fail on the unfixed code.
+"""
+
+import ast
+import math
+from pathlib import Path
+
+import pytest
+import torch
+
+
+def _has_real_cuda():
+ try:
+ torch.zeros(1).to("cuda")
+ return True
+ except Exception:
+ return False
+
+
+HAS_REAL_CUDA = _has_real_cuda()
+requires_cuda = pytest.mark.skipif(
+ not HAS_REAL_CUDA,
+ reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__",
+)
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py"
+
+CLASS_NAME = "LlamaRotaryEmbedding"
+
+# Llama-3.1-style rope_scaling.
+LLAMA3_ROPE_SCALING = {
+ "rope_type": "llama3",
+ "factor": 8.0,
+ "low_freq_factor": 1.0,
+ "high_freq_factor": 4.0,
+ "original_max_position_embeddings": 8192,
+}
+ROPE_THETA = 500000.0
+HEAD_DIM = 128
+MAX_POS = 131072
+
+
+# --- Layer 1: AST structural tripwire (stdlib only, no unsloth import) ---
+
+
+def _load_class_init():
+ tree = ast.parse(LLAMA_PY.read_text())
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ClassDef) and node.name == CLASS_NAME:
+ for sub in node.body:
+ if isinstance(sub, ast.FunctionDef) and sub.name == "__init__":
+ return sub
+ raise AssertionError(
+ f"{CLASS_NAME}.__init__ not found in {LLAMA_PY}; if it was renamed or "
+ "moved, update this guard so RoPE scaling stays protected (issue #2405)"
+ )
+
+
+def _config_branch(init_fn):
+ """The `if config is not None:` block at the top of __init__."""
+ for node in init_fn.body:
+ if isinstance(node, ast.If):
+ test = node.test
+ is_config_test = (
+ isinstance(test, ast.Compare)
+ and isinstance(test.left, ast.Name)
+ and test.left.id == "config"
+ )
+ if is_config_test:
+ return node
+ return None
+
+
+def test_config_path_inspects_rope_scaling():
+ init_fn = _load_class_init()
+ branch = _config_branch(init_fn)
+ assert branch is not None, (
+ f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` "
+ "branch; the config constructor path must read config.rope_scaling so "
+ "scaled models (llama3/linear/longrope) are not silently unscaled "
+ "(issue #2405)"
+ )
+
+ names = set()
+ for stmt in branch.body:
+ for sub in ast.walk(stmt):
+ if isinstance(sub, ast.Attribute):
+ names.add(sub.attr)
+ elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
+ names.add(sub.value)
+ assert "rope_scaling" in names, (
+ f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. "
+ "When a rotary class is built straight from a config (the path modern "
+ "transformers takes, since rotary moved to LlamaModel), the llama3 / "
+ "linear / longrope scaling must still be applied; otherwise long inputs "
+ "produce repeated-pattern gibberish (issue #2405)."
+ )
+
+ called = {
+ sub.func.id
+ for stmt in branch.body
+ for sub in ast.walk(stmt)
+ if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)
+ }
+ assert "_compute_config_rope_inv_freq" in called, (
+ f"{CLASS_NAME}.__init__ config path no longer calls "
+ "_compute_config_rope_inv_freq; the CPU behavioral tests below cover "
+ "that helper directly, so the constructor must stay wired to it or "
+ "scaled configs silently lose RoPE scaling again (issue #2405)."
+ )
+
+
+# --- Layer 2: CPU behavioral guard (pure helper, no instantiation) ---
+
+
+def _make_config(rope_scaling):
+ from transformers import LlamaConfig
+ return LlamaConfig(
+ hidden_size = 256,
+ num_attention_heads = 2,
+ num_key_value_heads = 2,
+ head_dim = HEAD_DIM,
+ rope_theta = ROPE_THETA,
+ max_position_embeddings = MAX_POS,
+ rope_scaling = rope_scaling,
+ )
+
+
+def _unsloth_rotary(config):
+ from unsloth.models import llama as llama_mod
+ return llama_mod.LlamaRotaryEmbedding(config = config)
+
+
+def _reference_inv_freq(config, rope_type):
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+ inv_freq, _attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, "cpu")
+ return inv_freq.float().cpu()
+
+
+def _vanilla_inv_freq():
+ return 1.0 / (
+ ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype = torch.int64).float() / HEAD_DIM)
+ )
+
+
+def _compute_helper(config, rope_scaling):
+ from unsloth.models.llama import _compute_config_rope_inv_freq
+ return _compute_config_rope_inv_freq(config, rope_scaling)
+
+
+def test_llama3_scaling_applied_to_inv_freq():
+ config = _make_config(LLAMA3_ROPE_SCALING)
+ got, attention_scaling = _compute_helper(config, config.rope_scaling)
+ expected = _reference_inv_freq(config, "llama3")
+ vanilla = _vanilla_inv_freq()
+
+ # Guard against a vacuous test.
+ assert not torch.allclose(
+ expected, vanilla, rtol = 1e-4
+ ), "test setup error: llama3-scaled inv_freq should differ from vanilla"
+ assert got is not None, (
+ "_compute_config_rope_inv_freq returned None for a llama3 config; the "
+ "config path is dropping config.rope_scaling, so long-context inference "
+ "degrades into repeated-pattern gibberish (issue #2405)."
+ )
+ got = got.float().cpu()
+ assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
+ "inv_freq for a llama3 config does not match transformers' llama3 RoPE "
+ "scaling (issue #2405).\n"
+ f"got[:6]={got[:6].tolist()}\nexpected[:6]={expected[:6].tolist()}"
+ )
+
+
+def test_default_rope_type_matches_vanilla_inv_freq():
+ config = _make_config(None)
+ got, attention_scaling = _compute_helper(config, {"rope_type": "default"})
+ assert got is not None
+ vanilla = _vanilla_inv_freq()
+ assert torch.allclose(got.float().cpu(), vanilla, rtol = 1e-4, atol = 1e-6), (
+ "default rope_type must equal the vanilla inv_freq; "
+ f"got[:6]={got[:6].tolist()} vanilla[:6]={vanilla[:6].tolist()}"
+ )
+
+
+def _cos_at_position(rot, position):
+ """cos row at one position, built like _set_cos_sin_cache but CPU-only."""
+ inv_freq = rot.inv_freq.float().cpu()
+ t = torch.tensor([position], dtype = torch.float32)
+ t = rot._apply_time_scaling(t.clone()) if hasattr(rot, "_apply_time_scaling") else t
+ freqs = torch.outer(t, inv_freq)
+ emb = torch.cat((freqs, freqs), dim = -1)
+ return emb.cos().squeeze(0)
+
+
+# --- Layer 3: CUDA behavioral guard (real instantiation needs a device) ---
+
+
+@requires_cuda
+def test_constructor_applies_llama3_scaling():
+ config = _make_config(LLAMA3_ROPE_SCALING)
+ rot = _unsloth_rotary(config)
+ got = rot.inv_freq.float().cpu()
+ expected = _reference_inv_freq(config, "llama3")
+ assert torch.allclose(
+ got, expected, rtol = 1e-4, atol = 1e-6
+ ), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)."
+
+
+@requires_cuda
+def test_constructor_unscaled_config_uses_vanilla_inv_freq():
+ rot = _unsloth_rotary(_make_config(None))
+ got = rot.inv_freq.float().cpu()
+ vanilla = _vanilla_inv_freq()
+ assert torch.allclose(
+ got, vanilla, rtol = 1e-4, atol = 1e-6
+ ), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq"
+
+
+@requires_cuda
+def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position():
+ scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
+ unscaled = _unsloth_rotary(_make_config(None))
+
+ pos = 10000
+ cos_scaled = _cos_at_position(scaled, pos)
+ cos_unscaled = _cos_at_position(unscaled, pos)
+ assert not torch.allclose(cos_scaled, cos_unscaled, rtol = 1e-4, atol = 1e-5), (
+ f"cos values at position {pos} are identical for a llama3-scaled and an "
+ "unscaled rotary embedding, which means scaling was dropped (issue "
+ "#2405). With correct llama3 scaling the low-frequency bands shrink by "
+ "up to 8x and must change the angles at long positions."
+ )
+
+
+@requires_cuda
+def test_extended_cache_keeps_scaling_after_growth():
+ scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
+ # Grow past the initial cache size (mirrors long-context decode).
+ dummy = torch.zeros(1, dtype = torch.float32)
+ scaled.extend_rope_embedding(dummy, seq_len = 40960)
+
+ config = _make_config(LLAMA3_ROPE_SCALING)
+ expected = _reference_inv_freq(config, "llama3")
+ got = scaled.inv_freq.float().cpu()
+ assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), (
+ "growing the RoPE cache (extend_rope_embedding) must preserve llama3 "
+ "scaling of inv_freq; long-context decode loses scaling otherwise "
+ "(issue #2405)."
+ )
+
+
+def test_object_style_rope_scaling_does_not_crash():
+ # Object-style rope_scaling must be normalized, not .get()'d directly.
+ from dataclasses import dataclass
+
+ from unsloth.models.llama import _compute_config_rope_inv_freq
+
+ @dataclass
+ class FakeRopeScalingConfig:
+ rope_type: str = "llama3"
+ factor: float = 8.0
+ low_freq_factor: float = 1.0
+ high_freq_factor: float = 4.0
+ original_max_position_embeddings: int = 8192
+
+ config = _make_config(LLAMA3_ROPE_SCALING)
+ inv_freq, attention_scaling = _compute_config_rope_inv_freq(config, FakeRopeScalingConfig())
+ assert inv_freq is not None, (
+ "object-style (non-dict) config.rope_scaling must be normalized, not "
+ "dropped; otherwise scaled models silently lose RoPE scaling again "
+ "(issue #2405)."
+ )
+ expected = _reference_inv_freq(config, "llama3")
+ assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)
+
+
+def test_object_style_rope_scaling_on_config_delegates_correctly():
+ # 'linear' has no inline fallback; only the normalized-config retry passes this.
+ from dataclasses import dataclass
+
+ from unsloth.models.llama import _compute_config_rope_inv_freq
+
+ @dataclass
+ class FakeLinearRopeScalingConfig:
+ rope_type: str = "linear"
+ factor: float = 4.0
+
+ dict_config = _make_config({"rope_type": "linear", "factor": 4.0})
+ expected = _reference_inv_freq(dict_config, "linear")
+
+ object_config = _make_config({"rope_type": "linear", "factor": 4.0})
+ object_config.rope_scaling = FakeLinearRopeScalingConfig()
+ inv_freq, attention_scaling = _compute_config_rope_inv_freq(
+ object_config, object_config.rope_scaling
+ )
+ assert inv_freq is not None, (
+ "linear rope_scaling exposed as a config object was silently dropped; "
+ "delegation must retry with a config copy carrying the normalized dict "
+ "(issue #2405)."
+ )
+ assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6)
diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py
index ede97ce68e..5da1e27a9d 100644
--- a/unsloth/_gpu_init.py
+++ b/unsloth/_gpu_init.py
@@ -30,6 +30,7 @@ from .import_fixes import (
disable_broken_causal_conv1d,
disable_broken_vllm,
configure_amdgpu_asic_id_table_path,
+ fix_bitsandbytes_rocm_arch_detection,
torchvision_compatibility_check,
fix_diffusers_warnings,
fix_huggingface_hub,
@@ -67,6 +68,8 @@ except Exception:
# Configure libdrm ids table path early so ROCm can resolve AMD GPU names.
configure_amdgpu_asic_id_table_path()
+# Must precede `import unsloth_zoo` below, which imports bnb on ROCm.
+fix_bitsandbytes_rocm_arch_detection()
disable_broken_causal_conv1d()
disable_broken_vllm()
fix_message_factory_issue()
@@ -75,6 +78,7 @@ torchvision_compatibility_check()
fix_diffusers_warnings()
fix_huggingface_hub()
del configure_amdgpu_asic_id_table_path
+del fix_bitsandbytes_rocm_arch_detection
del disable_broken_causal_conv1d
del disable_broken_vllm
del fix_message_factory_issue
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index b96f8cb7a0..fe078c5025 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -2341,13 +2341,15 @@ extra_eos_tokens = None,
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
"""
- # Strip only the left
+ # Strip only the left: trailing whitespace can be part of the repeated example
+ # (e.g. "{OUTPUT}\n"). Accidental trailing whitespace (#992) is retried on failure.
chat_template = chat_template.lstrip()
assert(tokenizer is not None)
if extra_eos_tokens is None: extra_eos_tokens = []
elif type(extra_eos_tokens) is str: extra_eos_tokens = [extra_eos_tokens,]
+ original_extra_eos_tokens = list(extra_eos_tokens)
vocab = tokenizer.get_vocab()
for extra_eos in extra_eos_tokens:
@@ -2454,6 +2456,20 @@ extra_eos_tokens = None,
f"{left_changed}"
)
except:
+ # Accidental trailing whitespace (#992) desyncs the two-example detection,
+ # so retry once without it. Templates that parse as-is are never altered.
+ rstripped_chat_template = chat_template.rstrip()
+ if rstripped_chat_template != chat_template:
+ try:
+ return construct_chat_template(
+ tokenizer = tokenizer,
+ chat_template = rstripped_chat_template,
+ default_system_message = default_system_message,
+ extra_eos_tokens = original_extra_eos_tokens,
+ )
+ except Exception:
+ pass
+
output_pos = chat_template.find("{OUTPUT}")
input_pos = chat_template.find("{INPUT}")
if output_pos == -1 or input_pos == -1:
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index 0622f73965..695a6a577a 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -1823,6 +1823,340 @@ def configure_amdgpu_asic_id_table_path():
return None
+# ---------------------------------------------------------------------------
+# bitsandbytes Windows ROCm fix: cextension.py runs get_rocm_gpu_arch()
+# (bnb >= 0.47) and get_rocm_warpsize() (0.49.x) at import, shelling out to
+# rocminfo / hipinfo.exe via PATH. Neither is on PATH on Windows (AMD torch
+# wheels put hipInfo.exe in venv Scripts), so every import logs ERROR +
+# WARNING, ROCM_GPU_ARCH becomes "unknown", and warp size defaults to 64:
+# wrong on RDNA (wave 32), breaking 4-bit blocksizes and
+# ALLOW_PREQUANTIZED_MODELS. Upstream fix unmerged (bitsandbytes#1969), so a
+# MetaPathFinder swaps both helpers for torch-device-props-first versions
+# right after bitsandbytes.cuda_specs executes, before cextension reads
+# them. Must run before `import unsloth_zoo` (imports bnb on ROCm).
+# ---------------------------------------------------------------------------
+
+_BNB_CUDA_SPECS_MODULE = "bitsandbytes.cuda_specs"
+_BNB_ROCM_FIX_FINDER_SENTINEL = "_unsloth_bnb_rocm_fix_finder"
+_BNB_ROCM_FIX_FUNCTION_FLAG = "__unsloth_bnb_rocm_fix__"
+
+
+def _torch_rocm_device_props():
+ """Device-0 props on a ROCm torch build with a visible GPU, else None.
+ Never raises; bnb's own import initializes the device context anyway."""
+ try:
+ import torch
+
+ if not getattr(getattr(torch, "version", None), "hip", None):
+ return None
+ if not torch.cuda.is_available():
+ return None
+ return torch.cuda.get_device_properties(0)
+ except Exception:
+ return None
+
+
+def _iter_hipinfo_paths():
+ """Yield existing hipInfo.exe paths: PATH, interpreter scripts dir (venv
+ and conda layouts), then HIP SDK / AMD installer locations."""
+ import shutil
+ import sysconfig
+
+ candidates = []
+ try:
+ resolved = shutil.which("hipinfo.exe")
+ if resolved:
+ candidates.append(resolved)
+ except Exception:
+ pass
+ try:
+ scripts_dir = sysconfig.get_path("scripts")
+ if scripts_dir:
+ candidates.append(os.path.join(scripts_dir, "hipInfo.exe"))
+ except Exception:
+ pass
+ executable_dir = os.path.dirname(sys.executable or "")
+ if executable_dir:
+ candidates.append(os.path.join(executable_dir, "hipInfo.exe"))
+ candidates.append(os.path.join(executable_dir, "Scripts", "hipInfo.exe"))
+ for env_key in ("HIP_PATH", "ROCM_PATH"):
+ root = os.environ.get(env_key, "").strip()
+ if root:
+ candidates.append(os.path.join(root, "bin", "hipInfo.exe"))
+ rocm_root = os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm")
+ try:
+ if os.path.isdir(rocm_root):
+ for version_dir in sorted(os.listdir(rocm_root), reverse = True):
+ candidates.append(os.path.join(rocm_root, version_dir, "bin", "hipInfo.exe"))
+ except Exception:
+ pass
+
+ seen = set()
+ for candidate in candidates:
+ try:
+ key = os.path.normcase(os.path.normpath(candidate))
+ if key in seen:
+ continue
+ seen.add(key)
+ if os.path.isfile(candidate):
+ yield candidate
+ except Exception:
+ continue
+
+
+def _run_hipinfo(hipinfo_path):
+ """Run hipInfo.exe and return its stdout, or "" on any failure."""
+ import subprocess
+ try:
+ result = subprocess.run(
+ [hipinfo_path],
+ capture_output = True,
+ text = True,
+ timeout = 15,
+ creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0),
+ )
+ return result.stdout or ""
+ except Exception as e:
+ _log_rocm_detection(f"Unsloth: `{hipinfo_path}` failed: {e}")
+ return ""
+
+
+def _unsloth_get_rocm_gpu_arch():
+ """Replaces bnb's get_rocm_gpu_arch: torch device props first (no
+ subprocess), then hipInfo.exe by absolute path, then a quiet "unknown"."""
+ try:
+ import torch
+ if not getattr(getattr(torch, "version", None), "hip", None):
+ return "unknown"
+ except Exception:
+ return "unknown"
+ props = _torch_rocm_device_props()
+ if props is not None:
+ try:
+ # gcnArchName may carry feature flags, e.g. "gfx90a:sramecc+:xnack-"
+ arch = str(props.gcnArchName).split(":")[0].strip()
+ if arch.startswith("gfx"):
+ return arch
+ except Exception:
+ pass
+ for hipinfo_path in _iter_hipinfo_paths():
+ match = re.search(r"gcnArchName:\s+gfx([a-zA-Z\d]+)", _run_hipinfo(hipinfo_path))
+ if match:
+ return "gfx" + match.group(1)
+ _log_rocm_detection(
+ "Unsloth: Could not detect the ROCm GPU architecture - bitsandbytes will see `unknown`."
+ )
+ return "unknown"
+
+
+def _unsloth_get_rocm_warpsize():
+ """Replaces bnb 0.49.x get_rocm_warpsize: upstream defaults to 64 when
+ rocminfo is missing, wrong on RDNA (wave 32)."""
+ try:
+ import torch
+ if not getattr(getattr(torch, "version", None), "hip", None):
+ return 32 # upstream behavior: NVIDIA warp size is always 32
+ except Exception:
+ return 64 # upstream behavior: default to 64 on failure
+ props = _torch_rocm_device_props()
+ if props is not None:
+ # torch 2.11 ROCm exposes warp_size; some builds used warpSize.
+ for attribute_name in ("warp_size", "warpSize"):
+ warp_size = getattr(props, attribute_name, None)
+ if isinstance(warp_size, int) and warp_size in (32, 64):
+ return warp_size
+ for hipinfo_path in _iter_hipinfo_paths():
+ match = re.search(r"^\s*warpSize:\s+(\d+)", _run_hipinfo(hipinfo_path), re.MULTILINE)
+ if match and int(match.group(1)) in (32, 64):
+ return int(match.group(1))
+ _log_rocm_detection(
+ "Unsloth: Could not detect the ROCm warp size - defaulting to 64 "
+ "(bitsandbytes' own default)."
+ )
+ return 64
+
+
+setattr(_unsloth_get_rocm_gpu_arch, _BNB_ROCM_FIX_FUNCTION_FLAG, True)
+setattr(_unsloth_get_rocm_warpsize, _BNB_ROCM_FIX_FUNCTION_FLAG, True)
+
+
+def _bnb_rocm_helper_is_broken(function):
+ """True only for upstream's subprocess-only detectors; co_names works
+ where getsource fails. Versions consulting torch props are untouched."""
+ if function is None or not callable(function):
+ return False
+ if getattr(function, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
+ return False # Already ours.
+ try:
+ function = inspect.unwrap(function)
+ except Exception:
+ pass
+ code = getattr(function, "__code__", None)
+ co_names = getattr(code, "co_names", ()) if code is not None else ()
+ if not co_names:
+ return False # C function or opaque wrapper -- do not touch.
+ if "get_device_properties" in co_names or "gcnArchName" in co_names:
+ return False # Fixed upstream -- no-op.
+ return "subprocess" in co_names
+
+
+def _patch_bnb_cuda_specs_module(module):
+ """Swap broken ROCm detection helpers on an executed cuda_specs module.
+ Returns True when the module ends up patched (now or previously)."""
+ patched = False
+ for attribute_name, replacement in (
+ ("get_rocm_gpu_arch", _unsloth_get_rocm_gpu_arch),
+ ("get_rocm_warpsize", _unsloth_get_rocm_warpsize),
+ ):
+ original = getattr(module, attribute_name, None)
+ if getattr(original, _BNB_ROCM_FIX_FUNCTION_FLAG, False):
+ patched = True # Already ours.
+ continue
+ if not _bnb_rocm_helper_is_broken(original):
+ continue
+ setattr(module, attribute_name, replacement)
+ patched = True
+ logger.info(
+ f"Unsloth: Patched bitsandbytes.cuda_specs.{attribute_name} - "
+ f"avoids PATH-dependent subprocess GPU detection on Windows ROCm."
+ )
+ return patched
+
+
+class _BnbCudaSpecsPatchLoader(importlib.abc.Loader):
+ __slots__ = ("_loader",)
+
+ def __init__(self, loader):
+ self._loader = loader
+
+ def create_module(self, spec):
+ create_module = getattr(self._loader, "create_module", None)
+ if create_module is None:
+ return None
+ return create_module(spec)
+
+ def exec_module(self, module):
+ self._loader.exec_module(module)
+ # Patch after the module body ran, before cextension calls it. The
+ # finder stays on sys.meta_path (same lifecycle as the blockers
+ # above) so importlib.reload(bitsandbytes.cuda_specs) re-patches.
+ try:
+ _patch_bnb_cuda_specs_module(module)
+ except Exception as e:
+ _log_rocm_detection(f"Unsloth: bitsandbytes ROCm detection patch failed: {e}")
+
+ def __getattr__(self, name):
+ # Delegate get_source / get_filename etc. so introspection works.
+ return getattr(self._loader, name)
+
+
+class _BnbCudaSpecsPatchFinder(importlib.abc.MetaPathFinder):
+ __slots__ = (_BNB_ROCM_FIX_FINDER_SENTINEL,)
+
+ def __init__(self):
+ setattr(self, _BNB_ROCM_FIX_FINDER_SENTINEL, True)
+
+ def find_spec(
+ self,
+ fullname,
+ path = None,
+ target = None,
+ ):
+ if fullname != _BNB_CUDA_SPECS_MODULE:
+ return None
+ # Delegate to remaining finders (editable installs, frozen apps)
+ # and wrap the loader that would actually be used.
+ spec = None
+ for finder in sys.meta_path:
+ if finder is self or getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
+ continue
+ finder_find_spec = getattr(finder, "find_spec", None)
+ if finder_find_spec is None:
+ continue
+ try:
+ spec = finder_find_spec(fullname, path, target)
+ except Exception:
+ spec = None
+ if spec is not None:
+ break
+ if spec is None or spec.loader is None:
+ return None
+ if not hasattr(spec.loader, "exec_module"):
+ return None # Legacy loader -- let the stock machinery handle it.
+ spec.loader = _BnbCudaSpecsPatchLoader(spec.loader)
+ return spec
+
+
+def _repair_imported_bitsandbytes_rocm_constants():
+ """bnb imported before unsloth: noise already fired, but fix detectors
+ and cached constants, incl. by-value ROCM_WARP_SIZE_64 copies."""
+ cuda_specs = sys.modules.get(_BNB_CUDA_SPECS_MODULE)
+ if cuda_specs is None:
+ return
+ if not _patch_bnb_cuda_specs_module(cuda_specs):
+ return
+
+ try:
+ arch = cuda_specs.get_rocm_gpu_arch()
+ except Exception:
+ arch = "unknown"
+ warp_size_64 = None
+ get_rocm_warpsize = getattr(cuda_specs, "get_rocm_warpsize", None)
+ if callable(get_rocm_warpsize):
+ try:
+ warp_size_64 = get_rocm_warpsize() == 64
+ except Exception:
+ warp_size_64 = None
+
+ for module_name, module in list(sys.modules.items()):
+ if module is None or module is cuda_specs:
+ continue
+ if module_name != "bitsandbytes" and not module_name.startswith("bitsandbytes."):
+ continue
+ try:
+ if arch != "unknown" and getattr(module, "ROCM_GPU_ARCH", None) == "unknown":
+ module.ROCM_GPU_ARCH = arch
+ if warp_size_64 is not None and isinstance(
+ getattr(module, "ROCM_WARP_SIZE_64", None), bool
+ ):
+ module.ROCM_WARP_SIZE_64 = warp_size_64
+ except Exception:
+ continue
+ logger.info("Unsloth: Repaired bitsandbytes ROCm arch / warp-size constants in place.")
+
+
+def fix_bitsandbytes_rocm_arch_detection():
+ """Fix bnb's import-time ROCm arch / warp-size detection on Windows
+ (see header above). No-op on non-Windows, non-ROCm, missing or
+ upstream-fixed bnb. Idempotent. Opt out: UNSLOTH_DISABLE_BNB_ROCM_FIX=1."""
+ if os.environ.get("UNSLOTH_DISABLE_BNB_ROCM_FIX", "0") == "1":
+ return
+ if sys.platform != "win32":
+ return
+ if not _is_rocm_torch_build():
+ return
+
+ # Already imported: prevention impossible, repair in place instead.
+ if _BNB_CUDA_SPECS_MODULE in sys.modules:
+ try:
+ _repair_imported_bitsandbytes_rocm_constants()
+ except Exception:
+ pass
+ return
+
+ try:
+ if importlib.util.find_spec("bitsandbytes") is None:
+ return
+ except Exception:
+ return
+
+ for finder in sys.meta_path:
+ if getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False):
+ return # Already installed -- idempotent.
+ sys.meta_path.insert(0, _BnbCudaSpecsPatchFinder())
+ _log_rocm_detection("Unsloth: Installed the bitsandbytes ROCm arch detection patch hook.")
+
+
def _is_causal_conv1d_name(module_name: str) -> bool:
return module_name == _CAUSAL_CONV1D_PREFIX or module_name.startswith(
_CAUSAL_CONV1D_PREFIX + "."
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 1a86ff8eb6..6baa9a1398 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-__version__ = "2026.6.2"
+__version__ = "2026.6.3"
__all__ = [
"SUPPORTS_BFLOAT16",
@@ -1264,6 +1264,18 @@ if is_openai_available():
from transformers import AutoTokenizer
from transformers.utils.import_utils import _is_package_available
+
+def _package_available(pkg_name: str) -> bool:
+ # transformers >= 5.x makes `_is_package_available` always return a
+ # `(exists, version)` tuple, which is truthy even when the package is
+ # absent; older versions returned a plain bool. Normalise to a bool so
+ # callers don't take "package present" branches for missing packages.
+ result = _is_package_available(pkg_name)
+ if isinstance(result, tuple):
+ return bool(result[0])
+ return bool(result)
+
+
SUPPORTS_BFLOAT16 = False
HAS_FLASH_ATTENTION = False
HAS_FLASH_ATTENTION_SOFTCAPPING = False
@@ -1274,7 +1286,7 @@ if DEVICE_TYPE == "cuda":
if major_version >= 8:
SUPPORTS_BFLOAT16 = True
- if _is_package_available("flash_attn"):
+ if _package_available("flash_attn"):
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
try:
try:
@@ -1319,7 +1331,7 @@ if DEVICE_TYPE == "cuda":
HAS_FLASH_ATTENTION = False
elif DEVICE_TYPE == "hip":
SUPPORTS_BFLOAT16 = True
- if _is_package_available("flash_attn"):
+ if _package_available("flash_attn"):
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
try:
try:
@@ -1981,7 +1993,7 @@ def is_bfloat16_supported():
def is_vLLM_available():
- return _is_package_available("vllm")
+ return _package_available("vllm")
# Patches models to add RoPE Scaling
diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py
new file mode 100644
index 0000000000..12596b432e
--- /dev/null
+++ b/unsloth/models/diffusion.py
@@ -0,0 +1,301 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+FastDiffusionModel: a transformers-only slow path for text-diffusion models (e.g. DiffusionGemma).
+
+These models use a block-diffusion sampling loop (custom generate) and a novel backbone, so we skip
+Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay
+bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the
+(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land.
+"""
+
+import os
+import torch
+from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+from ._utils import is_bfloat16_supported
+from .llama import logger
+
+__all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"]
+
+# transformers model_type strings routed to this slow path
+DIFFUSION_MODEL_TYPES = ("diffusion_gemma", "diffusion_gemma4")
+
+# Default LoRA targets: standard nn.Linear modules in the shared Gemma-4 backbone. The 128 MoE experts
+# are fused 3D Parameters (gate_up_proj/down_proj), not nn.Linear, so PEFT LoRA cannot target them.
+DIFFUSION_LORA_TARGETS = [
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ "o_proj", # attention
+ "gate_proj",
+ "up_proj",
+ "down_proj", # dense (non-expert) MLP
+]
+
+# Vision tower uses a custom Linear with the same suffix names; exclude it so only the text path is wrapped.
+DIFFUSION_LORA_EXCLUDE = r".*(vision_tower|embed_vision).*"
+
+
+def is_diffusion_model_type(model_types):
+ """model_types: str or iterable -> True if any is a known diffusion model_type."""
+ if isinstance(model_types, str):
+ model_types = (model_types,)
+ return any(mt in DIFFUSION_MODEL_TYPES for mt in model_types)
+
+
+def _resolve_diffusion_model_class(config):
+ """Resolve the HF model class for a diffusion checkpoint from config.architectures."""
+ import transformers
+
+ archs = getattr(config, "architectures", None) or []
+ for arch in archs:
+ cls = getattr(transformers, arch, None)
+ if cls is not None:
+ return cls
+ # Fallbacks across naming revisions.
+ for name in (
+ "DiffusionGemmaForBlockDiffusion",
+ "DiffusionGemma4ModelForBlockDiffusion",
+ "DiffusionGemma4ForBlockDiffusion",
+ ):
+ cls = getattr(transformers, name, None)
+ if cls is not None:
+ return cls
+ raise RuntimeError(
+ f"Unsloth: could not resolve a diffusion model class from architectures={archs}. "
+ "Ensure you have the transformers build that ships the DiffusionGemma implementation."
+ )
+
+
+def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only):
+ """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4``
+ classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the
+ type/arch names in-memory, and rebuild."""
+ try:
+ return AutoConfig.from_pretrained(
+ model_name,
+ token = token,
+ trust_remote_code = trust_remote_code,
+ revision = revision,
+ local_files_only = local_files_only,
+ )
+ except ValueError as e:
+ if "diffusion_gemma" not in str(e):
+ raise
+ import json
+ from transformers.utils import cached_file
+
+ cfg_path = cached_file(
+ model_name,
+ "config.json",
+ token = token,
+ revision = revision,
+ local_files_only = local_files_only,
+ )
+ with open(cfg_path, encoding = "utf-8") as f:
+ cd = json.load(f)
+ cd["model_type"] = "diffusion_gemma4"
+ cd.setdefault("architectures", ["DiffusionGemma4ModelForBlockDiffusion"])
+ if isinstance(cd.get("text_config"), dict):
+ cd["text_config"]["model_type"] = "diffusion_gemma4_text"
+ if isinstance(cd.get("vision_config"), dict):
+ cd["vision_config"]["model_type"] = "diffusion_gemma4_vision"
+ from transformers import DiffusionGemma4Config
+
+ return DiffusionGemma4Config.from_dict(cd)
+
+
+class FastDiffusionModel:
+ """transformers-only slow path for text-diffusion models."""
+
+ @staticmethod
+ def from_pretrained(
+ model_name = "google/diffusiongemma-26B-A4B-it",
+ max_seq_length = None, # API-compat; diffusion uses canvas_length
+ dtype = None,
+ load_in_4bit = False,
+ load_in_8bit = False,
+ load_in_16bit = False,
+ full_finetuning = False,
+ token = None,
+ device_map = "auto",
+ trust_remote_code = False,
+ attn_implementation = "eager", # exact match with the reference golden logits
+ revision = None,
+ return_tokenizer = True,
+ **kwargs,
+ ):
+ SUPPORTS_BFLOAT16 = is_bfloat16_supported()
+ if dtype is None:
+ dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
+ elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
+ logger.warning_once("Device does not support bfloat16. Will change to float16.")
+ dtype = torch.float16
+ assert dtype in (torch.float16, torch.bfloat16, torch.float32)
+
+ # Honor an explicit local_files_only; else fall back to the offline env vars.
+ local_files_only = kwargs.pop("local_files_only", None)
+ if local_files_only is None:
+ local_files_only = (
+ os.environ.get("HF_HUB_OFFLINE", "0") == "1"
+ or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1"
+ )
+ config = _load_diffusion_config(
+ model_name,
+ token,
+ trust_remote_code,
+ revision,
+ local_files_only,
+ )
+ model_type = getattr(config, "model_type", None)
+ if not is_diffusion_model_type(model_type):
+ raise RuntimeError(
+ f"Unsloth: FastDiffusionModel only supports diffusion model_types {DIFFUSION_MODEL_TYPES}, "
+ f"got '{model_type}'. Use FastModel/FastLanguageModel for autoregressive models."
+ )
+
+ model_cls = _resolve_diffusion_model_class(config)
+
+ load_kwargs = dict(
+ dtype = dtype,
+ device_map = device_map,
+ token = token,
+ trust_remote_code = trust_remote_code,
+ attn_implementation = attn_implementation,
+ revision = revision,
+ local_files_only = local_files_only,
+ )
+
+ # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips
+ # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision.
+ if load_in_4bit or load_in_8bit:
+ from transformers import BitsAndBytesConfig
+ if load_in_4bit:
+ qcfg = BitsAndBytesConfig(
+ load_in_4bit = True,
+ bnb_4bit_use_double_quant = True,
+ bnb_4bit_quant_type = "nf4",
+ bnb_4bit_compute_dtype = dtype,
+ llm_int8_skip_modules = [
+ "lm_head",
+ "embed_tokens",
+ "experts",
+ "self_conditioning",
+ "router",
+ ],
+ )
+ else:
+ qcfg = BitsAndBytesConfig(load_in_8bit = True)
+ load_kwargs["quantization_config"] = qcfg
+
+ print(f"==(( Unsloth: FastDiffusionModel (slow / transformers-only path) ))==")
+ print(f" Model: {model_name} | class: {model_cls.__name__} | model_type: {model_type}")
+ print(
+ f" dtype: {dtype} | 4bit: {load_in_4bit} | 8bit: {load_in_8bit} | attn: {attn_implementation}"
+ )
+
+ model = model_cls.from_pretrained(model_name, **load_kwargs).eval()
+ # Mark before any early return so get_peft_model/for_* route to the slow path.
+ model._unsloth_slow_diffusion = True
+
+ if not return_tokenizer:
+ return model, None
+
+ # Prefer the processor (chat template + tokenizer); fall back to a bare tokenizer. Returned as
+ # "tokenizer" to match the Unsloth (model, tokenizer) contract.
+ try:
+ tokenizer = AutoProcessor.from_pretrained(
+ model_name,
+ token = token,
+ trust_remote_code = trust_remote_code,
+ revision = revision,
+ local_files_only = local_files_only,
+ )
+ except Exception:
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_name,
+ token = token,
+ trust_remote_code = trust_remote_code,
+ revision = revision,
+ local_files_only = local_files_only,
+ )
+
+ return model, tokenizer
+
+ @staticmethod
+ def get_peft_model(
+ model,
+ r = 16,
+ target_modules = None,
+ lora_alpha = 16,
+ lora_dropout = 0.0,
+ bias = "none",
+ use_gradient_checkpointing = True,
+ random_state = 3407,
+ task_type = None,
+ **kwargs,
+ ):
+ """Attach a PEFT LoRA to the diffusion backbone (attention + dense MLP). No fused kernels."""
+ from peft import LoraConfig, get_peft_model as peft_get_peft_model
+
+ if target_modules is None:
+ target_modules = DIFFUSION_LORA_TARGETS
+
+ lora_kwargs = dict(
+ r = r,
+ lora_alpha = lora_alpha,
+ lora_dropout = lora_dropout,
+ bias = bias,
+ target_modules = target_modules,
+ task_type = task_type, # None: diffusion has no standard CAUSAL_LM head
+ **{k: v for k, v in kwargs.items() if k in ("modules_to_save", "init_lora_weights")},
+ )
+ # Exclude the vision tower's custom (non-Linear) modules that share suffix names.
+ exclude = kwargs.get("exclude_modules", DIFFUSION_LORA_EXCLUDE)
+ try:
+ lora_config = LoraConfig(exclude_modules = exclude, **lora_kwargs)
+ except TypeError:
+ # Older PEFT without exclude_modules: scope the target to the text decoder by regex.
+ lora_kwargs["target_modules"] = (
+ r".*model\.decoder\.layers\.\d+\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)"
+ )
+ lora_config = LoraConfig(**lora_kwargs)
+ if use_gradient_checkpointing:
+ model.gradient_checkpointing_enable()
+ if hasattr(model, "enable_input_require_grads"):
+ model.enable_input_require_grads()
+
+ model = peft_get_peft_model(model, lora_config)
+ model._unsloth_slow_diffusion = True
+ try:
+ model.print_trainable_parameters()
+ except Exception:
+ pass
+ return model
+
+ @staticmethod
+ def for_inference(model):
+ model.eval()
+ for _, m in model.named_modules():
+ if hasattr(m, "gradient_checkpointing"):
+ m.gradient_checkpointing = False
+ return model
+
+ @staticmethod
+ def for_training(model, use_gradient_checkpointing = True):
+ model.train()
+ if use_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"):
+ model.gradient_checkpointing_enable()
+ return model
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index d5883e4346..2d31f71ab1 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1622,6 +1622,93 @@ def _get_rope_theta(config, default = 10000.0):
return default
+def _rope_scaling_as_dict(rope_scaling):
+ """Normalize config.rope_scaling (dict or config object) to a dict; {} on failure."""
+ if isinstance(rope_scaling, dict):
+ return rope_scaling
+ for converter in ("to_dict", "dict"):
+ fn = getattr(rope_scaling, converter, None)
+ if callable(fn):
+ try:
+ d = fn()
+ if isinstance(d, dict):
+ return d
+ except Exception:
+ pass
+ try:
+ return {k: v for k, v in vars(rope_scaling).items() if not k.startswith("_")}
+ except TypeError:
+ return {}
+
+
+def _llama3_inv_freq_from_config(
+ config,
+ rope_scaling,
+ device = "cpu",
+):
+ """llama3 inv_freq with factors from config; fallback when modeling_rope_utils is missing."""
+ base = _get_rope_theta(config, default = 10000.0)
+ dim = getattr(config, "head_dim", None)
+ if dim is None:
+ dim = int(config.hidden_size // config.num_attention_heads)
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype = torch.int64, device = device).float() / dim)
+ )
+
+ scale_factor = rope_scaling.get("factor", 8.0)
+ low_freq_factor = rope_scaling.get("low_freq_factor", 1.0)
+ high_freq_factor = rope_scaling.get("high_freq_factor", 4.0)
+ old_context_len = rope_scaling.get("original_max_position_embeddings", 8192)
+
+ low_freq_wavelen = old_context_len / low_freq_factor
+ high_freq_wavelen = old_context_len / high_freq_factor
+ assert low_freq_wavelen != high_freq_wavelen
+
+ # Vectorized meta-llama bands: high freqs kept, low divided by factor, medium blended.
+ wavelen = 2 * math.pi / inv_freq
+ scaled = torch.where(wavelen > low_freq_wavelen, inv_freq / scale_factor, inv_freq)
+ smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
+ smoothed = (1 - smooth) * inv_freq / scale_factor + smooth * inv_freq
+ is_medium = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen)
+ return torch.where(is_medium, smoothed, scaled)
+
+
+def _compute_config_rope_inv_freq(config, rope_scaling):
+ """(inv_freq, attention_scaling) per config.rope_scaling via transformers'
+ ROPE_INIT_FUNCTIONS, with an inline llama3 fallback; (None, 1.0) on failure."""
+ original_rope_scaling = rope_scaling
+ rope_scaling = _rope_scaling_as_dict(rope_scaling)
+ rope_type = rope_scaling.get("rope_type", None) or rope_scaling.get("type", None)
+ try:
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+
+ rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type]
+ try:
+ inv_freq, attention_scaling = rope_init_fn(config, torch.device("cpu"))
+ except Exception:
+ # Object-style rope_scaling: retry with a config copy carrying the plain dict.
+ if isinstance(original_rope_scaling, dict):
+ raise
+ import copy as _copy
+
+ config_copy = _copy.copy(config)
+ config_copy.rope_scaling = rope_scaling
+ inv_freq, attention_scaling = rope_init_fn(config_copy, torch.device("cpu"))
+ return inv_freq.to(dtype = torch.float32, device = "cpu"), float(attention_scaling)
+ except Exception as exception:
+ if rope_type == "llama3":
+ try:
+ return _llama3_inv_freq_from_config(config, rope_scaling), 1.0
+ except Exception:
+ pass
+ logger.warning_once(
+ f"Unsloth: Could not apply RoPE scaling '{rope_type}' from config "
+ f"({type(exception).__name__}: {exception}); falling back to unscaled RoPE. "
+ "Long-context generation may degrade."
+ )
+ return None, 1.0
+
+
# Solves https://github.com/unslothai/unsloth/issues/168
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
@@ -1640,6 +1727,12 @@ class LlamaRotaryEmbedding(torch.nn.Module):
config = None, # [TODO] Hack to pass in config - need to remove later
):
super().__init__()
+ # cos/sin multiplier (1.0 except yarn / longrope); set before any cache build.
+ self.attention_scaling = 1.0
+ # Base-class-from-config path (modern transformers): derive inv_freq like
+ # transformers so config.rope_scaling is not dropped (#2405). Scaled
+ # subclasses are excluded to avoid double-scaling.
+ config_inv_freq = None
if config is not None:
# [TODO] Hack to pass in config - need to remove later
base = _get_rope_theta(config, default = base)
@@ -1652,6 +1745,13 @@ class LlamaRotaryEmbedding(torch.nn.Module):
device = DEVICE_TYPE_TORCH
max_position_embeddings = config.max_position_embeddings
+ rope_scaling = getattr(config, "rope_scaling", None)
+ if rope_scaling is not None and type(self) is LlamaRotaryEmbedding:
+ config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq(
+ config,
+ rope_scaling,
+ )
+
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
@@ -1660,12 +1760,17 @@ class LlamaRotaryEmbedding(torch.nn.Module):
self.multi_gpu_cos_cached = [None] * DEVICE_COUNT
self.multi_gpu_sin_cached = [None] * DEVICE_COUNT
- # Normal Llama-3 RoPE
- inv_freq = 1.0 / (
- self.base
- ** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim)
- )
- inv_freq = self._apply_inv_freq_scaling(inv_freq)
+ if config_inv_freq is not None:
+ inv_freq = config_inv_freq # already scaled; skip subclass scaling
+ else:
+ # Normal Llama-3 RoPE
+ inv_freq = 1.0 / (
+ self.base
+ ** (
+ torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim
+ )
+ )
+ inv_freq = self._apply_inv_freq_scaling(inv_freq)
self.register_buffer("inv_freq", inv_freq, persistent = False)
# Build here to make `torch.jit.trace` work.
@@ -1704,8 +1809,10 @@ class LlamaRotaryEmbedding(torch.nn.Module):
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim = -1)
- cos = emb.cos().to(dtype = dtype, device = device, non_blocking = True)
- sin = emb.sin().to(dtype = dtype, device = device, non_blocking = True)
+ # Applied here so attention_scaling survives extend_rope_embedding rebuilds;
+ # default 1.0 keeps unscaled paths bit-identical.
+ cos = (emb.cos() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
+ sin = (emb.sin() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True)
self.multi_gpu_cos_cached[device.index] = cos
self.multi_gpu_sin_cached[device.index] = sin
return cos, sin
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 0ee9e49ba1..4dc3046928 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -829,6 +829,7 @@ from ..kernels import (
post_patch_loss_function,
)
from .vision import FastBaseModel
+from .diffusion import FastDiffusionModel, is_diffusion_model_type
from transformers import (
AutoModelForCausalLM,
)
@@ -846,6 +847,25 @@ class FastModel(FastBaseModel):
model = _prepare_model_for_qat(model, qat_scheme)
return model
+ @staticmethod
+ def get_peft_model(model, *args, **kwargs):
+ # Route text-diffusion models (slow path) to the transformers-only PEFT helper.
+ if getattr(model, "_unsloth_slow_diffusion", False):
+ return FastDiffusionModel.get_peft_model(model, *args, **kwargs)
+ return FastBaseModel.get_peft_model(model, *args, **kwargs)
+
+ @staticmethod
+ def for_inference(model):
+ if getattr(model, "_unsloth_slow_diffusion", False):
+ return FastDiffusionModel.for_inference(model)
+ return FastBaseModel.for_inference(model)
+
+ @staticmethod
+ def for_training(model, use_gradient_checkpointing = True):
+ if getattr(model, "_unsloth_slow_diffusion", False):
+ return FastDiffusionModel.for_training(model, use_gradient_checkpointing)
+ return FastBaseModel.for_training(model, use_gradient_checkpointing)
+
@staticmethod
def from_pretrained(
model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
@@ -1065,6 +1085,24 @@ class FastModel(FastBaseModel):
local_files_only = True
kwargs["local_files_only"] = True
+ # Text-diffusion slow-path dispatch, factored so both the normal route (below) and the
+ # legacy-config fallback (in the AutoConfig except handler) share one call site.
+ def _dispatch_diffusion():
+ return FastDiffusionModel.from_pretrained(
+ model_name = model_name,
+ max_seq_length = max_seq_length,
+ dtype = dtype,
+ load_in_4bit = load_in_4bit,
+ load_in_8bit = load_in_8bit,
+ load_in_16bit = load_in_16bit,
+ full_finetuning = full_finetuning,
+ token = token,
+ device_map = device_map,
+ trust_remote_code = trust_remote_code,
+ revision = revision,
+ **kwargs,
+ )
+
try:
model_config = AutoConfig.from_pretrained(
model_name,
@@ -1078,6 +1116,12 @@ class FastModel(FastBaseModel):
raise
except Exception as error:
autoconfig_error = str(error)
+ # Legacy text-diffusion configs use model_type "diffusion_gemma", which current
+ # transformers does not register by name (it ships "diffusion_gemma4"). AutoConfig
+ # raises before we can dispatch; route straight to the diffusion slow path, whose
+ # loader aliases the legacy type to the gemma4 classes.
+ if "diffusion_gemma" in autoconfig_error and is_diffusion_model_type("diffusion_gemma"):
+ return _dispatch_diffusion()
if "architecture" in autoconfig_error:
if "qwen3_5" in autoconfig_error:
raise ImportError(
@@ -1126,6 +1170,13 @@ class FastModel(FastBaseModel):
)
model_types_all = ",".join(model_types) + ","
+ # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ----
+ # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's
+ # autoregressive kernel/compile patching and load the unmodified HF model (bit-identical to
+ # naive transformers), keeping only 4bit/8bit + PEFT LoRA conveniences.
+ if is_diffusion_model_type(model_types):
+ return _dispatch_diffusion()
+
# Save model types and loading method
lowered_model_name = model_name.lower()
string = os.environ.get("UNSLOTH_MODEL_NAME", "") + model_types_all
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 34336aa80f..a52ab359bd 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -283,6 +283,9 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
input_ids = kwargs["input"]
elif "input_features" in kwargs:
input_ids = kwargs["input_features"]
+ elif "inputs_embeds" in kwargs:
+ # canonical HF name for embedding inputs (e.g. multimodal generate)
+ input_ids = kwargs["inputs_embeds"]
elif "input_embeds" in kwargs:
input_ids = kwargs["input_embeds"]
elif "inputs" in kwargs:
@@ -1280,6 +1283,15 @@ class FastBaseModel:
tokenizer.padding_side = "left" # Force inference
if hasattr(tokenizer, "tokenizer"):
tokenizer.tokenizer.padding_side = "left" # Force inference
+ # Audio feature extractors must stay right padded: left (a text setting,
+ # forwarded by from_pretrained) shifts Whisper mels and desyncs Gemma 4
+ # audio token counts (crash on transformers < 5.10).
+ feature_extractor = getattr(tokenizer, "feature_extractor", None)
+ if (
+ feature_extractor is not None
+ and getattr(feature_extractor, "padding_side", None) == "left"
+ ):
+ feature_extractor.padding_side = "right"
m = model
while hasattr(m, "model"):
m.max_seq_length = max_seq_length
diff --git a/unsloth/save.py b/unsloth/save.py
index 4ff56cc477..629cbb9548 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -1092,7 +1092,6 @@ def unsloth_save_model(
gc.collect()
# Remove temporary location
- import shutil
shutil.rmtree(temporary_location, ignore_errors = True)
@@ -1224,7 +1223,6 @@ def install_llama_cpp_old(version = -10):
for i in range(30):
print(f"**[WARNING]** Deleting llama.cpp directory... {30-i} seconds left.")
time.sleep(1)
- import shutil
shutil.rmtree("llama.cpp", ignore_errors = True)
@@ -2494,7 +2492,6 @@ def unsloth_push_to_hub_gguf(
except Exception as e:
if cleanup_temp:
- import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
try:
shutil.rmtree(d)
@@ -2681,7 +2678,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
# Clean up temporary directory
if cleanup_temp:
print("Unsloth: Cleaning up temporary files...")
- import shutil
for d in [save_directory, f"{save_directory}_gguf"]:
if os.path.exists(d):
try:
diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py
index c8ec4c66c0..08d7e17e8c 100644
--- a/unsloth_cli/__init__.py
+++ b/unsloth_cli/__init__.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import os.path as _osp
+import os as _os
import sys as _sys
import typer
@@ -10,6 +10,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError
from unsloth_cli.commands.train import train
from unsloth_cli.commands.inference import inference
+from unsloth_cli.commands.chat import chat
from unsloth_cli.commands.export import export, list_checkpoints
from unsloth_cli.commands.studio import (
run as studio_run,
@@ -20,7 +21,7 @@ from unsloth_cli.commands.studio import (
# Canonicalise `-np` only under the `unsloth` console-script;
# third-party scripts that import unsloth_cli keep their argv intact.
-_entry_base = _osp.basename(_sys.argv[0]).lower() if _sys.argv else ""
+_entry_base = _os.path.basename(_sys.argv[0]).lower() if _sys.argv else ""
if _entry_base in {"unsloth", "unsloth.exe"}:
_expand_attached_np_short()
del _entry_base
@@ -53,11 +54,26 @@ def main(
help = "Show version and exit.",
),
):
- pass
+ if (
+ _sys.platform == "win32"
+ ): # this block catches unsloth running inside of System32 or any subdirs, this WILL cause errors if not prevented.
+ _cwd = _os.path.normcase(_os.path.normpath(_os.getcwd()))
+ _system32 = _os.path.normcase(
+ _os.path.normpath(_os.path.join(_os.environ.get("WINDIR", r"C:\Windows"), "System32"))
+ )
+ if _cwd == _system32 or _cwd.startswith(_system32 + _os.sep):
+ typer.secho(
+ "Refusing to run Unsloth inside System32 as it will lead to Errors.\n"
+ "cd to a normal working directory and try again.",
+ fg = "red",
+ err = True,
+ )
+ raise typer.Exit(code = 1)
app.command()(train)
app.command()(inference)
+app.command()(chat)
app.command()(export)
app.command("list-checkpoints")(list_checkpoints)
app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")
diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py
new file mode 100644
index 0000000000..5e734ee2a1
--- /dev/null
+++ b/unsloth_cli/_inference.py
@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Model loading and streaming shared by `inference` and `chat`."""
+
+import os
+import re
+import sys
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+_THINK_OPEN = ""
+_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL)
+
+
+def ensure_studio_backend_path() -> None:
+ backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend")
+ if backend_dir not in sys.path:
+ sys.path.insert(0, backend_dir)
+
+
+def configure_quiet_logging() -> None:
+ import logging
+
+ import structlog
+
+ # The CLI never configures structlog, so without this every backend INFO
+ # line prints. LOG_LEVEL is exported so the worker subprocess inherits it.
+ level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper()
+ level = getattr(logging, level_name, logging.WARNING)
+ structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level))
+ os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+
+
+def visible_text(text: str, show_thinking: bool) -> str:
+ if show_thinking:
+ return text
+ text = _THINK_BLOCK.sub("", text)
+ # Hold back an unclosed trailing so reasoning never leaks mid-stream.
+ open_idx = text.find(_THINK_OPEN)
+ if open_idx != -1:
+ text = text[:open_idx]
+ max_prefix = min(len(text), len(_THINK_OPEN) - 1)
+ for size in range(max_prefix, 0, -1):
+ if _THINK_OPEN.startswith(text[-size:]):
+ return text[:-size]
+ return text
+
+
+def stream_to_stdout(stream, show_thinking: bool) -> str:
+ # Backends yield the full text-so-far on each step (llama.cpp ends with a
+ # metadata dict, skipped); print the growing tail, return the raw text.
+ raw = ""
+ shown = ""
+ for chunk in stream:
+ if not isinstance(chunk, str):
+ continue
+ raw = chunk
+ rendered = visible_text(chunk, show_thinking)
+ delta = rendered[len(shown) :]
+ if delta:
+ sys.stdout.write(delta)
+ sys.stdout.flush()
+ shown = rendered
+ sys.stdout.write("\n")
+ sys.stdout.flush()
+ return raw
+
+
+def stream_markdown(stream, show_thinking: bool, *, console) -> str:
+ from rich.live import Live
+ from rich.markdown import Markdown
+ from rich.text import Text
+
+ raw = ""
+ with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live:
+ for chunk in stream:
+ if not isinstance(chunk, str):
+ continue
+ raw = chunk
+ visible = visible_text(chunk, show_thinking)
+ live.update(Markdown(visible) if visible.strip() else Text(""))
+ return raw
+
+
+def collect_stream(stream, show_thinking: bool) -> str:
+ raw = ""
+ for chunk in stream:
+ if isinstance(chunk, str):
+ raw = chunk
+ return visible_text(raw, show_thinking)
+
+
+def render_columns(
+ left_label: str,
+ left_text: str,
+ right_label: str,
+ right_text: str,
+ *,
+ console = None,
+) -> None:
+ from rich import box
+ from rich.console import Console
+ from rich.table import Table
+
+ table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False)
+ table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold")
+ table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold")
+ table.add_row(left_text or "", right_text or "")
+ (console or Console()).print(table)
+
+
+class ChatBackend:
+ """Uniform stream()/close() over the llama-server and Unsloth backends."""
+
+ def __init__(self, kind: str, backend) -> None:
+ self._kind = kind # "gguf" | "unsloth"
+ self._backend = backend
+
+ def stream(
+ self,
+ messages: list,
+ *,
+ system_prompt: str,
+ temperature: float,
+ top_p: float,
+ top_k: int,
+ max_new_tokens: int,
+ repetition_penalty: float,
+ enable_thinking: bool,
+ use_adapter: Optional[bool] = None,
+ ):
+ if self._kind == "gguf":
+ # llama-server takes the system prompt as the first message.
+ msgs = list(messages)
+ if system_prompt:
+ msgs = [{"role": "system", "content": system_prompt}, *msgs]
+ return self._backend.generate_chat_completion(
+ messages = msgs,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ max_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ enable_thinking = enable_thinking,
+ )
+ gen_kwargs = dict(
+ messages = messages,
+ system_prompt = system_prompt,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ enable_thinking = enable_thinking,
+ )
+ if use_adapter is not None:
+ return self._backend.generate_with_adapter_control(
+ use_adapter = use_adapter, **gen_kwargs
+ )
+ return self._backend.generate_chat_response(**gen_kwargs)
+
+ def close(self) -> None:
+ # Shut the worker down directly: the graceful unload_model waits for
+ # an ack that compare mode can swallow, hanging exit for minutes.
+ try:
+ if self._kind == "gguf":
+ self._backend.unload_model()
+ else:
+ self._backend._shutdown_subprocess(timeout = 2.0)
+ except Exception:
+ pass
+
+
+def resolve_model_config(model: str, *, hf_token: Optional[str]):
+ ensure_studio_backend_path()
+ from utils.models import ModelConfig
+
+ model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token)
+ if not model_config:
+ typer.echo("Could not resolve model config", err = True)
+ raise typer.Exit(code = 1)
+ return model_config
+
+
+def _load_gguf_backend(model_config, *, hf_token, max_seq_length):
+ ensure_studio_backend_path()
+ from core.inference.llama_cpp import LlamaCppBackend
+
+ llama_backend = LlamaCppBackend()
+ common = dict(
+ hf_variant = model_config.gguf_variant,
+ model_identifier = model_config.identifier,
+ is_vision = model_config.is_vision,
+ n_ctx = max_seq_length,
+ )
+ if model_config.gguf_hf_repo:
+ loaded = llama_backend.load_model(
+ hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common
+ )
+ else:
+ loaded = llama_backend.load_model(
+ gguf_path = model_config.gguf_file,
+ mmproj_path = model_config.gguf_mmproj_file,
+ mtp_draft_path = model_config.gguf_mtp_file,
+ **common,
+ )
+ if not loaded:
+ typer.echo("Model load failed", err = True)
+ raise typer.Exit(code = 1)
+ return ChatBackend("gguf", llama_backend)
+
+
+def load_chat_backend(
+ model: str,
+ *,
+ hf_token: Optional[str],
+ max_seq_length: int,
+ load_in_4bit: bool,
+ model_config = None,
+ fresh_backend: bool = False,
+):
+ """Load `model` in-process: GGUF via llama-server, else the orchestrator.
+
+ fresh_backend uses a private orchestrator so a second model (compare's
+ base column) can run alongside the main one.
+ """
+ if model_config is None:
+ model_config = resolve_model_config(model, hf_token = hf_token)
+
+ typer.echo(f"Loading {model}", err = True)
+
+ if model_config.is_gguf:
+ return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length)
+
+ if fresh_backend:
+ ensure_studio_backend_path()
+ from core.inference import InferenceOrchestrator
+ backend = InferenceOrchestrator()
+ else:
+ ensure_studio_backend_path()
+ from core.inference import get_inference_backend
+ backend = get_inference_backend()
+ if not backend.load_model(
+ config = model_config,
+ max_seq_length = max_seq_length,
+ load_in_4bit = load_in_4bit,
+ hf_token = hf_token,
+ ):
+ typer.echo("Model load failed", err = True)
+ raise typer.Exit(code = 1)
+ return ChatBackend("unsloth", backend)
+
+
+def find_studio_server(timeout: float = 0.4) -> Optional[str]:
+ import urllib.request
+ base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/")
+ try:
+ with urllib.request.urlopen(f"{base}/api/health", timeout = timeout):
+ return base
+ except Exception:
+ return None
+
+
+def _studio_token() -> Optional[str]:
+ """Self-issue a JWT: the CLI runs as the same OS user as the server, so it
+ signs with the same stored secret the server validates against."""
+ try:
+ import studio.backend.core # noqa: F401 puts studio/backend on sys.path
+
+ from studio.backend.auth import storage
+ from studio.backend.auth.authentication import create_access_token
+
+ row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone()
+ return create_access_token(row[0], desktop = True) if row else None
+ except Exception:
+ return None
+
+
+class HttpChatBackend:
+ """Chat against a running Studio server over its OpenAI-compatible API.
+
+ close() leaves the model loaded on purpose — the next session (or the
+ UI) starts instantly.
+ """
+
+ def __init__(self, base_url: str, token: str) -> None:
+ self._base = base_url
+ self._token = token
+
+ def _request(
+ self,
+ method: str,
+ path: str,
+ payload = None,
+ timeout = None,
+ ):
+ import json
+ import urllib.request
+
+ request = urllib.request.Request(
+ self._base + path,
+ data = None if payload is None else json.dumps(payload).encode(),
+ headers = {
+ "Authorization": f"Bearer {self._token}",
+ "Content-Type": "application/json",
+ },
+ method = method,
+ )
+ return urllib.request.urlopen(request, timeout = timeout)
+
+ def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None:
+ typer.echo(f"Loading {model} on the Studio server", err = True)
+ try:
+ self._request(
+ "POST",
+ "/api/inference/load",
+ {
+ "model_path": model,
+ "hf_token": hf_token,
+ "max_seq_length": max_seq_length,
+ "load_in_4bit": load_in_4bit,
+ },
+ ).close()
+ except Exception as exc:
+ typer.echo(f"Model load failed: {exc}", err = True)
+ raise typer.Exit(code = 1)
+
+ def stream(
+ self,
+ messages: list,
+ *,
+ system_prompt: str,
+ temperature: float,
+ top_p: float,
+ top_k: int,
+ max_new_tokens: int,
+ repetition_penalty: float,
+ enable_thinking: bool,
+ use_adapter: Optional[bool] = None,
+ ):
+ import json
+
+ msgs = list(messages)
+ if system_prompt:
+ msgs = [{"role": "system", "content": system_prompt}, *msgs]
+ resp = self._request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": "default",
+ "messages": msgs,
+ "stream": True,
+ "temperature": temperature,
+ "top_p": top_p,
+ "top_k": top_k,
+ "max_tokens": max_new_tokens,
+ "repetition_penalty": repetition_penalty,
+ "enable_thinking": enable_thinking,
+ },
+ )
+
+ def cumulative():
+ # Accumulate SSE deltas into the full-text-so-far convention the
+ # stream helpers expect.
+ text = ""
+ with resp:
+ for raw_line in resp:
+ line = raw_line.decode("utf-8", "replace").strip()
+ if not line.startswith("data:"):
+ continue
+ data = line[len("data:") :].strip()
+ if data == "[DONE]":
+ break
+ try:
+ parsed = json.loads(data)
+ except ValueError:
+ continue
+ if "error" in parsed:
+ raise RuntimeError(
+ f"Server error: {parsed['error'].get('message', 'Unknown server error')}"
+ )
+ try:
+ delta = parsed["choices"][0]["delta"].get("content")
+ except (KeyError, IndexError):
+ continue
+ if not delta:
+ continue
+ text += delta
+ # An emoji can arrive split across two deltas as lone
+ # surrogate halves: hold back a trailing half, merge pairs.
+ visible = text
+ if "\ud800" <= visible[-1] <= "\udbff":
+ visible = visible[:-1]
+ yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace")
+
+ return cumulative()
+
+ def close(self) -> None:
+ pass
+
+
+def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit):
+ """Backend on a running Studio server, or None (caller loads locally)."""
+ base_url = find_studio_server()
+ if not base_url:
+ return None
+ token = _studio_token()
+ if not token:
+ return None
+ backend = HttpChatBackend(base_url, token)
+ backend.ensure_loaded(
+ model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit
+ )
+ return backend
diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py
new file mode 100644
index 0000000000..c62916bc75
--- /dev/null
+++ b/unsloth_cli/commands/chat.py
@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from typing import Optional
+
+import typer
+from rich.console import Console
+
+from unsloth_cli._inference import (
+ collect_stream,
+ configure_quiet_logging,
+ connect_studio_server,
+ ensure_studio_backend_path,
+ load_chat_backend,
+ render_columns,
+ resolve_model_config,
+ stream_markdown,
+ visible_text,
+)
+
+_HELP = (
+ "Commands: /exit (quit), /reset (clear history), "
+ "/think (toggle reasoning), /compare (base vs tuned), /help"
+)
+
+
+def _you_prompt(colors: bool) -> str:
+ # The prompt must go through input(), not a separate print — readline
+ # redraws erase anything they didn't draw, eating the label. GNU readline
+ # wants colors wrapped in \001/\002; libedit (macOS) prints those
+ # literally, so it gets raw ANSI.
+ try:
+ import readline
+ except ImportError:
+ return "\n\x1b[1;36mYou: \x1b[0m" if colors else "\nYou: "
+ libedit = (
+ "libedit" in (readline.__doc__ or "") or getattr(readline, "backend", "") == "editline"
+ )
+ if not colors:
+ return "\nYou: "
+ if libedit:
+ return "\n\x1b[1;36mYou: \x1b[0m"
+ return "\n\001\x1b[1;36m\002You: \001\x1b[0m\002"
+
+
+def _compare_blocked_reason(model_config) -> Optional[str]:
+ if model_config.is_gguf:
+ return (
+ "GGUF models can't toggle adapters — load a LoRA fine-tune "
+ "(transformers backend) to compare base vs tuned."
+ )
+ if not model_config.is_lora:
+ return (
+ "this isn't a LoRA adapter — compare turns the adapter off for the "
+ "'base' column, so there's nothing to compare against."
+ )
+ return None
+
+
+def _get_base_load_in_4bit(model_config) -> bool:
+ """Determine load_in_4bit for base model based on tuned adapter precision."""
+ if not model_config.is_lora or not model_config.path:
+ # Fallback to default if not a LoRA or no path
+ return True
+
+ try:
+ import json
+ from pathlib import Path
+
+ adapter_cfg_path = Path(model_config.path) / "adapter_config.json"
+ if not adapter_cfg_path.exists():
+ return True
+
+ with open(adapter_cfg_path) as f:
+ adapter_cfg = json.load(f)
+
+ training_method = adapter_cfg.get("unsloth_training_method")
+ if training_method == "lora":
+ return False
+ elif training_method == "qlora":
+ return True
+ elif not training_method:
+ # Fallback: check base model name for -bnb-4bit suffix
+ if model_config.base_model and "-bnb-4bit" not in model_config.base_model.lower():
+ return False
+ return True
+ return True
+ except Exception:
+ return True
+
+
+def _compare_needs_second_model() -> bool:
+ # MLX can't toggle the adapter off, so compare loads the base separately.
+ # detect_hardware() would print into the chat (and import torch), so
+ # probe its MLX condition quietly: Apple Silicon with mlx installed.
+ try:
+ from studio.backend.utils.hardware import hardware as hw
+
+ if hw.DEVICE is not None:
+ return hw.DEVICE == hw.DeviceType.MLX
+ if not hw.is_apple_silicon():
+ return False
+ import mlx.core # noqa: F401
+
+ return True
+ except Exception:
+ return False
+
+
+def _pick_trained_model(console) -> str:
+ ensure_studio_backend_path()
+ from utils.models import scan_trained_models
+
+ trained = scan_trained_models()
+ if not trained:
+ typer.echo(
+ "No trained models found in your outputs folder. "
+ "Pass a model id or path: `unsloth chat `.",
+ err = True,
+ )
+ raise typer.Exit(code = 1)
+
+ console.print("Your trained models (newest first):", style = "bold")
+ for i, (display_name, _, model_type) in enumerate(trained, 1):
+ console.print(f" {i}. {display_name} ({model_type})", markup = False)
+
+ while True:
+ try:
+ raw = input(f"Chat with [1-{len(trained)}, Enter = 1]: ").strip()
+ except (EOFError, KeyboardInterrupt):
+ raise typer.Exit(code = 1)
+ if not raw:
+ return trained[0][1]
+ if raw.isdigit() and 1 <= int(raw) <= len(trained):
+ return trained[int(raw) - 1][1]
+ console.print(f"Pick a number between 1 and {len(trained)}.", style = "yellow")
+
+
+def chat(
+ model: Optional[str] = typer.Argument(
+ None, help = "HF model id or local path. Omit to pick one of your trained models."
+ ),
+ hf_token: Optional[str] = typer.Option(
+ None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed."
+ ),
+ temperature: float = typer.Option(0.7, "--temperature"),
+ top_p: float = typer.Option(0.9, "--top-p"),
+ top_k: int = typer.Option(40, "--top-k"),
+ max_new_tokens: int = typer.Option(512, "--max-new-tokens"),
+ repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"),
+ system_prompt: str = typer.Option(
+ "", "--system-prompt", help = "Optional system prompt for the conversation."
+ ),
+ max_seq_length: int = typer.Option(4096, "--max-seq-length"),
+ load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
+ think: bool = typer.Option(
+ False,
+ "--think/--no-think",
+ help = "Start with the model's reasoning shown. Toggle live with /think.",
+ ),
+ compare: bool = typer.Option(
+ False,
+ "--compare/--no-compare",
+ help = "Answer each prompt twice — base vs fine-tuned — side by side. "
+ "Needs a LoRA adapter. Toggle live with /compare.",
+ ),
+ verbose: bool = typer.Option(
+ False, "--verbose", "-v", help = "Show backend and llama-server logs."
+ ),
+ no_server: bool = typer.Option(
+ False,
+ "--no-server",
+ help = "Load the model in-process even if a Studio server is running.",
+ ),
+):
+ """Start an interactive chat with a model (loads once, stays warm)."""
+ if not verbose:
+ configure_quiet_logging()
+
+ console = Console()
+ err = Console(stderr = True)
+
+ if model is None:
+ model = _pick_trained_model(console)
+
+ # Resolve first so --compare can be rejected before the slow load.
+ model_config = resolve_model_config(model, hf_token = hf_token)
+ compare_blocked = _compare_blocked_reason(model_config)
+ if compare and compare_blocked:
+ err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
+ raise typer.Exit(code = 1)
+
+ load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
+
+ # Prefer a running Studio server: instant starts, model shared with the UI.
+ chat_backend = None if no_server else connect_studio_server(model, **load_opts)
+ server_mode = chat_backend is not None
+ if server_mode:
+ console.print(
+ "(Studio server connected — model stays warm after /exit)",
+ style = "bright_black",
+ )
+ else:
+ chat_backend = load_chat_backend(model, model_config = model_config, **load_opts)
+
+ name = model_config.display_name or model
+ show_thinking = think
+ compare_mode = compare
+ messages = []
+
+ # Compare's base column: server mode keeps the tuned model remote and
+ # loads the base locally; local MLX (no adapter toggle) does the same;
+ # local CUDA just toggles the adapter on the one loaded model.
+ dual_compare = compare_blocked is None and (server_mode or _compare_needs_second_model())
+ base_backend = None
+
+ def load_base_for_compare():
+ nonlocal base_backend
+ if base_backend is not None:
+ return True
+ base_id = model_config.base_model
+ if not base_id:
+ console.print(
+ "(compare unavailable: this adapter doesn't record its base model)",
+ style = "yellow",
+ )
+ return False
+ console.print(
+ f"(loading base model {base_id} for compare — keeps two models in memory)",
+ style = "bright_black",
+ markup = False,
+ )
+ try:
+ # Use the same precision as the tuned model for fair comparison
+ base_load_opts = dict(load_opts) # Copy original options
+ base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config)
+ base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts)
+ except Exception as exc:
+ err.print(f"(base model load failed: {exc})", style = "red", markup = False)
+ return False
+ return True
+
+ if compare and dual_compare and not load_base_for_compare():
+ raise typer.Exit(code = 1)
+
+ def generate(backend = None, use_adapter = None):
+ # Reads messages and show_thinking live, so /reset and /think apply.
+ return (backend or chat_backend).stream(
+ messages,
+ system_prompt = system_prompt,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ enable_thinking = show_thinking,
+ use_adapter = use_adapter,
+ )
+
+ console.print()
+ console.print(f"Chatting with {name}", style = "bold green", markup = False)
+ console.print(_HELP, style = "bright_black")
+
+ # legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage.
+ you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows)
+ assistant_label = "[bold magenta]Assistant:[/bold magenta]"
+
+ try:
+ while True:
+ try:
+ user = input(you_prompt).strip()
+ except (EOFError, KeyboardInterrupt):
+ console.print()
+ break
+
+ if not user:
+ continue
+ if user in ("/exit", "/quit"):
+ break
+ if user == "/reset":
+ messages = []
+ console.print("(history cleared)", style = "bright_black")
+ continue
+ if user == "/think":
+ show_thinking = not show_thinking
+ state = "on" if show_thinking else "off"
+ console.print(f"(thinking {state})", style = "bright_black")
+ continue
+ if user == "/compare":
+ if compare_blocked:
+ console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
+ continue
+ if not compare_mode and dual_compare and not load_base_for_compare():
+ continue
+ compare_mode = not compare_mode
+ state = "on" if compare_mode else "off"
+ console.print(f"(compare {state})", style = "bright_black")
+ continue
+ if user in ("/help", "/?"):
+ console.print(_HELP, style = "bright_black")
+ continue
+
+ messages.append({"role": "user", "content": user})
+
+ try:
+ if compare_mode:
+ console.print("(comparing base vs tuned…)", style = "bright_black")
+ if dual_compare:
+ base_text = collect_stream(generate(backend = base_backend), show_thinking)
+ tuned_text = collect_stream(generate(), show_thinking)
+ else:
+ base_text = collect_stream(generate(use_adapter = False), show_thinking)
+ tuned_text = collect_stream(generate(use_adapter = True), show_thinking)
+ console.print()
+ render_columns(
+ "base", base_text, f"{name} (tuned)", tuned_text, console = console
+ )
+ # History continues as the tuned model; base is just the reference.
+ answer = tuned_text
+ else:
+ console.print(assistant_label)
+ answer = stream_markdown(generate(), show_thinking, console = console)
+ except KeyboardInterrupt:
+ # Ctrl-C aborts this answer only; drop the unanswered turn.
+ console.print("\n(interrupted)", style = "bright_black")
+ messages.pop()
+ continue
+ except Exception as exc:
+ err.print(f"\n(error: {exc})", style = "red", markup = False)
+ messages.pop()
+ continue
+
+ messages.append(
+ {"role": "assistant", "content": visible_text(answer, show_thinking = False)}
+ )
+ finally:
+ chat_backend.close()
+ if base_backend is not None:
+ base_backend.close()
+ err.print("\nBye.", style = "bright_black")
diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py
index 18de68c4d9..5dbc32c7d2 100644
--- a/unsloth_cli/commands/inference.py
+++ b/unsloth_cli/commands/inference.py
@@ -1,11 +1,17 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import sys
from typing import Optional
import typer
+from unsloth_cli._inference import (
+ configure_quiet_logging,
+ connect_studio_server,
+ load_chat_backend,
+ stream_to_stdout,
+)
+
def inference(
model: str = typer.Argument(..., help = "HF model id or local path."),
@@ -25,45 +31,46 @@ def inference(
),
max_seq_length: int = typer.Option(2048, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
+ think: bool = typer.Option(
+ False,
+ "--think/--no-think",
+ help = "Show the model's reasoning. Off by default so reasoning "
+ "models answer directly instead of spending the token budget thinking.",
+ ),
+ verbose: bool = typer.Option(
+ False,
+ "--verbose",
+ "-v",
+ help = "Show backend and llama-server logs (otherwise only the answer).",
+ ),
+ no_server: bool = typer.Option(
+ False,
+ "--no-server",
+ help = "Load the model in-process even if a Studio server is running.",
+ ),
):
"""Run a single inference using the specified model."""
- from studio.backend.core import ModelConfig, get_inference_backend
+ if not verbose:
+ configure_quiet_logging()
- inference_backend = get_inference_backend()
- model_config = ModelConfig.from_ui_selection(
- dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False
- )
- if not model_config:
- typer.echo("Could not resolve model config", err = True)
- raise typer.Exit(code = 1)
-
- if not inference_backend.load_model(
- config = model_config,
- max_seq_length = max_seq_length,
- load_in_4bit = load_in_4bit,
- hf_token = hf_token,
- ):
- typer.echo("Model load failed", err = True)
- raise typer.Exit(code = 1)
-
- messages = [{"role": "user", "content": prompt}]
- stream = inference_backend.generate_chat_response(
- messages = messages,
- system_prompt = system_prompt,
- temperature = temperature,
- top_p = top_p,
- top_k = top_k,
- max_new_tokens = max_new_tokens,
- repetition_penalty = repetition_penalty,
- )
-
- typer.echo("Assistant:", nl = True)
- previous = ""
- for chunk in stream:
- delta = chunk[len(previous) :]
- if delta:
- sys.stdout.write(delta)
- sys.stdout.flush()
- previous = chunk
- sys.stdout.write("\n")
- sys.stdout.flush()
+ # A running Studio server keeps the model warm between runs, which is
+ # exactly what a one-shot command wants.
+ load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
+ chat_backend = None if no_server else connect_studio_server(model, **load_opts)
+ if chat_backend is None:
+ chat_backend = load_chat_backend(model, **load_opts)
+ try:
+ stream = chat_backend.stream(
+ [{"role": "user", "content": prompt}],
+ system_prompt = system_prompt,
+ temperature = temperature,
+ top_p = top_p,
+ top_k = top_k,
+ max_new_tokens = max_new_tokens,
+ repetition_penalty = repetition_penalty,
+ enable_thinking = think,
+ )
+ typer.echo("Assistant:")
+ stream_to_stdout(stream, show_thinking = think)
+ finally:
+ chat_backend.close()
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index b506f764d7..7493b87eee 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -597,6 +597,11 @@ def studio_default(
f"defaults to {_PARALLEL_DEFAULT_RUN}."
),
),
+ cloudflare: bool = typer.Option(
+ True,
+ "--cloudflare/--no-cloudflare",
+ help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
+ ),
):
"""Launch the Unsloth Studio server."""
# Runs before every subcommand (run/setup/update/...).
@@ -614,6 +619,16 @@ def studio_default(
err = True,
)
raise typer.Exit(2)
+ # Same for --no-cloudflare: it would not reach the subcommand.
+ if not cloudflare:
+ typer.echo(
+ f"Error: --no-cloudflare on `unsloth studio` applies to the "
+ f"plain-server path only. For `unsloth studio "
+ f"{ctx.invoked_subcommand}`, put it after the subcommand: "
+ f"`unsloth studio {ctx.invoked_subcommand} --no-cloudflare ...`",
+ err = True,
+ )
+ raise typer.Exit(2)
return
# Use the studio venv if it exists and we aren't already in it.
@@ -648,6 +663,8 @@ def studio_default(
args.append("--silent")
if api_only:
args.append("--api-only")
+ # Forward the explicit polarity (matches run.py's BooleanOptionalAction).
+ args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# On Windows os.execvp keeps the parent alive, so Ctrl+C
# would orphan the child; use Popen+wait instead.
if sys.platform == "win32":
@@ -689,6 +706,7 @@ def studio_default(
silent = silent,
api_only = api_only,
llama_parallel_slots = parallel,
+ cloudflare = cloudflare,
)
if frontend is not None:
run_kwargs["frontend_path"] = frontend
@@ -861,6 +879,11 @@ def run(
f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)."
),
),
+ cloudflare: bool = typer.Option(
+ True,
+ "--cloudflare/--no-cloudflare",
+ help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
+ ),
):
"""Start Studio, load a model, print an API key -- one-liner server.
@@ -980,6 +1003,8 @@ def run(
# Typer claims --parallel outside ctx.args; without this the
# child reverts to its default and silently drops the value.
args.extend(["--parallel", str(parallel)])
+ # Forward the explicit polarity (same rationale as --load-in-4bit above).
+ args.append("--cloudflare" if cloudflare else "--no-cloudflare")
# llama-server pass-through extras → child ctx.args → load payload.
if extra_llama_args:
args.extend(extra_llama_args)
@@ -997,7 +1022,13 @@ def run(
# ── 2. Start server (always suppress built-in banner) ─────────────
from studio.backend.run import run_server, _resolve_external_ip
- run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = parallel)
+ run_kwargs = dict(
+ host = host,
+ port = port,
+ silent = True,
+ llama_parallel_slots = parallel,
+ cloudflare = cloudflare,
+ )
if frontend is not None:
run_kwargs["frontend_path"] = frontend
app = run_server(**run_kwargs)
@@ -1011,32 +1042,41 @@ def run(
set_tool_policy(enable_tools)
- # 3. Wait for server health.
- if not silent:
- typer.echo("Starting Unsloth Studio...")
- if not _wait_for_server(actual_port):
- typer.echo("Error: server did not become healthy within 30 seconds.", err = True)
- raise typer.Exit(1)
+ # Steps 3-5 can abort (health timeout, model-load error, or Ctrl+C during the
+ # slow load); tear the server and its children (llama-server, cloudflared) down
+ # on any abort so they never orphan.
+ from studio.backend.run import _graceful_shutdown, _server
- # 4. Create API key in-process.
- api_key = _create_api_key_inprocess(api_key_name)
-
- # 5. Load model via HTTP.
- if not silent:
- typer.echo(f"Loading model: {model}...")
try:
- result = _load_model_via_http(
- port = actual_port,
- api_key = api_key,
- model = model,
- gguf_variant = gguf_variant,
- max_seq_length = max_seq_length,
- load_in_4bit = load_in_4bit,
- llama_extra_args = extra_llama_args,
- )
- except RuntimeError as exc:
- typer.echo(f"Error: {exc}", err = True)
- raise typer.Exit(1)
+ # 3. Wait for server health.
+ if not silent:
+ typer.echo("Starting Unsloth Studio...")
+ if not _wait_for_server(actual_port):
+ typer.echo("Error: server did not become healthy within 30 seconds.", err = True)
+ raise typer.Exit(1)
+
+ # 4. Create API key in-process.
+ api_key = _create_api_key_inprocess(api_key_name)
+
+ # 5. Load model via HTTP.
+ if not silent:
+ typer.echo(f"Loading model: {model}...")
+ try:
+ result = _load_model_via_http(
+ port = actual_port,
+ api_key = api_key,
+ model = model,
+ gguf_variant = gguf_variant,
+ max_seq_length = max_seq_length,
+ load_in_4bit = load_in_4bit,
+ llama_extra_args = extra_llama_args,
+ )
+ except RuntimeError as exc:
+ typer.echo(f"Error: {exc}", err = True)
+ raise typer.Exit(1)
+ except BaseException:
+ _graceful_shutdown(_server)
+ raise
loaded_model = result.get("model", model)
display_variant = f" ({gguf_variant})" if gguf_variant else ""
@@ -1045,6 +1085,8 @@ def run(
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
base_url = f"http://{display_host}:{actual_port}"
sdk_base_url = f"{base_url}/v1"
+ # run_server started the tunnel during the silent run above (0.0.0.0 only).
+ _cf_url = getattr(app.state, "cloudflare_url", None)
# Orange so the tool-policy notice stands out; printed under
# --silent / --yes too so the policy is never invisible.
@@ -1074,6 +1116,8 @@ def run(
typer.echo("")
typer.echo("=" * 56)
typer.echo(f" Unsloth Studio running at {base_url}")
+ if _cf_url:
+ typer.echo(f" Secure link access via Cloudflare: {_cf_url}")
typer.echo(f" Model loaded: {loaded_model}{display_variant}")
typer.echo(f" API Key: {api_key}")
typer.echo("")
@@ -1107,6 +1151,8 @@ def run(
else:
# Silent still prints URL + API key + tool-status policy.
typer.echo(f"URL: {base_url}")
+ if _cf_url:
+ typer.echo(f"Secure link access via Cloudflare: {_cf_url}")
typer.echo(f"API Key: {api_key}")
typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True)
diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py
new file mode 100644
index 0000000000..629cb46405
--- /dev/null
+++ b/unsloth_cli/tests/test_inference_chat.py
@@ -0,0 +1,401 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for the `unsloth chat` / `unsloth inference` CLI — fakes only, no model loads."""
+
+from __future__ import annotations
+
+import inspect
+import sys
+import types
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(_REPO_ROOT))
+
+
+import typer
+from rich.console import Console
+from typer.testing import CliRunner
+
+import unsloth_cli.commands.chat as chatmod
+from unsloth_cli._inference import (
+ ChatBackend,
+ HttpChatBackend,
+ collect_stream,
+ render_columns,
+ visible_text,
+)
+
+
+class _FakeConfig:
+ is_gguf = False
+ is_lora = True
+ display_name = "fake-model"
+ base_model = "fake/base"
+ path = None
+
+
+def _chat_app():
+ cli = typer.Typer()
+ cli.command()(chatmod.chat)
+ return cli
+
+
+def test_visible_text_passthrough_when_shown():
+ text = "reasoninganswer"
+ assert visible_text(text, show_thinking = True) == text
+
+
+def test_visible_text_strips_closed_think_block():
+ text = "step 1\nstep 2The answer is 42."
+ assert visible_text(text, show_thinking = False) == "The answer is 42."
+
+
+def test_visible_text_holds_unclosed_think():
+ # An open is held back so partial reasoning never leaks mid-stream.
+ assert visible_text("still thinking", show_thinking = False) == ""
+ assert visible_text("done.more thinking", show_thinking = False) == "done."
+
+
+def test_visible_text_holds_partial_think_prefix():
+ # Streams are cumulative, so the opening tag can arrive as "<", "". Hold possible tag prefixes until they are disambiguated.
+ assert visible_text("<", show_thinking = False) == ""
+ assert visible_text("rhel", "rhello"])
+ assert collect_stream(stream, show_thinking = False) == "hello"
+
+
+def test_render_columns_emits_both_answers_with_separator(capsys):
+ render_columns("base", "alpha", "tuned", "beta")
+ out = capsys.readouterr().out
+ assert "base" in out and "tuned" in out
+ assert "alpha" in out and "beta" in out
+ assert "│" in out
+
+
+def test_you_prompt_matches_readline_backend(monkeypatch):
+ gnu = types.ModuleType("readline")
+ gnu.__doc__ = "Importing this module enables command line editing using GNU readline."
+ monkeypatch.setitem(sys.modules, "readline", gnu)
+ prompt = chatmod._you_prompt(colors = True)
+ assert "You: " in prompt and "\001" in prompt
+
+ libedit = types.ModuleType("readline")
+ libedit.__doc__ = "Importing this module enables command line editing using libedit readline."
+ monkeypatch.setitem(sys.modules, "readline", libedit)
+ assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m"
+ assert chatmod._you_prompt(colors = False) == "\nYou: "
+
+ # Windows: no readline module at all; the console's own line editing
+ # handles backspace, so plain ANSI color (no markers) is safe.
+ monkeypatch.setitem(sys.modules, "readline", None)
+ assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m"
+ assert chatmod._you_prompt(colors = False) == "\nYou: "
+
+
+def test_chat_registered_on_app():
+ from unsloth_cli import app
+
+ # cmd.name is None until typer resolves it from the callback name.
+ names = {(cmd.name or cmd.callback.__name__) for cmd in app.registered_commands}
+ assert "chat" in names
+
+
+def test_chat_exits_cleanly_on_slash_exit(monkeypatch):
+ closed = []
+
+ class _FakeChatBackend:
+ def stream(self, *a, **k):
+ return iter(["hello"])
+
+ def close(self):
+ closed.append(True)
+
+ monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+ monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
+ monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
+
+ runner = CliRunner()
+ for args in (["fake-model"], ["fake-model", "--compare"]):
+ closed.clear()
+ result = runner.invoke(_chat_app(), args, input = "hi\n/exit\n")
+ assert result.exit_code == 0, result.output
+ assert closed == [True]
+ assert "Bye." in result.output
+ # The prompt must go through input() (readline-safe), not a print.
+ assert "You: " in result.output
+ assert "You: You:" not in result.output
+
+
+def test_pick_trained_model_lists_and_selects(monkeypatch):
+ fake_models = types.ModuleType("utils.models")
+ fake_models.scan_trained_models = lambda: [
+ ("run-new", "outputs/run-new", "lora"),
+ ("run-old", "outputs/run-old", "merged"),
+ ]
+ monkeypatch.setitem(sys.modules, "utils.models", fake_models)
+
+ monkeypatch.setattr("builtins.input", lambda prompt = "": "2")
+ assert chatmod._pick_trained_model(Console()) == "outputs/run-old"
+
+ monkeypatch.setattr("builtins.input", lambda prompt = "": "")
+ assert chatmod._pick_trained_model(Console()) == "outputs/run-new"
+
+
+def test_chat_no_arg_chats_with_picked_trained_model(monkeypatch):
+ class _FakeChatBackend:
+ def stream(self, *a, **k):
+ return iter(["hello"])
+
+ def close(self):
+ pass
+
+ resolved = []
+ monkeypatch.setattr(chatmod, "_pick_trained_model", lambda console: "outputs/run-42")
+ monkeypatch.setattr(
+ chatmod,
+ "resolve_model_config",
+ lambda model, **k: (resolved.append(model), _FakeConfig())[1],
+ )
+ monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
+ monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
+
+ result = CliRunner().invoke(_chat_app(), [], input = "/exit\n")
+ assert result.exit_code == 0, result.output
+ assert resolved == ["outputs/run-42"]
+
+
+def test_find_studio_server_none_when_not_running(monkeypatch):
+ import urllib.request
+
+ from unsloth_cli import _inference
+
+ def refuse(*a, **k):
+ raise OSError("connection refused")
+
+ monkeypatch.setattr(urllib.request, "urlopen", refuse)
+ assert _inference.find_studio_server() is None
+
+
+class _FakeSSEResponse:
+ def __init__(self, lines):
+ self._lines = lines
+
+ def __iter__(self):
+ return iter(self._lines)
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ return False
+
+
+def test_http_backend_streams_cumulative_text(monkeypatch):
+ backend = HttpChatBackend("http://localhost:8888", "token")
+ response = _FakeSSEResponse(
+ [
+ b'data: {"choices":[{"delta":{"content":"He"}}]}\n',
+ b"\n",
+ b'data: {"choices":[{"delta":{"content":"llo"}}]}\n',
+ b"data: [DONE]\n",
+ ]
+ )
+ monkeypatch.setattr(backend, "_request", lambda *a, **k: response)
+
+ out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS))
+ assert out == ["He", "Hello"]
+
+
+def test_http_backend_merges_emoji_split_across_deltas(monkeypatch):
+ backend = HttpChatBackend("http://localhost:8888", "token")
+ response = _FakeSSEResponse(
+ [
+ b'data: {"choices":[{"delta":{"content":"hi "}}]}\n',
+ b'data: {"choices":[{"delta":{"content":"\\ud83d"}}]}\n',
+ b'data: {"choices":[{"delta":{"content":"\\ude0a"}}]}\n',
+ b"data: [DONE]\n",
+ ]
+ )
+ monkeypatch.setattr(backend, "_request", lambda *a, **k: response)
+
+ out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS))
+ # The lone high surrogate is held back, then merged with its other half.
+ assert out == ["hi ", "hi ", "hi 😊"]
+
+
+def test_chat_prefers_running_studio_server(monkeypatch):
+ closed = []
+
+ class _FakeHttpBackend:
+ def stream(self, *a, **k):
+ return iter(["hello"])
+
+ def close(self):
+ closed.append("http")
+
+ local_loads = []
+ monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend())
+ monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: local_loads.append(1))
+ monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
+
+ result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
+
+ assert result.exit_code == 0, result.output
+ assert local_loads == []
+ assert "stays warm" in result.output
+ assert closed == ["http"]
+
+
+def test_chat_server_mode_compare_loads_base_locally(monkeypatch):
+ streamed, closed, base_loads = [], [], []
+
+ class _FakeHttpBackend:
+ def stream(self, *a, **k):
+ streamed.append("tuned")
+ return iter(["tuned-answer"])
+
+ def close(self):
+ closed.append("http")
+
+ class _FakeBaseBackend:
+ def stream(self, *a, **k):
+ streamed.append("base")
+ return iter(["base-answer"])
+
+ def close(self):
+ closed.append("base")
+
+ def fake_local_load(model, **kwargs):
+ base_loads.append((model, kwargs.get("fresh_backend", False)))
+ return _FakeBaseBackend()
+
+ monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend())
+ monkeypatch.setattr(chatmod, "load_chat_backend", fake_local_load)
+
+ result = CliRunner().invoke(_chat_app(), ["tuned-run"], input = "/compare\nhi\n/exit\n")
+
+ assert result.exit_code == 0, result.output
+ assert "(compare on)" in result.output
+ # Only the base model loaded locally, on its own private backend.
+ assert base_loads == [("fake/base", True)]
+ assert streamed == ["base", "tuned"]
+ assert set(closed) == {"http", "base"}
+
+
+def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
+ loads, streamed, closed = [], [], []
+
+ class _FakeLocalBackend:
+ def __init__(self, role):
+ self.role = role
+
+ def stream(self, *a, **k):
+ streamed.append((self.role, k.get("use_adapter")))
+ return iter([f"{self.role}-answer"])
+
+ def close(self):
+ closed.append(self.role)
+
+ def fake_load(model, **kwargs):
+ fresh = kwargs.get("fresh_backend", False)
+ loads.append((model, fresh))
+ return _FakeLocalBackend("base" if fresh else "tuned")
+
+ monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+ monkeypatch.setattr(chatmod, "load_chat_backend", fake_load)
+ monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: True)
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
+
+ result = CliRunner().invoke(_chat_app(), ["tuned-run", "--compare"], input = "hi\n/exit\n")
+
+ assert result.exit_code == 0, result.output
+ assert loads == [("tuned-run", False), ("fake/base", True)]
+ # Both models answered the turn, via plain generation (no adapter toggle).
+ assert ("base", None) in streamed and ("tuned", None) in streamed
+ assert set(closed) == {"tuned", "base"}
diff --git a/unsloth_cli/tests/test_studio_cloudflare_flag.py b/unsloth_cli/tests/test_studio_cloudflare_flag.py
new file mode 100644
index 0000000000..5c2a039547
--- /dev/null
+++ b/unsloth_cli/tests/test_studio_cloudflare_flag.py
@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for the `--cloudflare/--no-cloudflare` Studio flag.
+
+Pins the typer Option (default on) on both `unsloth studio` and
+`unsloth studio run`, and that the chosen polarity reaches the re-exec'd
+child and run_server. Modeled on test_studio_run_parallel_flag.py.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+from typer.testing import CliRunner
+
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+ sys.path.insert(0, str(_REPO_ROOT))
+
+
+def _studio():
+ from unsloth_cli.commands import studio as _studio_mod
+ return _studio_mod
+
+
+_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"]
+
+
+# ── option registration ──────────────────────────────────────────────
+
+
+def test_run_exposes_cloudflare_option_default_on():
+ import inspect
+
+ sig = inspect.signature(_studio().run)
+ assert "cloudflare" in sig.parameters
+ opt = sig.parameters["cloudflare"].default
+ decls = set(getattr(opt, "param_decls", []) or [])
+ assert "--cloudflare/--no-cloudflare" in decls
+ assert getattr(opt, "default", None) is True
+
+
+def test_studio_default_exposes_cloudflare_option_default_on():
+ import inspect
+
+ sig = inspect.signature(_studio().studio_default)
+ assert "cloudflare" in sig.parameters
+ opt = sig.parameters["cloudflare"].default
+ assert getattr(opt, "default", None) is True
+
+
+# ── re-exec forwarding: `unsloth studio run` ─────────────────────────
+
+
+class _ExecCaptured(SystemExit):
+ def __init__(self, argv):
+ super().__init__(0)
+ self.argv = list(argv)
+
+
+def _install_run_reexec_capture(monkeypatch, *, platform = "linux"):
+ studio_mod = _studio()
+ captured = []
+
+ monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+ fake_venv = Path("/fake/studio/venv/unsloth_studio")
+ monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+ fake_bin = fake_venv / "bin" / "unsloth"
+ real_is_file = Path.is_file
+ monkeypatch.setattr(
+ Path,
+ "is_file",
+ lambda self: True if str(self) == str(fake_bin) else real_is_file(self),
+ )
+ from unsloth_cli import _tool_policy as _tp_mod
+
+ monkeypatch.setattr(
+ _tp_mod,
+ "resolve_tool_policy",
+ lambda host, flag, yes, silent: False if flag is None else bool(flag),
+ )
+ monkeypatch.setattr(sys, "platform", platform)
+
+ def fake_execvp(file, argv):
+ captured.append(list(argv))
+ raise _ExecCaptured(argv)
+
+ monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+ return captured
+
+
+def _invoke_run(monkeypatch, args):
+ import typer as _typer
+
+ studio_mod = _studio()
+ captured = _install_run_reexec_capture(monkeypatch)
+ app = _typer.Typer()
+ app.command(
+ context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+ )(studio_mod.run)
+ CliRunner().invoke(app, args, catch_exceptions = True)
+ return captured
+
+
+@pytest.mark.parametrize(
+ "user_flag,expected,unexpected",
+ [
+ (None, "--cloudflare", "--no-cloudflare"), # default on
+ ("--cloudflare", "--cloudflare", "--no-cloudflare"),
+ ("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
+ ],
+)
+def test_run_reexec_forwards_cloudflare_polarity(monkeypatch, user_flag, expected, unexpected):
+ extras = [user_flag] if user_flag else []
+ captured = _invoke_run(monkeypatch, _BASE + extras)
+ assert len(captured) == 1, captured
+ argv = captured[0]
+ assert expected in argv, f"expected {expected} in child argv; got {argv}"
+ assert unexpected not in argv, f"unexpected {unexpected} in child argv; got {argv}"
+
+
+# ── re-exec forwarding: plain `unsloth studio` ───────────────────────
+
+
+def _invoke_studio_default(
+ monkeypatch,
+ args,
+ *,
+ platform = "linux",
+):
+ import typer as _typer
+
+ studio_mod = _studio()
+ captured = []
+
+ monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+ monkeypatch.setattr(studio_mod, "_ensure_studio_env_exported", lambda: None)
+ fake_venv = Path("/fake/studio/venv/unsloth_studio")
+ monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+ monkeypatch.setattr(studio_mod, "_find_run_py", lambda: Path("/fake/studio/run.py"))
+ monkeypatch.setattr(studio_mod, "_find_frontend_dist", lambda: None)
+ monkeypatch.setattr(sys, "platform", platform)
+
+ def fake_execvp(file, argv):
+ captured.append(list(argv))
+ raise _ExecCaptured(argv)
+
+ monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+
+ app = _typer.Typer()
+ app.command()(studio_mod.studio_default)
+ CliRunner().invoke(app, args, catch_exceptions = True)
+ return captured
+
+
+@pytest.mark.parametrize(
+ "user_flag,expected,unexpected",
+ [
+ (None, "--cloudflare", "--no-cloudflare"),
+ ("--no-cloudflare", "--no-cloudflare", "--cloudflare"),
+ ],
+)
+def test_studio_default_reexec_forwards_cloudflare(monkeypatch, user_flag, expected, unexpected):
+ extras = [user_flag] if user_flag else []
+ captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0"] + extras)
+ assert len(captured) == 1, captured
+ argv = captured[0]
+ assert expected in argv, f"expected {expected}; got {argv}"
+ assert unexpected not in argv, f"unexpected {unexpected}; got {argv}"
+
+
+# ── in-venv path forwards cloudflare into run_server ─────────────────
+
+
+class _RunServerCaptured(SystemExit):
+ def __init__(self, kwargs):
+ super().__init__(0)
+ self.kwargs = dict(kwargs)
+
+
+@pytest.mark.parametrize("user_flag,expected", [(None, True), ("--no-cloudflare", False)])
+def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, expected):
+ import types
+
+ studio_mod = _studio()
+ fake_venv = Path("/fake/studio/venv/unsloth_studio")
+ monkeypatch.setattr(sys, "prefix", str(fake_venv))
+ monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
+
+ from unsloth_cli import _tool_policy as _tp_mod
+
+ monkeypatch.setattr(
+ _tp_mod,
+ "resolve_tool_policy",
+ lambda host, flag, yes, silent: False if flag is None else bool(flag),
+ )
+
+ captured: dict = {}
+
+ def fake_run_server(**kwargs):
+ captured.update(kwargs)
+ raise _RunServerCaptured(kwargs)
+
+ fake_backend_run = sys.modules.setdefault(
+ "studio.backend.run", types.ModuleType("studio.backend.run")
+ )
+ fake_backend_run.run_server = fake_run_server
+ fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
+
+ import typer as _typer
+
+ app = _typer.Typer()
+ app.command(
+ context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+ )(studio_mod.run)
+ extras = [user_flag] if user_flag else []
+ CliRunner().invoke(app, _BASE + extras, catch_exceptions = True)
+
+ assert captured.get("cloudflare") is expected, captured
+
+
+# ── parent-level --no-cloudflare with a subcommand is rejected ───────
+
+
+def test_studio_default_rejects_no_cloudflare_with_subcommand(monkeypatch):
+ # `unsloth studio --no-cloudflare run ...` would not reach the subcommand,
+ # so it must error (mirrors --parallel) rather than silently still tunnel.
+ import typer as _typer
+
+ studio_mod = _studio()
+ app = _typer.Typer()
+ app.add_typer(studio_mod.studio_app, name = "studio")
+ result = CliRunner().invoke(app, ["studio", "--no-cloudflare", "run", "--model", "X"])
+ assert result.exit_code == 2, result.output
+ combined = (result.output or "") + (getattr(result, "stderr", "") or "")
+ assert "--no-cloudflare" in combined, combined
+
+
+# ── run() tears the server + tunnel down if startup aborts ───────────
+
+
+def test_run_in_venv_shuts_down_on_startup_abort(monkeypatch):
+ import types
+
+ studio_mod = _studio()
+ fake_venv = Path("/fake/studio/venv/unsloth_studio")
+ monkeypatch.setattr(sys, "prefix", str(fake_venv))
+ monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
+
+ from unsloth_cli import _tool_policy as _tp_mod
+
+ monkeypatch.setattr(
+ _tp_mod,
+ "resolve_tool_policy",
+ lambda host, flag, yes, silent: False if flag is None else bool(flag),
+ )
+
+ class _App:
+ class state:
+ server_port = 8888
+
+ shutdown_calls = []
+ backend = sys.modules.setdefault("studio.backend.run", types.ModuleType("studio.backend.run"))
+ backend.run_server = lambda **k: _App()
+ backend._resolve_external_ip = lambda: "1.2.3.4"
+ backend._server = object()
+ backend._shutdown_event = None
+ backend._graceful_shutdown = lambda server: shutdown_calls.append(server)
+
+ # set_tool_policy is imported as `from state.tool_policy import set_tool_policy`.
+ state_mod = sys.modules.setdefault("state", types.ModuleType("state"))
+ tp_mod = sys.modules.setdefault("state.tool_policy", types.ModuleType("state.tool_policy"))
+ tp_mod.set_tool_policy = lambda *a, **k: None
+ state_mod.tool_policy = tp_mod
+
+ # Force the health check to fail so startup aborts after run_server().
+ monkeypatch.setattr(studio_mod, "_wait_for_server", lambda *a, **k: False)
+
+ import typer as _typer
+
+ app = _typer.Typer()
+ app.command(
+ context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+ )(studio_mod.run)
+ result = CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0"], catch_exceptions = True)
+
+ assert result.exit_code == 1, result.output
+ assert len(shutdown_calls) == 1, "startup abort must call _graceful_shutdown"