diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 2457f08407..b63ac94b93 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh }" # Determinism (seed/temp) is applied at the server level by # serve-unsloth-run.sh --extra; agents inherit it through the API. TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" +# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex +# exec) it runs a full turn AND a separate small_model call to name the session, +# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared +# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it +# headroom (still well under the 40-min job budget); the fast agents keep the +# tight cap that still catches a real headless-TTY hang. +case "$AGENT" in + opencode) + # Double it, but only for a bare-integer seconds value. A GNU timeout(1) + # duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so + # the arithmetic never sees a non-number; timeout(1) parses it directly. + case "$TIMEOUT" in + *[!0-9]*) ;; + *) TIMEOUT=$(( TIMEOUT * 2 )) ;; + esac + ;; +esac # Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner # IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless diff --git a/install.sh b/install.sh index 5972379d26..6076721540 100755 --- a/install.sh +++ b/install.sh @@ -472,11 +472,13 @@ _on_install_exit() { _restore_studio_venv_replacement fi [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true exit "$_status" } -# Empty so an inherited value can never reach the trap's rm; only a temp dir -# this script creates below (Apple Silicon, spaced path) is ever removed. +# Empty so an inherited value never reaches the trap's rm; only temp paths this +# script creates below (spaced-path dir, torch-trio overrides) are removed. _UV_OVERRIDE_TMPDIR="" +_UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── @@ -1821,6 +1823,8 @@ tauri_log "STEP" "Creating virtual environment" mkdir -p "$STUDIO_HOME" _MIGRATED=false +# Empty so an inherited value can never masquerade as a probed torch version. +_PREV_TORCH_VER="" if [ -x "$VENV_DIR/bin/python" ]; then # why: matching guard to the .venv branch below -- in env-mode @@ -1838,6 +1842,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi + # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run + # rebuilds the venv for clean state, but must keep the torch release the user + # already has (see _previous_torch_pin below). Last line only: sitecustomize or + # import-hook noise on stdout must not corrupt the version. + _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ + "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" @@ -2187,6 +2197,68 @@ _torch_flavor_tag() { esac } +# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 +# ("torch>=A.B[.C],="*",<"*) ;; + *) echo "no"; return ;; + esac + _trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}" + _trw_ceil="${_trw_con##*,<}" + _v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}" + _f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}" + _c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}" + for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do + case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac + done + if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then + if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then + echo "yes" + return + fi + fi + echo "no" +} + +# Whether a re-run should keep the previous venv's torch: echo "torch==X.Y.Z" when the +# probed previous version ($1) has a flavor tag matching the freshly chosen cu*/cpu index +# leaf ($2) AND sits inside the active constraint window ($3), else "". Re-running +# `curl | sh` rebuilds the venv for clean state, but a healthy torch the user already +# validated must not be silently moved to a newer release (2.10 -> 2.11); a flavor +# change (cpu <-> cuda, cu126 -> cu130) still installs the correct new build, rocm +# leaves keep their floors (rocm7.2 must land 2.11 for the Strix _grouped_mm fix), and +# a release outside the window (2.3.x manual install, 2.12.x manual upgrade) is never +# kept: the installer's own bounds win. Opt out with UNSLOTH_TORCH_UPGRADE=1 to get +# the newest release. +_previous_torch_pin() { + _ptp_ver="$1" + _ptp_leaf="$2" + _ptp_con="$3" + [ -n "$_ptp_ver" ] || { echo ""; return; } + [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; } + case "$_ptp_leaf" in + cu[0-9]*|cpu) ;; + *) echo ""; return ;; + esac + _ptp_base="${_ptp_ver%%+*}" + # The base must look like a release (probe noise / garbage must never become a pin). + case "$_ptp_base" in + [0-9]*.[0-9]*) ;; + *) echo ""; return ;; + esac + [ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; } + if [ "$(_torch_flavor_tag "$_ptp_ver")" = "$_ptp_leaf" ]; then + echo "torch==$_ptp_base" + else + echo "" + fi +} + # Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. _expected_torch_flavor_tag() { @@ -2478,12 +2550,32 @@ case "$_torch_index_leaf" in *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. -# All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the +# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in +# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index +# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so +# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm +# leaf keeps the default <2.11.0. +case "$_torch_index_leaf" in + rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; + cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;; esac +# Re-run over an existing install: keep the previous venv's torch release instead of +# resolving the newest in range. The range stays in _PREV_FALLBACK_CONSTRAINT so the +# install can fall back when the exact release is not on the chosen index (custom +# mirrors may prune old wheels). Skipped for --no-torch (no previous probe runs). +_PREV_TORCH_PIN="" +_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT" +if [ "$SKIP_TORCH" = false ]; then + _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$_torch_index_leaf" "$TORCH_CONSTRAINT") + if [ -n "$_prev_pin" ]; then + _PREV_TORCH_PIN="$_prev_pin" + TORCH_CONSTRAINT="$_prev_pin" + substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + fi +fi + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". @@ -2705,6 +2797,43 @@ esac # ── Install unsloth directly into the venv (no activation needed) ── tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" + +# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares +# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio, +# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard +# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so +# freeze the trio via uv --overrides (overrides replace dependency requirements +# during resolution) while unsloth's other deps resolve normally. Sets +# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth +# install (migrated and fresh) must call this before resolving and rm it after. +_build_unsloth_torch_overrides() { + _UNSLOTH_TORCH_OVERRIDES="" + [ "$SKIP_TORCH" = false ] || return 0 + _torch_trio_pins=$("$_VENV_PY" -c " +from importlib.metadata import version, PackageNotFoundError +for _p in ('torch', 'torchvision', 'torchaudio'): + try: + print(_p + '==' + version(_p)) + except PackageNotFoundError: + pass +" 2>/dev/null) || _torch_trio_pins="" + case "$_torch_trio_pins" in + torch==*) + _UNSLOTH_TORCH_OVERRIDES=$(mktemp) + printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES" + # The CLI --overrides flag replaces any UV_OVERRIDE env file (same + # uv setting; macOS arm64 exports one here), so fold its pins in. + # awk, not cat: it drops inherited torch-trio lines (uv intersects + # duplicate overrides, so a conflicting pin would make resolution + # unsatisfiable) and newline-terminates the last line so an + # unterminated file cannot join two requirements into one. + for _ov_file in ${UV_OVERRIDE:-}; do + [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES" + done + ;; + esac +} + if [ "$_MIGRATED" = true ]; then # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state # in the new venv location, while preserving existing torch/CUDA @@ -2729,9 +2858,13 @@ if [ "$_MIGRATED" = true ]; then else # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. + _build_unsloth_torch_overrides run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-} + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2913,8 +3046,20 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + if [ -n "$_PREV_TORCH_PIN" ]; then + # Kept previous release: fall back to the supported range if the exact + # release is not resolvable from the chosen index (pruned mirror). + if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL"; then + substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN" + TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT" + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL" + fi + else + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL" + fi fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2927,9 +3072,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then ;; esac fi - # Fresh: Step 2 - install unsloth, preserving pre-installed torch + # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." + _build_unsloth_torch_overrides if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. @@ -2953,6 +3099,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps @@ -2962,8 +3109,11 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. if [ "$SKIP_TORCH" = false ]; then diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c9ab7eb83b..d7c7eed518 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ import atexit import contextlib import functools import json +import math import os import re import struct @@ -23,11 +24,22 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union +from typing import ( + Callable, + Collection, + Generator, + Iterable, + List, + Literal, + Mapping, + Optional, + Union, +) import httpx from core.inference.llama_server_args import ( + _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, @@ -234,8 +246,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# Plan-without-action re-prompt state (intent signal, caps, message) now lives -# in tool_call_parser, imported above under its old aliases. +# Plan-without-action re-prompt state now lives in tool_call_parser (imported above). # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -1431,7 +1442,10 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) -_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +# Layer-offload override detection. Single-sourced from llama_server_args, which +# also strips these (plus the MoE flags) from inherited extras; sharing the layer +# set keeps detection and stripping from drifting. +_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS _THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) @@ -1895,6 +1909,17 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # GPU memory strategy applied on the active load ("auto"/"manual"). + self._gpu_memory_mode: str = "auto" + # Manual-mode load options (echoed back so the UI round-trips them). + self._gpu_layers: int = -1 + # MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. + self._n_cpu_moe: int = 0 + # Relative model share per GPU (--tensor-split), in GPU order; None = + # default (llama.cpp splits by free VRAM). + self._tensor_split: Optional[List[float]] = None + # User-picked physical GPU indices (None = automatic selection). + self._gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -1909,6 +1934,11 @@ class LlamaCppBackend: self._spec_draft_n_max: Optional[int] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None + # MoE metadata (populated by _read_gguf_metadata): expert count (>0 = + # MoE) and leading dense-layer count (offsets --n-cpu-moe, which counts + # from layer 0). See the n_moe_layers property. + self._n_experts: Optional[int] = None + self._leading_dense_block_count: Optional[int] = None self._n_kv_heads: Optional[int] = None self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None @@ -2329,6 +2359,79 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def gpu_memory_mode(self) -> str: + """Active GPU memory strategy: 'auto' or 'manual' (gpu_layers < 0 = Auto/--fit, >= 0 = pinned).""" + return self._gpu_memory_mode + + @property + def gpu_layers(self) -> int: + """Requested --gpu-layers for manual mode (-1 when not manual).""" + return self._gpu_layers + + @property + def n_cpu_moe(self) -> int: + """MoE expert layers manual mode kept on CPU (--n-cpu-moe); 0 = none.""" + return self._n_cpu_moe + + @property + def tensor_split(self) -> Optional[List[float]]: + """Manual-mode relative model share per GPU (--tensor-split); None = + default (split by free VRAM).""" + return self._tensor_split + + @property + def gpu_ids(self) -> Optional[List[int]]: + """User-picked physical GPU indices, or None for automatic selection.""" + return self._gpu_ids + + @property + def n_layers(self) -> Optional[int]: + """Model layer count (GGUF block_count), or None if unknown.""" + return self._n_layers + + @property + def n_moe_layers(self) -> int: + """Number of MoE expert layers (the --n-cpu-moe ceiling), 0 if not MoE. + + block_count minus the leading dense layers (which carry no experts): + --n-cpu-moe counts from layer 0, so those dense layers are no-ops. + """ + if not self._n_experts or not self._n_layers: + return 0 + return max(0, self._n_layers - (self._leading_dense_block_count or 0)) + + @staticmethod + def _resolve_cpu_moe_flag( + n_cpu_moe: int, n_moe_layers: int, leading_dense: int + ) -> Optional[int]: + """The --n-cpu-moe value (absolute first-N layers), or None to omit it. + + Clamps the requested count to the model's MoE layers, then offsets past + the leading dense layers (--n-cpu-moe counts from layer 0). Returns None + for nothing-to-offload (0 requested) or a non-MoE model. + """ + if n_cpu_moe <= 0 or n_moe_layers <= 0: + return None + return leading_dense + min(n_cpu_moe, n_moe_layers) + + @staticmethod + def _sanitize_tensor_split(tensor_split: Optional[List[float]]) -> List[float]: + """Per-GPU shares with negative and non-finite entries clamped to 0. + + A direct caller's negative entry would launch a placement different + from the ratio the UI showed, and inf would pass a plain ``> 0`` total + gate and emit ``--tensor-split inf,...``. Returns [] for input that + can't be read as floats (the length gate at the call site then drops + the split). + """ + try: + return [ + x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split) + ] + except (TypeError, ValueError, OverflowError): + return [] + @property def layer_preserves_tensor_intent(self) -> bool: """True when a downgraded tensor request kept this layer load multi-GPU.""" @@ -2530,6 +2633,7 @@ class LlamaCppBackend: "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_fit_target": False, "supports_cache_ram": False, "supports_ctx_checkpoints": False, "supports_no_cache_prompt": False, @@ -2549,6 +2653,7 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_fit_target = False supports_cache_ram = False supports_ctx_checkpoints = False supports_no_cache_prompt = False @@ -2646,6 +2751,7 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_fit_target = _is_real("--fit-target") supports_cache_ram = _is_real("--cache-ram") supports_ctx_checkpoints = _is_real("--ctx-checkpoints") supports_no_cache_prompt = _is_real("--no-cache-prompt") @@ -2662,6 +2768,7 @@ class LlamaCppBackend: "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_fit_target": supports_fit_target, "supports_cache_ram": supports_cache_ram, "supports_ctx_checkpoints": supports_ctx_checkpoints, "supports_no_cache_prompt": supports_no_cache_prompt, @@ -2746,6 +2853,57 @@ class LlamaCppBackend: except ValueError: return None + @staticmethod + def _emit_child_gpu_visibility(env: dict, pinned: str) -> None: + """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on + ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child + seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP + mask at different layers, so the same indices apply twice -- ROCR reduces + and re-indexes from 0, then a non-zero HIP pin points out of range, HIP + enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone + narrows correctly; clear any inherited ROCR mask so it can't double up.""" + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env.pop("ROCR_VISIBLE_DEVICES", None) + except Exception as e: + logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + + @staticmethod + def _pin_visible_gpu_order_for_split(env: dict) -> None: + """Pin the child's GPU enumeration to the picker's order for a manual + ``--tensor-split`` across the whole visible set. CUDA's default + FASTEST_FIRST enumeration applies the shares to the wrong cards on + heterogeneous hosts (#5025), and CUDA_DEVICE_ORDER only fixes the + numbering base: an inherited numeric visibility mask ALSO defines + enumeration order, so a reordered parent mask (CUDA_VISIBLE_DEVICES=3,1) + would still hand the shares to the wrong cards. The UI built the split + positionally over get_backend_visible_gpu_info's device list (ascending + physical via nvidia-smi, inherited mask order on the torch fallback), so + re-emit the same set in that report order -- not an assumed ascending + sort. The visible set itself never changes. No mask, an empty mask, or a + UUID/MIG mask (which resolves to None) is left alone -- the multi-GPU + controls are hidden for the latter.""" + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + inherited = LlamaCppBackend._resolve_visible_physical_ids() + if not inherited: + return + order = None + try: + from utils.hardware import get_backend_visible_gpu_info + info = get_backend_visible_gpu_info() + if info.get("available") and info.get("index_kind") == "physical": + reported = [d["index"] for d in info.get("devices", [])] + if sorted(reported) == sorted(inherited): + order = reported + except Exception as e: + logger.debug("Could not read reported GPU order for split pin: %s", e) + if order is None: + order = sorted(inherited) + LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order)) + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where @@ -3262,6 +3420,20 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak + # from Studio's parent environment into llama-server and silently override + # the command assembled from the current request. Draft-model placement is + # intentionally separate and remains available to speculative decoding. + _MANUAL_PLACEMENT_ENV_VARS = ( + "LLAMA_ARG_CPU_MOE", + "LLAMA_ARG_N_CPU_MOE", + "LLAMA_ARG_N_GPU_LAYERS", + "LLAMA_ARG_TENSOR_SPLIT", + "LLAMA_ARG_FIT", + "LLAMA_ARG_FIT_TARGET", + "LLAMA_ARG_FIT_CTX", + ) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't # skip tensor for others; tensor is tried by default, recorded only on a real abort. @@ -3426,6 +3598,12 @@ class LlamaCppBackend: return env + @classmethod + def _clear_manual_placement_env(cls, env: dict[str, str]) -> None: + """Remove inherited main-model placement owned by Manual mode.""" + for name in cls._MANUAL_PLACEMENT_ENV_VARS: + env.pop(name, None) + @staticmethod def _select_gpus( model_size_bytes: int, @@ -4246,6 +4424,8 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -4335,6 +4515,8 @@ class LlamaCppBackend: arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", + f"{arch}.expert_count": "n_experts", + f"{arch}.leading_dense_block_count": "leading_dense_block_count", f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", @@ -4523,6 +4705,28 @@ class LlamaCppBackend: return None + @staticmethod + def _diffusion_gpu_arg(gpu_ids: Optional[List[int]], *, cpu_only: bool = False) -> str: + """Device token passed to the diffusion visual-server child. + + The visual engine replaces its child's CUDA visibility mask with this + token, so an unpinned load must carry forward the first token from the + parent's mask rather than turning a parent-relative ordinal into a new + physical selection. + """ + if gpu_ids: + return str(sorted(gpu_ids)[0]) + if cpu_only: + return "" + if "DG_GPU" in os.environ: + return os.environ["DG_GPU"] + parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES") + if parent_mask: + first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "") + if first and first != "-1": + return first + return "0" + def _start_diffusion_server( self, *, @@ -4533,6 +4737,7 @@ class LlamaCppBackend: model_identifier: str, n_ctx: int, extra_args: Optional[List[str]], + gpu_ids: Optional[List[int]] = None, ) -> bool: """Launch the OpenAI-compat diffusion shim (which drives the on-device visual decoder) and wait for health. Presents the same /v1 + /health @@ -4558,7 +4763,11 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. cpu_only = self._effective_gpu_count() == 0 - gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") + # Honor the GPU picker first: the diffusion runner takes a single device, + # so use the lowest selected GPU (matches the sorted set recorded below, so + # the device used == the echoed gpu_ids[0]). With no pick, fall back to the + # CPU-only mask, else DG_GPU / 0. + gpu = self._diffusion_gpu_arg(gpu_ids, cpu_only = cpu_only) cmd = list(shim_cmd) + [ "--gguf", @@ -4586,6 +4795,11 @@ class LlamaCppBackend: env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu + if gpu_ids: + # The visual server remasks via CUDA_VISIBLE_DEVICES=; pin PCI + # order (as the llama-server path does) so the picked physical id maps + # to the GPU the picker showed, not CUDA's default fastest-first order. + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) if extra_pythonpath: @@ -4631,6 +4845,23 @@ class LlamaCppBackend: self._model_identifier = model_identifier self._cache_type_kv = None self._gpu_offload_active = True + # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to + # defaults (the picked device is still recorded below) so /load, /status + # and reload dedup don't report a previous GGUF's manual settings. + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + # Diffusion is never tensor-parallel; clear any state left by a prior TP + # chat load (load_model phase 1 only kills the process, it doesn't run + # the unload reset) so /status doesn't misreport TP and an identical + # re-Apply doesn't reload against stale tensor-parallel state. + self._tensor_parallel = False + # Record only the single device the runner actually uses (the lowest + # selected GPU, chosen above) -- not the whole pick. The diffusion runner + # is single-device, so echoing a multi-GPU list would misreport placement + # in /status and let a re-Apply dedup against GPUs the runner never used. + self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -5721,6 +5952,11 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -5753,6 +5989,14 @@ class LlamaCppBackend: "speculative_type": speculative_type, "spec_draft_n_max": spec_draft_n_max, "tensor_parallel": tensor_parallel, + # GPU-memory placement: replayed on respawn so a server SIGKILL'd by + # GPU/RAM pressure reloads onto the same devices with the same + # offload, not the auto defaults. + "gpu_memory_mode": gpu_memory_mode, + "gpu_layers": gpu_layers, + "n_cpu_moe": n_cpu_moe, + "tensor_split": list(tensor_split) if tensor_split is not None else None, + "gpu_ids": list(gpu_ids) if gpu_ids is not None else None, "n_threads": n_threads, "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, @@ -5779,6 +6023,11 @@ class LlamaCppBackend: speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, tensor_parallel = tensor_parallel, + gpu_memory_mode = gpu_memory_mode, + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + gpu_ids = gpu_ids, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -5899,6 +6148,7 @@ class LlamaCppBackend: model_identifier = model_identifier, n_ctx = n_ctx, extra_args = extra_args, + gpu_ids = gpu_ids, ) if not binary: @@ -5960,6 +6210,59 @@ class LlamaCppBackend: # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) + # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or + # a per-GPU ratio still launches tensor mode -- and under the + # CPU-only mask below (no visible devices) that aborts the server + # instead of loading on CPU. Drop both here (nothing to split). + if gpu_memory_mode == "manual" and gpu_layers == 0: + if tensor_parallel or tensor_split: + logger.info( + "Manual gpu_layers=0: dropping tensor split/parallel " + "flags (nothing to split on the GPU)" + ) + tensor_parallel = False + tensor_split = None + # Record the requested strategy for /status and the load + # response. 'manual' has no fallback, so the request value is the + # value actually applied. + self._gpu_memory_mode = gpu_memory_mode + # The layer/MoE/split knobs apply only with an explicit offload + # (manual + gpu_layers >= 0); else record defaults so /status and + # /load don't report knobs the server never applied. + if gpu_memory_mode == "manual" and gpu_layers >= 0: + self._gpu_layers = gpu_layers + self._n_cpu_moe = n_cpu_moe + self._tensor_split = tensor_split + else: + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = sorted(gpu_ids) if gpu_ids else None + # Manual offload skips the TP planner but still emits --split-mode + # tensor at launch; drop it when fewer than 2 GPUs are in use -- + # tensor split is a no-op there and aborts on some architectures. + # Done before the cache-drop below so a quantized KV survives. + if ( + tensor_parallel + and gpu_memory_mode == "manual" + and gpu_layers >= 0 + and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2 + ): + logger.info( + "Tensor parallelism requested in manual mode but fewer " + "than 2 GPUs are in use; ignoring (needs >= 2)." + ) + tensor_parallel = False + # Drop TP for manual + Auto layers before the cache-drop below (like + # the <2-GPU guard above), so a requested quantized KV survives into + # the --fit load rather than being stripped for a tensor attempt. + if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0: + logger.info( + "Manual mode with Auto layers hands memory management to " + "llama.cpp --fit, which is incompatible with tensor " + "parallelism; ignoring the tensor split." + ) + tensor_parallel = False # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type # that would re-impose it when appended last). Layer split does @@ -6040,6 +6343,12 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + # Seed before the try: the except (GPU-selection failure -> + # --fit on) falls through to the launch which reads this, and the + # probe that assigns it may throw first. Captured before manual + # empty `gpus` so the speculative defaults stay GPU-aware and the + # CPU-fallback check still knows GPUs were present. + _detected_gpus: list[tuple[int, int]] = [] model_size = None # set in the fit try; used by the APU RAM guard # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). @@ -6057,6 +6366,18 @@ class LlamaCppBackend: _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + # GPU picker: restrict every mode to the chosen devices, so + # auto selection only considers them and manual mask to + # them (the env block below pins CUDA/HIP_VISIBLE_DEVICES). + if gpu_ids: + _picked = set(gpu_ids) + gpus = [g for g in gpus if g[0] in _picked] + + # GPUs the model will run on -- captured before manual + # empty `gpus` to bypass the planner. bool() drives the + # GPU-aware speculative defaults; the list feeds the + # CPU-fallback check. + _detected_gpus = list(gpus) def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): # Per-GPU usable budget for ranking: free - (1-frac)*total. @@ -6088,6 +6409,44 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx + # Manual + Auto layers (the Manual default): hand memory + # management to llama.cpp's --fit. Emptying the probed GPU set + # no-ops the selection/TP planning below, leaving gpu_indices + # None (an explicit gpu_ids pick still pins below) and use_fit + # True. An explicit context is honored (--fit optimizes around + # it); 0 lets --fit size it. + if gpu_memory_mode == "manual" and gpu_layers < 0: + # Tensor parallelism was already dropped above (before the + # cache-drop), so a quantized KV survives into this --fit load. + gpus = [] + effective_ctx = requested_ctx if requested_ctx > 0 else 0 + original_ctx = effective_ctx + # --fit aborts under --split-mode tensor; a raw extras + # --split-mode/--tensor-split (appended last) would + # otherwise reach llama-server. Strip it like the TP + # downgrade does. + extra_args = strip_split_mode_only(extra_args) + elif gpu_memory_mode == "manual": + # Manual offload (--gpu-layers + --fit off): no automatic + # device masking (a gpu_ids pick still pins below) or + # context cap -- the user owns both. tensor_parallel is + # honored but skips the memory-based planner (gpus = []); + # the toggle just emits --split-mode tensor (split by free + # VRAM, or by the Split ratio if set). + gpus = [] + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) + original_ctx = effective_ctx + # Strip the user --split-mode when the toggle owns the split + # (TP engaged -> Studio emits --split-mode tensor) or when the + # user asked for tensor (which aborts on a single GPU even if + # the manual <2-GPU guard downgraded TP). Otherwise keep their + # non-tensor mode (row/none/layer) -- the toggle can't express + # those. + if tensor_parallel or split_mode_override == "tensor": + extra_args = strip_split_mode_only(extra_args) + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always # engage; auto only on an MTP model >= 3B; ngram/off never. A @@ -6175,7 +6534,10 @@ class LlamaCppBackend: _extra_n_max = _extra_args_spec_draft_n_max(extra_args) _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max if _mtp_eff_n_max is None: - _mtp_eff_n_max = 2 if gpus else 3 + # _detected_gpus (not gpus) so manual -- which empty + # gpus to bypass the planner -- keep the GPU draft depth the + # launch flags also use, instead of the CPU default. + _mtp_eff_n_max = 2 if _detected_gpus else 3 # Separate-drafter weights live on GPU (an embedded head is # already in model_size). Size the drafter the launch loads, by # precedence: extras --model-draft (last-wins), else Unsloth's @@ -6313,7 +6675,8 @@ class LlamaCppBackend: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True - # Per-GPU weight proportions for tensor mode (None = even). + # Per-GPU weight proportions for tensor mode (None lets + # llama.cpp split by free VRAM). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # Flat MTP reserve fraction: used only as the fallback when the @@ -6388,7 +6751,12 @@ class LlamaCppBackend: # GPUs below that reserve from the set up front (gpu_indices # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus - if tensor_parallel: + # Manual mode owns the layer count and context, so it skips + # the memory-based planner; its toggle still emits + # --split-mode tensor below (split by free VRAM, or by the + # Split ratio if set). auto plans here. + plan_tp = tensor_parallel and gpu_memory_mode != "manual" + if plan_tp: # Deterministic per-device compute buffer (replicated on # every device in tensor mode); flat fallback when dims # are unavailable. _plan_tensor_parallel uses the same. @@ -6407,7 +6775,7 @@ class LlamaCppBackend: # free yet have no budget left. tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] - if tensor_parallel and len(tp_gpus) < 2: + if plan_tp and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only # or probe failed) it must not reach llama-server; and a @@ -6823,6 +7191,12 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # GPU picker: when no narrower subset was chosen (manual, or + # a failed/file-size selection), pin the whole picked set so the + # model can't spill onto an unpicked GPU. + if gpu_ids and gpu_indices is None: + gpu_indices = sorted(gpu_ids) + # Unified-memory APUs load weights into system RAM (under WSL the VM # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model @@ -6859,8 +7233,6 @@ class LlamaCppBackend: model_path, "--port", str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", str(n_parallel), "--flash-attn", @@ -6868,6 +7240,17 @@ class LlamaCppBackend: # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". "--no-context-shift", ] + # A positive context is always passed (in auto-fit, --fit then + # optimizes the gpu-layer offload around it). When auto-fit has + # no explicit context, omit -c so --fit sizes it to fit VRAM: + # "-c 0" would instead pin the FULL native context (llama.cpp's + # -c handler sets fit_params_min_ctx = UINT32_MAX on value 0, + # disabling --fit's reduction). See gpu_memory_mode. + auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0 + if effective_ctx > 0: + cmd.extend(["-c", str(effective_ctx)]) + elif not auto_fit: + cmd.extend(["-c", "0"]) # Report a clean public model id (matching GET /v1/models) rather # than the raw -m path in llama-server's own /v1/models and the @@ -6879,7 +7262,63 @@ class LlamaCppBackend: cmd.extend(["--alias", _alias]) fully_gpu_offloaded = False - if use_fit: + # Set when a positional --tensor-split is emitted, so the env block + # can pin CUDA to PCI order even without a GPU subset (see below). + manual_tensor_split_emitted = False + if gpu_memory_mode == "manual" and gpu_layers >= 0: + # Pin the user's layer count and disable auto-fit. --fit off + # also means _ctx_integrity_flags must not add --fit-ctx. + use_fit = False + cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"]) + # Keep the first n_cpu_moe MoE layers' experts on CPU. + moe_flag = self._resolve_cpu_moe_flag( + n_cpu_moe, + self.n_moe_layers, + self._leading_dense_block_count or 0, + ) + if moe_flag is not None: + cmd.extend(["--n-cpu-moe", str(moe_flag)]) + elif n_cpu_moe: + # Requested on a dense model: nothing was emitted, so + # don't report a count llama-server never received. + self._n_cpu_moe = 0 + # Distribute the model across GPUs by the user's per-GPU shares + # (--tensor-split). Works with layer split and tensor + # parallelism; --fit off means no fit/tensor abort. Only emit + # when >1 GPU is in use AND the list length matches that count: + # the field is hidden (not cleared) when the picker narrows to + # one, and a direct caller can send a stale ratio for a different + # GPU set. Studio drops any mismatch to the free-VRAM default + # (llama.cpp would silently zero-pad a short list, or abort past + # its 16-device cap). + _split_gpus = self._effective_gpu_count(gpu_indices) + if tensor_split and _split_gpus > 1: + # An all-zero/non-positive sanitized split assigns nothing + # anywhere, so fall through to the free-VRAM default in + # that case. + _sanitized_split = self._sanitize_tensor_split(tensor_split) + _split_total = sum(_sanitized_split) + if len(_sanitized_split) == _split_gpus and _split_total > 0: + cmd.extend( + ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)] + ) + self._tensor_split = _sanitized_split + manual_tensor_split_emitted = True + else: + logger.warning( + "Dropping manual --tensor-split (%d entries for " + "%d GPUs, sanitized total %s); llama.cpp's " + "free-VRAM split applies instead", + len(tensor_split), + _split_gpus, + _split_total, + ) + self._tensor_split = None + elif tensor_split: + # Single effective GPU: the split is never emitted, so + # don't report it as active via /status and /load. + self._tensor_split = None + elif use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- force all layers on GPU. --fit off is @@ -6897,6 +7336,7 @@ class LlamaCppBackend: self._ctx_integrity_flags( n_parallel, use_fit, + auto_fit, requested_ctx, effective_ctx, server_caps, @@ -6960,9 +7400,11 @@ class LlamaCppBackend: self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor - # rather than by layer. Multi-GPU only -- a no-op on a single - # GPU. Default (layer split) is left implicit by omitting the - # flag. See llama.cpp --split-mode. + # rather than by layer. The UI only offers it on multi-GPU; a + # direct single-GPU caller is redundant (supported archs no-op, + # unsupported ones abort and the /load path retries layer split). + # Default (layer split) is left implicit by omitting the flag. + # See llama.cpp --split-mode. if tensor_parallel: cmd.extend(["--split-mode", "tensor"]) if tp_tensor_split and len(tp_tensor_split) > 1: @@ -6994,7 +7436,7 @@ class LlamaCppBackend: extra_args = extra_args, model_identifier = model_identifier, model_path = model_path, - gpus = bool(gpus), + gpus = bool(_detected_gpus), binary = binary, mtp_draft_path = launch_mtp_draft_path, ) @@ -7112,6 +7554,8 @@ class LlamaCppBackend: # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) + if gpu_memory_mode == "manual": + self._clear_manual_placement_env(env) # Omitting --threads relies on llama.cpp's physical-core default, so # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the # arg handler and silently force hardware_concurrency(). #5692 @@ -7170,28 +7614,39 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device # (above), not here. - if gpu_indices is not None and not is_vulkan_backend: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - # Do NOT also set ROCR_VISIBLE_DEVICES to the same - # value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr - # layer and HIP_VISIBLE_DEVICES at the HIP layer, so - # setting both with the same physical indices applies - # the mask twice: ROCR reduces the visible set and - # re-indexes it from 0, then HIP indexes into the - # already-reduced set. A single non-zero pin (e.g. - # "1") then points out of range at the HIP layer, HIP - # enumerates 0 devices, and llama.cpp falls back to - # CPU ("ggml_cuda_init: no ROCm-capable device is - # detected"). The HIP mask alone narrows correctly; - # clear any inherited ROCR mask so it can't double up. - env.pop("ROCR_VISIBLE_DEVICES", None) - except Exception as e: - logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + # A deliberate zero-offload load with no GPU companions runs + # entirely on CPU, yet a visible CUDA device still costs the child + # ~0.5 GB (context + compute scratch) that the CPU-only + # classification below reports as free. Hide the GPUs so the load + # is exactly what it claims: zero VRAM (verified: GPU stays at idle + # baseline and generation runs). Companion loads keep the normal + # masking, and a user device pin (in extras or an inherited + # LLAMA_ARG_DEVICE) keeps control of its own devices -- the child + # aborts on a pin it can't see. The draft-device forms count too: + # llama-server parses them even with no drafter loaded. + _cpu_only_zero_offload = ( + gpu_memory_mode == "manual" + and gpu_layers == 0 + and not is_vulkan_backend + and not self._zero_offload_keeps_gpu_visible(cmd, env) + ) + if _cpu_only_zero_offload: + self._emit_child_gpu_visibility(env, "-1") + elif gpu_indices is not None and not is_vulkan_backend: + # When the user picked GPUs by index, align CUDA's ordering + # with the PCI-bus order the picker enumerated (nvidia-smi), + # so "GPU 1" in the UI is GPU 1 to llama.cpp -- not CUDA's + # default FASTEST_FIRST order (#5025). + if gpu_ids: + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + elif manual_tensor_split_emitted and not is_vulkan_backend: + # A manual per-GPU ratio across ALL GPUs (no explicit pick, so + # no CUDA_VISIBLE_DEVICES mask above): the UI built the + # --tensor-split list in ascending physical/PCI index order, + # so pin the child's enumeration to that order too. The whole + # visible set stays in use; only its ordering is fixed. + self._pin_visible_gpu_order_for_split(env) # Captured before any text-only fallback strips it from cmd. launched_with_mmproj = "--mmproj" in cmd @@ -7350,7 +7805,6 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) - self._reconcile_effective_ctx_with_server() self._max_context_length = ( max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) @@ -7535,6 +7989,10 @@ class LlamaCppBackend: "session; run 'unsloth studio update' to enable vision." ) cmd = self._strip_mmproj_args(_last_spawn_cmd) + # This retry bypasses _spawn_and_wait, so refresh the + # launched-argv snapshot itself -- the zero-offload + # classification below must not see the stripped --mmproj. + _last_spawn_cmd = list(cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) @@ -7566,6 +8024,13 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + # Server is up: adopt the real per-request context it allocated + # -- the length --fit chose, or a --parallel slot split -- so the + # reported context_length matches reality. (Querying /props + # before the spawn above always failed; the seeded value was the + # requested/native length.) + self._reconcile_effective_ctx_with_server() + # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] # clears, list sets. Source records hf_variant for the route's @@ -7580,11 +8045,24 @@ class LlamaCppBackend: self._mtp_runtime_fallback_active = _mtp_active_for_launched_server self._start_mtp_crash_watchdog() - # Catch silent CPU fallback when GPU was intended (#5106). - self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit, gpus or [] - ) - if self._gpu_offload_active is False: + # Catch silent CPU fallback when GPU was intended (#5106). Manual + # offload (no picker) leaves gpu_indices None and use_fit False, so + # include its GPU-layer intent; use the preserved probe since + # auto-layers/manual empty `gpus`. A deliberate zero-offload load + # classifies by its launched argv instead: the main model is + # CPU-only by construction and must read False (not None), or + # training needlessly unloads a server holding no VRAM. + _deliberate_cpu_only = gpu_memory_mode == "manual" and gpu_layers == 0 + if _deliberate_cpu_only: + self._gpu_offload_active = self._zero_offload_gpu_flag( + _last_spawn_cmd, _detected_gpus, env + ) + else: + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit or gpu_memory_mode == "manual", + _detected_gpus, + ) + if self._gpu_offload_active is False and not _deliberate_cpu_only: logger.warning( "llama-server appears to have loaded the model entirely " "on CPU even though Unsloth detected at least one GPU. " @@ -7947,6 +8425,11 @@ class LlamaCppBackend: gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, mtp_draft_path: Optional[str] = None, preserve_multi_gpu_on_layer: bool = False, ) -> bool: @@ -8003,6 +8486,38 @@ class LlamaCppBackend: ): return False + # The diffusion runner is mode-agnostic (always "auto", ignores the + # layer/MoE/split knobs), so a standing manual preference in the + # request must not force a needless reload -- only the GPU pick matters. + if not self._is_diffusion: + # A GPU-memory-mode flip (Unsloth / manual) must always reload. + if self._gpu_memory_mode != gpu_memory_mode: + return False + # Manual: a layer-count change always reloads (covers Auto(-1) <-> a + # pinned count); MoE/split only matter with an explicit offload. + if gpu_memory_mode == "manual" and ( + self._gpu_layers != gpu_layers + or ( + gpu_layers >= 0 + and ( + self._n_cpu_moe != n_cpu_moe + or (self._tensor_split or None) != (tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload (compare order-insensitively; None/[] + # both mean automatic). The diffusion runner collapses a multi-GPU pick + # to its single lowest device, so self._gpu_ids holds just that device; + # normalize the request the same way, or a multi-GPU pick that resolves + # to the same device needlessly reloads. + if self._is_diffusion: + requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None + else: + requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None + if (self._gpu_ids or None) != requested_gpu_pick: + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -8071,6 +8586,78 @@ class LlamaCppBackend: return None return classify_gpu_offload_lines(self._stdout_lines) + @staticmethod + def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the argv/env carries a GPU companion: any --mmproj form, or + a drafter (Studio's --model-draft, the extras aliases, or the + LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of + the main ``--gpu-layers``. A drafter explicitly forced to CPU + (--spec-draft-ngl 0 / --spec-draft-device cpu) doesn't count.""" + if any(str(a).startswith("--mmproj") for a in cmd): + return True + if _extra_args_mtp_draft_path(cmd, env) is None: + return False + return not _extra_args_draft_offloaded_to_cpu(cmd, env) + + @staticmethod + def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """Whether a zero-layer launch still has a reason to use visible GPUs. + + Keep this shared by child masking and post-launch residency bookkeeping: + a device pin, surviving tensor mode, mmproj, or GPU drafter prevents the + launch from being a confirmed zero-VRAM server. + """ + return ( + LlamaCppBackend._cmd_has_gpu_device_pin(cmd, env) + or _effective_tensor_parallel(cmd, False, env) + or LlamaCppBackend._cmd_has_gpu_companion(cmd, env) + ) + + @staticmethod + def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the effective main or draft ``--device`` pin names a GPU.""" + main_flags = {"--device", "-dev"} + draft_flags = {"--spec-draft-device", "-devd", "--device-draft"} + last_main: Optional[str] = None + last_draft: Optional[str] = None + args = [str(arg) for arg in cmd] + for index, raw in enumerate(args): + flag, equals, inline = raw.partition("=") + if flag not in main_flags and flag not in draft_flags: + continue + value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") + if flag in main_flags: + last_main = value + else: + last_draft = value + if last_main is None: + last_main = (env or {}).get("LLAMA_ARG_DEVICE") + + def _names_gpu(value: Optional[str]) -> bool: + if value is None: + return False + devices = [item.strip().lower() for item in value.split(",") if item.strip()] + return not devices or any(item not in ("cpu", "none") for item in devices) + + return _names_gpu(last_main) or _names_gpu(last_draft) + + @staticmethod + def _zero_offload_gpu_flag( + spawn_cmd: list, + detected_gpus: list, + env: Optional[Mapping[str, str]] = None, + ) -> Optional[bool]: + """GPU-residency flag for a deliberate manual zero-offload load. The + main model is CPU-only by construction, but device pins, tensor mode, + mmproj, and GPU drafters can still make the server hold VRAM. The counted + offload classifier cannot see those allocations. This uses the same + predicate as the launch-time zero-VRAM mask; None means no GPU signal.""" + if not detected_gpus: + return None + if LlamaCppBackend._is_vulkan_backend(): + return True + return LlamaCppBackend._zero_offload_keeps_gpu_visible(spawn_cmd, env) + def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not yet consumed by the next load_model. Lets the tensor->layer fallback @@ -8114,11 +8701,18 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -8181,6 +8775,10 @@ class LlamaCppBackend: # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False + # Reset to unknown so the training guard treats the next (still + # loading) server as VRAM-resident rather than reading the killed + # server's stale zero-offload flag until the health probe reclassifies. + self._gpu_offload_active = None # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() @@ -8785,7 +9383,12 @@ class LlamaCppBackend: @staticmethod def _ctx_integrity_flags( - n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + n_parallel: int, + use_fit: bool, + auto_fit: bool, + requested_ctx: int, + effective_ctx: int, + caps: dict, ) -> list[str]: """Flags that keep the per-request window equal to the advertised ctx. @@ -8793,14 +9396,28 @@ class LlamaCppBackend: ``--kv-unified`` default, silently splitting ``-c`` into per-slot windows of ``-c / N``; restore the shared pool so one request can use the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step - at an explicitly requested ctx (default floor is 4096) so it offloads - or fails instead of silently shrinking the window. + at an explicitly requested ctx so it offloads or fails instead of + silently shrinking the window. The 8192 auto-floor and the tighter + ``--fit-target`` margin apply only under Manual + Auto (``auto_fit``), + which omits ``-c``: on the legacy auto path ``-c 0`` already pins the + native window and ``--fit-ctx 8192`` would override it down to 8192. """ flags: list[str] = [] if n_parallel > 1 and caps.get("supports_kv_unified"): flags.append("--kv-unified") - if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): - flags.extend(["--fit-ctx", str(effective_ctx)]) + if use_fit and caps.get("supports_fit_ctx"): + if requested_ctx > 0 and effective_ctx > 0: + # Floor the fit step at the explicitly requested ctx. + flags.extend(["--fit-ctx", str(effective_ctx)]) + elif auto_fit: + # Manual + Auto omits -c, so floor at 8192 so --fit doesn't + # shrink the window below a usable size. + flags.extend(["--fit-ctx", "8192"]) + if use_fit and auto_fit and caps.get("supports_fit_target"): + # llama.cpp's --fit leaves 1 GiB free per device by default; + # tighten that to 512 MiB so it packs more of the model onto + # the GPU before spilling to system RAM. + flags.extend(["--fit-target", "512"]) return flags def _query_server_n_ctx(self) -> Optional[int]: diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 70d0dc774d..e72e10e071 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -186,12 +186,25 @@ _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) _TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS +# GPU-offload flags. Stripped only when the GPU Memory mode owns offload +# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's +# inherited -ngl is respected (the offload_overridden path), so this group is +# opt-in, not default. Layer flags are shared with llama_cpp's override +# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them). +_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( + {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"} +) +_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) +_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS + _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS ) # Shadowing flags that take no value -- strip the flag only, not the next token. -_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"} +) def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: @@ -424,6 +437,8 @@ def strip_shadowing_flags( strip_spec: bool = True, strip_template: bool = True, strip_split_mode: bool = True, + strip_tensor_split: bool = False, + strip_offload: bool = False, ) -> list[str]: """Strip flags that shadow first-class Unsloth settings. @@ -432,6 +447,12 @@ def strip_shadowing_flags( (same for cache / spec / template / split-mode). Each ``strip_*`` toggle controls one group; the route only strips groups whose first-class field the caller actually supplied. + + ``strip_split_mode`` removes both ``--split-mode`` and the coupled + ``--tensor-split`` (the Tensor Parallelism toggle owns the whole split). + ``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can + replace an inherited per-GPU ratio while leaving the user's ``--split-mode`` + row/none/layer choice intact. """ shadowing: set[str] = set() if strip_context: @@ -444,6 +465,10 @@ def strip_shadowing_flags( shadowing |= _TEMPLATE_FLAGS if strip_split_mode: shadowing |= _SPLIT_SHADOWING_FLAGS + if strip_tensor_split: + shadowing |= _TENSOR_SPLIT_FLAGS + if strip_offload: + shadowing |= _OFFLOAD_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py index 86ad8b9fd8..64ab38ec75 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -201,7 +201,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]: continue # Skip what Unsloth hides from its pickers (validation probe, RAG embed # weights): not chat models, so never an auto-switch target. - if _is_hidden_model(raw_id, getattr(info, "path", None)): + if _is_hidden_model( + raw_id, + getattr(info, "model_id", None), + getattr(info, "path", None), + ): continue # Advertise a client-facing alias, not an absolute filesystem path. loader_id = _advertised_loader_id(info) diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 2de32a68e4..f54d795731 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -87,6 +87,22 @@ def _names_gguf(model: str) -> bool: return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) +def gguf_repo_for_embedding_model(model: str) -> str: + """GGUF repo for ``model``, honoring an explicit companion override.""" + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + +def default_gguf_repo() -> str: + """GGUF companion for the env/default embedding model.""" + return gguf_repo_for_embedding_model(EMBEDDING_MODEL) + + def effective_gguf_repo() -> str: """GGUF repo for the llama-server backend, tracking the effective model. @@ -95,14 +111,7 @@ def effective_gguf_repo() -> str: ``-GGUF`` companion repo (the unsloth convention the default pair follows), or is used as-is when it already names a GGUF repo. """ - if "RAG_EMBED_GGUF_REPO" in os.environ: - return EMBED_GGUF_REPO - model = effective_embedding_model() - if model == DEFAULT_EMBEDDING_MODEL: - return EMBED_GGUF_REPO - if _names_gguf(model): - return model - return f"{model}-GGUF" + return gguf_repo_for_embedding_model(effective_embedding_model()) # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 1f38af9381..54a25482f2 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -37,6 +37,13 @@ from hub.services.models.common import ( _runtime_for_format, ) +# Imported at module scope (not inside the per-repo scan loop) so a broken +# import surfaces at startup instead of silently emptying the inventory: the +# scan loop swallows per-repo exceptions and would drop every repo. Lives under +# ``utils`` (not ``utils.models``) to avoid the eager model-config/checkpoint +# imports in ``utils/models/__init__.py``. +from utils.hidden_models import is_hidden_model + logger = get_logger(__name__) _repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = ( @@ -243,6 +250,13 @@ def invalidate_hf_cache_scans() -> None: hf_cache_scan.invalidate_hf_cache_scans() +def _is_hidden_infra_repo(*values: str | None) -> bool: + """True for infra-only repos (the RAG embedder and the llama.cpp install + validation probe) that are cached as a side effect of Studio itself and are + not usable chat models.""" + return is_hidden_model(*values) + + def _scan_cached_gguf() -> list[dict]: """Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread.""" cache_scans = all_hf_cache_scans() @@ -254,13 +268,24 @@ def _scan_cached_gguf() -> list[dict]: if str(repo_info.repo_type) != "model": continue repo_id = repo_info.repo_id + repo_path = Path(repo_info.repo_path) + snapshot_path = _cached_model_snapshot_path(repo_path) total_size = _repo_gguf_size_bytes(repo_info) has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id) + is_hidden_infra = _is_hidden_infra_repo( + repo_id, + str(repo_path), + str(snapshot_path) if snapshot_path is not None else None, + ) + # Hide infra repos unless the user downloaded a variant via + # the Hub; variant state only exists for user downloads. + if is_hidden_infra and not has_variant_state: + continue if total_size == 0 and not has_variant_state: continue partial = hf_cache_scan.is_gguf_repo_partial( repo_id, - Path(repo_info.repo_path), + repo_path, ) if total_size == 0 and not partial: continue @@ -283,6 +308,9 @@ def _scan_cached_gguf() -> list[dict]: requires_variant = True, ) ) + # Visible infra variants remain management-only. + if is_hidden_infra: + row["capabilities"]["can_chat"] = False if _prefer_cache_row(row, existing): seen_lower[key] = row except Exception as e: @@ -475,6 +503,15 @@ def _scan_cached_models() -> list[dict]: if str(repo_info.repo_type) != "model": continue repo_id = repo_info.repo_id + repo_path = Path(repo_info.repo_path) + snapshot_path = _cached_model_snapshot_path(repo_path) + # The non-GGUF embedder has no variant downloads; always hide. + if _is_hidden_infra_repo( + repo_id, + str(repo_path), + str(snapshot_path) if snapshot_path is not None else None, + ): + continue has_main_gguf = _repo_has_gguf_files(repo_info) payload = _repo_non_gguf_model_payload(repo_info) if payload.size_bytes == 0: @@ -486,7 +523,6 @@ def _scan_cached_models() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) - repo_path = Path(repo_info.repo_path) snapshot_partial = hf_cache_scan.is_snapshot_partial( "model", repo_id, diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py index a3782efead..b34532fa35 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -36,6 +36,7 @@ from hub.utils.paths import ( ) from hub.services.models import common as model_common from hub.services.models.ollama import scan_ollama_dir +from utils.hidden_models import is_hidden_model logger = get_logger(__name__) _MAX_MODELS_PER_CUSTOM_FOLDER = 200 @@ -623,6 +624,20 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI ) +def _filter_hidden_models(local_models: List[LocalModelInfo]) -> list[LocalModelInfo]: + """Remove infrastructure-only models from the shared local inventory.""" + visible: list[LocalModelInfo] = [] + for model in local_models: + resolved_cache_path = ( + hf_cache_scan.resolve_hf_cache_realpath(Path(model.path)) + if model.source == "hf_cache" + else None + ) + if not is_hidden_model(model.id, model.model_id, model.path, resolved_cache_path): + visible.append(model) + return visible + + async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse: """List local model candidates from every supported on-device source.""" hf_cache_dir = _resolve_hf_cache_dir() @@ -653,7 +668,7 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel ollama_dirs, ) local_models += await _collect_models_from_custom_folders() - models = _dedupe_local_models(local_models) + models = _dedupe_local_models(_filter_hidden_models(local_models)) return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 2c33e09b2b..693d945ee1 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -439,6 +439,287 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa assert row["capabilities"]["requires_variant"] is True +def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path): + probe = _repo( + "ggml-org/models", + [_file("tinyllamas/stories260K.gguf", 1_200_000)], + tmp_path / "probe", + ) + embedder = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)], + tmp_path / "embedder", + ) + chat = _repo("Org/Chat-GGUF", [_file("Q4_K_M.gguf", 100)], tmp_path / "chat") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [probe, embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + + assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"] + + +def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + embedder = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [ + _file("bge-small-en-v1.5-f16.gguf", 60_000_000), + _file("bge-small-en-v1.5-Q8_0.gguf", 35_000_000), + ], + tmp_path / "embedder", + ) + # Variant manifests only exist for user Hub downloads, not auto-downloads. + assert download_manifest.write_manifest( + "model", + "unsloth/bge-small-en-v1.5-GGUF", + "Q8_0", + [download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)], + "http", + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + + assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"] + assert result["cached"][0]["capabilities"]["can_chat"] is False + + +def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path): + embedder_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5" + embedder_path.mkdir(parents = True) + embedder = _repo( + "unsloth/bge-small-en-v1.5", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + embedder_path, + ) + chat_path = tmp_path / "hub" / "models--Org--Chat" + chat_path.mkdir(parents = True) + chat = _repo( + "Org/Chat", + [_file("config.json", 12), _file("model.safetensors", 100)], + chat_path, + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_models()} + + assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat"] + + +def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + gguf_path = tmp_path / "hub" / "models--Org--PathEmbedder-GGUF" + gguf_path.mkdir(parents = True) + gguf = _repo( + "Org/PathEmbedder-GGUF", + [_file("model-F16.gguf", 60_000_000)], + gguf_path, + ) + model_path = tmp_path / "hub" / "models--Org--PathEmbedder" + model_path.mkdir(parents = True) + model = _repo( + "Org/PathEmbedder", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + model_path, + ) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(model_path), + ) + monkeypatch.setattr( + rag_config, + "effective_gguf_repo", + lambda: str(gguf_path), + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, model])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + +def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + gguf_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder-GGUF" + gguf_snapshot = gguf_path / "snapshots" / "gguf-revision" + gguf_snapshot.mkdir(parents = True) + gguf = _repo( + "Org/SnapshotEmbedder-GGUF", + [_file("model-F16.gguf", 60_000_000)], + gguf_path, + ) + model_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder" + model_snapshot = model_path / "snapshots" / "model-revision" + model_snapshot.mkdir(parents = True) + model = _repo( + "Org/SnapshotEmbedder", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + model_path, + ) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(model_snapshot), + ) + monkeypatch.setattr( + rag_config, + "effective_gguf_repo", + lambda: str(gguf_snapshot), + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, model])], + ) + + def _resolve_snapshot(repo_path): + return str( + { + gguf_path: gguf_snapshot, + model_path: model_snapshot, + }.get(Path(repo_path), Path(repo_path)) + ) + + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + _resolve_snapshot, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + +def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder( + monkeypatch, tmp_path +): + # A custom embedder with a generic basename ("org/model") must be hidden by + # EXACT repo-id match only. An unrelated cached chat model whose id merely + # contains "model" (e.g. "user/model-chat") must stay on device: substring + # basename matching used to drop real chat models from the inventory. + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + def _model_repo(repo_id: str): + path = tmp_path / "hub" / f"models--{repo_id.replace('/', '--')}" + path.mkdir(parents = True) + return _repo( + repo_id, + [_file("config.json", 12), _file("model.safetensors", 100)], + path, + ) + + embedder = _model_repo("org/model") + chat = _model_repo("user/model-chat") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_models()} + + assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"] + + +def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path): + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + gguf = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)], + tmp_path / "default-gguf", + ) + weights_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5" + weights_path.mkdir(parents = True) + weights = _repo( + "unsloth/bge-small-en-v1.5", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + weights_path, + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, weights])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj(): requirements = gguf_variants._build_gguf_variant_requirements( [ @@ -1610,6 +1891,63 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_ assert rows[0].capabilities.requires_variant is True +def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_path): + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + + def _row(repo_id: str): + repo_path = tmp_path / f"models--{repo_id.replace('/', '--')}" + return model_common._local_model_info( + scan_path = repo_path, + load_path = repo_path, + source = "hf_cache", + model_format = "safetensors", + model_id = repo_id, + ) + + rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")]) + + assert [row.model_id for row in rows] == ["org/chat-model"] + + +def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + embedder_path = tmp_path / "hub" / "models--org--embedder" + embedder_snapshot = embedder_path / "snapshots" / "revision" + embedder_snapshot.mkdir(parents = True) + chat_path = tmp_path / "hub" / "models--org--chat-model" + chat_path.mkdir(parents = True) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(embedder_snapshot), + ) + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path), + ) + + def _row(repo_id: str, repo_path: Path): + return model_common._local_model_info( + scan_path = repo_path, + load_path = repo_path, + source = "hf_cache", + model_format = "safetensors", + model_id = repo_id, + ) + + rows = local_inventory._filter_hidden_models( + [_row("org/embedder", embedder_path), _row("org/chat-model", chat_path)] + ) + + assert [row.model_id for row in rows] == ["org/chat-model"] + + def test_model_download_job_helpers_preserve_idle_shape(): key = downloads._download_job_key("Org/Model", None) status = downloads._job_status(key) diff --git a/studio/backend/main.py b/studio/backend/main.py index 4797764ce7..81d4c16e52 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1156,9 +1156,23 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") enriched_devices.append(enriched_dev) + # Whether GGUF loads accept an explicit gpu_ids pick: /load and + # /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu + # ordinals) and on Vulkan-only builds (--device pins ggml's own + # ordinals), so the picker must not offer them. + try: + from core.inference.llama_cpp import LlamaCppBackend + from utils.hardware import DeviceType, get_device + gpu_ids_supported = ( + get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend() + ) + except Exception as e: + logger.debug(f"Could not resolve gpu_ids support: {e}") + gpu_ids_supported = True gpu_info = { "available": visibility_info.get("available", False), "devices": enriched_devices, + "gguf_gpu_ids_supported": gpu_ids_supported, } _system_gpu_cache = (time.monotonic(), gpu_info) return gpu_info diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f3ae0f70df..d51d35189b 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -64,7 +64,7 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", ) speculative_type: Optional[str] = Field( None, @@ -100,6 +100,66 @@ class LoadRequest(BaseModel): "No effect on a single GPU. Ignored for non-GGUF models." ), ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GPU memory strategy for GGUF models. 'auto' (default): Unsloth " + "selects GPUs and caps context to fit VRAM. 'manual': you own the " + "offload. Leave gpu_layers at -1 (Auto) to hand memory management to " + "llama.cpp's --fit (no device masking, no context auto-reduce, no " + "gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers " + "and n_cpu_moe yourself (--fit off), with tensor_parallel still " + "applying (split by free VRAM unless tensor_split is set, no planner). " + "Ignored for non-GGUF." + ), + ) + gpu_layers: int = Field( + -1, + ge = -1, + description = ( + "Manual mode only: number of layers to offload to the GPU " + "(--gpu-layers, with --fit off). A value >= the model's layer count " + "offloads all of them. -1 = Auto: hand layer + context sizing to " + "llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'." + ), + ) + n_cpu_moe: int = Field( + 0, + ge = 0, + description = ( + "Manual mode only: keep the first N MoE expert layers on the CPU " + "(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of " + "MoE layers offloaded (the backend offsets past any leading dense " + "layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + tensor_split: Optional[List[float]] = Field( + None, + description = ( + "Manual mode only: relative share of the model per GPU (--tensor-split), " + "in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let " + "llama.cpp use its default, which splits by free VRAM. Any list given is " + "passed through as-is, so send [1, 1] to force an even split. Ignored " + "unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + + @field_validator("tensor_split") + @classmethod + def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]: + # A negative / non-finite / all-zero split is silently dropped at launch + # (stored as None) yet still compared raw in the reload dedupe, so an + # identical Apply reloads forever. Reject it up front; [] = no split. + if not value: + return value + import math + + if any((not math.isfinite(v)) or v < 0 for v in value): + raise ValueError("tensor_split entries must be finite and non-negative") + if sum(value) <= 0: + raise ValueError("tensor_split must have a positive total") + return value + llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -133,6 +193,14 @@ class ValidateModelRequest(BaseModel): max_seq_length: int = Field(0, ge = 0, le = 1048576) load_in_4bit: bool = Field(True) gpu_ids: Optional[List[int]] = Field(None) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GGUF GPU-memory strategy intended for the follow-up load. Manual " + "placement bypasses the training coexistence estimate: Auto layers " + "delegate fitting to llama.cpp, while explicit layers are user-owned." + ), + ) include_context_length: bool = Field( False, description = "Also read the native context length from the local GGUF header. " @@ -188,6 +256,16 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + layer_count: Optional[int] = Field( + None, + description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read " + "from the header alongside context_length; None when not read.", + ) + moe_layer_count: Optional[int] = Field( + None, + description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF " + "header alongside context_length; 0 for dense models, None when not read.", + ) # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. requires_transformers_upgrade: bool = Field( False, @@ -333,6 +411,34 @@ class LoadResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) class UnloadResponse(BaseModel): @@ -461,6 +567,42 @@ class InferenceStatusResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + requested_context_length: Optional[int] = Field( + None, + description = ( + "The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the " + "UI re-seed a Manual + Auto-layers context pin on hydration, where " + "context_length only exposes the resolved value. None for non-GGUF." + ), + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3d527bf317..136e4f7645 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, Callable, List, Optional, Union +from typing import Any, Callable, List, Literal, Optional, Union import json import httpx from loggers import get_logger @@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: - """Whether an inherited --split-mode should be stripped on reload. + """Whether an inherited --split-mode (and its coupled --tensor-split) should + be stripped on reload. The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ layer modes, so only strip when the toggle overrides it: tensor being turned on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes - survive. Shared by the inheritance strip and the already-loaded stale check - so they agree on what reload would do. + survive. A manual per-GPU ratio is handled by _should_strip_tensor_split, + which strips only --tensor-split so the inherited mode is kept. Shared by the + inheritance strip and the already-loaded stale check so they agree on what + reload would do. """ fields_set = getattr(request, "model_fields_set", set()) return "tensor_parallel" in fields_set and ( @@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _should_strip_tensor_split(request: LoadRequest) -> bool: + """Whether an inherited --tensor-split alone should be stripped on reload. + + Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio + it emits its own --tensor-split (an inherited one, appended last, would + override it), and with the ratio cleared it wants llama.cpp's default + free-VRAM split. Either way an inherited --tensor-split must go, else the + cleared case silently keeps the stale ratio while status reports None. + Unlike _should_strip_split_mode this leaves --split-mode untouched, so a + user's row/none/layer mode survives a Studio split-ratio edit. When the + Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode + (called alongside this at every site) strips --split-mode anyway. + """ + return ( + getattr(request, "gpu_memory_mode", "auto") == "manual" + and getattr(request, "gpu_layers", -1) >= 0 + ) + + def _carry_preserved_tensor_intent( *, preserved: bool, same_model: bool, explicit_drop: bool ) -> bool: @@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings( else strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) ) if not _tensor_parallel_matches_loaded( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # The diffusion runner is mode-agnostic (it always reports "auto" and ignores + # the layer/MoE/split knobs), so a standing manual preference in the request + # must not force a needless reload -- only the GPU pick matters. + if not llama_backend.is_diffusion: + if request.gpu_memory_mode != llama_backend.gpu_memory_mode: + return False + # Manual: a layer-count change always reloads; MoE/split only matter with + # an explicit offload (gpu_layers >= 0), so a leftover value under Auto + # must not force one. Mirrors LlamaCppBackend._already_in_target_state. + if request.gpu_memory_mode == "manual" and ( + request.gpu_layers != llama_backend.gpu_layers + or ( + request.gpu_layers >= 0 + and ( + request.n_cpu_moe != llama_backend.n_cpu_moe + or (request.tensor_split or None) != (llama_backend.tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU + # request to its single lowest device (it drives one device only), so the + # backend records just that device; compare the request the same way, or a + # multi-GPU pick that resolves to the same device needlessly reloads. + if llama_backend.is_diffusion: + _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None + else: + _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None + if _req_gpu_ids != llama_backend.gpu_ids: + return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement # re-selects instead of keeping the all-GPU mask (#6659). The effective check @@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings( # contain any shadow flag, so the reload path strips them rather than # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - # Mirror the reload's conditional split-mode strip, so a preserved - # non-tensor mode (row/none/layer) isn't seen as stale and doesn't - # trigger a needless reload of a healthy server. + # Mirror the reload's conditional strips, so a preserved non-tensor mode + # (row/none/layer) isn't seen as stale and doesn't trigger a needless + # reload of a healthy server, while an inherited offload/ratio flag that + # the reload *would* strip is correctly seen as stale. if ( backend_extra and strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) != backend_extra ): @@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb( return None +def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: + """Classify a GGUF as diffusion, normal, or unknown before it is loaded. + + ``None`` is important here: a remote GGUF whose header is not cached can + still be routed to the single-GPU diffusion runner after download. Treating + that case as normal would let Manual mode skip the training guard even + though the runner ignores Manual's llama-server placement controls. + """ + identity = " ".join( + str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") + ).lower() + if "diffusion" in identity: + return True + + try: + main = getattr(config, "gguf_file", None) + if not (main and Path(main).is_file()): + repo = getattr(config, "gguf_hf_repo", None) + variant = getattr(config, "gguf_variant", None) + if repo and variant: + from hub.utils.gguf import resolve_local_gguf_path + main = resolve_local_gguf_path(repo, variant) + if not main or not Path(main).is_file(): + return None + + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A successfully decoded architecture proves that this is a normal + # llama-server GGUF. No architecture means the lightweight probe could + # not establish the routing decision, so preserve the unknown state. + if getattr(probe, "_architecture", None): + return False + return None + except Exception as e: + logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + return None + + def _guard_chat_load_against_training( config: ModelConfig, *, @@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training( requested_gpu_ids: Optional[List[int]], llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + gpu_memory_mode: Literal["auto", "manual"] = "auto", ) -> None: - """Refuse loading a local chat model that would OOM an active training run. + """Protect active training from automatically placed chat-model loads. + No-op when training is inactive or unknown. `load_in_4bit` must be the - effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the - model would not fit alongside training.""" + effective quantization (see _effective_load_in_4bit). Manual chat-GGUF + placement is an explicit override: Auto layers delegate fitting to + llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is + estimated here. Diffusion is still guarded because its mode-agnostic runner + ignores those controls and uses one GPU. An unclassified GGUF is guarded as + potentially diffusion until its local header proves otherwise. Other loads + raise HTTP 409 when they would not fit beside training. + """ from core.training import get_training_backend from routes.training_vram import can_load_chat_during_training @@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training( return is_gguf = bool(getattr(config, "is_gguf", False)) + diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False + if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: + return + + diffusion_gpu = None + if is_gguf and diffusion_kind is not False: + # Use the same token selection as the runner: an explicit pick wins, + # followed by DG_GPU, the first parent-visible token, then GPU 0. + diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( + requested_gpu_ids, + cpu_only = LlamaCppBackend._effective_gpu_count() == 0, + ) + required_override_gb = ( _estimate_gguf_required_gb( config, @@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training( requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, required_override_gb = required_override_gb, + single_device_gpu = diffusion_gpu, ) if ok: return @@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training( raise HTTPException(status_code = 409, detail = detail) +def _resolve_inherited_extra_args( + request, + config: ModelConfig, + model_identifier: str, + extra_llama_args: Optional[list[str]], + effective_chat_template_override: Optional[str] = None, +) -> Optional[list[str]]: + """Effective pass-through extras for a GGUF request that omitted the field: + the previous same-model load's extras, shadow-stripped, so a settings-Apply + reload (which does not round-trip the extras field) keeps them (#5401).""" + if getattr(request, "llama_extra_args", None) is not None: + return extra_llama_args + if not getattr(config, "is_gguf", False): + return extra_llama_args + llama_backend = get_llama_cpp_backend() + if not llama_backend.extra_args: + return extra_llama_args + # Inherit the previous load's extras (the chat-settings Apply path doesn't + # round-trip them; an explicit [] still clears). Gated on (model_identifier, + # hf_variant) to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins CLI + # parse against a freshly-supplied first-class field. + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request field: callers + # commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick + # flows. ``config.gguf_variant`` is the variant load_model was actually + # invoked with, so both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower()) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend doesn't + # inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field was set by the caller, so + # an inherited --chat-template-file survives an Apply that omits + # chat_template_override. A bundled family template (e.g. gemma-4) counts as + # a first-class template even when the request omits chat_template_override, + # so strip the inherited --chat-template-file then too -- else the stale arg + # (appended last) shadows the bundled template while Studio reports its caps. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set), + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args), + # manual + per-GPU ratio emits its own --tensor-split; drop + # an inherited one (appended last would override it) while + # keeping the user's --split-mode row/none/layer choice. + strip_tensor_split = _should_strip_tensor_split(request), + # manual emits its own --fit/--gpu-layers, so an inherited offload flag + # must not last-wins-override it. auto leaves a user's inherited -ngl + # alone. getattr: a validate request reuses this resolver, no offload fields. + strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual", + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Shouldn't happen on already-validated args; degrade to + # no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + return extra_llama_args + + def _model_json_response(model, status_code: int = 200) -> Response: """Serialize a pydantic response once via pydantic-core. @@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre None if request.llama_extra_args is None else extra_llama_args ) + # Manual mode owns the offload flags: strip them from EXPLICIT extras + # too (the inherited path already does), or a last-wins --gpu-layers / + # --fit in extras re-enables GPU offload on a load status reports as + # CPU-only. Manual + per-GPU ratio owns --tensor-split the same way. + if request.gpu_memory_mode == "manual" and extra_llama_args: + _stripped_explicit = strip_shadowing_flags( + extra_llama_args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = True, + ) + if _stripped_explicit != extra_llama_args: + logger.info( + "Manual GPU memory owns the offload flags; stripping them " + "from explicit llama_extra_args: %s -> %s", + extra_llama_args, + _stripped_explicit, + ) + extra_llama_args = _stripped_explicit + + # Keep every downstream consumer on the normalized explicit list. In + # particular, the already-loaded comparator must not compare the raw + # request's managed offload flags against the stripped launch state. + request = request.model_copy(update = {"llama_extra_args": extra_llama_args}) + model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") ) @@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) else: if ( @@ -4187,12 +4434,41 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409. + # GGUF supports gpu_ids: validate the pick up front (before the training + # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects + # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts + # are rejected outright: the picker's indices are torch-xpu ordinals neither + # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin + # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + # Same reasoning for a Vulkan-only build: --device pins ggml's own + # Vulkan ordinals, so a physical pick can land on the wrong card on + # masked or non-contiguous hosts. + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre "architectures)" ) - # Refuse a load that would OOM active training, before the unload step below - # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. + # Inherit the previous same-model load's pass-through extras when this + # request omits the field (a settings-Apply reload doesn't round-trip + # them); shadow-stripped so an inherited flag can't override a + # first-class field the caller did set (#5401). + extra_llama_args = _resolve_inherited_extra_args( + request, + config, + model_identifier, + extra_llama_args, + effective_chat_template_override, + ) + + # Apply the training coexistence policy before the unload step below + # frees the resident model. Off-loop: the default-mode guard does sync work. await asyncio.to_thread( _guard_chat_load_against_training, config, @@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre requested_gpu_ids = effective_gpu_ids, llama_extra_args = extra_llama_args, n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1), + gpu_memory_mode = request.gpu_memory_mode, ) # ── GGUF path: load via llama-server ────────────────────── @@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre from core.inference.llama_cpp import gguf_load_in_flight gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) - # Inherit llama_extra_args from the previous load when the request - # omits the field (the chat-settings Apply path doesn't round-trip - # them; explicit [] still clears). Gated on (model_identifier, - # hf_variant) to refuse cross-model pickup, and shadowing flags are - # stripped so an inherited override can't win the last-wins CLI - # parse against a freshly-supplied first-class field. - if request.llama_extra_args is None and llama_backend.extra_args: - source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request - # field: callers commonly omit gguf_variant for local - # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ - # variant`` is the variant load_model was actually - # invoked with (see the HF / local branches below), so - # both sides of the comparison key off the same string. - resolved_variant = (config.gguf_variant or "").lower() - request_variant = (request.gguf_variant or "").lower() - stored_variant = (source[1] or "").lower() if source else "" - same_model = bool( - source and source[0] and source[0].lower() == model_identifier.lower() - ) - if request.gguf_variant: - variant_mismatch = request_variant != stored_variant - else: - variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) - same_source = same_model and not variant_mismatch - if not same_source: - logger.info( - "Not inheriting llama_extra_args: stored args came from %s, loading %s", - source, - (model_identifier, resolved_variant), - ) - # Cross-model: clear explicitly so the backend doesn't - # inherit via "no opinion" semantics. - extra_llama_args = [] - else: - # Strip only the groups whose first-class field was set by - # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. A bundled family - # template (e.g. the gemma-4 override) is an effective - # first-class template setting even when the raw request - # omits chat_template_override, so strip the inherited - # --chat-template-file in that case too -- otherwise the stale - # extra arg (appended last) shadows the bundled template while - # Unsloth reports the bundled template's capabilities. - fields_set = getattr(request, "model_fields_set", set()) - stripped = strip_shadowing_flags( - llama_backend.extra_args, - strip_context = "max_seq_length" in fields_set, - strip_cache = "cache_type_kv" in fields_set, - strip_spec = ( - "speculative_type" in fields_set or "spec_draft_n_max" in fields_set - ), - strip_template = ( - "chat_template_override" in fields_set - or effective_chat_template_override is not None - ), - strip_split_mode = _should_strip_split_mode( - request, llama_backend.extra_args - ), - ) - try: - extra_llama_args = validate_extra_args(stripped) - except ValueError: - # Shouldn't happen on already-validated args; degrade to - # no-extras rather than 400 if managed flags changed. - logger.warning( - "Stored llama_extra_args failed revalidation; loading without them: %s", - stripped, - ) - extra_llama_args = [] - else: - if extra_llama_args: - logger.info( - "Inheriting llama_extra_args from previous " - "load (same model, shadow-stripped): %s", - extra_llama_args, - ) - # Block cache writes that would race the download manager. This runs # after pass-through argument inheritance so a carried --no-mmproj # changes the companion requirement exactly as it does for the load. @@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, spec_draft_n_max = request.spec_draft_n_max, + gpu_memory_mode = request.gpu_memory_mode, + gpu_layers = request.gpu_layers, + n_cpu_moe = request.n_cpu_moe, + tensor_split = request.tensor_split, + gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, ) if config.gguf_hf_repo: @@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -4795,7 +5018,9 @@ def _requires_security_review_for_model( @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( - request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) + request: ValidateModelRequest, + fastapi_request: Request = None, + current_subject: str = Depends(get_current_subject), ): """ Lightweight validation endpoint for model identifiers. @@ -4823,15 +5048,39 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) - # Refuse early (before the frontend unloads to load this) if it can't fit - # alongside training, using the same settings /load uses so they agree. + # Apply the same training coexistence policy as /load before the frontend + # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400. + # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is + # a clean 400) before the guard sizes the model against training VRAM. + # XPU-host picks are rejected like /load (no defined mapping from the + # picker's torch-xpu ordinals to the launcher's device spaces). if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -4895,16 +5144,32 @@ async def validate_model( latest_tier_active_for, config.identifier, request.hf_token ): effective_load_in_4bit = False - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) + # A metadata-only probe just reads the GGUF header and allocates no VRAM, + # so it must not be refused by the training guard. Real loads validate + # without include_context_length and /load applies the guard again. + if not request.include_context_length: + # Match /load's inherited llama.cpp extras and parallel slot count so + # validation cannot pass a smaller estimate than the subsequent load. + effective_extra_args = _resolve_inherited_extra_args( + request, config, model_identifier, None + ) + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + llama_extra_args = effective_extra_args, + n_parallel = ( + getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if fastapi_request is not None + else 1 + ), + gpu_memory_mode = request.gpu_memory_mode, + ) # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a # mixed repo are inert for this load, so gating on them is a false positive. Only @@ -4918,10 +5183,15 @@ async def validate_model( # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. + # Staged header dims (one read): native context, total layer count, and + # MoE expert-layer count -- let the staged flow size the context, GPU- + # layers and manual --n-cpu-moe sliders before the load. context_length: Optional[int] = None + layer_count: Optional[int] = None + moe_layer_count: Optional[int] = None if request.include_context_length and is_gguf: from hub.utils.gguf import resolve_local_gguf_path - from utils.models.gguf_metadata import read_gguf_context_length + from utils.models.gguf_metadata import read_gguf_staged_dims # Best-effort: a header-read failure must never fail validation of an # otherwise-valid model (the outer except turns it into a 400). @@ -4937,9 +5207,15 @@ async def validate_model( model_identifier, request.gguf_variant ) if local_gguf: - context_length = read_gguf_context_length(local_gguf) + # Header walk reads tokenizer arrays for dense models (tens of + # ms); keep it off the event loop. + dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) + if dims: + context_length = dims["context_length"] + layer_count = dims["layer_count"] + moe_layer_count = dims["moe_layer_count"] except Exception as e: - logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + logger.debug("Header probe failed for %s: %s", model_log_label, e) return ValidateModelResponse( valid = True, @@ -4954,6 +5230,8 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + layer_count = layer_count, + moe_layer_count = moe_layer_count, requires_transformers_upgrade = transformers_upgrade is not None, transformers_upgrade = transformers_upgrade, ) @@ -5593,6 +5871,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)): speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + requested_context_length = llama_backend.requested_n_ctx, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 742ecde3ba..0806c2f513 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -59,59 +59,12 @@ def _safe_is_dir(path) -> bool: return False -# Hub repo id shape ("owner/name", no leading separator); anything else is -# treated as a local filesystem path. -_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") - - -def _is_hidden_model(*values: str | None) -> bool: - """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or - EMBED_GGUF_REPO basename) or the llama.cpp install validation probe - (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). - None are usable chat models; the probe can be cached as a side effect of - installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected. A local-path embedder is matched by exact resolved - path only: a generic basename like "model" must not substring-hide - unrelated chat models.""" - from core.rag import config as rag_config - - needles = [ - # The validation probe's repo (matches the cached repo id) and its exact - # filename (matches the on-disk path). The filename carries the .gguf so - # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. - "ggml-org/models", - "stories260k.gguf", - ] - exact_paths: list[str] = [] - for model in ( - rag_config.effective_embedding_model(), - rag_config.effective_gguf_repo(), - ): - if _HF_REPO_ID_RE.match(model): - needles.append(model.split("/")[-1].lower()) - else: - resolved = _safe_resolve(Path(model).expanduser()) - if resolved: - exact_paths.append(resolved.lower()) - for v in values: - if not v: - continue - low = v.lower() - if any(n in low for n in needles): - return True - if exact_paths: - resolved = _safe_resolve(Path(v).expanduser()) - if resolved and resolved.lower() in exact_paths: - return True - return False - - -def _safe_resolve(path: Path) -> Optional[str]: - """resolve() to a string, or None when the path is inaccessible.""" - try: - return str(path.resolve()) - except OSError: - return None +# Shared with the hub inventory scans; keep the private aliases so existing +# importers (core.inference.local_model_resolver, tests) stay valid. +from utils.hidden_models import ( + _safe_resolve, + is_hidden_model as _is_hidden_model, +) backend_path = Path(__file__).parent.parent.parent @@ -853,7 +806,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: key = lambda item: (item.updated_at or 0), reverse = True, ) - return [m for m in models if not _is_hidden_model(m.id, m.path)] + return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)] @router.get("/local", response_model = LocalModelListResponse) @@ -2778,7 +2731,11 @@ async def get_gguf_variants( ], has_vision = response.has_vision, default_variant = response.default_variant, - context_length = _read_native_context_length(repo_id, is_local = local), + # The header walk reads tokenizer arrays on dense models (tens of + # ms per uncached file); keep it off the event loop. + context_length = await asyncio.to_thread( + _read_native_context_length, repo_id, is_local = local + ), ) except HTTPException: raise diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 1ddfc0eacb..ab0fd2fd99 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject from auth.storage import rotate_preview_link_secret +from core.rag.config import default_gguf_repo, effective_gguf_repo from loggers import get_logger from utils.utils import safe_error_detail, log_and_http_error from utils.personalization_settings import ( @@ -263,14 +264,18 @@ class EmbeddingModelPayload(BaseModel): class EmbeddingModelResponse(BaseModel): embedding_model: str + embedding_gguf_repo: str default_embedding_model: str + default_embedding_gguf_repo: str is_custom: bool def _embedding_model_response() -> EmbeddingModelResponse: return EmbeddingModelResponse( embedding_model = get_rag_embedding_model(), + embedding_gguf_repo = effective_gguf_repo(), default_embedding_model = default_embedding_model(), + default_embedding_gguf_repo = default_gguf_repo(), is_custom = get_stored_embedding_model() is not None, ) diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index fb361d3359..fd96fe2175 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -197,15 +197,18 @@ def can_load_chat_during_training( requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, required_override_gb: Optional[float] = None, + single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: """Decide if a NEW chat model can load without OOMing active training (inverse of can_keep_chat_during_training: training is already resident, so size the chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. `load_in_4bit` must be effective - (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any - CUDA case it can't size, so a load never OOMs training.""" + required_override_gb over the visible pool. ``single_device_gpu`` is the + exact physical device token selected by a single-device runner. + `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA + allows the load; default-deny on any CUDA case it can't size, so a load never + OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -251,26 +254,49 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. + if single_device_gpu is not None: + mode = "single_device" + elif is_gguf: + mode = "gguf" + else: + mode = "explicit" required_gb = required_override_gb if required_gb is None: required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs) if required_gb is None: - mode = "explicit" if requested_gpu_ids else "gguf" return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if requested_gpu_ids: + if single_device_gpu is not None: + token = str(single_device_gpu).strip() + if not token: + # Empty token = a CPU-only single-device runner (e.g. a CPU + # diffusion GGUF): it uses no GPU VRAM, so it never threatens + # active training and can always load. + return True, {"mode": "single_device", "reason": "cpu_only"} + try: + selected_gpu = int(token) + if selected_gpu < 0: + raise ValueError + except (TypeError, ValueError): + # A non-numeric device token (e.g. a CUDA UUID / MIG handle) + # can't be mapped to a free-VRAM index, but the runner still + # drives ONE device. Size against the worst-case visible device + # (min free), never the aggregate pool, so a single-device load + # is never OK'd on capacity it can't use and OOMs training. + free_vals = [min(free_by_index.values())] if free_by_index else [] + else: + free_vals = [free_by_index.get(selected_gpu, 0.0)] + elif requested_gpu_ids: # Invalid ids -> load_model 400s first, so don't block; missing id = 0. try: resolved = resolve_requested_gpu_ids(requested_gpu_ids) except ValueError: - return True, {"mode": "explicit", "reason": "invalid_gpu_ids"} + return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] - mode = "explicit" else: # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. free_vals = list(free_by_index.values()) - mode = "gguf" if not free_vals: return False, {"mode": mode, "reason": "no_visible_gpus"} diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index d4a7cae208..b3e6255d55 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -120,12 +120,151 @@ def test_is_hidden_model_hides_validation_probe_everywhere(): assert models_route._is_hidden_model( None, "/hf/models--ggml-org--models/snapshots/abc/tinyllamas/stories260K.gguf" ) + # A Windows-style snapshot path must match too, even on a POSIX interpreter + # (the filename check splits on both separators). + assert models_route._is_hidden_model( + r"C:\Users\u\.cache\huggingface\hub\models--ggml-org--models\snapshots\abc\tinyllamas\stories260K.gguf" + ) assert not models_route._is_hidden_model("unsloth/gemma-3-270m-it-GGUF") # The exact-filename needle must not hide a real repo that merely # references stories260K in its name. assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF") +def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch): + """A custom embedder with a generic basename is hidden by EXACT repo-id + match only, so unrelated cached repos that merely contain the basename stay + visible. Regression: substring basename matching hid real chat models like + ``user/model-chat`` from the On Device inventory.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + # The exact embedder repo and its GGUF companion are hidden. + assert models_route._is_hidden_model("org/model") + assert models_route._is_hidden_model("org/model-GGUF") + # Unrelated repos that merely contain "model" must NOT be hidden. + assert not models_route._is_hidden_model("user/model-chat") + assert not models_route._is_hidden_model("org/model-instruct") + assert not models_route._is_hidden_model("acme/remodelled-chat") + # The validation probe stays hidden regardless of embedder config. + assert models_route._is_hidden_model("ggml-org/models") + + +def test_is_hidden_model_matches_repo_derived_local_paths(monkeypatch): + """Match exact repo-derived cache and LM Studio paths.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + assert models_route._is_hidden_model( + "/cache/models--org--model/snapshots/abc/model.safetensors" + ) + assert models_route._is_hidden_model( + r"C:\Users\u\.cache\huggingface\hub\models--org--model-GGUF\snapshots\abc" + ) + assert models_route._is_hidden_model("/lm-studio/org/model-GGUF/model-Q8_0.gguf") + assert not models_route._is_hidden_model("/lm-studio/user/model-chat/model-Q8_0.gguf") + assert not models_route._is_hidden_model("/cache/models--org--model-instruct") + + +def test_is_hidden_model_prefers_existing_relative_path(monkeypatch, tmp_path): + """Prefer an existing relative path over repo-id syntax.""" + from core.rag import config as rag_config + + embedder = tmp_path / "models" / "embedder" + embedder.mkdir(parents = True) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + + assert models_route._is_hidden_model(str(embedder)) + + +def test_is_hidden_model_keeps_stale_default_embedder_hidden(monkeypatch): + """Keep default embedders hidden after a settings change.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5") + assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5-GGUF") + assert models_route._is_hidden_model("/models/bge-small-en-v1.5") + assert models_route._is_hidden_model("/models/bge-small-en-v1.5-F16.gguf") + assert models_route._is_hidden_model(r"C:\models\bge-small-en-v1.5-Q8_0.gguf") + # Repo IDs still use exact matching, and similar local basenames must have + # a real separator after the static default name. + assert not models_route._is_hidden_model("user/bge-small-en-v1.5-chat") + assert not models_route._is_hidden_model("/models/bge-small-en-v1.50") + + +def test_is_hidden_model_keeps_env_default_hidden_after_override(monkeypatch): + """A persisted override must not expose the deployment's env default.""" + from core.rag import config as rag_config + + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default") + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + assert models_route._is_hidden_model("org/env-default") + assert models_route._is_hidden_model("org/env-default-GGUF") + assert models_route._is_hidden_model("org/custom") + assert models_route._is_hidden_model("org/custom-GGUF") + assert not models_route._is_hidden_model("org/env-default-chat") + + +def test_hidden_models_importable_without_heavy_model_stack(): + """The hub cache scanner imports ``is_hidden_model`` at module scope, so it + must not drag in ``utils/models/__init__`` (the model-config + checkpoint + stack). Verify in a clean interpreter that importing the helper touches + neither ``utils.models`` nor those heavy submodules, and still classifies + the probe.""" + import os + import subprocess + import textwrap + + backend = Path(__file__).resolve().parents[1] + code = textwrap.dedent( + """ + import sys + + class _Blocker: + _blocked = ( + "utils.models", + "utils.models.model_config", + "utils.models.checkpoints", + ) + + def find_spec(self, name, path=None, target=None): + if name in self._blocked: + raise ImportError("blocked heavy import: " + name) + return None + + sys.meta_path.insert(0, _Blocker()) + from utils.hidden_models import is_hidden_model + + loaded = sorted(m for m in sys.modules if m.startswith("utils.models")) + assert not loaded, loaded + assert is_hidden_model("ggml-org/models") is True + assert is_hidden_model("unsloth/gemma-3-270m-it-GGUF") is False + print("HIDDEN_MODELS_IMPORT_OK") + """ + ) + env = dict(os.environ, PYTHONPATH = str(backend)) + proc = subprocess.run( + [sys.executable, "-c", code], + capture_output = True, + text = True, + env = env, + ) + assert proc.returncode == 0, proc.stderr + assert "HIDDEN_MODELS_IMPORT_OK" in proc.stdout + + def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path): """The ggml-org/models / stories260K install validation probe can land in the HF cache as a side effect of installing the prebuilt llama-server. diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 63dba8579c..7daa4224aa 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): devices, required_override = None, estimate = None, + single_device_gpu = None, + gpu_ids = None, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})), patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): ok, info = tv.can_load_chat_during_training( @@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = gpu_ids, is_gguf = True, required_override_gb = required_override, + single_device_gpu = single_device_gpu, ) return ok, info, auto_mock @@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0) self.assertTrue(ok) + def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self): + # gpu_ids narrows llama.cpp's candidate pool but does not turn its + # self-placement into HF device_map="balanced". The uneven selected + # pair therefore keeps the aggregate GGUF check without an even-share + # floor on the nearly-full card. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)), + required_override = 20.0, + gpu_ids = [0, 1], + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf") + + def test_single_device_uses_selected_gpu(self): + # The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an + # unrelated training-heavy GPU 1 has only 10 GB free. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "0", + ) + self.assertTrue(ok) + self.assertEqual(info["usable_gb"], 45.0) + + blocked, blocked_info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "1", + ) + self.assertFalse(blocked) + self.assertEqual(blocked_info["usable_gb"], 10.0) + + def test_single_device_unresolved_token_sizes_against_worst_device(self): + # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a + # free-VRAM index. The runner still drives ONE device, so size against the + # worst-case visible device (min free), not the aggregate pool: one GPU + # with 80 GB free vs a 20 GB model -> allow. + ok, info, _ = self._run( + devices = _devices((0, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "single_device") + self.assertNotIn("reason", info) + + def test_single_device_unresolved_token_refuses_when_worst_device_full(self): + # Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB + # model) -> refuse (default-deny), not on an unresolved-token technicality. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("reason"), "unresolved_gpu_id") + + def test_single_device_unresolved_token_uses_min_free_not_aggregate(self): + # The single-device runner uses ONE device but we can't tell which from a + # UUID token. Sizing against the aggregate pool would let a 20 GB model + # "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing + # training. Min-free (2 GB) is the safe worst case -> refuse. + ok, info, _ = self._run( + devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "single_device") + + def test_single_device_cpu_token_allows(self): + # An empty device token = a CPU-only single-device runner (CPU diffusion + # GGUF): it uses no GPU VRAM, so it never threatens training -> allow + # regardless of how full the GPUs are. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "", + ) + self.assertTrue(ok) + self.assertEqual(info["reason"], "cpu_only") + def test_estimate_unavailable_refuses(self): # No override and the estimator can't size it -> default-deny. ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None) @@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): captured = None, training_active, decision, + gpu_memory_mode = "auto", + requested_gpu_ids = None, ): config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None) with _stub_guard_deps( @@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = requested_gpu_ids, + gpu_memory_mode = gpu_memory_mode, ) def test_noop_when_training_inactive(self): @@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_allows_when_fits(self): self._guard(training_active = True, decision = (True, {"mode": "auto"})) + def test_diffusion_detection_uses_name_before_download(self): + config = SimpleNamespace( + identifier = "unsloth/DiffusionGemma-GGUF", + gguf_hf_repo = "unsloth/DiffusionGemma-GGUF", + gguf_file = None, + ) + self.assertTrue(self.route._classify_diffusion_gguf(config)) + + def test_uncached_gguf_classification_remains_unknown(self): + config = SimpleNamespace( + identifier = "owner/renamed-model", + gguf_hf_repo = "owner/renamed-model", + gguf_variant = "Q4_K_M", + gguf_file = None, + ) + self.assertIsNone(self.route._classify_diffusion_gguf(config)) + + def test_diffusion_detection_reuses_loader_metadata_probe(self): + import tempfile + + seen = [] + + class _Probe: + is_diffusion = False + _architecture = None + + def _read_gguf_metadata(self, path): + seen.append(path) + self.is_diffusion = True + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertTrue(self.route._classify_diffusion_gguf(config)) + self.assertEqual(seen, [str(model)]) + + def test_local_chat_gguf_classification_is_definitive(self): + import tempfile + class _Probe: + is_diffusion = False + _architecture = "llama" + + def _read_gguf_metadata(self, _path): + pass + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertFalse(self.route._classify_diffusion_gguf(config)) + + def test_manual_known_normal_gguf_bypasses_training_estimate(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (False, {"reason": "must not run"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(captured, []) + + def test_manual_unknown_gguf_keeps_single_device_training_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = None), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "2", + ), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "2") + + def test_manual_diffusion_uses_single_device_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "gguf"}), + gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) + + def test_unpinned_diffusion_uses_runner_default_gpu(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_effective_gpu_count", + return_value = 2, + ), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "3", + ) as gpu_arg, + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + gpu_arg.assert_called_once_with(None, cpu_only = False) + self.assertEqual(captured[0]["single_device_gpu"], "3") + def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} with self.assertRaises(HTTPException) as exc: @@ -467,36 +691,115 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): self.assertEqual(captured[0]["load_in_4bit"], False) self.assertEqual(captured[0]["max_seq_length"], 4096) - def test_rejects_gguf_with_gpu_ids_before_guard(self): - # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard. + def test_validate_forwards_manual_gpu_memory_mode_to_guard(self): from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0]) + request = ValidateModelRequest( + model_path = "unsloth/model-GGUF", + gguf_variant = "Q4_K_M", + gpu_memory_mode = "manual", + ) cfg = SimpleNamespace( - identifier = "x.gguf", - display_name = "x", + identifier = "unsloth/model-GGUF", + display_name = "model-GGUF", is_gguf = True, is_lora = False, is_vision = False, path = None, base_model = None, ) - captured = [] + captured = {} with ( patch.object( self.route, "_resolve_model_identifier_for_request", - return_value = ("x.gguf", "x.gguf", False), + return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False), ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), ): - with self.assertRaises(HTTPException) as exc: - asyncio.run(self.route.validate_model(request, current_subject = "u")) - self.assertEqual(exc.exception.status_code, 400) - self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail) - self.assertEqual(captured, []) # guard never reached + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("gpu_memory_mode"), "manual") + + def test_validate_forwards_inherited_extras_and_parallel_to_guard(self): + # Regression: /load resolves inherited same-model extras and passes the + # real slot count to the guard; validate must do the same, else it sizes + # a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and + # /load then 409s after the frontend has already unloaded. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + captured = {} + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"]) + self.assertIn("n_parallel", captured) + + def test_metadata_probe_skips_training_guard(self): + # A header-only probe (include_context_length) allocates no VRAM, so the + # training guard must not run -- else the staging GPU-layers / MoE sliders + # it feeds are hidden exactly when a during-training user needs them. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", + max_seq_length = 4096, + include_context_length = True, + ) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(guard_called, []) # ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ────── diff --git a/studio/backend/tests/test_cuda_torch_spec.py b/studio/backend/tests/test_cuda_torch_spec.py new file mode 100644 index 0000000000..928cef787e --- /dev/null +++ b/studio/backend/tests/test_cuda_torch_spec.py @@ -0,0 +1,73 @@ +# 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 _CUDA_TORCH_PKG_SPEC in install_python_stack.py. + +The CUDA repair path installs the torch trio from an exclusive --index-url (no +PyPI fallback), so these pinned ranges decide which torch the venv gets. The +upper bound is locked to the 2.11.x family to match the base image and rocm7.2 +spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from packaging.requirements import Requirement + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + + +def _load_module(monkeypatch): + """(Re-)import and return install_python_stack (mirrors test_torchao_select).""" + sys.modules.pop("install_python_stack", None) + monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent)) + import install_python_stack + + return install_python_stack + + +def _spec_of(pkg_spec: str): + """Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet.""" + return Requirement(pkg_spec).specifier + + +@pytest.mark.parametrize( + "index, allowed, rejected", + [ + # torch: 2.11.x allowed (matches base image); 2.12.x excluded. + (0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]), + # torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out. + (1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]), + # torchaudio: same 2.11.x window as torch. + (2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]), + ], +) +def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected): + mod = _load_module(monkeypatch) + spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index]) + for v in allowed: + assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}" + for v in rejected: + assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}" + + +def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch): + """CUDA and rocm7.2 target the same torch 2.11.x family, so their upper + bounds must stay in lockstep (bump both together at 2.12.x).""" + mod = _load_module(monkeypatch) + rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def _upper(pkg_spec: str) -> str: + for clause in _spec_of(pkg_spec): + if clause.operator == "<": + return clause.version + raise AssertionError(f"no upper bound in {pkg_spec!r}") + + for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True): + assert _upper(cuda_pkg) == _upper( + rocm_pkg + ), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}" diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 940b35d7ba..b3fa98b604 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -52,6 +52,16 @@ def client(monkeypatch): monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, + "effective_gguf_repo", + lambda: f"{saved.get('model', 'unsloth/default-embed')}-GGUF", + ) + monkeypatch.setattr( + settings, + "default_gguf_repo", + lambda: "unsloth/default-embed-GGUF", + ) app = FastAPI() app.include_router(settings.router) @@ -257,6 +267,13 @@ def test_clean_repo_saves_under_force(client, monkeypatch): r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) assert r.status_code == 200 assert saved.get("model") == "acme/clean-embed" + assert r.json() == { + "embedding_model": "acme/clean-embed", + "embedding_gguf_repo": "acme/clean-embed-GGUF", + "default_embedding_model": "unsloth/default-embed", + "default_embedding_gguf_repo": "unsloth/default-embed-GGUF", + "is_custom": True, + } def test_load_sink_refuses_flagged_model(monkeypatch): diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py index 3be4af0e32..bcf3ded71c 100644 --- a/studio/backend/tests/test_embedding_model_settings.py +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -53,3 +53,10 @@ def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeyp assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL assert ems.get_stored_embedding_model() is None + + +def test_env_default_derives_its_gguf_companion(monkeypatch): + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default-embedder") + + assert rag_config.default_gguf_repo() == "org/env-default-embedder-GGUF" diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index a5be07f8e3..ec0330ce05 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -15,6 +15,7 @@ from utils.models.gguf_metadata import ( pairing_score, read_gguf_context_length, read_gguf_general_metadata, + read_gguf_staged_dims, read_mmproj_audio_capability, ) @@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path): assert read_gguf_context_length(str(p)) is None +# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ---- + + +def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path): + assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf at all") + assert read_gguf_staged_dims(str(p)) is None + + +def test_staged_dims_moe_with_leading_dense(tmp_path: Path): + # GLM-4.7-Flash shape: context + total layers + MoE layers in one read. + p = _write_synthetic_gguf( + tmp_path / "glm.gguf", + {"general.architecture": "deepseek2"}, + extra_uint32 = { + "deepseek2.context_length": 202752, + "deepseek2.block_count": 47, + "deepseek2.expert_count": 64, + "deepseek2.leading_dense_block_count": 1, + }, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 202752, + "layer_count": 47, + "moe_layer_count": 46, + } + + +def test_staged_dims_dense_model(tmp_path: Path): + # Dense: layer_count present, moe_layer_count 0 (slider hidden). + p = _write_synthetic_gguf( + tmp_path / "dense.gguf", + {"general.architecture": "qwen3"}, + extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 40960, + "layer_count": 36, + "moe_layer_count": 0, + } + + +def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path): + # Experts present, no leading_dense key -> every block is a MoE layer. + p = _write_synthetic_gguf( + tmp_path / "moe.gguf", + {"general.architecture": "qwen35moe"}, + extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 40, + "moe_layer_count": 40, + } + + +def test_staged_dims_uint64_block_count(tmp_path: Path): + # block_count stored as uint64 (vtype 10) still parses; moe == block_count. + p = _write_synthetic_gguf( + tmp_path / "moe64.gguf", + {"general.architecture": "gpt-oss"}, + extra_uint32 = {"gpt-oss.expert_count": 32}, + extra_uint64 = {"gpt-oss.block_count": 24}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 24, + "moe_layer_count": 24, + } + + def test_context_length_read_from_uint64(tmp_path: Path): # Some models store context_length as a uint64 (vtype 10). p = _write_synthetic_gguf( diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py new file mode 100644 index 0000000000..b17274197f --- /dev/null +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GPU Memory mode dropdown. + +The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the +chat UI through the load request. "manual" lets the user own the offload: with +``gpu_layers < 0`` (Auto, the default) it hands all memory management to +llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no +gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers +and MoE offload itself (``--fit off``). These tests pin: + + * the pydantic request/response/status contract (snake_case key, default + "auto", unknown values rejected), + * the backend ``gpu_memory_mode`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that the manual + Auto-layers branch in ``load_model`` empties the probed + GPU set and drops tensor parallelism so the selection below no-ops, while + the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``. +""" + +from __future__ import annotations + +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx is a real, installed backend dependency: import it so the genuine module +# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since +# setdefault installs it before real httpx loads, would poison a combined pytest +# run -- routes/inference references httpx.Response (and other attrs) at def time. +import httpx # noqa: F401 + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default "auto") ─────────────── + + +def test_load_request_defaults_gpu_memory_mode_auto(): + assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto" + + +def test_load_request_round_trips_json_key(): + req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"}) + assert req.gpu_memory_mode == "manual" + assert req.model_dump()["gpu_memory_mode"] == "manual" + + +def test_load_request_rejects_unknown_mode(): + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus") + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_memory_mode(model_cls): + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + manual = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + ) + else: + default = model_cls() + manual = model_cls(gpu_memory_mode = "manual") + assert default.model_dump()["gpu_memory_mode"] == "auto" + assert manual.model_dump()["gpu_memory_mode"] == "manual" + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_gpu_memory_mode_property_defaults_auto(): + assert LlamaCppBackend().gpu_memory_mode == "auto" + + +def test_gpu_memory_mode_property_reflects_field(): + backend = LlamaCppBackend() + backend._gpu_memory_mode = "manual" + assert backend.gpu_memory_mode == "manual" + + +def test_unload_resets_gpu_memory_mode(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._gpu_memory_mode = "manual" + backend.unload_model() + assert backend.gpu_memory_mode == "auto" + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._gpu_memory_mode = gpu_memory_mode + return backend + + +def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = gpu_memory_mode, + ) + + +@pytest.mark.parametrize("mode", ["auto", "manual"]) +def test_already_in_target_state_matches_same_mode(mode): + assert _target_state(_loaded_backend(mode), mode) is True + + +@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")]) +def test_already_in_target_state_reloads_on_mode_change(loaded, requested): + # Flipping the dropdown either direction must force a reload so the command + # is rebuilt with/without the Unsloth GPU masking. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_ignores_mode_for_diffusion(): + # The diffusion runner is mode-agnostic (always "auto"), so a standing manual + # preference must not force a needless reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + assert _target_state(backend, "manual") is True + + +# ── load_model: manual + Auto layers bypasses Unsloth GPU management ── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel(): + # Emptying the probed set makes the selection / TP planning below no-op, so + # gpu_indices stays None and use_fit True (--fit on). + src = _load_model_source() + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:') + assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)" + block = src[gate : gate + 1400] + assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set" + # --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped. + assert "strip_split_mode_only(extra_args)" in block + assert "requested_ctx if requested_ctx > 0 else 0" in block + # The branch sits before GPU selection assigns gpu_indices; --fit on is its emission. + assert gate < src.find("gpu_indices, use_fit = None, True") + assert 'cmd.extend(["--fit", "on"])' in src + # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so + # a requested quantized cache survives into the --fit load. + tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:') + assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel" + assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400] + cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache") + assert cache_drop != -1 + assert ( + tp_drop < cache_drop + ), "TP must drop before the cache-drop so a quantized KV survives --fit" + + +def test_auto_layers_never_sends_ctx_size_zero(): + # Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning + # the full native context and disabling --fit's reduction. So the base cmd + # must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit) + # case, and a positive context is passed through (which --fit optimizes + # layers around). + src = _load_model_source() + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] + assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list" + assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c" + assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src + zero = src.find('cmd.extend(["-c", "0"])') + assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case' + guard = src.rfind("elif not auto_fit:", 0, zero) + assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard' + + +def test_manual_mode_clears_inherited_main_model_placement_env(): + env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS} + env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7" + env["UNRELATED"] = "kept" + + LlamaCppBackend._clear_manual_placement_env(env) + + assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS)) + assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7" + assert env["UNRELATED"] == "kept" + + +def test_load_model_sanitizes_manual_env_after_building_child_env(): + src = _load_model_source() + env_build = src.find("env = self._llama_server_env_for_binary(binary)") + env_clear = src.find("self._clear_manual_placement_env(env)", env_build) + launch = src.find("subprocess.Popen", env_build) + assert env_build != -1 + assert env_build < env_clear < launch + + +# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ─────────── + + +def test_load_request_accepts_manual(): + req = LoadRequest( + model_path = "owner/repo", + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + ) + assert req.gpu_memory_mode == "manual" + assert req.gpu_layers == 20 + assert req.n_cpu_moe == 8 + assert req.tensor_split == [2, 1] + + +def test_load_request_manual_defaults(): + req = LoadRequest(model_path = "owner/repo") + assert req.gpu_layers == -1 + assert req.n_cpu_moe == 0 + assert req.tensor_split is None + + +@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]]) +def test_load_request_rejects_degenerate_tensor_split(bad): + # A negative/non-finite/all-zero split is dropped at launch but compared raw + # in the reload dedupe, so it would reload forever -- reject it up front. + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", tensor_split = bad) + + +@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None]) +def test_load_request_accepts_valid_tensor_split(good): + assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good + + +def test_route_normalizes_explicit_extras_before_reload_dedupe(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + load_impl = route_src[route_src.index("async def _load_model_impl") :] + strip = load_impl.index("_stripped_explicit = strip_shadowing_flags") + normalize = load_impl.index( + 'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})' + ) + dedupe = load_impl.index("and _request_matches_loaded_settings(") + assert strip < normalize < dedupe + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_manual_fields(model_cls): + if model_cls is LoadResponse: + obj = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + else: + obj = model_cls( + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + dumped = obj.model_dump() + assert dumped["gpu_memory_mode"] == "manual" + assert dumped["gpu_layers"] == 20 + assert dumped["n_cpu_moe"] == 8 + assert dumped["tensor_split"] == [2, 1] + assert dumped["n_layers"] == 32 + assert dumped["n_moe_layers"] == 32 + + +def test_manual_properties_default_and_reflect_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + backend._gpu_layers = 20 + backend._n_cpu_moe = 8 + backend._tensor_split = [2, 1] + assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8 + assert backend.tensor_split == [2, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + + +def test_n_moe_layers_property(): + # 0 for a dense model (hides the slider); block_count for all-MoE; + # block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46). + b = LlamaCppBackend() + b._n_layers = 36 + b._n_experts = None + assert b.n_moe_layers == 0 + b._n_experts = 128 + b._leading_dense_block_count = None + assert b.n_moe_layers == 36 + b._n_layers = 47 + b._leading_dense_block_count = 1 + assert b.n_moe_layers == 46 + + +def _target_state_manual( + backend, + *, + gpu_layers, + n_cpu_moe, + tensor_split = None, +): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = "manual", + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + ) + + +def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change(): + backend = _loaded_backend("manual") + backend._gpu_layers = 20 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same knobs -> no reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True + # Changed layer count -> reload. + assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False + # Changed MoE offload -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False + # Added a GPU split -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False + # Same GPU split -> no reload. + backend._tensor_split = [2, 1] + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True + + +def test_auto_layers_reload_tracks_only_gpu_layers(): + # Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover + # request value must not reload -- only a gpu_layers change (Auto -> pinned) does. + backend = _loaded_backend("manual") + backend._gpu_layers = -1 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same Auto, leftover MoE/split in the request -> still no reload. + assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True + # Auto -> explicit offload reloads. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False + + +def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe(): + src = _load_model_source() + gate = src.find('elif gpu_memory_mode == "manual":') + assert gate != -1, "load_model must have an explicit-offload manual branch" + block = src[gate : gate + 700] + # Empties the probed set (skips the planner) but keeps the user's TP choice + # (only the Auto-layers branch above drops TP). + assert "gpus = []" in block + assert "tensor_parallel = False" not in block + # The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0. + assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src + assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src + # MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below). + assert "_resolve_cpu_moe_flag(" in src + assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src + # A count requested on a dense model is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a count + # llama-server never received (same rule as the tensor-split drop below). + moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])') + assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300] + assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300] + # The offload path forces use_fit False so --fit-ctx is never added under --fit off. + emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])') + assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80] + + +def test_status_reports_requested_context_length(): + # The hydration path re-seeds a Manual+Auto context pin from the REQUESTED + # n_ctx (0 = Auto); context_length only exposes the resolved value. + assert "requested_context_length" in InferenceStatusResponse.model_fields + s = InferenceStatusResponse(requested_context_length = 8192) + assert s.model_dump()["requested_context_length"] == 8192 + assert InferenceStatusResponse().model_dump()["requested_context_length"] is None + # The /status route must actually wire it from the backend (a declared-but- + # never-populated field would leave hydration silently reverting the pin). + from pathlib import Path as _P + + route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert "requested_context_length = llama_backend.requested_n_ctx" in route_src + + +def test_manual_offload_emits_tensor_split(): + # The offload path emits --tensor-split from the per-GPU shares, only when + # provided, with >1 GPU in use, AND matching that count (a stale ratio on a + # narrowed picker or a mismatched direct-API list must not emit -- llama- + # server aborts on a split/GPU-count mismatch). + src = _load_model_source() + assert "if tensor_split and _split_gpus > 1:" in src + # Emit only on a length match AND a positive sanitized total: a mismatched + # or all-zero split aborts llama-server / assigns nothing, so it's dropped. + # The emitted list is the sanitized one (clamping tested behaviorally below). + assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src + assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src + assert '"--tensor-split"' in src + # Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch. + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:') + nxt = src.find("elif use_fit:", gate) + assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt] + # A split with a single effective GPU is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a ratio + # llama-server never received and the dedupe baseline preserves it. + assert "elif tensor_split:" in src[gate:nxt] + drop = src.find("elif tensor_split:", gate, nxt) + assert "self._tensor_split = None" in src[drop : drop + 250] + + +def test_sanitize_tensor_split_clamps_negative_and_non_finite(): + # Negative entries would launch a placement different from the ratio the + # UI showed; inf passes a plain > 0 total gate and would emit + # "--tensor-split inf,..." (llama.cpp normalizes shares by the running + # total, so an inf poisons the shares from that entry on). Both clamp to 0. + sanitize = LlamaCppBackend._sanitize_tensor_split + assert sanitize([2, 1]) == [2.0, 1.0] + assert sanitize([-1, 2]) == [0.0, 2.0] + assert sanitize([float("inf"), 1]) == [0.0, 1.0] + assert sanitize([float("nan"), 1]) == [0.0, 1.0] + # All-zero survives sanitization; the call site's total gate drops it. + assert sanitize([0, 0]) == [0.0, 0.0] + # Unreadable input -> []; the call site's length gate drops it. + assert sanitize(["x", 1]) == [] + assert sanitize([10**400, 1]) == [] + + +def test_zero_offload_mask_honors_device_pin_spellings(): + # A user device pin must keep the GPUs visible: llama-server aborts on a + # pin it can't see ('error: invalid device'). The pin can arrive as + # --device or its -dev alias, as the draft forms (parsed even with no + # drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var. + load_src = _load_model_source() + assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src + block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin) + for flag in ( + '"--device"', + '"-dev"', + '"--spec-draft-device"', + '"-devd"', + '"--device-draft"', + ): + assert flag in block + assert '"LLAMA_ARG_DEVICE"' in block + + +def test_resolve_cpu_moe_flag(): + # Clamp the requested MoE-layer count to the model's MoE layers, then offset + # past leading dense layers (--n-cpu-moe counts from layer 0). + R = LlamaCppBackend._resolve_cpu_moe_flag + assert R(0, 40, 0) is None # nothing requested + assert R(8, 0, 0) is None # dense model (no MoE layers) + assert R(8, 40, 0) == 8 # all-MoE: direct + assert R(100, 40, 0) == 40 # clamp to the MoE layer count + # GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46. + assert R(5, 46, 1) == 6 # offset past the 1 dense layer + assert R(46, 46, 1) == 47 # all MoE on CPU == block_count + + +def test_manual_allows_tensor_parallel_via_split_mode(): + # Manual offload keeps the user's TP choice but skips the memory-based planner + # (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The + # --split-mode tensor emission gates on tensor_parallel alone, so manual + # reaches it -- with tp_tensor_split None it's an even split (no + # --tensor-split). --fit off means no fit/tensor abort. + src = _load_model_source() + assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src + assert "if plan_tp:" in src + assert "if plan_tp and len(tp_gpus) < 2:" in src + sm = src.find('cmd.extend(["--split-mode", "tensor"])') + assert sm != -1, "TP must emit --split-mode tensor" + guard = src.rfind("if tensor_parallel:", 0, sm) + assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel" + # The tensor-split is only emitted for a planned (non-even) split, which + # manual never produces, so manual stays an even split. + assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src + + +def test_fit_sets_target_margin(): + # Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB. + caps = {"supports_fit_target": True} + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps) + assert flags[flags.index("--fit-target") + 1] == "512" + # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins + # native there, so the tighter margin must not ride along. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps) + # Not emitted when fit is off. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps) + # Not emitted when the binary lacks support. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, True, 0, 0, {"supports_fit_target": False} + ) + + +# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ───────────────────── + + +def test_load_request_accepts_gpu_ids(): + req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0]) + assert req.gpu_ids == [1, 0] + assert LoadRequest(model_path = "owner/repo").gpu_ids is None + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_ids(model_cls): + if model_cls is LoadResponse: + obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + else: + obj = model_cls(gpu_ids = [1]) + assert obj.model_dump()["gpu_ids"] == [1] + + +def test_gpu_ids_property_default_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_ids is None + backend._gpu_ids = [0, 1] + assert backend.gpu_ids == [0, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_ids is None + + +def _target_state_gpu_ids(backend, gpu_ids): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_ids = gpu_ids, + ) + + +def test_gpu_ids_reload_detection_is_order_insensitive(): + backend = _loaded_backend("auto") + backend._gpu_ids = [0, 1] + # Same set, different order -> no reload. + assert _target_state_gpu_ids(backend, [1, 0]) is True + # Different set -> reload. + assert _target_state_gpu_ids(backend, [0]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): + # The diffusion runner drives only its single lowest device, so the backend + # records [lowest]. A later multi-GPU request that still resolves to that + # same lowest device must dedupe (no needless reload); a request whose lowest + # device moves, or that drops the pick, must reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick + assert _target_state_gpu_ids(backend, [3, 1]) is True + assert _target_state_gpu_ids(backend, [1]) is True + # Lowest device changes (2, not 1) -> reload. + assert _target_state_gpu_ids(backend, [3, 2]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_start_diffusion_server_resets_tensor_parallel(): + # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model + # phase 1 only kills the process, it skips the unload reset). Diffusion is never + # TP, so startup must clear it -- else /status misreports TP and an identical + # diffusion re-Apply reloads against stale tensor-parallel state. + src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) + assert "self._tensor_parallel = False" in src + + +def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): + # The route-level reload dedupe mirrors the backend: for a loaded diffusion + # model it compares the request against the single recorded device, not the + # full requested list, or a same-device multi-GPU pick reloads needlessly. + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] + guard = match_impl.index("if llama_backend.is_diffusion:") + collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") + compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") + assert guard < collapse < compare + + +# ── Manual tensor split: child enumeration pinned to the picker's order ────── + + +def _patch_split_pin_env(monkeypatch, *, inherited, reported): + """Point the pin helper at a fake inherited mask and picker report. + ``reported`` None = enumeration unavailable (falls back to ascending).""" + import utils.hardware as hw + + monkeypatch.setattr( + LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited) + ) + info = ( + {"available": False} + if reported is None + else { + "available": True, + "index_kind": "physical", + "devices": [{"index": i} for i in reported], + } + ) + monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info) + + +def test_split_pin_reorders_inherited_numeric_mask(monkeypatch): + # Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but + # nvidia-smi reported the picker's list ascending -- the mask must be + # re-emitted in that order or the per-GPU shares land on the wrong cards. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch): + # Torch-fallback enumeration (no nvidia-smi) reports devices in inherited + # mask order, so the picker's split list follows the mask -- the pin must + # keep that order, not re-sort it into a mismatch. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "3,1" + + +def test_split_pin_falls_back_to_ascending_without_report(monkeypatch): + # Enumeration unavailable: ascending physical is the best guess (it matches + # the dominant nvidia-smi report order). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_without_mask_only_sets_pci_order(monkeypatch): + # No inherited mask (or a UUID/MIG one resolving to None): enumeration order + # is fully fixed by CUDA_DEVICE_ORDER, so no mask is written. + _patch_split_pin_env(monkeypatch, inherited = None, reported = None) + env = {} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"} + + +def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch): + # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR + # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP + # would index into the already-reduced set). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + assert env["HIP_VISIBLE_DEVICES"] == "1,3" + assert "ROCR_VISIBLE_DEVICES" not in env + + +# ── Diffusion single-device selection ─────────────────────────────────────── + + +def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + monkeypatch.setenv("DG_GPU", "7") + assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1" + + +def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch): + monkeypatch.delenv("DG_GPU", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "3" + + +def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch): + monkeypatch.setenv("DG_GPU", "GPU-abc") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc" + assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == "" + + +# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ───────── + + +def test_zero_offload_flag_false_without_companions(): + # CPU-only by construction: False lets training skip unloading a server that + # holds no VRAM. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False + + +@pytest.mark.parametrize( + "companion", + ["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"], +) +def test_zero_offload_flag_true_with_companion(companion): + # mmproj / a drafter offload to GPU regardless of --gpu-layers, so the + # server still holds VRAM and training must unload it. Drafter detection + # reuses the extras parser, so pass-through aliases count too. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_inline_companion_forms(): + cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_drafter(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + "device_args", + [ + ["--device", "CUDA0"], + ["--device=CUDA0"], + ["-dev", "CUDA0"], + ["--spec-draft-device", "CUDA0"], + ["--device-draft=CUDA0"], + ], +) +def test_zero_offload_flag_true_with_device_pin(device_args): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_device_pin(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_DEVICE": "CUDA0"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + ("device_args", "env"), + [ + (["--device", "cpu"], {}), + (["--device=none"], {}), + (["--spec-draft-device", "cpu"], {}), + ([], {"LLAMA_ARG_DEVICE": "none"}), + (["--device", "CUDA0", "--device", "cpu"], {}), + ], +) +def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False + + +def test_zero_offload_flag_true_with_surviving_tensor_mode(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True)) + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_none_without_gpus(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None + + +def test_cmd_has_gpu_companion_detection(): + # The env mask for CPU-only zero-offload loads keys off this scan: any + # --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible. + has = LlamaCppBackend._cmd_has_gpu_companion + assert has(["llama-server", "-m", "m.gguf"], {}) is False + assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True + assert has(["llama-server", "--mmproj=p.gguf"], {}) is True + assert has(["llama-server", "-md", "d.gguf"], {}) is True + assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True + + +def test_cmd_companion_ignores_cpu_forced_drafter(): + # A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs + # and training may leave the server alone. + has = LlamaCppBackend._cmd_has_gpu_companion + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"] + assert has(cmd, {}) is False + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"] + assert has(cmd, {}) is False + # mmproj still counts even alongside a CPU drafter. + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"] + assert has(cmd, {}) is True diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 69ad560788..d4f2fbe993 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase): self.assertIn("only supported on CUDA devices", str(exc_info.exception)) - def test_inference_route_rejects_gpu_ids_for_gguf(self): + def test_inference_route_validates_gpu_ids_for_gguf(self): + # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still + # validated: a rejected pick surfaces as a clean 400, not the old + # "not supported for GGUF" rejection. Patch the validator so the test + # is deterministic regardless of the host's (or a prior test's) GPU env. + import utils.hardware.hardware as hardware_mod + inference_route = _load_route_module( "inference_route_module_for_gguf_gpu_ids_test", "routes/inference.py", @@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase): ) ) + # The validator's ValueError becomes a clean 400 (not the removed + # "not supported for GGUF" rejection). self.assertEqual(exc_info.exception.status_code, 400) - self.assertIn("GGUF", exc_info.exception.detail) + self.assertIn("gpu_ids", exc_info.exception.detail.lower()) + self.assertNotIn("not supported", exc_info.exception.detail.lower()) def test_training_route_returns_400_for_invalid_gpu_ids(self): training_route = _load_route_module( diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py index f320d29a02..662c918305 100644 --- a/studio/backend/tests/test_llama_cpp_no_context_shift.py +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list(): "conditional branch -- otherwise some code paths would still " "run with silent context shift enabled." ) - # Pin that it sits next to -c / --ctx so the grouping makes sense. - assert '"-c"' in block assert '"--flash-attn"' in block + # -c is emitted in the conditional right after the base list, not inside + # it: auto-fit (--fit on with no pinned context) must omit -c entirely, + # because "-c 0" pins the full native context and disables --fit's + # VRAM-based sizing. Pin that it still sits next to the base block so the + # context grouping stays intact. + after = rest[end_rel : end_rel + 1000] + assert '"-c"' in after, ( + "-c must still be emitted in the conditional immediately after the " + "base cmd list (omitted only in auto-fit, where --fit sizes context)." + ) def _iter_lines_with_offset(text: str): diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index 488645ee5a..fe1e67edad 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -225,31 +225,46 @@ def test_kv_unified_added_for_multi_slot(): """Explicit --parallel N disables llama-server's auto-slots kv-unified default, splitting -c into per-slot windows of -c/N; Unsloth must restore the shared pool so one request can use the full advertised context.""" - flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL) assert "--kv-unified" in flags def test_kv_unified_skipped_for_single_slot_or_old_build(): assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 4, False, 98304, 98304, _CAPS_NONE + 4, False, False, 98304, 98304, _CAPS_NONE ) def test_fit_ctx_floors_explicit_request_under_fit(): - flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + # An explicit requested ctx floors --fit-ctx at that value on any --fit + # path, including legacy auto (auto_fit False). + flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL) assert flags[flags.index("--fit-ctx") + 1] == "98304" -def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): +def test_fit_ctx_skipped_without_fit_or_support(): + # No --fit on -> no --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) - assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + # --fit on but the binary doesn't support --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, True, 98304, 98304, _CAPS_NONE + 1, True, True, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit(): + # Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't + # shrink it to a tiny size. + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "8192" + # Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the + # 8192 floor must NOT ride along and override that pin. + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 0, 262144, _CAPS_ALL ) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index ba52afad1c..c6d16363f8 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -747,6 +747,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too(): assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] +def test_strip_offload_is_opt_in_and_covers_moe(): + base = dict( + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) + # Default: offload (incl. MoE) flags are NOT stripped. + assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [ + "--n-cpu-moe", + "8", + "--top-k", + "20", + ] + # Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest. + assert strip_shadowing_flags( + ["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"], + **base, + strip_offload = True, + ) == ["--top-k", "20"] + # Boolean --cpu-moe drops the flag only, not the following value. + assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [ + "--seed", + "-1", + ] + + @pytest.mark.parametrize( "args", [ @@ -796,6 +824,23 @@ def test_strip_split_mode_only_drops_tensor_split_too(): assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] +def test_strip_tensor_split_alone_preserves_split_mode(): + # Manual mode emits its own --tensor-split, so an inherited ratio is dropped + # -- but the user's --split-mode row/none/layer choice (which the manual + # ratio toggle can't express) must survive. strip_tensor_split removes only + # the ratio, unlike strip_split_mode which removes the whole group. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = True, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 90fbd19297..c4c0ce15c9 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -697,6 +697,10 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): normal.write_bytes(b"x" * 32) probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe probe.write_bytes(b"x" * 32) + embedder = tmp_path / "embedding-Q8_0.gguf" + embedder.write_bytes(b"x" * 32) + local_default_embedder = tmp_path / "bge-small-en-v1.5-F16.gguf" + local_default_embedder.write_bytes(b"x" * 32) def _info(mid, path): return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) @@ -704,7 +708,22 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): monkeypatch.setattr( models_route, "_scan_models_dir", - lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + lambda *a, **k: [ + _info("org/Normal-GGUF", normal), + _info("ggml-org/models", probe), + SimpleNamespace( + id = str(embedder), + path = str(embedder), + model_id = "unsloth/bge-small-en-v1.5-GGUF", + display_name = "embedding-Q8_0", + ), + SimpleNamespace( + id = str(local_default_embedder), + path = str(local_default_embedder), + model_id = None, + display_name = local_default_embedder.name, + ), + ], ) monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) @@ -713,6 +732,8 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): index = resolver._index() assert "org/normal-gguf" in index # keys are normalized to lowercase assert "ggml-org/models" not in index + assert "unsloth/bge-small-en-v1.5-gguf" not in index + assert str(local_default_embedder).lower() not in index # And the hidden probe cannot be auto-switched to by name. resolver._scan = (0.0, {}) assert resolver.resolve_local_gguf("ggml-org/models") is None @@ -1729,6 +1750,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): # host path in /v1/models, yet the model stays resolvable by that path too. from types import SimpleNamespace import routes.models as models_route + from storage import studio_db + import utils.paths as paths gguf = tmp_path / "model-Q4_K_M.gguf" gguf.write_bytes(b"x" * 32) @@ -1742,6 +1765,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: []) + monkeypatch.setattr(studio_db, "list_scan_folders", lambda: []) resolver._scan = (0.0, {}) # The advertised id is the alias, never the absolute path. diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 06f72d3b9f..00c7aeac69 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): src = _load_model_source() assert '"--tensor-split"' in src gate = src.find("if tensor_parallel:") - ts = src.find('"--tensor-split"') + # Find the TP block's emission (after the gate); manual mode emits its own + # --tensor-split earlier in the source from the user's per-GPU shares. + ts = src.find('"--tensor-split"', gate) nxt_else = src.find("self._tensor_parallel = False") assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + assert "tp_tensor_split" in src[gate:nxt_else] def test_mtp_decode_probe_wired_under_tensor_parallel(): diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index fb0989b306..d1372ca415 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -126,10 +126,21 @@ _ALLOWED_TP_DROP_GUARDS = { # Capability: --split-mode tensor aborted for this (binary, model) (#6415). # Self-healing -- tried by default, skipped only after a real abort (vs #6416). "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", - # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. - "tensor_parallel and len(tp_gpus) < 2", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated + # on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414). + "plan_tp and len(tp_gpus) < 2", # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. "_tp_weight_budget_mib <= _tp_required_mib", + # Manual mode, Auto layers: --fit owns memory and is incompatible with a + # tensor split, so TP is dropped (surfaced via logger.info) before the + # cache-drop, so a quantized KV survives into the --fit load (#6414). + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)", + # Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use. + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers >= 0) and (self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2)", + # Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode + # launch under the CPU-only GPU mask (no visible devices) aborts the server + # instead of the intended CPU-only load (#6414). + "gpu_memory_mode == 'manual' and gpu_layers == 0", } @@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): full GPU set too, so it is symmetric with the budget/geometry downgrades and doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" src = inspect.getsource(LlamaCppBackend.load_model) - gate = src.find("tensor_parallel and len(tp_gpus) < 2") + gate = src.find("plan_tp and len(tp_gpus) < 2") assert gate != -1 # Bound to exactly this block: from its gate to the next (budget) downgrade. nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py new file mode 100644 index 0000000000..4263b012eb --- /dev/null +++ b/studio/backend/tests/test_training_config_popover_source.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Source-level regression guards for the Training Config popover data source +(#6853). + +The live Training Progress popover used to read the editable form store +(useTrainingConfigStore) while a run was active, so it showed stale/static +values whenever the user touched the form after starting the run; only the +History view read the run's saved config snapshot. These guards pin the fixed +wiring: both views feed ProgressSection a config override mapped from +GET /api/train/runs/{id}, and ProgressSection prefers that override whenever +one is present -- not only for historical views. +""" + +from __future__ import annotations + +from pathlib import Path + +_STUDIO_FRONTEND = Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio" + + +def _read(rel: str) -> str: + return (_STUDIO_FRONTEND / rel).read_text(encoding = "utf-8") + + +def test_progress_section_prefers_override_over_form_store(): + src = _read("sections/progress-section.tsx") + # Fields key on the override's presence, not isHistorical: a live view passing + # an override wins over the store; without one, live keeps the store while + # History shows blanks rather than unrelated live form values. + assert "const cfg = configOverride ?? (isHistorical ? undefined : config)" in src + assert "const cfgEpochs = cfg?.epochs" in src + assert "isHistorical ? configOverride?.epochs" not in src + + +def test_live_view_fetches_the_active_run_config(): + src = _read("live-training-view.tsx") + # Live view resolves the run's saved config snapshot by job id... + assert "getTrainingRun(" in src + assert "mapRunConfigToOverride(" in src + # ...and hands it to the popover. + assert "configOverride={runConfigOverride}" in src + + +def test_live_view_fetches_as_soon_as_the_job_id_exists(): + # start_training() inserts the run row BEFORE the pump consumes any event, so + # the saved config is available during configuring/loading/downloading. The + # job id is therefore the whole readiness condition: gating on a first step + # or a terminal phase would show the wrong config for the entire pre-step + # window of a long load, or for a run adopted from another client. + src = _read("live-training-view.tsx") + assert "if (!runtime.jobId) {" in src + assert "[runtime.jobId, fetchedRunConfig, fetchAttempt]" in src + # No step/phase readiness gate may creep back in. + assert "runRowReady" not in src + + +def test_live_view_retries_the_transient_row_miss(): + # start_training() creates the row before the pump, but a lookup racing that + # commit can still 404. Nothing else in the effect deps changes on failure, so + # the retry must be explicit and bounded, else a genuinely absent row would + # poll forever instead of falling back to the form store. + src = _read("live-training-view.tsx") + assert "RUN_CONFIG_FETCH_RETRIES" in src + assert "RUN_CONFIG_FETCH_RETRY_MS" in src + assert "setFetchAttempt(" in src + assert "attempts >= RUN_CONFIG_FETCH_RETRIES" in src + # The budget is keyed by job so a new run always starts fresh. + assert "fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0" in src + # The pending retry must be cancelled with the effect. + assert "clearTimeout(retryTimer)" in src + + +def test_live_view_prefers_saved_training_method(): + # The method label / LoRA-row visibility must come from the run snapshot, + # not the editable form (which may have changed since the run started). + src = _read("live-training-view.tsx") + assert "runConfigOverride?.trainingMethod ?? config.trainingMethod" in src + + +def test_history_view_uses_the_shared_mapper(): + src = _read("historical-training-view.tsx") + # Shared mapper, not a re-inlined field-by-field copy that could drift. + assert "mapRunConfigToOverride(detail.config)" in src + assert "num_epochs" not in src + + +def test_shared_mapper_matches_backend_config_keys(): + src = _read("sections/run-config-override.ts") + # The mapper reads the run config JSON the backend snapshots at job start; + # keep the key set pinned so a silent rename breaks loudly here. + for key in ( + "training_type", + "load_in_4bit", + "num_epochs", + "batch_size", + "learning_rate", + "max_steps", + "max_seq_length", + "warmup_steps", + "optim", + "lora_r", + "lora_alpha", + "lora_dropout", + "use_rslora", + "use_loftq", + ): + assert key in src, f"run-config mapper lost backend key {key}" diff --git a/studio/backend/utils/hidden_models.py b/studio/backend/utils/hidden_models.py new file mode 100644 index 0000000000..20d0bb966e --- /dev/null +++ b/studio/backend/utils/hidden_models.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Infra-only model detection shared by the model routes and the hub +inventory. Lives directly under ``utils`` (not ``utils.models``) so the hub +cache scanner can import it without pulling in ``utils/models/__init__.py``, +which eagerly loads the model-config/checkpoint stack, and without importing +``routes.models`` (import-time side effects, would cycle).""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Optional + +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + +# The llama.cpp install-validation probe repo. Always hidden. +_PROBE_REPO_ID = "ggml-org/models" +# The probe's on-disk filename. Carries the ".gguf" so it stays specific and +# does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. +_PROBE_FILENAME = "stories260k.gguf" +# Keep previously cached defaults hidden after settings changes. +_DEFAULT_EMBEDDING_REPO_IDS = { + "unsloth/bge-small-en-v1.5", + "unsloth/bge-small-en-v1.5-GGUF", +} +# Local copies do not always retain the repo id. Keep a narrow basename +# fallback for Studio's static default embedder only; configured custom repos +# remain exact-match-only. +_DEFAULT_EMBEDDING_PATH_BASENAMES = {"bge-small-en-v1.5"} + + +def _safe_resolve(path: Path) -> Optional[str]: + """resolve() to a string, or None when the path is inaccessible.""" + try: + return str(path.resolve()) + except OSError: + return None + + +def _existing_resolved_path(value: str) -> Optional[str]: + """Resolve an existing local path.""" + path = Path(value).expanduser() + try: + if not path.exists(): + return None + except OSError: + return None + return _safe_resolve(path) + + +def _path_contains_repo_id(value: str, repo_ids: set[str]) -> bool: + """Match exact repo-derived path segments.""" + parts = [part for part in value.lower().replace("\\", "/").split("/") if part] + for repo_id in repo_ids: + owner, name = repo_id.split("/", 1) + if f"models--{owner}--{name}" in parts: + return True + if any( + parts[index] == owner and parts[index + 1] == name for index in range(len(parts) - 1) + ): + return True + return False + + +def _path_basename_is_default_embedder(value: str) -> bool: + """Match a default embedder folder or a suffixed local weight filename.""" + normalized = value.lower().replace("\\", "/").rstrip("/") + basename = normalized.rsplit("/", 1)[-1] + return any( + basename == needle + or any(basename.startswith(f"{needle}{separator}") for separator in ("-", "_", ".")) + for needle in _DEFAULT_EMBEDDING_PATH_BASENAMES + ) + + +def is_hidden_model(*values: str | None) -> bool: + """True if any id/path is the RAG embedding model (the effective embedder + or its GGUF companion repo) or the llama.cpp install validation probe + (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). + None are usable chat models; the probe can be cached as a side effect of + installing the prebuilt llama-server and otherwise sorts smallest, so it + would be auto-selected. + + Hub repo ids are matched EXACTLY (case-insensitive full "owner/name"), so a + custom embedder with a generic basename like "org/model" cannot substring + hide unrelated cached repos such as "user/model-chat" or "org/model-GGUF". + Existing paths take precedence over the identical ``owner/name`` repo + shape. Cache and LM Studio paths use exact repo-derived segments. Local + copies of the static default embedder also use a boundary-aware basename + fallback; configured custom repos never do.""" + from core.rag import config as rag_config + + hidden_repo_ids = { + _PROBE_REPO_ID.lower(), + *(repo_id.lower() for repo_id in _DEFAULT_EMBEDDING_REPO_IDS), + } + exact_paths: list[str] = [] + for model in { + rag_config.EMBEDDING_MODEL, + rag_config.default_gguf_repo(), + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + }: + existing_path = _existing_resolved_path(model) + if existing_path: + exact_paths.append(existing_path.lower()) + elif _HF_REPO_ID_RE.match(model): + hidden_repo_ids.add(model.lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if _HF_REPO_ID_RE.match(v): + # A repo id ("owner/name"): match the hidden set exactly. It is + # never a filesystem path, so skip the path/filename checks. + if low in hidden_repo_ids: + return True + continue + # Anything else is treated as a filesystem path (the cached snapshot + # path, or a local model id). Match the probe by its exact filename and + # any configured local-path embedder by exact resolved path. Split on + # both separators so a Windows-style path ("...\\stories260K.gguf") is + # matched even when this runs on a POSIX interpreter (and vice versa). + if low.replace("\\", "/").rsplit("/", 1)[-1] == _PROBE_FILENAME: + return True + if _path_basename_is_default_embedder(v): + return True + if _path_contains_repo_id(v, hidden_repo_ids): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index c24ec28e1d..50b3cd3513 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,9 +50,11 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} -# Native training context length (``{arch}.context_length``). None = absent / -# unreadable. Lets the UI show the real context ceiling before a model loads. -_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} +# GGUF header dims for the staged/deferred-load UI: context_length, layer_count +# (block_count), and moe_layer_count (block_count minus leading dense layers; 0 +# if not MoE). One cached pass fills all three so the staged sheet can size every +# slider before the model loads. None = unreadable / not a GGUF. +_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {} def _cache_key(path: str) -> Optional[_CacheKey]: @@ -142,32 +144,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out -def read_gguf_context_length(path: str) -> Optional[int]: - """Return the GGUF's native training context length (``{arch}.context_length``), - or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). - Lets the UI populate the context slider before the model is loaded.""" +def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]: + """GGUF header dims for the staged-load UI in one cached pass: + ``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None + when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not + a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size + the context, GPU-layers and MoE sliders before the model loads.""" key = _cache_key(path) if key is None: return None with _CACHE_LOCK: - if key in _CONTEXT_CACHE: - return _CONTEXT_CACHE[key] - result = _parse_gguf_context_length(path) + if key in _DIMS_CACHE: + return _DIMS_CACHE[key] + result = _parse_gguf_staged_dims(path) with _CACHE_LOCK: - while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES: try: - _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + _DIMS_CACHE.pop(next(iter(_DIMS_CACHE))) except StopIteration: break - _CONTEXT_CACHE[key] = result + _DIMS_CACHE[key] = result return result -def _parse_gguf_context_length(path: str) -> Optional[int]: - # The context key is architecture-namespaced (``llama.context_length`` etc.), - # so we learn the key only after reading ``general.architecture``. GGUF writes - # general.* before arch.* keys, matching the loader's own parser. - ctx_key: Optional[str] = None +def read_gguf_context_length(path: str) -> Optional[int]: + """Native training context length (``{arch}.context_length``), or ``None``. + Thin accessor over read_gguf_staged_dims.""" + dims = read_gguf_staged_dims(path) + return dims["context_length"] if dims else None + + +def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]: + """Walk a GGUF header once and return the requested architecture-namespaced + uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are + ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF + writes general.* before arch.* keys, matching the loader's own parser). + Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or + partial when some keys are absent).""" + arch: Optional[str] = None + found: Dict[str, int] = {} try: with open(path, "rb") as f: head = f.read(24) @@ -204,28 +219,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]: sbytes = f.read(slen) if len(sbytes) < slen: break - ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" - elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + arch = sbytes.decode("utf-8", "replace") + elif ( + arch is not None + and vtype in (4, 10) + and key.startswith(f"{arch}.") + and key[len(arch) + 1 :] in wanted_suffixes + ): width = 4 if vtype == 4 else 8 n_bytes = f.read(width) if len(n_bytes) < width: break - value = struct.unpack(" 0 else None + found[key[len(arch) + 1 :]] = struct.unpack( + " Optional[Dict[str, Optional[int]]]: + vals = _parse_gguf_arch_uints( + path, + frozenset( + { + "context_length", + "block_count", + "expert_count", + "leading_dense_block_count", + } + ), + ) + if vals is None: + return None + ctx = vals.get("context_length") + block = vals.get("block_count") + # A real context/layer count is positive; treat 0/garbage as absent so the + # UI never builds a slider with max < min. + context_length = ctx if ctx and ctx > 0 else None + layer_count = block if block and block > 0 else None + # MoE layer count = block_count - leading dense layers, only when experts + # exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in + # core/inference/llama_cpp.py. + if not vals.get("expert_count") or not block: + moe_layer_count: Optional[int] = 0 + else: + moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0)) + return { + "context_length": context_length, + "layer_count": layer_count, + "moe_layer_count": moe_layer_count, + } # Strings (8) and arrays (9) are handled inline. diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts index ec75b17f20..08492ab480 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts @@ -2,7 +2,9 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 // Per-model pre-load inference settings, persisted in localStorage so the load -// dialog can offer "Remember settings for ". +// dialog can offer "Remember settings for ". GGUF picks only: every +// field is a llama.cpp load knob, so all save/restore call sites gate on +// GGUF-ness (a non-GGUF blob would only snapshot leftover standing values). const KEY = "unsloth_load_settings"; @@ -12,14 +14,22 @@ export interface RememberedLoadSettings { speculativeType: string | null; specDraftNMax: number | null; tensorParallel: boolean; + // GPU Memory controls. Optional so an older blob (which lacked them) still + // parses, leaving the live knobs untouched on apply. The mode is kept with the + // manual knobs (gpuLayers/nCpuMoe are ignored outside Manual mode). A null + // selectedGpuIds is meaningful (all GPUs), so it's distinguished from absent. + // The per-GPU split ratio is deliberately NOT remembered: it's positionally + // bound to the exact GPU set/order and unvalidated, so it would mismatch. + gpuMemoryMode?: "auto" | "manual"; + gpuLayers?: number; + nCpuMoe?: number; + selectedGpuIds?: number[] | null; } -// Storage key for a pick's remembered settings. The remembered knobs are -// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the -// right values differ per quant. An HF repo collapses all its GGUF variants into -// one `id`, so fold the variant in to scope settings per quant. Local .gguf -// paths key by their file path (already file-specific); native drag-drop files -// key by display label, so same-named files in different folders share an entry. +// Storage key for a pick's remembered settings, scoped per quant (the VRAM-budget +// knobs differ per quant). An HF repo collapses its GGUF variants into one `id`, +// so fold the variant in. Local .gguf paths are already file-specific; native +// drag-drop files key by display label, so same-named files share an entry. export function rememberedLoadSettingsKey(selection: { id: string; ggufVariant?: string | null; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0bf46e7343..7083f02288 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -45,12 +45,18 @@ import { import { type PendingImageEditReference, type RagAutoInject, + GPU_LAYERS_AUTO, + loadedGpuMemoryFieldsUnlessStaged, + reconcilePersistedGpuIds, resolveLoadedSpeculativeSettings, resolveSpeculativeSettingsForLoad, + persistGpuMemoryModeOnLoad, resolveToolsEnabledOnLoad, saveSpeculativeType, useChatRuntimeStore, } from "../stores/chat-runtime-store"; +import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "../presets/preset-policy"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { useExternalProvidersStore } from "../stores/external-providers-store"; import { shouldPreserveFullOutput, @@ -1489,6 +1495,13 @@ async function autoLoadSmallestModel(): Promise<{ max_seq_length: number; is_lora: boolean; gguf_variant?: string | null; + // GGUF-only: scopes the training guard to the same placement policy /load + // will use. Manual mode must match because it makes placement user-owned. + // The layer/MoE/split/KV/spec knobs are deliberately not sent: Auto mode's + // guard sizes conservatively, while Manual mode bypasses that estimate. + // The safetensors fallback omits both fields and uses HF auto-placement. + gpu_ids?: number[]; + gpu_memory_mode?: "auto" | "manual"; }): Promise { const validation = await validateModel({ ...payload, @@ -1520,12 +1533,18 @@ async function autoLoadSmallestModel(): Promise<{ return false; } const currentStore = useChatRuntimeStore.getState(); - const remembered = loadRememberedLoadSettings( - rememberedLoadSettingsKey({ - id: candidate.id, - ggufVariant: candidate.ggufVariant, - }), - ); + // Blobs are saved for GGUF picks only (the sheet gates on it), so don't + // let a legacy non-GGUF blob feed a stale context/spec choice into a + // safetensors auto-load. + const remembered = + candidate.kind === "gguf" + ? loadRememberedLoadSettings( + rememberedLoadSettingsKey({ + id: candidate.id, + ggufVariant: candidate.ggufVariant, + }), + ) + : null; const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId: candidate.id, ggufVariant: candidate.ggufVariant, @@ -1537,6 +1556,38 @@ async function autoLoadSmallestModel(): Promise<{ maxSeqLength: candidate.maxSeqLength, presetSource: currentStore.activePresetSource, }); + // The GPU knobs are per-model, so read them from the same remembered + // settings that fed effectiveMaxSeqLength -- on a background auto-load the + // live store holds session defaults, not the saved Manual mode / layer pin / + // GPU pick. Absent fields fall back like applyRememberedLoadSettings: the + // mode to the store (a persisted standing preference), the per-model knobs to + // their defaults. The saved GPU pick is reconciled against the GPUs present + // now, like the interactive restore. + const effectiveGpuMemoryMode = + remembered?.gpuMemoryMode ?? currentStore.gpuMemoryMode; + const effectiveGpuLayers = remembered?.gpuLayers ?? GPU_LAYERS_AUTO; + const effectiveNCpuMoe = remembered?.nCpuMoe ?? 0; + if (remembered?.selectedGpuIds != null) { + // Warm the device cache first: on a cold cache the reconcile passes the + // saved pick through unvalidated, and a stale cross-host pick then fails + // the load with the picker hidden. + await ensureGpuDeviceCache(); + } + const effectiveGpuIds = + remembered?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(remembered.selectedGpuIds) + : null; + // Under Manual GPU memory + Auto layers, llama.cpp's --fit owns context + // sizing, so send 0 (or the pinned length). GGUF-only; a no-op otherwise. + // The context pin is per-model too, so it comes from remembered settings, + // not the live store. + const fitMaxSeqLength = resolveFitMaxSeqLength( + candidate.kind === "gguf", + effectiveGpuMemoryMode, + effectiveGpuLayers, + remembered?.contextLength ?? null, + effectiveMaxSeqLength, + ); const effectiveSpeculativeType = remembered?.speculativeType ?? specSettings.speculativeType; const effectiveSpecDraftNMax = @@ -1544,9 +1595,16 @@ async function autoLoadSmallestModel(): Promise<{ if ( !(await canAutoLoad({ model_path: candidate.id, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: fitMaxSeqLength, is_lora: false, gguf_variant: candidate.ggufVariant, + // The same remembered-derived GPU pick the load below sends. + ...(candidate.kind === "gguf" + ? { + gpu_ids: effectiveGpuIds ?? undefined, + gpu_memory_mode: effectiveGpuMemoryMode, + } + : {}), })) ) { skippedAutoLoadCandidates.add( @@ -1558,7 +1616,7 @@ async function autoLoadSmallestModel(): Promise<{ const loadResp = await loadModel({ model_path: candidate.id, hf_token: hfToken, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: fitMaxSeqLength, load_in_4bit: true, is_lora: false, gguf_variant: candidate.ggufVariant, @@ -1567,8 +1625,22 @@ async function autoLoadSmallestModel(): Promise<{ speculative_type: effectiveSpeculativeType, spec_draft_n_max: effectiveSpecDraftNMax, tensor_parallel: remembered?.tensorParallel ?? false, + // GGUF-only: the safetensors fallback loads via HF auto-placement (no + // explicit pins). The split ratio is deliberately never remembered + // (positionally bound to an exact GPU set), so auto-load leaves llama.cpp's + // free-VRAM default in charge rather than sending a stale store value. + ...(candidate.kind === "gguf" + ? { + gpu_memory_mode: effectiveGpuMemoryMode, + gpu_layers: effectiveGpuLayers, + n_cpu_moe: effectiveNCpuMoe, + gpu_ids: effectiveGpuIds ?? undefined, + } + : {}), }); saveSpeculativeType(effectiveSpeculativeType); + // Self-gates on is_gguf (skips diffusion), so persists only for a real GGUF load. + persistGpuMemoryModeOnLoad(loadResp, effectiveGpuMemoryMode); useChatRuntimeStore .getState() .setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined); @@ -1597,6 +1669,15 @@ async function autoLoadSmallestModel(): Promise<{ store.setModels([...store.models, autoModel]); } if (candidate.kind === "gguf") { + // Keep an explicit Manual+Auto context pin the load just applied (so a + // later Apply doesn't silently revert it to auto-fit sizing), mirroring + // the interactive path's keepCustomCtx; other cases baseline on + // ggufContextLength. + const keepCustomCtx = resolveManualAutoCtxPin( + effectiveGpuMemoryMode, + effectiveGpuLayers, + remembered?.contextLength ?? null, + ); useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, ggufMaxContextLength: @@ -1613,6 +1694,10 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + ...loadedGpuMemoryFieldsUnlessStaged(loadResp, { + customContextLength: keepCustomCtx, + }), + loadedCustomContextLength: keepCustomCtx, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1633,6 +1718,9 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + // Non-GGUF response: clears any stale GPU baseline a prior manual-GPU + // GGUF load left, matching the interactive/status sibling load paths. + ...loadedGpuMemoryFieldsUnlessStaged(loadResp), defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1820,12 +1908,17 @@ async function autoLoadSmallestModel(): Promise<{ duration: 30000, }); try { + const rt = useChatRuntimeStore.getState(); if ( !(await canAutoLoad({ model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", max_seq_length: 0, is_lora: false, gguf_variant: "UD-Q4_K_XL", + // The same live-store GPU pick the load below sends (a fresh default + // model has no remembered settings to prefer). + gpu_ids: rt.selectedGpuIds ?? undefined, + gpu_memory_mode: rt.gpuMemoryMode, })) ) { toast.dismiss(toastId); @@ -1835,6 +1928,9 @@ async function autoLoadSmallestModel(): Promise<{ const loadResp = await loadModel({ model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", hf_token: hfToken, + // Model default under both modes: Auto layers + no pin means + // resolveFitMaxSeqLength returns 0 for every mode (the canAutoLoad + // preflight above sends the same). max_seq_length: 0, load_in_4bit: true, is_lora: false, @@ -1842,8 +1938,20 @@ async function autoLoadSmallestModel(): Promise<{ trust_remote_code: trustRemoteCode, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, + // GPU Memory mode is a standing preference, so honor it on auto-load. + // The layer/MoE/split knobs and the context pin are per-model: the live + // store may hold edits drafted for a staged pick, and a fresh default + // model has no remembered settings, so those stay at their defaults like + // the cached-candidate path. The GPU pick deliberately differs (it's the + // picker's current on-screen selection, which the canAutoLoad preflight + // above already committed to). + gpu_memory_mode: rt.gpuMemoryMode, + gpu_layers: GPU_LAYERS_AUTO, + n_cpu_moe: 0, + gpu_ids: rt.selectedGpuIds ?? undefined, }); saveSpeculativeType(specSettings.speculativeType); + persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode); useChatRuntimeStore .getState() .setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL"); @@ -1880,6 +1988,10 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + ...loadedGpuMemoryFieldsUnlessStaged(loadResp), + // Drives the GPU Memory controls' diffusion gate; set alongside the + // GPU fields on every load path so the gate can't read stale. + loadedIsDiffusion: loadResp.is_diffusion ?? false, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index ebf9461172..0f6af38033 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -127,28 +127,38 @@ export async function validateModel( native_path_lease: payload.nativePathLease ?? null, hf_token: payload.hf_token, gguf_variant: payload.gguf_variant ?? null, - // Send the intended load settings so validate's VRAM check matches the - // follow-up /load and doesn't unload for a load /load would then reject. + // Intended load settings so validate's preflight matches the follow-up + // /load. Default placement is sized against the selected GPUs. max_seq_length: payload.max_seq_length, load_in_4bit: payload.load_in_4bit, + gpu_ids: payload.gpu_ids, + // Manual placement is an explicit override: Auto layers use llama.cpp + // --fit, while a pinned layer count is owned by the user. Tell validate + // so it applies the same training-guard policy as /load. + gpu_memory_mode: payload.gpu_memory_mode, }), }); return parseJsonOrThrow(response); } /** - * Read a GGUF's native context length from its local header (no GPU load, no - * download). Returns null when the file isn't downloaded yet, the model isn't a - * GGUF, or it's gated. For a native (drag-drop / picked) file, pass - * `nativePathToken` so the backend reads the granted local path. Used by the - * deferred-load staging flow to fill the context slider before the single load. + * Read a GGUF's header dims (native context length, total layer count, MoE + * expert-layer count) from its local file (no GPU load, no download). All are + * null when the file isn't downloaded yet, the model isn't a GGUF, or it's + * gated. For a native (drag-drop / picked) file, pass `nativePathToken` so the + * backend reads the granted local path. Used by the deferred-load staging flow + * to size the context, GPU-layers and MoE sliders before the single load. */ -export async function fetchGgufContextLength(payload: { +export async function fetchGgufStagedMetadata(payload: { model_path: string; gguf_variant?: string | null; hf_token?: string | null; nativePathToken?: string | null; -}): Promise { +}): Promise<{ + contextLength: number | null; + layerCount: number | null; + moeLayerCount: number | null; +}> { let nativePathLease: string | null = null; if (payload.nativePathToken) { try { @@ -156,8 +166,8 @@ export async function fetchGgufContextLength(payload: { await consumeNativePathToken(payload.nativePathToken, "validate-model") ).nativePathLease; } catch { - // Lease expired / revoked: degrade to no context (the load can re-mint). - return null; + // Lease expired / revoked: degrade to no metadata (the load can re-mint). + return { contextLength: null, layerCount: null, moeLayerCount: null }; } } const response = await authFetch("/api/inference/validate", { @@ -172,7 +182,11 @@ export async function fetchGgufContextLength(payload: { }), }); const res = await parseJsonOrThrow(response); - return res.context_length ?? null; + return { + contextLength: res.context_length ?? null, + layerCount: res.layer_count ?? null, + moeLayerCount: res.moe_layer_count ?? null, + }; } export async function unloadModel(payload: UnloadModelRequest): Promise { diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index ec0ad977bf..217eaf8b6d 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1445,9 +1445,11 @@ export function ChatPage({ // were already seeded on stage, so keepSpeculative only when a config was // saved -- otherwise the standing speculative preference should win. autoLoadStagedRef.current = (pending) => { - const remembered = loadRememberedLoadSettings( - rememberedLoadSettingsKey(pending), - ); + // Blobs are saved for GGUF picks only (the sheet gates on it), so don't + // let a legacy non-GGUF blob claim a seeded config here. + const remembered = hasGgufSource(pending) + ? loadRememberedLoadSettings(rememberedLoadSettingsKey(pending)) + : null; void selectModel({ ...pending, isDownloaded: true, @@ -2813,6 +2815,11 @@ export function ChatPage({ selectModel({ id: state.params.checkpoint, ggufVariant: state.activeGgufVariant ?? undefined, + // A native (drag-drop / picked) GGUF's checkpoint is only a display + // label, so the reload needs its path token to re-mint a lease -- + // else applying the now-exposed GPU/context controls can't resolve + // the file. Null for non-native loads, which reload by id as before. + nativePathToken: state.activeNativePathToken ?? undefined, forceReload: true, isDownloaded: true, loadingDescription: "Reloading with updated chat template.", diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cedd298ecf..b368a811fa 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -55,6 +55,7 @@ import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; +import { useGpuDevices } from "@/hooks/use-gpu-info"; import { useIsMobile } from "@/hooks/use-mobile"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { cn } from "@/lib/utils"; @@ -99,8 +100,11 @@ import { providerSupportsFastMode, } from "./provider-capabilities"; import { + GPU_LAYERS_AUTO, + distributeByWeight, isPendingGguf, pendingSelectionMatches, + rebalanceSplit, useChatRuntimeStore, } from "./stores/chat-runtime-store"; import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section"; @@ -250,6 +254,7 @@ function ParamSlider({ displayValue, info, valueSize, + disabled, }: { label: string; value: number; @@ -260,6 +265,7 @@ function ParamSlider({ displayValue?: string; info?: ReactNode; valueSize?: number; + disabled?: boolean; }) { return (
@@ -279,6 +285,7 @@ function ParamSlider({ displayValue={displayValue} ariaLabel={label} size={valueSize ?? 4} + disabled={disabled} />
onChange(snapToStep(v, step, min, max))} className="panel-slider" + disabled={disabled} /> ); @@ -540,8 +548,17 @@ export function ChatSettingsPanel({ const base = slash >= 0 ? id.slice(slash + 1) : id; return base || id; })(); + const activeNativePathToken = useChatRuntimeStore( + (s) => s.activeNativePathToken, + ); + const loadedGgufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + // A GGUF loaded from a native path / direct .gguf has no HF variant, so key + // off the same signal the status hydration uses -- variant OR native token OR + // a GGUF context -- else the GPU Memory controls hide for a loaded local GGUF. const isLoadedGguf = - useChatRuntimeStore((s) => s.activeGgufVariant) != null; + useChatRuntimeStore((s) => s.activeGgufVariant) != null || + activeNativePathToken != null || + loadedGgufContextLength != null; // While a pick is staged the sheet configures *that* model, so its GGUF-ness // (not the currently loaded model's) decides whether the GGUF-only controls // show. Otherwise a staged non-GGUF Hub repo would inherit the loaded GGUF's @@ -607,6 +624,25 @@ export function ChatSettingsPanel({ const loadedTensorParallel = useChatRuntimeStore( (s) => s.loadedTensorParallel, ); + const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode); + const setGpuMemoryMode = useChatRuntimeStore((s) => s.setGpuMemoryMode); + const loadedGpuMemoryMode = useChatRuntimeStore((s) => s.loadedGpuMemoryMode); + const loadedIsDiffusion = useChatRuntimeStore((s) => s.loadedIsDiffusion); + const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers); + const setGpuLayers = useChatRuntimeStore((s) => s.setGpuLayers); + const loadedGpuLayers = useChatRuntimeStore((s) => s.loadedGpuLayers); + const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe); + const setNCpuMoe = useChatRuntimeStore((s) => s.setNCpuMoe); + const loadedNCpuMoe = useChatRuntimeStore((s) => s.loadedNCpuMoe); + const splitRatio = useChatRuntimeStore((s) => s.splitRatio); + const setSplitRatio = useChatRuntimeStore((s) => s.setSplitRatio); + const loadedSplitRatio = useChatRuntimeStore((s) => s.loadedSplitRatio); + const ggufLayerCount = useChatRuntimeStore((s) => s.ggufLayerCount); + const moeLayerCount = useChatRuntimeStore((s) => s.moeLayerCount); + const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds); + const setSelectedGpuIds = useChatRuntimeStore((s) => s.setSelectedGpuIds); + const loadedGpuIds = useChatRuntimeStore((s) => s.loadedGpuIds); + const gpuDevices = useGpuDevices(); const chatTemplateOverride = useChatRuntimeStore( (s) => s.chatTemplateOverride, ); @@ -614,6 +650,9 @@ export function ChatSettingsPanel({ (s) => s.loadedChatTemplateOverride, ); const customContextLength = useChatRuntimeStore((s) => s.customContextLength); + const loadedCustomContextLength = useChatRuntimeStore( + (s) => s.loadedCustomContextLength, + ); const setCustomContextLength = useChatRuntimeStore( (s) => s.setCustomContextLength, ); @@ -641,10 +680,14 @@ export function ChatSettingsPanel({ : null; useEffect(() => { if (!pendingKey) return; - const saved = loadRememberedLoadSettings(pendingKey); + // GGUF-only, like the stageOrLoad / Hub restore paths: every remembered + // field is a llama.cpp knob, so a non-GGUF pick has nothing to restore -- + // and applying its blob would clobber the standing gpuMemoryMode with a + // stale snapshot (the save on Load below is gated the same way). + const saved = pendingIsGguf ? loadRememberedLoadSettings(pendingKey) : null; setRemember(saved != null); if (saved) applyRememberedLoadSettings(saved); - }, [pendingKey, applyRememberedLoadSettings]); + }, [pendingKey, pendingIsGguf, applyRememberedLoadSettings]); // While staging, the sheet reflects the STAGED model, so its header context // takes precedence over the loaded model's (which may differ or be larger). const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength; @@ -661,15 +704,132 @@ export function ChatSettingsPanel({ const ctxDisplayValue = customContextLength ?? baseContext ?? ""; const ctxMaxValue = baseNativeContext ?? baseContext ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; - const ctxDirty = customContextLength !== null; + const ctxDirty = customContextLength !== loadedCustomContextLength; const specDirty = speculativeType !== loadedSpeculativeType; const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax; const tpDirty = tensorParallel !== (loadedTensorParallel ?? false); + // A loaded diffusion GGUF runs mode-agnostic (pins all layers on one GPU, + // ignores --fit/--gpu-layers), so the GPU Memory mode + manual controls don't + // apply -- hide them and don't let the preserved standing mode read as dirty. + // The GPU picker still applies (diffusion pins the chosen device). A staged pick + // keeps the controls (a pending pick's diffusion-ness isn't known until load). + const gpuModeApplies = + isGguf && (pendingSelection != null || !loadedIsDiffusion); + const gpuDirty = + gpuModeApplies && gpuMemoryMode !== (loadedGpuMemoryMode ?? "auto"); + const isManual = gpuModeApplies && gpuMemoryMode === "manual"; + // Manual with the GPU Layers slider at "Auto" (leftmost): --fit owns the whole + // layout, so the offload knobs (MoE, split, TP) don't apply. + const autoLayers = isManual && gpuLayers < 0; + // GPUs actually in use: the picked subset, or all visible when none picked. + const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index); + // TP is off with fewer than 2 GPUs in use (single GPU, or the picker narrowed + // to one): tensor split is a no-op there and aborts on some archs. Mirrors the + // multi-GPU gate on the GPU picker / Split ratio. (Under Auto layers the whole + // TP control is hidden -- llama.cpp's --fit aborts under --split-mode tensor.) + const tpDisabled = gpusInUse.length <= 1; + // Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback): + // llama.cpp counts the output layer as one more offloadable layer past the + // repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so + // the slider max must reach it or full offload is unreachable. While staging, + // use the staged model's layer count (read from its header). + const stagedLayerCount = pendingSelection?.layerCount ?? null; + const modelLayerCount = pendingIsGguf ? stagedLayerCount : ggufLayerCount; + const gpuLayersMax = modelLayerCount != null ? modelLayerCount + 1 : 256; + // MoE-offload slider: shown only for MoE models, capped at their MoE-layer + // count. While staging, use the staged model's count (read from its header); + // otherwise the loaded model's. + const stagedMoeLayerCount = pendingSelection?.moeLayerCount ?? null; + const moeLayersMax = pendingIsGguf + ? (stagedMoeLayerCount ?? 0) + : (moeLayerCount ?? 0); + const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0; + // gpuLayers always counts; MoE only with an explicit layer count (see above). + const manualDirty = + isManual && + (gpuLayers !== loadedGpuLayers || + (!autoLayers && nCpuMoe !== (loadedNCpuMoe ?? 0))); + // GPU picker: only meaningful on multi-GPU, and only when the reported + // indices are physical (relative ordinals from a parent CUDA_VISIBLE_DEVICES + // mask can't be mapped back to pin a device). null = use all (auto). + const showGpuPicker = + isGguf && + gpuDevices.length > 1 && + gpuDevices.every((d) => d.physicalIndex); + const isGpuChecked = (index: number) => + selectedGpuIds === null || selectedGpuIds.includes(index); + const toggleGpu = (index: number) => { + const all = gpuDevices.map((d) => d.index); + const current = selectedGpuIds ?? all; + const next = current.includes(index) + ? current.filter((i) => i !== index) + : [...current, index].sort((a, b) => a - b); + if (next.length === 0) return; // keep at least one GPU selected + setSelectedGpuIds(next.length === all.length ? null : next); + // The per-GPU split is positional, so any change to the set of GPUs in use + // invalidates it: drop it (the sliders fall back to the VRAM-weighted + // default). TP needs 2+ GPUs, so disable it when only one remains. + setSplitRatio(null); + if (next.length <= 1) { + setTensorParallel(false); + } + }; + const gpuIdsKey = (ids: number[] | null) => (ids === null ? "auto" : ids.join(",")); + const gpuIdsDirty = gpuIdsKey(selectedGpuIds) !== gpuIdsKey(loadedGpuIds); + // Per-GPU layer split (--tensor-split): manual + 2+ GPUs in use. One slider + // per GPU, each a layer count; together they sum to the GPU Layers total. + const showSplitRatio = + isManual && !autoLayers && showGpuPicker && gpusInUse.length > 1; + // The total the per-GPU counts sum to (the GPU Layers slider value); 0 under + // Auto, where the split is hidden. The devices behind the GPUs in use, for + // labels + the VRAM-weighted default. + const splitTotal = Math.max(0, Math.min(gpuLayers, gpuLayersMax)); + const gpusInUseDevices = gpusInUse.map( + (i) => gpuDevices.find((d) => d.index === i) ?? null, + ); + // Displayed per-GPU counts. splitRatio is a stable reference balance (only a + // slider edit changes it), rescaled to the current total; deriving rather than + // mutating it on GPU Layers changes keeps the balance intact when the total + // passes through low values or Auto. No saved split: free-VRAM-weighted default + // (llama.cpp's unset default splits by free VRAM, so the first edit starts from + // the default's placement, not a total-VRAM ratio that can land layers on a + // busy GPU). A genuine 0 (a full GPU) is a real weight, not missing data: the + // probe's no-data case degrades to the total server-side, and an all-zero list + // falls back to an even split in distributeByWeight. Not yet sent. + const splitCounts = + splitRatio && splitRatio.length === gpusInUse.length + ? distributeByWeight(splitTotal, splitRatio) + : distributeByWeight( + splitTotal, + gpusInUseDevices.map((d) => d?.memoryFreeGb ?? d?.memoryTotalGb ?? 1), + ); + const setSplitCount = (k: number, v: number) => + setSplitRatio(rebalanceSplit(splitTotal, splitCounts, k, v)); + const splitRatioDirty = + isManual && + !autoLayers && + JSON.stringify(splitRatio ?? null) !== JSON.stringify(loadedSplitRatio ?? null); + // Auto-fit context (Manual + Auto layers): <= 0 means "Auto" (--fit sizes it); + // a positive value pins it. Surface the length --fit chose once it's loaded. + const fitCtxAuto = autoLayers && (customContextLength ?? 0) <= 0; + const loadedAutoLayers = + loadedGpuMemoryMode === "manual" && (loadedGpuLayers ?? GPU_LAYERS_AUTO) < 0; + const fitResolvedCtx = + fitCtxAuto && loadedAutoLayers ? ggufContextLength : null; // A saved chat-template override is a reload-time setting too, so surface // Apply for a template-only edit (otherwise it could never be applied). const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride; const modelSettingsDirty = - kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty || templateDirty; + kvDirty || + ctxDirty || + specDirty || + specDraftDirty || + tpDirty || + gpuDirty || + manualDirty || + gpuIdsDirty || + splitRatioDirty || + templateDirty; const [presetNameInput, setPresetNameInput] = useState(activePreset); const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); const [systemPromptDraft, setSystemPromptDraft] = useState(""); @@ -980,7 +1140,64 @@ export function ChatSettingsPanel({ )} {isGguf && ( <> - {showContextControl && ( + {showContextControl && (autoLayers ? ( +
+
+
+ + Context Length + + + Auto: llama.cpp's --fit sizes the context to fit VRAM. + Set a length to pin it instead -- --fit then optimizes + GPU layer offload around it. The length --fit chose + shows here after loading. + +
+ { + setCustomContextLength(v > 0 ? v : null); + }} + ariaLabel="Context Length" + size={8} + disabled={modelControlsDisabled} + /> +
+ { + // Far-left snaps to Auto; otherwise to the nearest 1024. + if (v < 512) { + setCustomContextLength(null); + } else { + setCustomContextLength(Math.round(v / 1024) * 1024); + } + }} + className="panel-slider" + disabled={modelControlsDisabled} + /> + {fitResolvedCtx != null && ( +

+ llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens. +

+ )} +
+ ) : (
@@ -1036,7 +1253,7 @@ export function ChatSettingsPanel({

)}
- )} + ))}
@@ -1191,6 +1408,163 @@ export function ChatSettingsPanel({ )} )} + {gpuModeApplies && ( +
+
+ + GPU Memory + + +
+
+ Default: Unsloth + fits the model and context to your GPUs. +
+
+ Manual: set GPU + Layers yourself. Leave it on Auto to let llama.cpp size + the context and offload overflow (including MoE experts) + to RAM. +
+
+ +
+
+ +
+
+ )} + {isManual && ( + <> + + Layers to keep on the GPU (--gpu-layers); the rest run + on CPU. Auto lets llama.cpp size the split (and the + context) to fit VRAM. At the maximum, the whole model + is on the GPU. + + } + /> + {showMoeSlider && ( + + Keep the experts of this many MoE layers on the CPU + (--n-cpu-moe) to save VRAM. 0 = all experts on the + GPU; at the maximum, all are on the CPU. + + } + /> + )} + {showSplitRatio && ( +
+
+ + Layers per GPU + + + Splits GPU Layers across GPUs (--tensor-split). + Without Tensor Parallelism each value is the layer + count on that GPU; with it, every GPU holds a slice + of each layer, so the values are only a ratio. + +
+ {gpusInUseDevices.map((d, k) => ( + setSplitCount(k, v)} + valueSize={6} + disabled={modelControlsDisabled} + /> + ))} +
+ )} + + )} + {showGpuPicker && ( +
+
+ + GPUs + + + Which GPUs this model may use. Unchecked GPUs are hidden + from llama.cpp (CUDA_VISIBLE_DEVICES, or + HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use + every GPU. + +
+
+ {gpuDevices.map((d) => ( +
+ + GPU {d.index}: {d.name} + {d.memoryTotalGb + ? ` · ${Math.round(d.memoryTotalGb)} GB` + : ""} + + toggleGpu(d.index)} + data-test-id={`gpu-pick-${d.index}`} + disabled={modelControlsDisabled} + /> +
+ ))} +
+
+ )} + {gpuModeApplies && !autoLayers && (
@@ -1206,10 +1580,11 @@ export function ChatSettingsPanel({ className="panel-switch shrink-0" checked={tensorParallel} onCheckedChange={setTensorParallel} - disabled={modelControlsDisabled} + disabled={tpDisabled || modelControlsDisabled} data-test-id="tensor-parallel-switch" />
+ )} )} {/* No persistent "enable custom code" toggle: it is consented per model @@ -1228,14 +1603,21 @@ export function ChatSettingsPanel({ {Math.round((stagedDownloadFraction ?? 0) * 100)}%

)} - + {/* GGUF picks only: a non-GGUF pick shows none of the load + knobs the blob captures, so there is nothing to remember. */} + {pendingIsGguf && ( + + )} {stagedLoading ? ( // Mid-load: nothing to load or abandon until it settles, so disable.
) : null} - + {/* The template override is a load-time knob too (applied on the next + reload) and the in-flight load already snapshotted it, so lock its + editors like the sibling controls -- a mid-load save would be + silently clobbered by the load response despite its toast. */} +
)} @@ -2086,7 +2477,7 @@ function BypassPermissionsToggle() { ); } -function ChatTemplateFields() { +function ChatTemplateFields({ disabled = false }: { disabled?: boolean }) { const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate); const override = useChatRuntimeStore((s) => s.chatTemplateOverride); const setOverride = useChatRuntimeStore((s) => s.setChatTemplateOverride); @@ -2120,7 +2511,8 @@ function ChatTemplateFields() { @@ -2131,7 +2523,8 @@ function ChatTemplateFields() { -
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 10e0904e4f..3003b52230 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -29,9 +29,14 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { + GPU_LAYERS_AUTO, isLocalModelPath, + loadedGpuMemoryFields, + loadedGpuMemoryFieldsUnlessStaged, pendingSelectionMatches, + persistGpuMemoryModeOnLoad, readPersistedSpeculativeType, + reconcilePersistedGpuIds, resolveToolsEnabledOnLoad, saveSpeculativeType, useChatRuntimeStore, @@ -46,9 +51,12 @@ import { } from "../lib/apply-inference-status-to-store"; import { mergeBackendRecommendedInference, + resolveFitMaxSeqLength, resolveLoadMaxSeqLength, + resolveManualAutoCtxPin, } from "../presets/preset-policy"; import { recordLastLocalModelLoad } from "../utils/last-local-model-load"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { isMultimodalResponse, } from "../types/api"; @@ -291,9 +299,12 @@ async function syncInferenceStatusToStore(options?: { if (statusRes.active_model && !isExternalSelectionActive) { const checkpointId = resolveInferenceCheckpointId(statusRes); if (checkpointId) { + const previousGgufVariant = + useChatRuntimeStore.getState().activeGgufVariant; setCheckpoint(checkpointId, statusRes.gguf_variant); applyActiveModelStatusToStore(statusRes, { previousCheckpoint: selectedCheckpoint, + previousGgufVariant, }); // setModels(listRes...) above used catalog data, which omits audio // capability. Re-apply live status so attach gates survive a refresh. @@ -511,7 +522,11 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.isDownloaded ?? false; const model = models.find((entry) => entry.id === modelId); const lora = loras.find((entry) => entry.id === modelId); - const isGguf = explicitIsGguf ?? model?.isGguf ?? false; + // A native path-token selection is a local GGUF by construction (the + // native model intents only grant .gguf files), but its id is a display + // label that need not end in ".gguf" -- without this, Manual + Auto + // layers would pin the UI context instead of letting --fit size it. + const isGguf = explicitIsGguf ?? model?.isGguf ?? nativePathToken != null; const loraIsAdapter = lora?.exportType === "lora"; const isLora = explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false; @@ -578,18 +593,27 @@ export function useChatModelRuntime() { let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; let approvedRemoteCodeFingerprint: string | null = null; const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const previousActiveNativePathToken = + stateBeforeUnload.activeNativePathToken; const previousIsGguf = previousModel?.isGguf === true || previousVariant != null + || previousActiveNativePathToken != null || (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); - const rollbackMaxSeqLength = previousIsGguf - ? (stateBeforeUnload.ggufContextLength ?? 0) - : maxSeqLength; + // Respect the rolled-back model's auto-layers mode: a Manual+Auto model + // with an unpinned (auto) context must reload with 0 (so --fit + // re-auto-sizes), not the positive context it happened to pick (which + // the backend would treat as a pin). + const rollbackMaxSeqLength = resolveFitMaxSeqLength( + previousIsGguf, + stateBeforeUnload.loadedGpuMemoryMode ?? "auto", + stateBeforeUnload.loadedGpuLayers ?? GPU_LAYERS_AUTO, + stateBeforeUnload.loadedCustomContextLength, + previousIsGguf ? (stateBeforeUnload.ggufContextLength ?? 0) : maxSeqLength, + ); const hfToken = stateBeforeUnload.hfToken || null; const previousModelRequiresTrustRemoteCode = stateBeforeUnload.modelRequiresTrustRemoteCode; - const previousActiveNativePathToken = - stateBeforeUnload.activeNativePathToken; // Snapshot the load settings at click time, before the awaits below // (validation, the trust dialog, unload). For a staged Load these knobs // stay editable and a sheet-close revert (abandonStagedModel) can fire @@ -598,11 +622,29 @@ export function useChatModelRuntime() { // updates this snapshot in lock-step so non-staged loads are unchanged. const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; - const loadCustomContextLength = stateBeforeUnload.customContextLength; + // gpuMemoryMode is a standing preference (kept across a model switch); + // the rest are per-model knobs the reset below clears, so they are + // re-baselined there in lock-step with the store. + let loadCustomContextLength = stateBeforeUnload.customContextLength; const loadGgufContextLength = stateBeforeUnload.ggufContextLength; const loadTensorParallel = stateBeforeUnload.tensorParallel; const loadActivePresetSource = stateBeforeUnload.activePresetSource; const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; + const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode; + let loadGpuLayers = stateBeforeUnload.gpuLayers; + let loadNCpuMoe = stateBeforeUnload.nCpuMoe; + let loadSplitRatio = stateBeforeUnload.splitRatio; + // Reconcile the persisted pick against the GPUs present now, so a stale + // cross-host / now-hidden pick is dropped before /load rather than + // rejected there. Warm the device cache first: load-on-selection can + // run before any GPU hook mounted, and a cold cache would pass the + // pick through unvalidated. validateGpuIds derives from this too. + if (stateBeforeUnload.selectedGpuIds != null) { + await ensureGpuDeviceCache(); + } + let loadSelectedGpuIds = reconcilePersistedGpuIds( + stateBeforeUnload.selectedGpuIds, + ); let loadSpeculativeType = stateBeforeUnload.speculativeType; let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; try { @@ -615,16 +657,47 @@ export function useChatModelRuntime() { // context can exceed maxSeqLength, so sizing on raw maxSeqLength could // pass, unload, then have /load refuse it. Uses the click-time // snapshot (same values loadModel uses below), so the two agree. - const validateMaxSeqLength = resolveLoadMaxSeqLength({ - modelId, - ggufVariant, - customContextLength: loadCustomContextLength, - ggufContextLength: loadGgufContextLength, - currentCheckpoint, - activeGgufVariant: loadActiveGgufVariant, - maxSeqLength, - presetSource: loadActivePresetSource, - }); + // Mirror what /load does on a cross-model switch: the reset below + // clears the per-model Auto-layers context pin + GPU pick, and + // Manual+Auto sizes context through resolveFitMaxSeqLength. + // gpuMemoryMode is a standing preference, kept across the switch. + // A same-repo quant switch (same checkpoint, different gguf_variant) + // is a different model for per-model knobs: the pinned context, + // gpuLayers, GPU pick, and MoE offload are scoped per variant, so + // treat a variant change like a model switch and re-baseline them. + const switchingModelOrVariant = + currentCheckpoint !== modelId || + (loadActiveGgufVariant ?? null) !== (ggufVariant ?? null); + const resetsPerModelSettings = Boolean( + currentCheckpoint && switchingModelOrVariant && !keepSpeculative, + ); + const validateCustomContextLength = resetsPerModelSettings + ? null + : loadCustomContextLength; + const validateGpuIds = resetsPerModelSettings + ? null + : loadSelectedGpuIds; + // The reset below re-baselines gpuLayers to Auto; mirror it here. + const validateGpuLayers = resetsPerModelSettings + ? GPU_LAYERS_AUTO + : loadGpuLayers; + const validateMaxSeqLength = resolveFitMaxSeqLength( + isGguf, + loadGpuMemoryMode, + validateGpuLayers, + validateCustomContextLength, + resolveLoadMaxSeqLength({ + modelId, + ggufVariant, + isGguf, + customContextLength: validateCustomContextLength, + ggufContextLength: loadGgufContextLength, + currentCheckpoint, + activeGgufVariant: loadActiveGgufVariant, + maxSeqLength, + presetSource: loadActivePresetSource, + }), + ); const validation = await validateModel({ model_path: modelId, nativePathLease: validateNativePathLease, @@ -633,6 +706,8 @@ export function useChatModelRuntime() { load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, + gpu_ids: validateGpuIds ?? undefined, + ...(isGguf ? { gpu_memory_mode: loadGpuMemoryMode } : {}), }); // Upgrade consent runs before the security dialogs; Accept installs and the load continues. if (validation.requires_transformers_upgrade) { @@ -697,18 +772,52 @@ export function useChatModelRuntime() { // keepSpeculative skips this for a staged Load: the user picked the // mode for this model on the sidebar, so honor it (the backend still // falls back at runtime if the model has no MTP head). - if (currentCheckpoint && currentCheckpoint !== modelId && !keepSpeculative) { + if (resetsPerModelSettings) { const persistedSpeculativeType = readPersistedSpeculativeType(); useChatRuntimeStore.setState({ speculativeType: persistedSpeculativeType, loadedSpeculativeType: persistedSpeculativeType, specDraftNMax: null, loadedSpecDraftNMax: null, + // Per-model GPU knobs must not follow onto a different model + // (gpuMemoryMode is a standing preference and is kept). + selectedGpuIds: null, + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + // A Manual+Auto context pin is per-model; clear it so a different + // model loads at Auto/native, not the previous model's pin. + customContextLength: null, }); loadSpeculativeType = persistedSpeculativeType; loadSpecDraftNMax = null; + // Keep the click-time snapshot in lock-step with the store reset so + // the load below sizes against the cleared per-model knobs, not the + // previous model's (gpuMemoryMode is standing, so left as captured). + loadCustomContextLength = null; + loadSelectedGpuIds = null; + loadGpuLayers = GPU_LAYERS_AUTO; + loadNCpuMoe = 0; + loadSplitRatio = null; } + // Pinning layers on the SAME model keeps the currently resolved + // context: with no explicit pin, a manual+pinned reload would send 0, + // which the backend's --fit off branch treats as the NATIVE context -- + // far larger than the sheet shows when the load was fit-sized (Default + // or Manual + Auto layers may auto-reduce context to fit VRAM), a + // likely OOM. ggufContextLength is that resolved value; a model already + // at native reloads unchanged, so this is safe for any prior mode. + if ( + isGguf && + !switchingModelOrVariant && + loadGpuMemoryMode === "manual" && + loadGpuLayers >= 0 && + loadCustomContextLength == null && + (loadGgufContextLength ?? 0) > 0 + ) { + loadCustomContextLength = loadGgufContextLength; + } const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId, ggufVariant, @@ -720,13 +829,20 @@ export function useChatModelRuntime() { maxSeqLength, presetSource: loadActivePresetSource, }); + const loadMaxSeqLength = resolveFitMaxSeqLength( + isGguf, + loadGpuMemoryMode, + loadGpuLayers, + loadCustomContextLength, + effectiveMaxSeqLength, + ); const effectiveChatTemplateOverride = loadChatTemplateOverride?.trim() ? loadChatTemplateOverride : null; const loadResponse = await loadModel({ model_path: modelId, nativePathLease: loadNativePathLease, hf_token: hfToken, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: loadMaxSeqLength, load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, @@ -737,6 +853,11 @@ export function useChatModelRuntime() { speculative_type: loadSpeculativeType, spec_draft_n_max: loadSpecDraftNMax, tensor_parallel: loadTensorParallel, + gpu_memory_mode: loadGpuMemoryMode, + gpu_layers: loadGpuLayers, + n_cpu_moe: loadNCpuMoe, + tensor_split: loadSplitRatio ?? undefined, + gpu_ids: loadSelectedGpuIds ?? undefined, }); // If cancelled while loading, don't update UI to show @@ -747,6 +868,9 @@ export function useChatModelRuntime() { // preference now (the requested intent, not the resolved echo; // saveSpeculativeType keeps only the universal auto/ngram/off). saveSpeculativeType(loadSpeculativeType); + // Persist the GPU Memory mode only on a successful load (not on + // dropdown change), so an abandoned selection doesn't stick. + persistGpuMemoryModeOnLoad(loadResponse, loadGpuMemoryMode); const currentParams = useChatRuntimeStore.getState().params; setParams( @@ -782,9 +906,13 @@ export function useChatModelRuntime() { const reportedNativeCtx = loadResponse.is_gguf ? (loadResponse.native_context_length ?? null) : null; - // A successful reload has applied settings, so clear pending custom - // context state and display the backend-reported effective context. - const keepCustomCtx = null; + // Keep an explicit Manual+Auto context pin (so a later Apply doesn't + // revert it to Auto); other cases baseline on ggufContextLength. + const keepCustomCtx = resolveManualAutoCtxPin( + loadGpuMemoryMode, + loadGpuLayers, + loadCustomContextLength, + ); const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false; const reasoningStyle = loadResponse.reasoning_style ?? "enable_thinking"; const supportsReasoning = loadResponse.supports_reasoning ?? false; @@ -837,11 +965,13 @@ export function useChatModelRuntime() { loadedKvCacheDtype: loadedKv, tensorParallel: loadedTp, loadedTensorParallel: loadedTp, + ...loadedGpuMemoryFields(loadResponse), speculativeType: loadedSpec, loadedSpeculativeType: loadedSpec, specDraftNMax: loadResponse.spec_draft_n_max ?? null, loadedSpecDraftNMax: loadResponse.spec_draft_n_max ?? null, customContextLength: keepCustomCtx, + loadedCustomContextLength: keepCustomCtx, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: effectiveChatTemplateOverride, loadedChatTemplateOverride: effectiveChatTemplateOverride, @@ -938,7 +1068,7 @@ export function useChatModelRuntime() { } } try { - await loadModel({ + const rollbackResponse = await loadModel({ model_path: previousCheckpoint, nativePathLease: rollbackNativePathLease, hf_token: hfToken, @@ -951,14 +1081,51 @@ export function useChatModelRuntime() { // Resend the previous model's pinned approval so restoring it is not re-blocked. approved_remote_code_fingerprint: approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null, + chat_template_override: + stateBeforeUnload.loadedChatTemplateOverride, + cache_type_kv: stateBeforeUnload.loadedKvCacheDtype, + speculative_type: + stateBeforeUnload.loadedSpeculativeType, + spec_draft_n_max: + stateBeforeUnload.loadedSpecDraftNMax, // Restore the previous model in the split mode it was running, // not the default layer split. tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false, + gpu_memory_mode: stateBeforeUnload.loadedGpuMemoryMode ?? "auto", + gpu_layers: stateBeforeUnload.loadedGpuLayers ?? -1, + n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0, + tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined, + gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined, }); + const rollbackSpeculativeType = normalizeSpeculativeType( + rollbackResponse.speculative_type, + ); useChatRuntimeStore.setState({ activeNativePathToken: previousActiveNativePathToken ?? null, - loadedSpeculativeType: null, - loadedSpecDraftNMax: null, + loadedSpeculativeType: rollbackSpeculativeType, + loadedSpecDraftNMax: + rollbackResponse.spec_draft_n_max ?? null, + loadedKvCacheDtype: rollbackResponse.cache_type_kv ?? null, + loadedChatTemplateOverride: + stateBeforeUnload.loadedChatTemplateOverride, + // Re-baseline the GPU knobs from the rolled-back load's own + // response (the shared seeding every load path uses): the + // refresh() below can't do it, since the status reseed is + // gated off while modelLoading is still true. A failed staged + // Load stays staged for retry, so the staged hold applies. + ...loadedGpuMemoryFieldsUnlessStaged(rollbackResponse, { + tensorParallel: rollbackResponse.tensor_parallel ?? false, + loadedTensorParallel: + rollbackResponse.tensor_parallel ?? false, + // refresh() is held while modelLoading remains true, so + // restore the rolled-back model's context pin directly. + customContextLength: + stateBeforeUnload.loadedCustomContextLength, + }), + loadedTensorParallel: + rollbackResponse.tensor_parallel ?? false, + loadedCustomContextLength: + stateBeforeUnload.loadedCustomContextLength, }); await refresh(); } catch { diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts index d8076c720b..a3e7a2d264 100644 --- a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts +++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts @@ -7,7 +7,7 @@ import { useRepoDownload } from "@/features/hub/download-manager/use-repo-downlo import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download"; import { useLatestRef } from "@/features/hub/hooks/use-latest-ref"; -import { fetchGgufContextLength } from "../api/chat-api"; +import { fetchGgufStagedMetadata } from "../api/chat-api"; import { isPendingGguf, pendingSelectionMatches, @@ -46,8 +46,16 @@ export function useStagedModelPreparation(opts?: { const pendingDownloaded = useChatRuntimeStore( (s) => s.pendingSelection?.isDownloaded ?? false, ); - const pendingHasContext = useChatRuntimeStore( - (s) => s.pendingSelection?.contextLength != null, + // "Already probed" must key off layerCount / moeLayerCount, which only the + // full header probe fills (it sets all three together, so either is a + // reliable marker). contextLength alone can be list-seeded from + // /gguf-variants, which returns no layer/MoE counts -- treating it as + // complete would skip the probe and leave the GPU Layers slider at its 256 + // fallback and the MoE slider hidden until the model loads. + const pendingHasMetadata = useChatRuntimeStore( + (s) => + s.pendingSelection?.layerCount != null || + s.pendingSelection?.moeLayerCount != null, ); const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection); const onAutoLoadRef = useLatestRef(opts?.onAutoLoad); @@ -69,25 +77,31 @@ export function useStagedModelPreparation(opts?: { if (!current?.id || !isPendingGguf(current)) return; const { id, ggufVariant, nativePathToken } = current; try { - const contextLength = await fetchGgufContextLength({ - model_path: id, - gguf_variant: ggufVariant, - hf_token: useChatRuntimeStore.getState().hfToken || null, - nativePathToken, - }); + const { contextLength, layerCount, moeLayerCount } = + await fetchGgufStagedMetadata({ + model_path: id, + gguf_variant: ggufVariant, + hf_token: useChatRuntimeStore.getState().hfToken || null, + nativePathToken, + }); // Apply only if the same model is still staged (the user may have switched // picks or loaded/cancelled while the request was in flight). const latest = useChatRuntimeStore.getState().pendingSelection; if ( latest && - contextLength != null && - pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) + pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) && + (contextLength != null || layerCount != null || moeLayerCount != null) ) { - setPendingSelection({ ...latest, contextLength }); + setPendingSelection({ + ...latest, + contextLength, + layerCount, + moeLayerCount, + }); } } catch { - // Leave contextLength null: the context slider stays hidden and the user - // can still load (context fills in from the load response afterwards). + // Leave metadata null: the context/MoE sliders stay hidden and the user + // can still load (they fill in from the load response afterwards). } }, [setPendingSelection]); @@ -125,7 +139,7 @@ export function useStagedModelPreparation(opts?: { if ( !pendingId || (!pendingIsGguf && !pendingIsHubRepo) || - pendingHasContext + pendingHasMetadata ) { return; } @@ -146,7 +160,7 @@ export function useStagedModelPreparation(opts?: { pendingIsGguf, pendingIsHubRepo, pendingDownloaded, - pendingHasContext, + pendingHasMetadata, startDownloadRef, fetchMetadataRef, ]); diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index a4b5f848e2..69bb38bbbe 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -2,13 +2,17 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getInferenceStatus } from "../api/chat-api"; -import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { + mergeBackendRecommendedInference, + resolveManualAutoCtxPin, +} from "../presets/preset-policy"; import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, type ReasoningEffort, type ReasoningStyle, loadOptionalBool, + loadedGpuMemoryFields, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; @@ -20,6 +24,10 @@ import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; +function sameArray(a: T[] | null, b: T[] | null): boolean { + return JSON.stringify(a) === JSON.stringify(b); +} + // Canonicalises backend / persisted speculative mode values onto the UI modes. export function normalizeSpeculativeType( v: string | null | undefined, @@ -119,6 +127,10 @@ function ensureActiveModelInStoreList( export type ApplyInferenceStatusOptions = { previousCheckpoint?: string; + /** activeGgufVariant BEFORE the caller's setCheckpoint synced it to the + * status -- without it a variant-only switch underneath the tab reads as + * steady state and the hydration reseed keeps the old quant's baselines. */ + previousGgufVariant?: string | null; }; /** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */ @@ -144,9 +156,13 @@ export function applyActiveModelStatusToStore( ); } + const previousGgufVariant = + options.previousGgufVariant !== undefined + ? options.previousGgufVariant + : store.activeGgufVariant; const hydratingExistingModel = previousCheckpoint !== checkpointId || - store.activeGgufVariant !== (status.gguf_variant ?? null); + previousGgufVariant !== (status.gguf_variant ?? null); const supportsReasoning = status.supports_reasoning ?? false; const reasoningAlwaysOn = status.reasoning_always_on ?? false; const reasoningStyle = status.reasoning_style ?? "enable_thinking"; @@ -185,6 +201,66 @@ export function applyActiveModelStatusToStore( // While a load is in flight, performLoad owns the load params. Seeding them // from a stale poll here would clobber the values the load dialog just set. const seedLoadParams = !prevState.modelLoading; + // A Manual + Auto-layers load sent its positive context pin as max_seq_length, + // and status only exposes the RESOLVED context; re-seed the pin from the + // requested value (parity with the load paths' keepCustomCtx). Baselines + // unconditionally: anything but an applicable pin is null, so a previous + // model's pin can't survive a model change underneath and reload at the old length. + const gpuPin = status.is_gguf + ? resolveManualAutoCtxPin( + status.gpu_memory_mode ?? "auto", + status.gpu_layers ?? -1, + status.requested_context_length ?? null, + ) + : null; + const incomingGpuMode = status.is_gguf + ? (status.gpu_memory_mode ?? "auto") + : null; + const incomingGpuLayers = + incomingGpuMode === "manual" ? (status.gpu_layers ?? null) : null; + const incomingNCpuMoe = + incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null; + const incomingSplit = + incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null; + const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null; + const gpuStatusChanged = + prevState.loadedGpuMemoryMode !== incomingGpuMode || + prevState.loadedGpuLayers !== incomingGpuLayers || + prevState.loadedNCpuMoe !== incomingNCpuMoe || + !sameArray(prevState.loadedSplitRatio, incomingSplit) || + !sameArray(prevState.loadedGpuIds, incomingGpuIds) || + prevState.loadedCustomContextLength !== gpuPin; + const gpuMemoryEditsPending = + (prevState.loadedGpuMemoryMode !== null && + prevState.gpuMemoryMode !== prevState.loadedGpuMemoryMode) || + (prevState.loadedGpuMemoryMode === "manual" && + (prevState.gpuLayers !== prevState.loadedGpuLayers || + prevState.nCpuMoe !== prevState.loadedNCpuMoe || + !sameArray(prevState.splitRatio, prevState.loadedSplitRatio))) || + prevState.customContextLength !== prevState.loadedCustomContextLength; + const gpuIdsEditPending = !sameArray( + prevState.selectedGpuIds, + prevState.loadedGpuIds, + ); + const incomingGpuFields = loadedGpuMemoryFields(status); + // A same-model reload from another client advances every loaded baseline. + // Preserve each editable group only when this tab has an unapplied change. + const preserveSameModelEdits = gpuStatusChanged && !hydratingExistingModel; + const gpuStatusFields = { + ...incomingGpuFields, + customContextLength: gpuPin, + loadedCustomContextLength: gpuPin, + ...(preserveSameModelEdits && + gpuMemoryEditsPending && { + gpuMemoryMode: prevState.gpuMemoryMode, + gpuLayers: prevState.gpuLayers, + nCpuMoe: prevState.nCpuMoe, + splitRatio: prevState.splitRatio, + customContextLength: prevState.customContextLength, + }), + ...(preserveSameModelEdits && + gpuIdsEditPending && { selectedGpuIds: prevState.selectedGpuIds }), + }; useChatRuntimeStore.setState({ supportsReasoning, @@ -215,30 +291,51 @@ export function applyActiveModelStatusToStore( loadedIsMultimodal: isMultimodalResponse(status), loadedIsDiffusion: status.is_diffusion ?? false, specFallbackReason: status.spec_fallback_reason ?? null, + // The spec / KV seeds share the GPU-fields reseed mechanism below: a + // non-GGUF status leaves their loaded baselines null, so the "unseeded" + // guard re-fires every refresh -- hold them too while a staged pick's + // settings are being edited, or the refresh resets the staged edit. + // hydratingExistingModel reopens every load-param seed: when the active + // model changed underneath this tab (auto-switch, another client), the + // old model's baselines are stale and must adopt the new status. ...(seedLoadParams && - prevState.loadedSpeculativeType === null && { + prevState.pendingSelection == null && + (prevState.loadedSpeculativeType === null || hydratingExistingModel) && { speculativeType: currentSpecType, loadedSpeculativeType: currentSpecType, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.spec_draft_n_max !== undefined && - prevState.loadedSpecDraftNMax === null && - prevState.specDraftNMax === null && { + (hydratingExistingModel || + (prevState.loadedSpecDraftNMax === null && + prevState.specDraftNMax === null)) && { specDraftNMax: status.spec_draft_n_max ?? null, loadedSpecDraftNMax: status.spec_draft_n_max ?? null, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.cache_type_kv !== undefined && - prevState.loadedKvCacheDtype === null && { + (prevState.loadedKvCacheDtype === null || hydratingExistingModel) && { kvCacheDtype: status.cache_type_kv, loadedKvCacheDtype: status.cache_type_kv, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.tensor_parallel !== undefined && - prevState.loadedTensorParallel === null && { + (prevState.loadedTensorParallel === null || hydratingExistingModel) && { tensorParallel: status.tensor_parallel, loadedTensorParallel: status.tensor_parallel, }), + // Re-seed on first hydration, model/variant changes, or a same-model backend + // placement change. gpuStatusFields preserves dirty local edits in the last + // case while advancing their loaded baselines. + ...(seedLoadParams && + prevState.pendingSelection == null && + (prevState.loadedGpuMemoryMode === null || + hydratingExistingModel || + gpuStatusChanged) && + gpuStatusFields), ...(status.chat_template_override !== undefined && prevState.loadedChatTemplateOverride === null && prevState.chatTemplateOverride === null && { @@ -298,7 +395,11 @@ export async function tryAdoptServerActiveModel(): Promise { if (previousCheckpoint) { return true; } + const previousGgufVariant = useChatRuntimeStore.getState().activeGgufVariant; store.setCheckpoint(checkpointId, status.gguf_variant); - applyActiveModelStatusToStore(status, { previousCheckpoint }); + applyActiveModelStatusToStore(status, { + previousCheckpoint, + previousGgufVariant, + }); return true; } diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index f96ee91f1b..23d79a35e1 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -339,3 +339,34 @@ export function resolveLoadMaxSeqLength({ } return maxSeqLength; } + +/** + * Adjust a resolved max-seq-length for the GPU Memory mode. Under Manual + Auto + * layers (GGUF, gpuLayers < 0) llama.cpp's --fit owns context sizing, so send 0 + * (the backend omits -c) unless the user pinned a length; every other case keeps + * the resolved fallback. Shared by every GGUF load path so they can't drift. + */ +export function resolveFitMaxSeqLength( + isGguf: boolean | null | undefined, + gpuMemoryMode: "auto" | "manual", + gpuLayers: number, + customContextLength: number | null, + fallback: number, +): number { + if (!isGguf || gpuMemoryMode !== "manual" || gpuLayers >= 0) return fallback; + return customContextLength && customContextLength > 0 ? customContextLength : 0; +} + +// A Manual + Auto-layers load sends its positive context pin as max_seq_length; +// keep it across a status reseed/Apply so the model isn't reverted to auto-fit +// sizing. Anything else (Auto mode, pinned layers, no pin) baselines to null. +// The caller keeps its own isGguf/targetIsGguf guard inline. +export function resolveManualAutoCtxPin( + gpuMemoryMode: "auto" | "manual", + gpuLayers: number, + customContextLength: number | null, +): number | null { + return gpuMemoryMode === "manual" && gpuLayers < 0 && (customContextLength ?? 0) > 0 + ? customContextLength + : null; +} diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a0813fe27b..31e50ee60c 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -84,6 +84,8 @@ import { useTransformersUpgradeDialogStore, } from "@/features/transformers-upgrade"; import { loadModel, validateModel } from "./api/chat-api"; +import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "./presets/preset-policy"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { parseExternalModelId, providerTypeSupportsVision, @@ -95,8 +97,11 @@ import { usePlusMenuPrefsStore, } from "./stores/plus-menu-prefs-store"; import { + loadedGpuMemoryFieldsUnlessStaged, type ReasoningEffort, + reconcilePersistedGpuIds, resolveLoadedSpeculativeSettings, + persistGpuMemoryModeOnLoad, resolveSpeculativeSettingsForLoad, saveSpeculativeType, useChatRuntimeStore, @@ -1037,10 +1042,32 @@ export function SharedComposer({ return parts[parts.length - 1] || id; } + // Warm the device cache before the snapshot below reconciles the GPU + // pick: on a cold cache the reconcile passes a stale pick through. + if (store.selectedGpuIds != null) { + await ensureGpuDeviceCache(); + } + // The GPU/offload knobs both compare loads must use, snapshotted at Send. + // ensureModelLoaded runs sequentially and the first load's response echo + // (loadedGpuMemoryFields) rewrites the live store -- a non-GGUF or Auto + // first model resets gpuLayers/nCpuMoe/split/pick to defaults -- so + // reading the store per load would hand model 2 the first model's echoed + // defaults instead of the settings the user pressed Send with. + const compareLoadKnobs = { + gpuMemoryMode: store.gpuMemoryMode, + gpuLayers: store.gpuLayers, + nCpuMoe: store.nCpuMoe, + splitRatio: store.splitRatio, + // Reconcile the pick against the GPUs present now, like the model-switch + // path: an early remember-restore can hold a stale cross-host pick that + // /load would reject (the device cache is populated by send time). + selectedGpuIds: reconcilePersistedGpuIds(store.selectedGpuIds), + tensorParallel: store.tensorParallel, + customContextLength: store.customContextLength, + }; // Set when an accepted transformers install unloaded the active model // server-side; a later failure must then clear the stale checkpoint. let upgradeUnloadedActive = false; - // Helper: load a model and update store checkpoint async function ensureModelLoaded( sel: CompareModelSelection, @@ -1057,15 +1084,35 @@ export function SharedComposer({ if (isAlreadyActive) { return "ready"; } + const targetIsGguf = + sel.id.toLowerCase().endsWith(".gguf") || sel.ggufVariant != null; + // Size validation exactly as the load below, so the training-guard + // preflight checks the footprint that actually loads (under Manual + Auto + // layers the load sends 0 / the pinned context, not raw maxSeqLength). + const compareMaxSeqLength = resolveFitMaxSeqLength( + targetIsGguf, + compareLoadKnobs.gpuMemoryMode, + compareLoadKnobs.gpuLayers, + compareLoadKnobs.customContextLength, + maxSeqLength, + ); const validation = await validateModel({ model_path: sel.id, hf_token: currentStore.hfToken || null, - max_seq_length: maxSeqLength, + max_seq_length: compareMaxSeqLength, load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, trust_remote_code: loadTrustRemoteCode, chat_template_override: effectiveChatTemplateOverride, + // Scope the validate to the picked GPUs. GGUF-only, like the load + // below: a non-GGUF target must not inherit a hidden GGUF GPU pick. + ...(targetIsGguf + ? { + gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined, + gpu_memory_mode: compareLoadKnobs.gpuMemoryMode, + } + : {}), }); // Upgrade dialog first (mirrors the primary load path). if (validation.requires_transformers_upgrade) { @@ -1114,7 +1161,7 @@ export function SharedComposer({ const resp = await loadModel({ model_path: sel.id, hf_token: useChatRuntimeStore.getState().hfToken || null, - max_seq_length: maxSeqLength, + max_seq_length: compareMaxSeqLength, load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, @@ -1123,10 +1170,25 @@ export function SharedComposer({ chat_template_override: effectiveChatTemplateOverride, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, - // Honor the Tensor Parallelism toggle on compare loads too. - tensor_parallel: currentStore.tensorParallel, + // Honor the Tensor Parallelism + GPU Memory choices on compare loads. + // GGUF-only, like the auto-load path: the picker is a GGUF control, + // so a non-GGUF target loads via HF auto-placement instead of being + // pinned to a leftover GGUF pick it can't even show. + tensor_parallel: compareLoadKnobs.tensorParallel, + ...(targetIsGguf + ? { + gpu_memory_mode: compareLoadKnobs.gpuMemoryMode, + gpu_layers: compareLoadKnobs.gpuLayers, + n_cpu_moe: compareLoadKnobs.nCpuMoe, + tensor_split: compareLoadKnobs.splitRatio ?? undefined, + gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined, + } + : {}), }); saveSpeculativeType(specSettings.speculativeType); + // Persist the GPU Memory mode on a non-diffusion GGUF compare-load too, + // so an applied manual choice survives a restart. + persistGpuMemoryModeOnLoad(resp, compareLoadKnobs.gpuMemoryMode); upgradeUnloadedActive = false; const store = useChatRuntimeStore.getState(); store.setCheckpoint( @@ -1136,6 +1198,17 @@ export function SharedComposer({ store.setModelRequiresTrustRemoteCode( resp.requires_trust_remote_code ?? false, ); + // Keep an explicit Manual+Auto context pin the load just applied (so a + // later Apply/Reset doesn't silently revert the model to auto-fit + // sizing), mirroring the interactive path's keepCustomCtx. Non-GGUF + // compare loads don't send the pin, so their baseline clears. + const keepCustomCtx = targetIsGguf + ? resolveManualAutoCtxPin( + compareLoadKnobs.gpuMemoryMode, + compareLoadKnobs.gpuLayers, + compareLoadKnobs.customContextLength, + ) + : null; useChatRuntimeStore.setState({ supportsReasoning: resp.supports_reasoning ?? false, reasoningAlwaysOn: resp.reasoning_always_on ?? false, @@ -1144,6 +1217,32 @@ export function SharedComposer({ supportsTools: resp.supports_tools ?? false, tensorParallel: resp.tensor_parallel ?? false, loadedTensorParallel: resp.tensor_parallel ?? false, + customContextLength: keepCustomCtx, + loadedCustomContextLength: keepCustomCtx, + // Seed the loaded GGUF context (interactive/auto-load parity): the + // settings sheet keys the GGUF GPU controls off it for a direct .gguf + // with no variant, and a later Apply reads it as the resolved context. + ...(targetIsGguf + ? { + ggufContextLength: resp.context_length ?? 131072, + ggufMaxContextLength: + resp.max_context_length ?? resp.context_length ?? 131072, + ggufNativeContextLength: resp.native_context_length ?? null, + } + : { ggufContextLength: null }), + // Compare loads resolve by id (HF repo / local path), never through a + // native-path lease, so a token left by a previously loaded native + // GGUF is stale here -- isLoadedGguf keys off it, and a stale token + // would dress a non-GGUF compare load in GGUF controls. Mirror the + // interactive path, which writes it on every load success. + activeNativePathToken: null, + // Held under an open staged pick: setCheckpoint preserves a stage on + // the empty->active transition, so a compare load can complete with + // staged GPU edits still on screen. + ...loadedGpuMemoryFieldsUnlessStaged(resp), + // Drives the GPU Memory controls' diffusion gate; set alongside the + // GPU fields on every load path so the gate can't read stale. + loadedIsDiffusion: resp.is_diffusion ?? false, loadedIsMultimodal: isMultimodalResponse(resp), ...resolveLoadedSpeculativeSettings(resp), }); diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 192ce1ec69..5786947118 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -7,6 +7,10 @@ import { mirrorHfTokenInto, useHfTokenStore, } from "@/features/hub"; +import { + cachedPinnableGpuIndices, + ensureGpuDeviceCache, +} from "@/hooks/use-gpu-info"; import { toast } from "@/lib/toast"; import { create } from "zustand"; import { isExternalModelId, parseExternalModelId } from "../external-providers"; @@ -74,6 +78,7 @@ export const CHAT_RAG_AUTOINJECT_MIN_SCORE_KEY = export const CHAT_RAG_OCR_KEY = "unsloth_chat_rag_ocr_scanned"; export const CHAT_RAG_CAPTION_KEY = "unsloth_chat_rag_caption_figures"; export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type"; +export const CHAT_GPU_MEMORY_MODE_KEY = "unsloth_chat_gpu_memory_mode"; // Persist only the model-agnostic intents (auto/ngram/off). MTP modes // (mtp/mtp+ngram) and spec_draft_n_max stay session-only: a persisted MTP @@ -497,6 +502,213 @@ export function saveSpeculativeType(value: string | null): void { } } +// GPU Memory strategy is a standing preference (like speculative type), not a +// per-model setting: a "manual" choice persists across model switches and reloads. +export function readPersistedGpuMemoryMode(): "auto" | "manual" { + return loadString(CHAT_GPU_MEMORY_MODE_KEY, "auto") === "manual" ? "manual" : "auto"; +} + +export function saveGpuMemoryMode(value: "auto" | "manual"): void { + saveString(CHAT_GPU_MEMORY_MODE_KEY, value); +} + +/** Persist the GPU Memory mode after a load, but only for a non-diffusion GGUF: + * non-GGUF has no such mode, and diffusion runs mode-agnostic (reports "auto"), + * so neither must clobber the standing manual preference. */ +export function persistGpuMemoryModeOnLoad( + resp: { is_gguf?: boolean; is_diffusion?: boolean }, + mode: "auto" | "manual", +): void { + if (resp.is_gguf && !resp.is_diffusion) saveGpuMemoryMode(mode); +} + +// Manual-mode gpu_layers sentinel: -1 = Auto (hand layer + context sizing to +// llama.cpp's --fit). The Manual default; "all on GPU" is the slider's max. +export const GPU_LAYERS_AUTO = -1; + +// Round real-valued shares to integers summing exactly to `total`, giving the +// leftover units to the largest fractional parts (largest-remainder method). +function largestRemainder(shares: number[], total: number): number[] { + const out = shares.map((x) => Math.floor(x)); + let rem = total - out.reduce((a, b) => a + b, 0); + const byFrac = shares + .map((x, i) => ({ i, frac: x - Math.floor(x) })) + .sort((a, b) => b.frac - a.frac); + for (let k = 0; rem > 0 && k < byFrac.length; k++, rem--) out[byFrac[k].i] += 1; + return out; +} + +// Spread `total` layers across GPUs in proportion to `weights` (e.g. per-GPU +// VRAM), as integers summing exactly to `total`; even split for all-zero/empty +// weights. Default per-GPU layer split before the user edits it (mirrors +// llama.cpp's free-VRAM default). +export function distributeByWeight(total: number, weights: number[]): number[] { + if (weights.length === 0) return []; + const t = Math.max(0, Math.floor(total)); + const sum = weights.reduce((a, b) => a + b, 0); + const w = sum > 0 ? weights : weights.map(() => 1); + const wSum = w.reduce((a, b) => a + b, 0); + return largestRemainder( + w.map((x) => (t * x) / wSum), + t, + ); +} + +// Set GPU `index` to `value` and rebalance the rest so per-GPU counts still sum +// to `total`; others absorb the remainder in proportion to their counts (evenly +// if all zero). The --tensor-split editor: counts are sent verbatim, and +// llama.cpp gives each GPU exactly its count when gpu_layers == sum(counts). +export function rebalanceSplit( + total: number, + counts: number[], + index: number, + value: number, +): number[] { + const v = Math.max(0, Math.min(value, total)); + const out = counts.slice(); + const otherIdx = counts.map((_, i) => i).filter((i) => i !== index); + // No other GPU to absorb the remainder: this one holds everything. + if (otherIdx.length === 0) { + out[index] = total; + return out; + } + out[index] = v; + const dist = distributeByWeight( + total - v, + otherIdx.map((i) => counts[i]), + ); + otherIdx.forEach((i, k) => (out[i] = dist[k])); + return out; +} + +// Validate a persisted gpu_ids pick against the GPUs present right now, before +// restoring it from remembered settings. Returns null (= automatic) when the +// pick is stale (none of the saved ids exist, or the host can't pin a multi-GPU +// set), so a saved [1] on a now-1-GPU host doesn't get sent and rejected with no +// way to clear it. A null pick (= automatic) passes through unchanged, and an +// unpopulated device cache leaves the pick alone (the backend still guards). +export function reconcilePersistedGpuIds( + ids: number[] | null, +): number[] | null { + if (ids == null) return ids; + const pinnable = cachedPinnableGpuIndices(); + if (pinnable === null) return ids; // cache not ready: can't validate, keep it + const kept = ids.filter((i) => pinnable.includes(i)); + return kept.length > 0 ? kept : null; +} + +// Store fields derived from a load/status response's GPU-memory settings. +// Shared by every load path so the manual-knob round-trip can't drift. +export function loadedGpuMemoryFields(resp: { + is_gguf?: boolean; + is_diffusion?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + n_layers?: number | null; + n_moe_layers?: number; + gpu_ids?: number[] | null; +}) { + // GPU-memory state is meaningful only for a GGUF chat load. A non-GGUF response + // still carries gpu_memory_mode (its default "auto" is serialized), so gate on + // the authoritative is_gguf flag, not the field's presence -- otherwise loading + // a transformers model would reset the standing manual preference. + if (!resp.is_gguf) { + // Clear the GPU pick / offload baseline a prior GGUF load may have left, so it + // reflects the non-GGUF model (no pin) -- else a stale loadedGpuIds reads as + // dirty (gpuIdsDirty is ungated) and Reset restores it while the picker is + // hidden. gpuMemoryMode (the standing preference) is kept, but its loaded + // baseline clears to null so Reset preserves the preference, not a stale mode. + return { + selectedGpuIds: null, + loadedGpuIds: null, + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + }; + } + const mode = resp.gpu_memory_mode ?? "auto"; + const gpuIds = resp.gpu_ids ?? null; + // Layer/MoE/split knobs apply (and are reported) only in manual mode; in auto + // the server ignores them, so don't seed the loaded baseline or the editable + // knobs with values it never applied. In manual, the server reports gpu_layers + // = -1 under Auto, which round-trips the slider back to its Auto position. + const manualKnobs = + mode === "manual" + ? { + loadedGpuLayers: resp.gpu_layers ?? null, + loadedNCpuMoe: resp.n_cpu_moe ?? null, + loadedSplitRatio: resp.tensor_split ?? null, + gpuLayers: resp.gpu_layers ?? GPU_LAYERS_AUTO, + nCpuMoe: resp.n_cpu_moe ?? 0, + splitRatio: resp.tensor_split ?? null, + } + : { + loadedGpuLayers: null, + loadedNCpuMoe: null, + loadedSplitRatio: null, + // Auto ignores these, so reset the editable knobs too (not just the + // loaded baseline) -- else a later switch back to Manual would snapshot + // and send a previous model's stale gpuLayers/nCpuMoe/split that this + // load never applied. Mirrors the non-GGUF branch above. + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + }; + return { + // A diffusion GGUF runs mode-agnostic (pins all layers on one GPU, reports + // "auto"), so adopt everything a chat GGUF does EXCEPT the live standing + // preference -- the next chat load must still honor the user's manual choice. + // The loaded baseline is still "auto", but the UI hides mode controls for a + // loaded diffusion model so it can't read as dirty against the preference. + ...(resp.is_diffusion ? {} : { gpuMemoryMode: mode }), + loadedGpuMemoryMode: mode, + ggufLayerCount: resp.n_layers ?? null, + // MoE expert-layer count: the n_cpu_moe slider max, and 0 hides the slider. + moeLayerCount: resp.n_moe_layers ?? null, + // The picker reflects what loaded (the request sent the user's pick). + selectedGpuIds: gpuIds, + loadedGpuIds: gpuIds, + ...manualKnobs, + }; +} + +/** loadedGpuMemoryFields (plus any seedExtras), unless a staged pick is open. + * + * With a staged pick open (the load fired mid-staging), preserve its editable + * GPU knobs and seedExtras, but still advance every loaded baseline. Otherwise + * cancelling the stage restores its edits onto the newly loaded model. The + * status reseed cannot repair that while pendingSelection holds it off. + */ +export function loadedGpuMemoryFieldsUnlessStaged( + resp: Parameters[0], + seedExtras?: T, +) { + const fields = loadedGpuMemoryFields(resp); + if (useChatRuntimeStore.getState().pendingSelection != null) { + return { + loadedGpuMemoryMode: fields.loadedGpuMemoryMode, + loadedGpuLayers: fields.loadedGpuLayers, + loadedNCpuMoe: fields.loadedNCpuMoe, + loadedSplitRatio: fields.loadedSplitRatio, + loadedGpuIds: fields.loadedGpuIds, + // These are metadata ceilings for the model that actually loaded, not + // editable values from the open stage. Advance them with the baselines + // so abandoning the stage cannot expose the previous model's limits. + ggufLayerCount: fields.ggufLayerCount, + moeLayerCount: fields.moeLayerCount, + }; + } + return { ...fields, ...seedExtras }; +} + /** A local model staged for a deferred load (see `pendingSelection`). Shape is * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ export type PendingModelSelection = { @@ -515,6 +727,13 @@ export type PendingModelSelection = { * Scoped here (not the shared `ggufContextLength`) so a staged model's * metadata never pollutes the currently-loaded model's context display. */ contextLength?: number | null; + /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (llama.cpp counts the output layer as offloadable too); + * scoped here like contextLength. */ + layerCount?: number | null; + /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling); + * 0 for dense models, scoped here like contextLength. */ + moeLayerCount?: number | null; /** "Load on selection" on + un-cached GGUF: download via the manager (global * indicator) without opening the sheet, then load once the download finishes. */ autoLoad?: boolean; @@ -743,6 +962,32 @@ type ChatRuntimeStore = { tensorParallel: boolean; /** Backend-reported tensor-parallel state; null until first hydrated. */ loadedTensorParallel: boolean | null; + /** GPU memory strategy for GGUF loads. "auto" = Unsloth picks GPUs and context + * to fit; "manual" = you own the offload (gpuLayers < 0 = Auto/--fit, >= 0 + * pins layers + nCpuMoe). */ + gpuMemoryMode: "auto" | "manual"; + /** Backend-reported gpu memory mode; null until first hydrated. */ + loadedGpuMemoryMode: "auto" | "manual" | null; + /** Manual mode: layers to offload to GPU. -1 = Auto (--fit); >= model layer + * count = all. */ + gpuLayers: number; + loadedGpuLayers: number | null; + /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */ + nCpuMoe: number; + loadedNCpuMoe: number | null; + /** Manual mode: per-GPU layer counts (--tensor-split), in GPU-in-use order; + * null = unset (llama.cpp splits by free VRAM). */ + splitRatio: number[] | null; + /** Backend-reported per-GPU split ratio (--tensor-split); null = unset. */ + loadedSplitRatio: number[] | null; + /** Model layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (the output layer is offloadable too). */ + ggufLayerCount: number | null; + /** MoE expert-layer count: the nCpuMoe slider max; 0/null hides the slider. */ + moeLayerCount: number | null; + /** Picked physical GPU indices (null = use all / automatic). */ + selectedGpuIds: number[] | null; + loadedGpuIds: number[] | null; /** Persisted: when false, picking a local model stages it as * `pendingSelection` (and opens settings) instead of loading immediately, * so load settings can be set before the single load. */ @@ -766,6 +1011,9 @@ type ChatRuntimeStore = { * per step, cleared when the run ends, never persisted into the transcript. */ activeDiffusionCanvas: DiffusionCanvasFrame | null; customContextLength: number | null; + /** The pinned context the loaded model used (null = Auto), so dirty-tracking + * and a later fit Apply can tell an explicit pin apart from Auto. */ + loadedCustomContextLength: number | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; loadedChatTemplateOverride: string | null; @@ -884,6 +1132,11 @@ type ChatRuntimeStore = { * which skip the sheet but must still honor a saved config. */ applyRememberedLoadSettings: (settings: RememberedLoadSettings) => void; setTensorParallel: (value: boolean) => void; + setGpuMemoryMode: (mode: "auto" | "manual") => void; + setGpuLayers: (value: number) => void; + setNCpuMoe: (value: number) => void; + setSplitRatio: (value: number[] | null) => void; + setSelectedGpuIds: (ids: number[] | null) => void; setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; @@ -1101,11 +1354,12 @@ function setScalarSettingVersion( /** The "revert to the loaded model" baseline for the editable load knobs. * Shared by resetModelSettingsToLoaded (full revert) and stageModel (which - * overrides speculative to start a fresh pick from the standing default). */ + * overrides speculative and the per-model GPU knobs to start a fresh pick). */ function loadedBaselineSettings(s: ChatRuntimeStore) { const hasLoadedModel = Boolean(s.params.checkpoint); return { - customContextLength: null, + // Revert to the loaded model's pin (null = Auto), not a blanket Auto. + customContextLength: s.loadedCustomContextLength, kvCacheDtype: s.loadedKvCacheDtype, tensorParallel: s.loadedTensorParallel ?? false, speculativeType: hasLoadedModel @@ -1113,6 +1367,20 @@ function loadedBaselineSettings(s: ChatRuntimeStore) { : readPersistedSpeculativeType(), specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null, chatTemplateOverride: s.loadedChatTemplateOverride, + // GPU memory mode is a standing preference; revert to the loaded model's + // mode (or the persisted default when nothing is loaded). Manual knobs and + // the GPU pick are per-model and revert to their loaded baseline. A loaded + // model with no applicable mode -- diffusion ("auto" baseline) or non-GGUF + // (null baseline) -- keeps the live preference so Reset can't drop it. + gpuMemoryMode: !hasLoadedModel + ? readPersistedGpuMemoryMode() + : s.loadedIsDiffusion + ? s.gpuMemoryMode + : (s.loadedGpuMemoryMode ?? s.gpuMemoryMode), + gpuLayers: s.loadedGpuLayers ?? GPU_LAYERS_AUTO, + nCpuMoe: s.loadedNCpuMoe ?? 0, + splitRatio: s.loadedSplitRatio ?? null, + selectedGpuIds: s.loadedGpuIds, }; } @@ -1213,6 +1481,18 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedSpecDraftNMax: null, tensorParallel: false, loadedTensorParallel: null, + gpuMemoryMode: readPersistedGpuMemoryMode(), + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + selectedGpuIds: null, + loadedGpuIds: null, loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), @@ -1221,6 +1501,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedIsMultimodal: false, loadedIsDiffusion: false, customContextLength: null, + loadedCustomContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1455,9 +1736,23 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedSpecDraftNMax: null, tensorParallel: false, loadedTensorParallel: null, + // Standing preference: survives unload, unlike the per-model knobs above. + gpuMemoryMode: readPersistedGpuMemoryMode(), + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + selectedGpuIds: null, + loadedGpuIds: null, loadedIsMultimodal: false, loadedIsDiffusion: false, customContextLength: null, + loadedCustomContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1753,17 +2048,67 @@ export const useChatRuntimeStore = create((set, get) => ({ setSpeculativeType: (speculativeType) => set({ speculativeType }), setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }), setTensorParallel: (tensorParallel) => set({ tensorParallel }), + // Standing preference, but persisted only on a successful load (see + // use-chat-model-runtime), not on selection -- so an unapplied pick the user + // resets/abandons doesn't stick to the next session. + setGpuMemoryMode: (gpuMemoryMode) => set({ gpuMemoryMode }), + setGpuLayers: (gpuLayers) => set({ gpuLayers }), + setNCpuMoe: (nCpuMoe) => set({ nCpuMoe }), + setSplitRatio: (splitRatio) => set({ splitRatio }), + setSelectedGpuIds: (selectedGpuIds) => set({ selectedGpuIds }), resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)), - applyRememberedLoadSettings: (settings) => + applyRememberedLoadSettings: (settings) => { + const gpuCacheWasCold = cachedPinnableGpuIndices() === null; + const restoredGpuIds = + settings.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(settings.selectedGpuIds) + : undefined; // Coalesce every field: a blob persisted by an older/newer build can omit // keys, and a raw spread would push `undefined` into fields typed non-null. + // The GPU knobs are spread only when present, but first reset the per-model + // ones to defaults: this path (load-on-selection) starts from the loaded + // model's baseline and skips the model-switch reset, so a blob omitting + // gpuLayers/nCpuMoe/selectedGpuIds (older build) or splitRatio (never + // remembered) must not inherit the previous model's placement. gpuMemoryMode + // (standing preference) is NOT reset, only applied when the blob carries it; + // selectedGpuIds keeps a meaningful null (all GPUs), so it keys off undefined. set({ + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + selectedGpuIds: null, customContextLength: settings.contextLength ?? null, kvCacheDtype: settings.kvCacheDtype ?? null, speculativeType: settings.speculativeType ?? "auto", specDraftNMax: settings.specDraftNMax ?? null, tensorParallel: settings.tensorParallel ?? false, - }), + ...(settings.gpuMemoryMode != null && { + gpuMemoryMode: settings.gpuMemoryMode, + }), + ...(settings.gpuLayers != null && { gpuLayers: settings.gpuLayers }), + ...(settings.nCpuMoe != null && { nCpuMoe: settings.nCpuMoe }), + ...(restoredGpuIds !== undefined && { + // Reconcile against the GPUs present now (see reconcilePersistedGpuIds): + // a saved [1] on a 1-GPU host (or under relative/UUID visibility) would + // hide the picker yet still send gpu_ids, which the backend rejects. + selectedGpuIds: restoredGpuIds, + }), + }); + // A cold cache makes the synchronous restore provisional. Reconcile again + // when the shared fetch completes, but only if this exact restored array is + // still current so a user edit, stage change, or load cannot be overwritten. + if (gpuCacheWasCold && restoredGpuIds != null) { + void ensureGpuDeviceCache().then(() => { + set((state) => { + if (state.selectedGpuIds !== restoredGpuIds) return state; + const reconciled = reconcilePersistedGpuIds(restoredGpuIds); + return reconciled === restoredGpuIds + ? state + : { selectedGpuIds: reconciled }; + }); + }); + } + }, setLoadOnSelection: (loadOnSelection) => { saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection); set({ loadOnSelection }); @@ -1798,6 +2143,22 @@ export const useChatRuntimeStore = create((set, get) => ({ // Load's keepSpeculative) a forced MTP mode onto a model that may lack it. speculativeType: readPersistedSpeculativeType(), specDraftNMax: null, + // Keep the on-screen GPU Memory selection (loadedBaselineSettings would + // otherwise revert it to the loaded model's mode, dropping a Manual choice + // just made). Use the live store value, not the persisted one, which can + // lag a mode hydrated from an out-of-band load. + gpuMemoryMode: s.gpuMemoryMode, + // Per-model GPU knobs start from defaults too so a fresh pick doesn't + // inherit the loaded model's layer/MoE/split/GPU choices, matching the + // immediate-switch reset. + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + selectedGpuIds: null, + // Fresh pick starts at Auto context (loadedBaselineSettings would + // otherwise restore the current model's pin). Leaves the baseline + // intact, like the GPU knobs, so abandoning restores the loaded pin. + customContextLength: null, }; }); }, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index d72c406fdd..c24ddde5f5 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -65,6 +65,18 @@ export interface LoadModelRequest { * of by layer for GGUF models. Multi-GPU only; no effect on a single GPU. */ tensor_parallel?: boolean | null; + /** GPU memory strategy for GGUF models. "auto" (default): Unsloth selects GPUs + * and caps context to fit VRAM. "manual": you own the offload -- gpu_layers + * -1 (Auto) hands sizing to llama.cpp's --fit, >= 0 pins layers/n_cpu_moe. */ + gpu_memory_mode?: "auto" | "manual"; + /** Manual mode: layers to offload to GPU (--gpu-layers, --fit off); -1 = Auto (--fit). */ + gpu_layers?: number; + /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */ + n_cpu_moe?: number; + /** Manual mode: relative model share per GPU (--tensor-split), in GPU order. */ + tensor_split?: number[] | null; + /** Picked physical GPU indices (omit/empty = automatic). */ + gpu_ids?: number[]; } export interface ValidateModelResponse { @@ -80,6 +92,13 @@ export interface ValidateModelResponse { requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; + /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (llama.cpp counts the output layer as offloadable too); null + * until downloaded. */ + layer_count?: number | null; + /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling); + * 0 for dense models, null until downloaded. */ + moe_layer_count?: number | null; /** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */ requires_transformers_upgrade?: boolean; /** Set only when requires_transformers_upgrade. */ @@ -159,6 +178,14 @@ export interface LoadModelResponse { spec_draft_n_max?: number | null; /** Whether tensor-parallel split (--split-mode tensor) is active. */ tensor_parallel?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + n_layers?: number | null; + /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ + n_moe_layers?: number; + gpu_ids?: number[] | null; } export interface UnloadModelRequest { @@ -203,6 +230,17 @@ export interface InferenceStatusResponse { spec_draft_n_max?: number | null; /** Whether tensor-parallel split (--split-mode tensor) is active. */ tensor_parallel?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a + * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */ + requested_context_length?: number | null; + gpu_ids?: number[] | null; + n_layers?: number | null; + /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ + n_moe_layers?: number; /** * Why MTP was disabled on the loaded model despite being requested. * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable diff --git a/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts b/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts new file mode 100644 index 0000000000..f78679310f --- /dev/null +++ b/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts @@ -0,0 +1,46 @@ +// 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 { loadEmbeddingModelSettings } from "@/features/settings"; +import { useEffect, useState } from "react"; +import { useInventoryVersion } from "../stores/inventory-events"; + +/** Backend-resolved embedding repos that optimistic inventory rows must hide. */ +export function useHiddenEmbeddingModelIds( + enabled: boolean, +): ReadonlySet { + const inventoryVersion = useInventoryVersion(); + const [hiddenIds, setHiddenIds] = useState>( + () => new Set(), + ); + + // biome-ignore lint/correctness/useExhaustiveDependencies: inventory invalidation must reload backend-resolved embedder ids + useEffect(() => { + if (!enabled) { + return; + } + let cancelled = false; + loadEmbeddingModelSettings() + .then((settings) => { + if (cancelled) { + return; + } + setHiddenIds( + new Set( + [ + settings.embeddingModel, + settings.embeddingGgufRepo, + settings.defaultEmbeddingModel, + settings.defaultEmbeddingGgufRepo, + ].map((value) => value.trim().toLowerCase()), + ), + ); + }) + .catch(() => undefined); + return () => { + cancelled = true; + }; + }, [enabled, inventoryVersion]); + + return hiddenIds; +} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index 630daa48ad..d57f9636fe 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -63,6 +63,7 @@ import { useDiscoverSearch } from "./hooks/use-discover-search"; import { useFeedWriteBack } from "./hooks/use-feed-write-back"; import { useHubFeed } from "./hooks/use-hub-feed"; import { useHubModelVram } from "./hooks/use-hub-model-vram"; +import { useHiddenEmbeddingModelIds } from "./hooks/use-hidden-embedding-models"; import { useModelsSelection } from "./hooks/use-models-selection"; import { CHANNEL_TO_SECTION, @@ -73,7 +74,10 @@ import { SECTION_TO_CHANNEL, findChannel, } from "./lib/channels"; -import { isHiddenModelId } from "./lib/hidden-models"; +import { + isConfiguredHiddenModelId, + isHiddenModelId, +} from "./lib/hidden-models"; import { inventoryRowMatches, tokenizeQuery } from "./lib/inventory-search"; import { resolveOwnerProviderLogo } from "./lib/provider-logos"; import { @@ -386,6 +390,7 @@ export function ModelsPage() { useState("all"); const isDiscoverTab = tab === "discover"; const isDatasetMode = resourceType === "datasets"; + const hiddenEmbeddingModelIds = useHiddenEmbeddingModelIds(!isDatasetMode); const urlSection = hubSearch.section ?? null; const isModelDiscover = isDiscoverTab && !isDatasetMode; const sectionChannelId: ChannelId | null = urlSection @@ -700,6 +705,7 @@ export function ModelsPage() { return discoverRows.filter( (row) => !isHiddenModelId(row.id) && + !isConfiguredHiddenModelId(hiddenEmbeddingModelIds, row.id) && // The default feed only shows models with a provider logo. (!isFeedMode || resolveOwnerProviderLogo(row.owner, row.repo) !== null) && @@ -714,6 +720,7 @@ export function ModelsPage() { ); }, [ discoverRows, + hiddenEmbeddingModelIds, isDatasetMode, isFeedMode, effectiveDiscoverFormat, @@ -739,7 +746,11 @@ export function ModelsPage() { effectiveCachedRows, effectiveLocalRows, ) - .filter((row) => !isHiddenModelId(row.id)) + .filter( + (row) => + !isHiddenModelId(row.id) && + !isConfiguredHiddenModelId(hiddenEmbeddingModelIds, row.id), + ) .filter((row) => matchesFormat(row.result.isGguf, "gguf")) // Same fit filter as the main Discover list, so the feed carousel // honors the toggle too. @@ -751,6 +762,7 @@ export function ModelsPage() { ), [ hubFeed.trending.results, + hiddenEmbeddingModelIds, modelDiscoveryInventorySignature, fitOnDeviceOnly, gpu, @@ -778,22 +790,29 @@ export function ModelsPage() { () => (isDiscoverTab ? [] : tokenizeQuery(deferredDebouncedQuery)), [isDiscoverTab, deferredDebouncedQuery], ); - // Hide infra models (e.g. the RAG embedder bge-small-en-v1.5) from the On - // Device list like Discover, but reveal a row when a query matches it so the - // user can confirm it is already downloaded. + // Server cache rows already apply variant-aware infra hiding. Optimistic + // rows are not server-confirmed, so apply the client filter first. const isVisibleInventoryRow = useCallback( - (row: CachedInventoryRow | LocalInventoryRow) => - // Local rows can have a null repoId and an id that is a hash rather than - // the file path/name, so also check path/title (the backend's - // _is_hidden_model checks the on-disk path for the same reason). - !isHiddenModelId( - row.id, - row.repoId, - row.kind !== "cache" ? row.path : undefined, - row.kind !== "cache" ? row.title : undefined, - ) || - (inventoryTokens.length > 0 && inventoryRowMatches(row, inventoryTokens)), - [inventoryTokens], + (row: CachedInventoryRow | LocalInventoryRow) => { + if (row.kind === "cache") { + return ( + !row.optimistic || + (!isHiddenModelId(row.id, row.repoId, row.cachePath) && + !isConfiguredHiddenModelId( + hiddenEmbeddingModelIds, + row.id, + row.repoId, + row.cachePath, + )) + ); + } + // Local rows may lack a repo id, so also check path and title. + return ( + !isHiddenModelId(row.id, row.repoId, row.path, row.title) || + (inventoryTokens.length > 0 && inventoryRowMatches(row, inventoryTokens)) + ); + }, + [hiddenEmbeddingModelIds, inventoryTokens], ); // Format filter is a deliberate scope narrowing, so hard-filter it out. The // text query instead drives dim-not-filter on On Device (see ModelsCatalog) so diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts index 3515f6ca76..5d4151e87d 100644 --- a/studio/frontend/src/features/hub/index.ts +++ b/studio/frontend/src/features/hub/index.ts @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { cancelStagedModelDownload } from "./download-manager"; +export { bumpInventoryVersion } from "./stores/inventory-events"; export { getHfToken, mirrorHfTokenInto, diff --git a/studio/frontend/src/features/hub/inventory/inventory-hints.ts b/studio/frontend/src/features/hub/inventory/inventory-hints.ts index af9f254ab3..5e202e3150 100644 --- a/studio/frontend/src/features/hub/inventory/inventory-hints.ts +++ b/studio/frontend/src/features/hub/inventory/inventory-hints.ts @@ -12,6 +12,7 @@ export type InventoryHintRow = { repo_id: string; size_bytes: number; partial?: boolean; + optimistic?: boolean; }; export type InventoryHintReconciliation = { @@ -41,6 +42,7 @@ function optimisticRow(hint: InventoryHint): InventoryHintRow { repo_id: hint.repoId, size_bytes: hint.bytes ?? 0, partial: false, + optimistic: true, }; } @@ -101,9 +103,14 @@ function mergeInventoryHint( if (idx === -1) { return [...rows, seed]; } + const serverRow = rows[idx]; const merged = { - ...rows[idx], - ...seed, + ...serverRow, + // A completed hint may arrive before a partial server scan catches up. In + // that case keep the synthetic row non-runnable. A complete server row is + // already authoritative even when its runnable-weight size is smaller than + // the hint's full-snapshot byte count, so do not mark that merge optimistic. + ...(serverRow.partial ? seed : { optimistic: false }), size_bytes: Math.max(rowSizeBytes(rows[idx]), rowSizeBytes(seed)), }; return [...rows.slice(0, idx), merged, ...rows.slice(idx + 1)]; diff --git a/studio/frontend/src/features/hub/inventory/types.ts b/studio/frontend/src/features/hub/inventory/types.ts index c86ffb1d86..6f65a56037 100644 --- a/studio/frontend/src/features/hub/inventory/types.ts +++ b/studio/frontend/src/features/hub/inventory/types.ts @@ -54,6 +54,7 @@ export interface CachedInventoryRow { libraryName?: string | null; quantMethod?: string | null; liveDownload?: boolean; + optimistic?: boolean; } export interface LocalInventoryRow { diff --git a/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts b/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts index a7dc7ae3f3..fea7b3d331 100644 --- a/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts +++ b/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts @@ -204,6 +204,7 @@ function liveDownloadInventoryRows( size_bytes: job.displayBytes, partial: true, partial_transport: null, + optimistic: true, }, modelFormat, ), diff --git a/studio/frontend/src/features/hub/inventory/view-models.ts b/studio/frontend/src/features/hub/inventory/view-models.ts index 63d70418be..334050fab4 100644 --- a/studio/frontend/src/features/hub/inventory/view-models.ts +++ b/studio/frontend/src/features/hub/inventory/view-models.ts @@ -176,6 +176,7 @@ export function buildCachedInventoryRow( runtime?: string | null; format_variant?: string | null; capabilities?: BackendModelCapabilities | null; + optimistic?: boolean; }, fallbackFormat: ModelInventoryFormat, ): CachedInventoryRow { @@ -185,6 +186,15 @@ export function buildCachedInventoryRow( const inferredFromEndpoint = rawModelFormat === "unknown" && modelFormat !== "unknown"; const requiresVariant = modelFormat === "gguf"; + const capabilities = normalizeCapabilities( + inferredFromEndpoint ? null : row.capabilities, + modelFormat, + row.partial ?? false, + requiresVariant, + ); + if (row.optimistic) { + capabilities.canChat = false; + } return { kind: "cache", id: @@ -202,12 +212,7 @@ export function buildCachedInventoryRow( modelFormat, ), formatVariant: row.format_variant ?? null, - capabilities: normalizeCapabilities( - inferredFromEndpoint ? null : row.capabilities, - modelFormat, - row.partial ?? false, - requiresVariant, - ), + capabilities, bytes: row.size_bytes, cachePath: row.cache_path ?? null, partial: row.partial ?? false, @@ -216,6 +221,7 @@ export function buildCachedInventoryRow( tags: row.tags, libraryName: row.library_name ?? null, quantMethod: row.quant_method ?? null, + optimistic: row.optimistic, }; } diff --git a/studio/frontend/src/features/hub/lib/hidden-models.ts b/studio/frontend/src/features/hub/lib/hidden-models.ts index 2dbe257947..634a061e0c 100644 --- a/studio/frontend/src/features/hub/lib/hidden-models.ts +++ b/studio/frontend/src/features/hub/lib/hidden-models.ts @@ -1,11 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -// Infra models hidden from every browse/preview list (Hub discover and the chat -// model selector). Mirrors the backend `_is_hidden_model`: the RAG embedding -// model and the llama.cpp validation probe are not usable chat models. Per-repo -// file/download views are NOT filtered, so a reinstall still shows the model as -// already downloaded. +// Infra models hidden from browse/preview lists (Hub Discover, the chat model +// selector, and local on-device rows). Mirrors the backend +// `utils.hidden_models`: the RAG embedding model and the llama.cpp validation +// probe are not usable chat models. Server-confirmed cache rows are trusted +// because the backend applies variant-aware filtering. Optimistic cache rows +// still use these needles until the server confirms them. Per-repo views are +// not filtered, so reinstall flows still show downloaded files. const HIDDEN_NEEDLES = [ "bge-small-en-v1.5", // RAG embedder: unsloth/bge-small-en-v1.5[-GGUF] "ggml-org/models", // llama.cpp validation probe repo @@ -17,8 +19,20 @@ export function isHiddenModelId( ...values: (string | null | undefined)[] ): boolean { return values.some((v) => { - if (!v) return false; + if (!v) { + return false; + } const lower = v.toLowerCase(); return HIDDEN_NEEDLES.some((needle) => lower.includes(needle)); }); } + +/** Exact-match configured infra repos without hiding similarly named models. */ +export function isConfiguredHiddenModelId( + configuredIds: ReadonlySet, + ...values: (string | null | undefined)[] +): boolean { + return values.some( + (value) => value != null && configuredIds.has(value.trim().toLowerCase()), + ); +} diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 9a61142f73..cc21559f38 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -2,11 +2,14 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { authFetch } from "@/features/auth"; +import { bumpInventoryVersion } from "@/features/hub"; import { readFastApiError } from "@/lib/format-fastapi-error"; export type EmbeddingModelSettings = { embeddingModel: string; + embeddingGgufRepo: string; defaultEmbeddingModel: string; + defaultEmbeddingGgufRepo: string; isCustom: boolean; }; @@ -14,8 +17,12 @@ type ApiEmbeddingModelSettings = { // biome-ignore lint/style/useNamingConvention: API schema embedding_model: string; // biome-ignore lint/style/useNamingConvention: API schema + embedding_gguf_repo: string; + // biome-ignore lint/style/useNamingConvention: API schema default_embedding_model: string; // biome-ignore lint/style/useNamingConvention: API schema + default_embedding_gguf_repo: string; + // biome-ignore lint/style/useNamingConvention: API schema is_custom: boolean; }; @@ -30,7 +37,9 @@ export class EmbeddingModelBlockedError extends Error {} function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { return { embeddingModel: settings.embedding_model, + embeddingGgufRepo: settings.embedding_gguf_repo, defaultEmbeddingModel: settings.default_embedding_model, + defaultEmbeddingGgufRepo: settings.default_embedding_gguf_repo, isCustom: settings.is_custom, }; } @@ -75,7 +84,9 @@ export async function updateEmbeddingModelSettings( await readFastApiError(res, "Failed to save embedding model"), ); } - return fromApi(await res.json()); + const settings = fromApi(await res.json()); + bumpInventoryVersion(); + return settings; } export async function resetEmbeddingModelSettings(): Promise { @@ -87,5 +98,7 @@ export async function resetEmbeddingModelSettings(): Promise; @@ -147,25 +148,7 @@ export function HistoricalTrainingView({ } const viewData = mapToViewData(detail, t); - const configOverride = detail.config - ? { - epochs: detail.config.num_epochs as number | undefined, - batchSize: detail.config.batch_size as number | undefined, - learningRate: detail.config.learning_rate as string | undefined, - maxSteps: detail.config.max_steps as number | undefined, - contextLength: detail.config.max_seq_length as number | undefined, - warmupSteps: detail.config.warmup_steps as number | undefined, - optimizerType: detail.config.optim as string | undefined, - loraRank: detail.config.lora_r as number | undefined, - loraAlpha: detail.config.lora_alpha as number | undefined, - loraDropout: detail.config.lora_dropout as number | undefined, - loraVariant: detail.config.use_rslora - ? "rslora" - : detail.config.use_loftq - ? "loftq" - : "lora", - } - : undefined; + const configOverride = mapRunConfigToOverride(detail.config); return (
diff --git a/studio/frontend/src/features/studio/live-training-view.tsx b/studio/frontend/src/features/studio/live-training-view.tsx index cce39adbf4..0aecc7030e 100644 --- a/studio/frontend/src/features/studio/live-training-view.tsx +++ b/studio/frontend/src/features/studio/live-training-view.tsx @@ -1,18 +1,42 @@ // 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 { cn } from "@/lib/utils"; import { + getTrainingRun, useTrainingConfigStore, useTrainingRuntimeStore, } from "@/features/training"; import type { TrainingViewData } from "@/features/training"; +import { cn } from "@/lib/utils"; import type { ReactElement } from "react"; +import { useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { ChartsSection } from "./sections/charts-section"; import { ProgressSection } from "./sections/progress-section"; +import { + type RunConfigOverride, + mapRunConfigToOverride, +} from "./sections/run-config-override"; import { TrainingStartOverlay } from "./training-start-overlay"; +/** Retry budget for the run-config lookup. The row is inserted at + * start_training(), but a lookup issued in the same instant can still miss it; + * a few short retries cover that without polling a genuinely absent row. */ +const RUN_CONFIG_FETCH_RETRIES = 5; +const RUN_CONFIG_FETCH_RETRY_MS = 1000; + +/** The fetched run config only applies while it belongs to the active job; + * a stale record from a previous run falls back to the form store. */ +function activeRunOverride( + fetched: { jobId: string; override: RunConfigOverride | undefined } | null, + jobId: string | null, +): RunConfigOverride | undefined { + if (fetched === null || fetched.jobId !== jobId) { + return undefined; + } + return fetched.override; +} + export function LiveTrainingView(): ReactElement { const runtime = useTrainingRuntimeStore( useShallow((state) => ({ @@ -52,6 +76,59 @@ export function LiveTrainingView(): ReactElement { })), ); + // Show the ACTIVE run's saved config, not the editable form store the user may + // have changed since starting (#6853). start_training() commits the run row + // before the pump, so the job id alone gates the fetch; the bounded retry below + // covers the narrow uncommitted window, and until it loads ProgressSection falls + // back to the form store. The result is keyed by job id and filtered at render. + const [fetchedRunConfig, setFetchedRunConfig] = useState<{ + jobId: string; + override: RunConfigOverride | undefined; + } | null>(null); + // Retry budget for the transient 404 below, keyed by job so a new run always + // starts with a fresh budget. + const [fetchAttempt, setFetchAttempt] = useState<{ + jobId: string; + count: number; + } | null>(null); + useEffect(() => { + if (!runtime.jobId) { + return; + } + const jobId = runtime.jobId; + if (fetchedRunConfig !== null && fetchedRunConfig.jobId === jobId) { + return; // already resolved for this job + } + const attempts = fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0; + const controller = new AbortController(); + let retryTimer: ReturnType | undefined; + getTrainingRun(jobId, controller.signal) + .then((detail) => { + setFetchedRunConfig({ + jobId, + override: mapRunConfigToOverride(detail.config), + }); + }) + .catch(() => { + // A lookup racing the row commit can miss transiently; nothing else in + // the deps changes on failure, so retry explicitly. Bounded so a genuinely + // absent row falls back to the form store instead of polling forever. + if (controller.signal.aborted || attempts >= RUN_CONFIG_FETCH_RETRIES) { + return; + } + retryTimer = setTimeout(() => { + setFetchAttempt({ jobId, count: attempts + 1 }); + }, RUN_CONFIG_FETCH_RETRY_MS); + }); + return () => { + controller.abort(); + if (retryTimer !== undefined) { + clearTimeout(retryTimer); + } + }; + }, [runtime.jobId, fetchedRunConfig, fetchAttempt]); + const runConfigOverride = activeRunOverride(fetchedRunConfig, runtime.jobId); + const activeProjectName = runtime.startProjectName !== null ? runtime.startProjectName.trim() || null @@ -76,7 +153,11 @@ export function LiveTrainingView(): ReactElement { isTrainingRunning: runtime.isTrainingRunning, modelName: runtime.startModelName ?? config.selectedModel ?? "", projectName: activeProjectName, - trainingMethod: config.trainingMethod ?? "", + // Prefer the saved run's method: the form may have been edited (e.g. LoRA + // -> Full) after the run started, which would relabel the run and hide its + // saved LoRA rows in the popover. + trainingMethod: + runConfigOverride?.trainingMethod ?? config.trainingMethod ?? "", lossHistory: runtime.lossHistory, lrHistory: runtime.lrHistory, gradNormHistory: runtime.gradNormHistory, @@ -105,7 +186,11 @@ export function LiveTrainingView(): ReactElement { )} >
- +
o.value === cfgOptimizerType)?.label ?? diff --git a/studio/frontend/src/features/studio/sections/run-config-override.ts b/studio/frontend/src/features/studio/sections/run-config-override.ts new file mode 100644 index 0000000000..a1272bfeb0 --- /dev/null +++ b/studio/frontend/src/features/studio/sections/run-config-override.ts @@ -0,0 +1,54 @@ +// 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 { parseBackendTrainingMethod } from "@/features/training"; + +/** Shape of the Training Config popover's data when it is driven by a saved + * run snapshot instead of the editable form store. */ +export interface RunConfigOverride { + trainingMethod?: string; + epochs?: number; + batchSize?: number; + learningRate?: string; + maxSteps?: number; + contextLength?: number; + warmupSteps?: number; + optimizerType?: string; + loraRank?: number; + loraAlpha?: number; + loraDropout?: number; + loraVariant?: string; +} + +/** Map a saved run's config (GET /api/train/runs/{id} `detail.config`) into the + * Training Config popover's override shape. Shared by the History view and the + * live Current Run view so both read the same authoritative run snapshot + * instead of the editable form store (#6853). */ +export function mapRunConfigToOverride( + config: Record | null | undefined, +): RunConfigOverride | undefined { + if (!config) { + return undefined; + } + return { + trainingMethod: parseBackendTrainingMethod( + config.training_type, + config.load_in_4bit, + ), + epochs: config.num_epochs as number | undefined, + batchSize: config.batch_size as number | undefined, + learningRate: config.learning_rate as string | undefined, + maxSteps: config.max_steps as number | undefined, + contextLength: config.max_seq_length as number | undefined, + warmupSteps: config.warmup_steps as number | undefined, + optimizerType: config.optim as string | undefined, + loraRank: config.lora_r as number | undefined, + loraAlpha: config.lora_alpha as number | undefined, + loraDropout: config.lora_dropout as number | undefined, + loraVariant: config.use_rslora + ? "rslora" + : config.use_loftq + ? "loftq" + : "lora", + }; +} diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts index 553dcc2af5..a0d249ff1b 100644 --- a/studio/frontend/src/features/training/index.ts +++ b/studio/frontend/src/features/training/index.ts @@ -12,6 +12,7 @@ export { getTrainingRunDisplayTitle, getTrainingRunModelSubtitle, } from "./lib/run-display"; +export { parseBackendTrainingMethod } from "./lib/training-methods"; export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar"; export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle"; export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch"; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index 1e313acdf3..db2cc021be 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -15,6 +15,19 @@ export interface GpuInfo { systemRamTotalGb: number } +export interface SystemGpuDevice { + index: number; + name: string; + memoryTotalGb: number; + /** Free VRAM at fetch time. Degrades to the total when the utilization + * probe had no usage data; 0 only when the total is unknown too. */ + memoryFreeGb: number; + /** "physical" = `index` is a stable physical/PCI id safe to pin via gpu_ids; + * "relative" = an ordinal into a parent CUDA_VISIBLE_DEVICES mask, which the + * backend can't map back, so the picker must not offer it. */ + physicalIndex: boolean; +} + const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", @@ -25,70 +38,135 @@ const DEFAULT_GPU: GpuInfo = { systemRamTotalGb: 0 }; -// Module-level cache so multiple components share one fetch. -let cachedGpu: GpuInfo | null = null; -let fetchPromise: Promise | null = null; +// One module-level cache so every GPU hook shares a single /api/system fetch. +let cachedSystem: SystemInfoResponse | null = null; +let systemPromise: Promise | null = null; -async function fetchGpuOnce(): Promise { - if (cachedGpu) return cachedGpu; - if (fetchPromise) return fetchPromise; - - fetchPromise = (async () => { +async function fetchSystemOnce(): Promise { + if (cachedSystem) return cachedSystem; + if (systemPromise) return systemPromise; + systemPromise = (async () => { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - - const data = await res.json() as SystemInfoResponse; - const gpuData = data?.gpu; - - // CPU/RAM exist even on hosts without a GPU, so populate them on every path. - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const base = { - cpuCore: data?.cpu?.physical_count ?? 0, - cpuThread: data?.cpu?.logical_count ?? 0, - systemRamAvailableGb: data?.memory?.available_gb ?? 0, - systemRamTotalGb: data?.memory?.total_gb ?? 0, - }; - - const devices = gpuData?.devices ?? []; - const info: GpuInfo = - gpuData?.available && devices.length - ? { - ...base, - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), - } - : { ...DEFAULT_GPU, ...base }; - cachedGpu = info; - return info; + cachedSystem = (await res.json()) as SystemInfoResponse; + return cachedSystem; } catch { - // Reset promise so subsequent calls retry (e.g. backend wasn't ready) - fetchPromise = null; - return DEFAULT_GPU; + systemPromise = null; // reset so a later call retries (backend not ready) + return null; } })(); + return systemPromise; +} - return fetchPromise; +function toGpuInfo(data: SystemInfoResponse | null): GpuInfo { + // CPU/RAM exist even on GPU-less hosts (e.g. Mac), so populate them on every + // path: unified-memory math still needs a RAM budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, + }; + const gpuData = data?.gpu; + const devices = gpuData?.devices ?? []; + if (!gpuData?.available || !devices.length) { + return { ...DEFAULT_GPU, ...base }; + } + return { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + }; +} + +function toGpuDevices(data: SystemInfoResponse | null): SystemGpuDevice[] { + // Unpinnable configurations must hide every pick surface: XPU indices are + // torch-xpu ordinals no applicator speaks, and Vulkan-only builds pin ggml's + // own ordinals -- /load and /validate 400 picks on both, so the backend + // reports gpu.gguf_gpu_ids_supported and every gate keyed on physicalIndex + // (picker, persisted-pick reconcile) follows it. The device flavor lives on + // the TOP-LEVEL device_backend field; absent support info defaults to + // pinnable (older backend). + const pinnableBackend = + data?.device_backend !== "xpu" && + data?.gpu?.gguf_gpu_ids_supported !== false; + return (data?.gpu?.devices ?? []) + .filter((d) => typeof d.index === "number") + .map((d) => ({ + index: d.index as number, + name: d.name ?? `GPU ${d.index}`, + memoryTotalGb: d.memory_total_gb ?? 0, + memoryFreeGb: d.vram_free_gb ?? 0, + physicalIndex: pinnableBackend && d.index_kind === "physical", + })); +} + +/** Aggregate GPU info from /api/system; shares one module-level fetch across all GPU hooks. */ +export function useGpuInfo(): GpuInfo { + const [gpu, setGpu] = useState( + cachedSystem ? toGpuInfo(cachedSystem) : DEFAULT_GPU, + ); + useEffect(() => { + // No early return on cachedSystem: a consumer mounting as the cache fills + // (between render and effect) would otherwise stay stuck at the default. + let cancelled = false; + fetchSystemOnce().then((d) => { + if (!cancelled) setGpu(toGpuInfo(d)); + }); + return () => { + cancelled = true; + }; + }, []); + return gpu; +} + +/** All backend-visible GPUs (index, name, total VRAM); shares the same fetch. */ +export function useGpuDevices(): SystemGpuDevice[] { + const [devices, setDevices] = useState( + cachedSystem ? toGpuDevices(cachedSystem) : [], + ); + useEffect(() => { + // No early return on cachedSystem: a consumer mounting as the cache fills + // (between render and effect) would otherwise stay stuck at the default. + let cancelled = false; + fetchSystemOnce().then((d) => { + if (!cancelled) setDevices(toGpuDevices(d)); + }); + return () => { + cancelled = true; + }; + }, []); + return devices; } /** - * Fetch GPU info from /api/system. Cached at module level, so only one request - * is made no matter how many components call this hook. + * Await the shared /api/system fetch so cachedPinnableGpuIndices (and the + * store's reconcilePersistedGpuIds) can validate a persisted pick before a + * load path sends it -- on a cold cache the reconcile passes ids through + * unvalidated, and a stale cross-host pick then fails /load with the picker + * hidden. Resolves immediately once the module cache is warm; a failed fetch + * keeps the cache cold, preserving the "can't validate, backend guards" + * degradation. */ -export function useGpuInfo(): GpuInfo { - const [gpu, setGpu] = useState(cachedGpu ?? DEFAULT_GPU); +export async function ensureGpuDeviceCache(): Promise { + await fetchSystemOnce(); +} - useEffect(() => { - if (cachedGpu) return; - - let cancelled = false; - fetchGpuOnce().then((info) => { - if (!cancelled) setGpu(info); - }); - return () => { cancelled = true; }; - }, []); - - return gpu; -} \ No newline at end of file +/** + * Pinnable physical GPU indices from the already-fetched /api/system cache, for + * non-React code (the store) that needs to validate a persisted `gpu_ids` pick + * without triggering a fetch. Returns: + * - `null` when the cache isn't populated yet (caller can't validate, so keep + * the pick and let the backend guard reject a truly bad one); + * - `[]` when the host has no pinnable multi-GPU set (single GPU, or relative/ + * UUID-masked indices) -- the picker is hidden, so any saved pick is stale; + * - the physical indices otherwise. + */ +export function cachedPinnableGpuIndices(): number[] | null { + if (!cachedSystem) return null; + const physical = toGpuDevices(cachedSystem).filter((d) => d.physicalIndex); + // Mirrors the sheet's showGpuPicker gate: only a 2+ physical-GPU host can pin. + return physical.length > 1 ? physical.map((d) => d.index) : []; +} diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts index a135cce86e..8cfe2bace4 100644 --- a/studio/frontend/src/hooks/use-system.ts +++ b/studio/frontend/src/hooks/use-system.ts @@ -40,6 +40,9 @@ export interface SystemInfoResponse { gpu: { available: boolean; backend?: string; + /** Whether GGUF loads accept an explicit gpu_ids pick (false on XPU hosts + * and Vulkan-only builds, where /load and /validate 400 picks). */ + gguf_gpu_ids_supported?: boolean; backend_cuda_visible_devices?: string | null; parent_visible_gpu_ids?: number[]; index_kind?: string; diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py index 66d3d80920..53d281d435 100644 --- a/tests/python/test_construct_chat_template_validation.py +++ b/tests/python/test_construct_chat_template_validation.py @@ -104,3 +104,93 @@ def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(chat_t ) assert "{INPUT}" not in jinja_template assert "{OUTPUT}" not in jinja_template + + +_SYSTEM_CHAT_TEMPLATE = ( + "{SYSTEM}\n" + "### User: {INPUT}\n### Assistant: {OUTPUT}" + "### User: {INPUT}\n### Assistant: {OUTPUT}" +) + + +def _render(jinja_template, messages): + from jinja2.sandbox import ImmutableSandboxedEnvironment + + env = ImmutableSandboxedEnvironment() + env.globals["raise_exception"] = lambda message: (_ for _ in ()).throw(RuntimeError(message)) + return env.from_string(jinja_template).render( + messages = messages, + bos_token = "", + eos_token = "", + add_generation_prompt = False, + ) + + +@pytest.mark.parametrize("default_system_message", [None, "You are helpful."]) +def test_system_message_is_consumed_by_the_system_part(default_system_message): + """A caller-supplied system message must be rendered by the system part and + skipped by the message loop, whatever `default_system_message` is. + + With `default_system_message = None` the generated template used to bind + `loop_messages` only inside the `{% if %}` arm. The `Fix missing + loop_messages` step then saw no unconditional binding, rewrote the loop back + to `messages`, and the system message reached the loop and tripped + `raise_exception`. + """ + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _SYSTEM_CHAT_TEMPLATE, + default_system_message = default_system_message, + extra_eos_tokens = [""], + ) + rendered = _render( + jinja_template, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Hi"}, + ], + ) + assert rendered.count("Be terse.") == 1 + assert rendered.count("Hi") == 1 + # A caller system message overrides the default; the default must not leak in. + if default_system_message is not None: + assert default_system_message not in rendered + + +def test_absent_system_message_still_renders_without_default(): + """`default_system_message = None` with no system message in the input must + keep working -- the `{% else %}` arm has to bind `loop_messages = messages`.""" + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _SYSTEM_CHAT_TEMPLATE, + default_system_message = None, + extra_eos_tokens = [""], + ) + rendered = _render(jinja_template, [{"role": "user", "content": "Hi"}]) + assert "Hi" in rendered + + +_NO_SYSTEM_CHAT_TEMPLATE = ( + "PREAMBLE\n" + "### User: {INPUT}\n### Assistant: {OUTPUT}" + "### User: {INPUT}\n### Assistant: {OUTPUT}" +) + + +def test_static_prefix_without_system_still_rejects_system_message(): + """A template with a static prefix but no {SYSTEM} placeholder cannot render a + caller system message, so it must still raise rather than silently drop it.""" + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _NO_SYSTEM_CHAT_TEMPLATE, + default_system_message = None, + extra_eos_tokens = [""], + ) + with pytest.raises(RuntimeError, match = "Only user and assistant roles are supported!"): + _render( + jinja_template, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Hi"}, + ], + ) diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 4322f0c7d6..c58808689b 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -69,6 +69,21 @@ class TestStructuralTorchConstraint: def test_tightened_assignment_exists(self): assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh + def test_cuda_constraint_widened_to_2_12(self): + """A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x + land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC); + without it cu128/cu130 resolves torch 2.10.x.""" + assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh + + def test_cuda_case_widens_via_index_leaf(self): + """The cu* branch of the _torch_index_leaf case sets the widened + constraint (parallel to rocm7.2), anchored on the leaf.""" + m = re.search( + r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"', + self._sh, + ) + assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found" + def test_variable_used_in_pip_install(self): """$TORCH_CONSTRAINT must appear in a uv pip install line.""" assert '"$TORCH_CONSTRAINT"' in self._sh @@ -384,6 +399,71 @@ class TestTorchConstraintShell: logged = log_file.read_text() assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}" + # Mirrors the _torch_index_leaf case in install.sh: rocm7.2 -> 2.11.x floor, + # CUDA -> widened <2.12.0 ceiling, else (CPU/older ROCm) -> default. Anchored + # on the final path segment, so a mirror base path containing cu*/rocm7.2 but + # ending in a cpu/older-rocm leaf keeps the default. + _INDEX_SNIPPET = textwrap.dedent(r""" + #!/bin/bash + set -e + TORCH_INDEX_URL="{index_url}" + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + _torch_index_leaf="${TORCH_INDEX_URL%/}" + _torch_index_leaf="${_torch_index_leaf##*/}" + case "$_torch_index_leaf" in + rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; + cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;; + esac + echo "$TORCH_CONSTRAINT" + """).strip() + + def _resolve_index(self, tmp_path: pathlib.Path, index_url: str) -> str: + script_file = tmp_path / "index_snippet.sh" + script_file.write_text(self._INDEX_SNIPPET.replace("{index_url}", index_url)) + script_file.chmod(0o755) + result = subprocess.run( + ["bash", str(script_file)], + capture_output = True, + text = True, + timeout = 10, + ) + assert result.returncode == 0, f"Script failed: {result.stderr}" + return result.stdout.strip() + + @pytest.mark.parametrize("leaf", ["cu118", "cu124", "cu126", "cu128", "cu130"]) + def test_cuda_index_widens_to_2_12(self, tmp_path, leaf): + url = f"https://download.pytorch.org/whl/{leaf}" + assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0" + + def test_rocm72_index_uses_211_floor(self, tmp_path): + url = "https://download.pytorch.org/whl/rocm7.2" + assert self._resolve_index(tmp_path, url) == "torch>=2.11.0,<2.12.0" + + def test_cpu_index_keeps_default(self, tmp_path): + # /cpu must NOT match the */cu[0-9]* branch. + url = "https://download.pytorch.org/whl/cpu" + assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0" + + def test_older_rocm_index_keeps_default(self, tmp_path): + url = "https://download.pytorch.org/whl/rocm7.1" + assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0" + + def test_cuda_index_custom_mirror_widens(self, tmp_path): + url = "https://internal.example.com/pytorch/cu128" + assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0" + + @pytest.mark.parametrize( + "url", + [ + "https://internal.example.com/pytorch/cu128/cpu", + "https://internal.example.com/cu128/whl/rocm7.1", + ], + ) + def test_cuda_in_mirror_path_but_noncuda_leaf_keeps_default(self, tmp_path, url): + # A cu128 in the mirror base path must not widen when the leaf is cpu / + # older ROCm: the case anchors on _torch_index_leaf, not the whole URL. + assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0" + # Group 3 -- E2E tokenizers fix (requires network, ~2-5 min) @pytest.mark.e2e diff --git a/tests/sh/test_previous_torch_pin.sh b/tests/sh/test_previous_torch_pin.sh new file mode 100644 index 0000000000..253ede8a27 --- /dev/null +++ b/tests/sh/test_previous_torch_pin.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for install.sh's _previous_torch_pin, which keeps the previous +# venv's torch release on a re-run (curl | sh over an existing install) instead +# of silently moving the user to a newer release. Helpers are extracted from +# install.sh and sourced. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +# Extract _previous_torch_pin and its dependencies _torch_flavor_tag and +# _torch_release_in_window. +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_torch_release_in_window()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_previous_torch_pin()/,/^}/p' "$INSTALL_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +unset UNSLOTH_TORCH_UPGRADE + +echo "=== _previous_torch_pin: matching flavor keeps the release ===" +assert_eq "cu126 wheel on cu126 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "cu130 wheel on cu130 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "untagged wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'cu128' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: flavor change installs the new build ===" +assert_eq "cu126 wheel on cu130 leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu wheel on cu126 leaf" "" "$(_previous_torch_pin '2.10.0+cpu' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "cu126 wheel on cpu leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cpu' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: rocm and unknown leaves never pin ===" +assert_eq "rocm7.2 leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'rocm7.2' 'torch>=2.4,<2.12.0')" +assert_eq "gfx leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'gfx120X-all' 'torch>=2.4,<2.12.0')" +assert_eq "unknown mirror leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'simple' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: probe noise never becomes a pin ===" +assert_eq "empty version" "" "$(_previous_torch_pin '' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'cpu' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: out-of-window releases never pin ===" +assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'cu118' 'torch>=2.4,<2.12.0')" +assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'cpu' 'torch>=2.4,<2.11.0')" +assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'cpu' 'torch>=2.6,<2.11.0')" +assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch')" +assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' '')" + +echo "=== _torch_release_in_window ===" +assert_eq "in window" "yes" "$(_torch_release_in_window '2.10.0' 'torch>=2.4,<2.12.0')" +assert_eq "at floor" "yes" "$(_torch_release_in_window '2.4.0' 'torch>=2.4,<2.12.0')" +assert_eq "below floor" "no" "$(_torch_release_in_window '2.3.1' 'torch>=2.4,<2.12.0')" +assert_eq "at ceiling" "no" "$(_torch_release_in_window '2.12.0' 'torch>=2.4,<2.12.0')" +assert_eq "next major" "no" "$(_torch_release_in_window '3.0.0' 'torch>=2.4,<2.12.0')" +assert_eq "patch-level floor" "yes" "$(_torch_release_in_window '2.11.5' 'torch>=2.11.0,<2.12.0')" +assert_eq "no ceiling -> no" "no" "$(_torch_release_in_window '2.10.0' 'torch>=2.4')" +assert_eq "garbage minor -> no" "no" "$(_torch_release_in_window '2.x' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: UNSLOTH_TORCH_UPGRADE=1 opts out ===" +assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" + +echo "=== install.sh wiring ===" +# The probe must run against the OLD venv, before it is moved aside for rollback. +_probe_line=$(grep -n '_PREV_TORCH_VER=\$(' "$INSTALL_SH" | head -1 | cut -d: -f1) +_move_line=$(grep -n '_start_studio_venv_replacement "\$VENV_DIR"' "$INSTALL_SH" | head -1 | cut -d: -f1) +assert_eq "probe exists" "yes" "$([ -n "$_probe_line" ] && echo yes)" +assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n "$_move_line" ] && [ "$_probe_line" -lt "$_move_line" ] && echo yes)" +# A kept release that vanished from the index must fall back to the supported range. +assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)" +assert_eq "pin gated on SKIP_TORCH" "yes" "$(grep -q 'if \[ "\$SKIP_TORCH" = false \]; then' "$INSTALL_SH" && echo yes)" + +echo "" +if [ "$FAIL" -gt 0 ]; then + echo "$FAIL check(s) FAILED" + exit 1 +fi +echo "All $PASS checks passed" diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index 293a709360..d60dfc9f90 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -108,6 +108,20 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" _hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded" +# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch +# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC). +_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) +assert_eq "CUDA TORCH_CONSTRAINT widened to <2.12.0" "1" "$_cuda_widen" + +# Widening keys off the final leaf (_torch_index_leaf), not the full URL, so a +# mirror base path with cu*/rocm7.2 but a cpu/older-rocm leaf is not mis-widened. +_cuda_case=$(grep -c 'cu\[0-9\]\*)' "$INSTALL_SH" || true) +_has_cuda_case=$([ "$_cuda_case" -ge 1 ] && echo "yes" || echo "no") +assert_eq "cu* index case adjusts TORCH_CONSTRAINT" "yes" "$_has_cuda_case" +_leaf_case=$(grep -c 'case "\$_torch_index_leaf" in' "$INSTALL_SH" || true) +_has_leaf_constraint=$([ "$_leaf_case" -ge 2 ] && echo "yes" || echo "no") +assert_eq "constraint case anchors on _torch_index_leaf" "yes" "$_has_leaf_constraint" + echo "" echo "=== Structural: tokenizers in no-torch-runtime.txt ===" diff --git a/tests/sh/test_unsloth_torch_override.sh b/tests/sh/test_unsloth_torch_override.sh new file mode 100644 index 0000000000..7e8e3f5b5b --- /dev/null +++ b/tests/sh/test_unsloth_torch_override.sh @@ -0,0 +1,131 @@ +#!/bin/bash +# 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 torch-trio --overrides guard on the Step-2 unsloth installs in +# install.sh. A released unsloth wheel can pin an older torch (2026.7.2 declares +# torch<2.11.0); without the overrides file a with-deps PyPI resolve downgrades +# the trio Step 1 installed, and the flavor guard misses it (PyPI's torch 2.10 +# default is itself cu128-flavored). Same assertion pattern as test_torch_constraint.sh. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +assert_true() { + _label="$1"; _ok="$2" + if [ "$_ok" = "0" ]; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label" + FAIL=$((FAIL + 1)) + fi +} + +echo "=== test_unsloth_torch_override ===" + +# 1. Every with-deps unsloth install carries the overrides expansion (local, +# generic, migrated); the --no-deps no-torch paths need no guard. +_local_block=$(grep -A2 '"install unsloth (local)"' "$INSTALL_SH") +printf '%s' "$_local_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "local (with-deps) unsloth install passes --overrides" "$?" + +_generic_block=$(grep -A2 '"install unsloth" uv pip install' "$INSTALL_SH") +printf '%s' "$_generic_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "generic (with-deps) unsloth install passes --overrides" "$?" + +_migrated_block=$(grep -A3 '"install unsloth (migrated)"' "$INSTALL_SH") +printf '%s' "$_migrated_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "migrated (with-deps) unsloth install passes --overrides" "$?" + +_no_torch_block=$(grep -A2 '"install unsloth (no-torch)"' "$INSTALL_SH") +if printf '%s' "$_no_torch_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi +assert_true "no-torch (--no-deps) unsloth install has no overrides" "$_rc" + +_migrated_nt_block=$(grep -A2 '"install unsloth (migrated no-torch)"' "$INSTALL_SH") +if printf '%s' "$_migrated_nt_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi +assert_true "migrated no-torch (--no-deps) unsloth install has no overrides" "$_rc" + +# 2. The overrides file is only built when SKIP_TORCH=false. +grep -B2 '_torch_trio_pins=\$(' "$INSTALL_SH" | grep -q 'SKIP_TORCH" = false' +assert_true "overrides file build is gated on SKIP_TORCH=false" "$?" + +# 3. The pin-collection snippet emits exact ==pins for the installed trio (run +# the embedded python against this test's interpreter). +_snippet=$(sed -n '/_torch_trio_pins=\$("\$_VENV_PY" -c "/,/^" 2>\/dev\/null)/p' "$INSTALL_SH" \ + | sed '1s/.*-c "//' | sed '$d') +_out=$(python3 -c "$_snippet" 2>&1) || true +# torch may or may not be importable on the test host; the snippet must not +# crash and every line it does emit must be an exact pkg==version pin. +if [ -n "$_out" ]; then + printf '%s\n' "$_out" | grep -vqE '^(torch|torchvision|torchaudio)==.+$' && _rc=1 || _rc=0 +else + _rc=0 +fi +assert_true "pin snippet emits only exact trio ==pins (or nothing)" "$_rc" + +# 4. The temp overrides file is cleaned up after Step 2. +grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' "$INSTALL_SH" +assert_true "overrides temp file is removed after the unsloth installs" "$?" + +# 5. Any UV_OVERRIDE env file is folded in (the CLI --overrides flag would +# otherwise replace it, dropping e.g. the macOS arm64 darwin overrides). +grep -q 'for _ov_file in \${UV_OVERRIDE:-}' "$INSTALL_SH" +assert_true "UV_OVERRIDE env files are merged into the overrides file" "$?" + +# 6. The EXIT trap also removes the overrides file, so a failed Step 2 (set -e +# fires before the normal-path rm) cannot leak it. +sed -n '/_on_install_exit() {/,/^}/p' "$INSTALL_SH" \ + | grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "EXIT trap removes the overrides temp file on failure" "$?" + +# 7. The UV_OVERRIDE fold filters inherited files instead of cat-ing them (run +# the extracted awk program on sample files): (a) inherited torch-trio lines +# are dropped so the generated exact pins win (uv intersects duplicates); +# (b) every line is newline-terminated so an unterminated file cannot join +# two requirements into one. +_awk_prog=$(sed -n "s/.*awk '\(.*\)' \"\$_ov_file\".*/\1/p" "$INSTALL_SH") +[ -n "$_awk_prog" ] +assert_true "UV_OVERRIDE fold uses the trio-filtering awk program" "$?" + +_ov_dir=$(mktemp -d) +printf '%s' 'transformers>=4.57.6' > "$_ov_dir/ov1.txt" # no trailing newline +cat > "$_ov_dir/ov2.txt" <<'EOF' +# comment survives +torch<2.11.0 +torchvision==0.25.0 +torchaudio!=2.11.0 +torchmetrics==1.0 +anyio<4.14.0 +EOF +_merged="$_ov_dir/merged.txt" +printf '%s\n' 'torch==2.11.0+cu128' > "$_merged" +for _f in "$_ov_dir/ov1.txt" "$_ov_dir/ov2.txt"; do + awk "$_awk_prog" "$_f" >> "$_merged" +done + +grep -qx 'transformers>=4.57.6' "$_merged" +assert_true "no-trailing-newline override stays a separate requirement line" "$?" + +if grep -qx 'torchmetrics==1.0' "$_merged" && grep -qx 'anyio<4.14.0' "$_merged"; then + _rc=0 +else + _rc=1 +fi +assert_true "unrelated inherited overrides are preserved" "$_rc" + +if grep -qE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged" \ + && [ "$(grep -cE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged")" != "1" ]; then + _rc=1 +else + _rc=0 +fi +grep -qx 'torch==2.11.0+cu128' "$_merged" || _rc=1 +assert_true "inherited torch-trio lines are dropped; generated pin wins" "$_rc" +rm -rf "$_ov_dir" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] || exit 1 diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 283b099107..76f462c5cf 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -1,5 +1,8 @@ """Register each model set and check the registered ids exist on the HF Hub.""" +import os +import subprocess +import sys from dataclasses import dataclass import pytest @@ -77,3 +80,85 @@ def test_quant_type(): assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models) quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH] assert all(quant_tag in m.model_path for m in dynamic_quant_models) + + +def _run_registry_child(body: str) -> subprocess.CompletedProcess: + """Run ``body`` in a fresh interpreter that first imports this directory's + ``conftest`` so it inherits the same GPU-free harness the pytest session + uses (device_type stubs plus torch.cuda probe patches). Without it, + ``import unsloth.registry`` raises ``NotImplementedError`` from + ``unsloth_zoo.device_type`` on no-accelerator CI runners, so the child + would exit non-zero and the test would fail even though the registry code + is correct. A fresh process also keeps each check independent of any + ``register_models()`` calls other tests make on the shared registry. + """ + tests_dir = os.path.dirname(os.path.abspath(__file__)) + prelude = ( + f"import sys; sys.path.insert(0, {tests_dir!r})\n" + "try:\n" + " import conftest # noqa: F401 GPU-free harness on no-accelerator runners\n" + "except Exception:\n" + " pass\n" + ) + return subprocess.run( + [sys.executable, "-c", prelude + body], + capture_output = True, + text = True, + check = False, + ) + + +def test_importing_registry_does_not_register_models(): + """Importing the registry must not populate MODEL_REGISTRY on its own. + + ``_deepseek`` used to call ``register_deepseek_models(...)`` at module + scope, so merely importing ``unsloth.registry`` registered models as an + import side effect, unlike every other family which only registers on + demand. + """ + result = _run_registry_child( + "import unsloth.registry\n" + "from unsloth.registry.registry import MODEL_REGISTRY\n" + "print('REGISTRY_SIZE', len(MODEL_REGISTRY))" + ) + assert result.returncode == 0, ( + f"registry import subprocess exited {result.returncode}\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + size_lines = [line for line in result.stdout.splitlines() if line.startswith("REGISTRY_SIZE")] + assert size_lines == ["REGISTRY_SIZE 0"], result.stdout + result.stderr + + +def test_register_models_registers_no_upstream_originals(): + """``register_models()`` must register each family's ``unsloth``-org models + and must NOT leak upstream vendor "original" models. + + Before the fix, ``_deepseek``'s import-time + ``register_deepseek_models(include_original_model = True)`` set the + ``_IS_DEEPSEEK_*_REGISTERED`` guards, so the later default + ``register_models()`` early-returned for deepseek and its 10 ``deepseek-ai`` + originals leaked permanently (129 -> 139). This asserts the whole registry + is ``unsloth``-org after ``register_models()`` while deepseek is still + registered via the normal path. Runs in a fresh interpreter so it is + independent of other tests' registry mutations. + """ + result = _run_registry_child( + "import unsloth.registry\n" + "from unsloth.registry import register_models\n" + "from unsloth.registry.registry import MODEL_REGISTRY\n" + "register_models()\n" + "orgs = sorted({m.org for m in MODEL_REGISTRY.values()})\n" + "deepseek = [k for k in MODEL_REGISTRY if 'deepseek' in k.lower()]\n" + "print('ORGS', orgs)\n" + "print('NUM_DEEPSEEK', len(deepseek))" + ) + assert result.returncode == 0, ( + f"register_models subprocess exited {result.returncode}\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + out = result.stdout + # Every registered model is unsloth-org: no upstream "original" leaked. + assert "ORGS ['unsloth']" in out, out + result.stderr + # Deepseek is still registered via the normal path, just without originals. + deepseek_lines = [line for line in out.splitlines() if line.startswith("NUM_DEEPSEEK")] + assert deepseek_lines and int(deepseek_lines[0].split()[1]) > 0, out + result.stderr diff --git a/tests/version_compat/test_trl_fake_train_cpu.py b/tests/version_compat/test_trl_fake_train_cpu.py index 4dae696282..2f7b469428 100644 --- a/tests/version_compat/test_trl_fake_train_cpu.py +++ b/tests/version_compat/test_trl_fake_train_cpu.py @@ -158,6 +158,37 @@ except Exception: _MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM" +def _guard_finite_logits(model): + """Keep the LM head logits finite so GRPO sampling can't crash. + + ``test_grpo_trains_on_cpu`` samples completions from a tiny, *untrained* + random model on CPU. Driven autoregressively -- and nudged by the fake + reward's optimizer step between the two train steps -- such a model can emit + non-finite logits, so ``torch.multinomial`` inside ``generate()`` + intermittently raises "probability tensor contains either `inf`, `nan` or + element < 0". That is a well-known nondeterministic sampling failure, not an + Unsloth/TRL regression: the Trainer already fixes the seed, but CPU reduction + order is not bit-reproducible, so the blow-up still surfaces every so often. + + Sanitize the logits to a finite, bounded range (out of place, so autograd + stays valid) before they reach the sampler. This test asserts the train loop + runs end to end, not the (deliberately meaningless) numerics, so bounding the + logits changes nothing it checks while making the run reliable. + """ + + def _finite_logits_hook(_module, _inputs, output): + logits = getattr(output, "logits", None) + if logits is None: + return output + # nan_to_num maps nan -> 0 and the infinities to large finite values; + # clamp then bounds everything to [-30, 30]. + output.logits = torch.nan_to_num(logits).clamp(-30.0, 30.0) + return output + + model.register_forward_hook(_finite_logits_hook) + return model + + def _load_plain(): """Tiny plain HF model + tokenizer on CPU. Skips (not fails) if the model cannot be fetched -- that is a network/hub issue, not an unsloth regression.""" @@ -233,6 +264,11 @@ def test_grpo_trains_on_cpu(tmp_path): assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" model, tok = _load_plain() + # GRPO is the only canary that autoregressively samples completions, so it is + # the only one that can hit the non-finite-logits multinomial crash. Install + # the guard here (not in _load_plain) so the SFT/DPO canaries keep asserting + # against the model's true, unclamped outputs. + _guard_finite_logits(model) ds = Dataset.from_list([{"prompt": "hi there"}] * 4) cfg = GRPOConfig( output_dir = str(tmp_path / "ci_grpo"), diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index f47c78ba80..b857c34bcb 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2652,6 +2652,12 @@ extra_eos_tokens = None, "{{ '" + full_system + "' }}"\ "{% set loop_messages = messages %}"\ "{% endif %}" + elif "{SYSTEM}" in system_part: + # Only bind loop_messages when the template can render a caller system + # message. A static prefix with no {SYSTEM} must still raise, not drop it. + partial_system += "{% else %}"\ + "{% set loop_messages = messages %}"\ + "{% endif %}" else: partial_system += "{% endif %}" diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index e29190f0f2..618453fb82 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -171,8 +171,6 @@ def _list_deepseek_r1_distill_models(): return distill_models -register_deepseek_models(include_original_model = True) - if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info