ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into fix/studio-export-multi-gpu-device-map

Browse source
This commit is contained in:
Hakan Baysal 2026-07-19 16:33:57 +03:00 • committed by GitHub
commit a94677c31f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
70 changed files with 6313 additions and 554 deletions
Download patch file Download diff file Expand all files Collapse all files

17
.github/scripts/agent-guides-drive.sh vendored
View file

@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh <mode> <agent>}"
# 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

168
install.sh
View file

@ -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],<D.E.F"). Compares at major.minor granularity, which is exact
# for the windows this script uses (ceilings are always X.Y.0); a non-.0 ceiling
# would only make this conservative (excludes the whole ceiling minor). Anything
# unparseable answers "no" so the caller fails toward the supported range.
_torch_release_in_window() {
_trw_con="$2"
case "$_trw_con" in
"torch>="*",<"*) ;;
*) 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

715
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

27
studio/backend/core/inference/llama_server_args.py
View file

@ -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] = []

6
studio/backend/core/inference/local_model_resolver.py
View file

@ -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)

25
studio/backend/core/rag/config.py
View file

@ -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

40
studio/backend/hub/services/models/cache_inventory.py
View file

@ -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,

17
studio/backend/hub/services/models/local_inventory.py
View file

@ -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),

338
studio/backend/hub/tests/test_model_services.py
View file

@ -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)

14
studio/backend/main.py
View file

@ -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

144
studio/backend/models/inference.py
View file

@ -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 = (

518
studio/backend/routes/inference.py
View file

@ -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,

67
studio/backend/routes/models.py
View file

@ -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

5
studio/backend/routes/settings.py
View file

@ -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,
)

42
studio/backend/routes/training_vram.py
View file

@ -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"}

139
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -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.

333
studio/backend/tests/test_chat_load_during_training.py
View file

@ -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) ──────

73
studio/backend/tests/test_cuda_torch_spec.py Normal file
View file

@ -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}"

17
studio/backend/tests/test_embedding_model_security_gate.py
View file

@ -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):

7
studio/backend/tests/test_embedding_model_settings.py
View file

@ -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"

73
studio/backend/tests/test_gguf_metadata.py
View file

@ -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(

879
studio/backend/tests/test_gpu_memory_mode.py Normal file
View file

@ -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

18
studio/backend/tests/test_gpu_selection.py
View file

@ -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(

12
studio/backend/tests/test_llama_cpp_no_context_shift.py
View file

@ -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):

31
studio/backend/tests/test_llama_cpp_props_readback.py
View file

@ -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
)

45
studio/backend/tests/test_llama_server_args.py
View file

@ -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"],

27
studio/backend/tests/test_openai_auto_switch.py
View file

@ -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.

5
studio/backend/tests/test_tensor_parallel.py
View file

@ -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():

17
studio/backend/tests/test_tp_vision_regression.py
View file

@ -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)

109
studio/backend/tests/test_training_config_popover_source.py Normal file
View file

@ -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}"

142
studio/backend/utils/hidden_models.py Normal file
View file

@ -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

109
studio/backend/utils/models/gguf_metadata.py
View file

@ -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}.<suffix>``; 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("<I" if vtype == 4 else "<Q", n_bytes)[0]
# A real context length is positive; treat 0/garbage as
# absent so the UI never builds a slider with max < min.
return value if value > 0 else None
found[key[len(arch) + 1 :]] = struct.unpack(
"<I" if vtype == 4 else "<Q", n_bytes
)[0]
if len(found) == len(wanted_suffixes):
break
else:
if not _skip_gguf_value(f, vtype):
break
except (struct.error, UnicodeDecodeError):
break
except OSError as e:
logger.debug(f"read_gguf_context_length: cannot open {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: cannot open {path}: {e}")
return None
except Exception as e:
logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: parse failure on {path}: {e}")
return None
return None
return found
def _parse_gguf_staged_dims(path: str) -> 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.

24
studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
View file

@ -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 <model>".
// dialog can offer "Remember settings for <model>". 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;

128
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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<boolean> {
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),

38
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -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<ValidateModelResponse>(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<number | null> {
}): 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<ValidateModelResponse>(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<void> {

13
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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.",

450
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -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 (
<div className="space-y-3.5">
@ -279,6 +285,7 @@ function ParamSlider({
displayValue={displayValue}
ariaLabel={label}
size={valueSize ?? 4}
disabled={disabled}
/>
</div>
<Slider
@ -288,6 +295,7 @@ function ParamSlider({
value={[value]}
onValueChange={([v]) => onChange(snapToStep(v, step, min, max))}
className="panel-slider"
disabled={disabled}
/>
</div>
);
@ -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 ? (
<div className="space-y-3.5">
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Context Length
</span>
<InfoHint>
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.
</InfoHint>
</div>
<NumericValueInput
value={fitCtxAuto ? 0 : (customContextLength ?? 0)}
displayValue={fitCtxAuto ? "Auto" : undefined}
min={0}
max={ctxMaxValue ?? undefined}
step={1}
onChange={(v) => {
setCustomContextLength(v > 0 ? v : null);
}}
ariaLabel="Context Length"
size={8}
disabled={modelControlsDisabled}
/>
</div>
<Slider
min={0}
max={ctxMaxValue ?? 4096}
step={1024}
value={[
fitCtxAuto
? 0
: Math.min(
customContextLength ?? 0,
ctxMaxValue ?? 4096,
),
]}
onValueChange={([v]) => {
// 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 && (
<p className="text-[11px] text-nav-fg/40">
llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens.
</p>
)}
</div>
) : (
<div className="space-y-3.5">
<div className="flex items-center justify-between gap-3">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -1036,7 +1253,7 @@ export function ChatSettingsPanel({
</p>
)}
</div>
)}
))}
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -1191,6 +1408,163 @@ export function ChatSettingsPanel({
)}
</>
)}
{gpuModeApplies && (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
GPU Memory
</span>
<InfoHint>
<div className="flex flex-col gap-1.5">
<div>
<span className="font-medium">Default:</span> Unsloth
fits the model and context to your GPUs.
</div>
<div>
<span className="font-medium">Manual:</span> set GPU
Layers yourself. Leave it on Auto to let llama.cpp size
the context and offload overflow (including MoE experts)
to RAM.
</div>
</div>
</InfoHint>
</div>
<div className="flex shrink-0 items-center gap-1.5">
<Select
value={gpuMemoryMode}
onValueChange={(v) => {
setGpuMemoryMode(v as "auto" | "manual");
}}
// An in-flight staged load already snapshotted its
// settings, so edits here could not apply -- disable like
// the sibling context/KV/spec controls.
disabled={modelControlsDisabled}
>
<SelectTrigger
animateRadius={false}
icon={ChevronDownStandardIcon}
iconClassName="size-3.5"
className="grid h-7 w-[136px] min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 rounded-full border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.1] pl-3 pr-2 py-0 text-[13px]! font-medium text-nav-fg focus-visible:ring-0 focus-visible:border-transparent [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0"
data-test-id="gpu-memory-mode-select"
>
<SelectValue />
</SelectTrigger>
<SelectContent className="menu-soft-surface ring-0 border-0 rounded-lg">
<SelectItem value="auto">Default</SelectItem>
<SelectItem value="manual">Manual</SelectItem>
</SelectContent>
</Select>
</div>
</div>
)}
{isManual && (
<>
<ParamSlider
label="GPU Layers"
value={Math.max(GPU_LAYERS_AUTO, Math.min(gpuLayers, gpuLayersMax))}
min={GPU_LAYERS_AUTO}
max={gpuLayersMax}
step={1}
onChange={setGpuLayers}
disabled={modelControlsDisabled}
displayValue={autoLayers ? "Auto" : undefined}
valueSize={6}
info={
<>
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 && (
<ParamSlider
label="MoE Layers on CPU"
value={Math.min(nCpuMoe, moeLayersMax)}
min={0}
max={moeLayersMax}
step={1}
onChange={setNCpuMoe}
disabled={modelControlsDisabled}
valueSize={6}
info={
<>
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 && (
<div className="space-y-3.5">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Layers per GPU
</span>
<InfoHint>
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.
</InfoHint>
</div>
{gpusInUseDevices.map((d, k) => (
<ParamSlider
key={d?.index ?? k}
label={`GPU ${d?.index ?? k}`}
value={Math.min(splitCounts[k] ?? 0, splitTotal)}
min={0}
max={splitTotal}
step={1}
onChange={(v) => setSplitCount(k, v)}
valueSize={6}
disabled={modelControlsDisabled}
/>
))}
</div>
)}
</>
)}
{showGpuPicker && (
<div className="space-y-2">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
GPUs
</span>
<InfoHint>
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.
</InfoHint>
</div>
<div className="flex flex-col gap-2">
{gpuDevices.map((d) => (
<div
key={d.index}
className="flex items-center justify-between gap-3"
>
<span className="min-w-0 truncate text-[12px] text-nav-fg/80">
GPU {d.index}: {d.name}
{d.memoryTotalGb
? ` · ${Math.round(d.memoryTotalGb)} GB`
: ""}
</span>
<Switch
className="panel-switch shrink-0"
checked={isGpuChecked(d.index)}
onCheckedChange={() => toggleGpu(d.index)}
data-test-id={`gpu-pick-${d.index}`}
disabled={modelControlsDisabled}
/>
</div>
))}
</div>
</div>
)}
{gpuModeApplies && !autoLayers && (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -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"
/>
</div>
)}
</>
)}
{/* No persistent "enable custom code" toggle: it is consented per model
@ -1228,14 +1603,21 @@ export function ChatSettingsPanel({
{Math.round((stagedDownloadFraction ?? 0) * 100)}%
</p>
)}
<label className="flex cursor-pointer items-center gap-2 pb-1.5 text-[12px] text-muted-foreground">
<Checkbox
className="size-3.5 rounded-full [&_[data-slot=checkbox-indicator]_svg]:size-2.5"
checked={remember}
onCheckedChange={(v) => setRemember(v === true)}
/>
Remember settings next time
</label>
{/* GGUF picks only: a non-GGUF pick shows none of the load
knobs the blob captures, so there is nothing to remember. */}
{pendingIsGguf && (
<label className="flex cursor-pointer items-center gap-2 pb-1.5 text-[12px] text-muted-foreground">
<Checkbox
className="size-3.5 rounded-full [&_[data-slot=checkbox-indicator]_svg]:size-2.5"
checked={remember}
onCheckedChange={(v) => setRemember(v === true)}
// The save/clear already ran in the Load click handler, so
// a mid-load toggle could not apply -- lock it like the knobs.
disabled={modelControlsDisabled}
/>
Remember settings next time
</label>
)}
{stagedLoading ? (
// Mid-load: nothing to load or abandon until it settles, so disable.
<Button
@ -1255,9 +1637,10 @@ export function ChatSettingsPanel({
// Persist (or clear) this model's load knobs before loading.
// Context is stored as the override (null = auto), never the
// resolved native value, so restoring can't force an OOM.
const pid = pendingSelection
? rememberedLoadSettingsKey(pendingSelection)
: null;
// GGUF-only, like the restore effect: saving for a
// non-GGUF pick would snapshot leftover standing values
// its hidden controls never showed.
const pid = pendingIsGguf ? pendingKey : null;
if (pid) {
if (remember) {
saveRememberedLoadSettings(pid, {
@ -1266,6 +1649,10 @@ export function ChatSettingsPanel({
speculativeType,
specDraftNMax,
tensorParallel,
gpuMemoryMode,
gpuLayers,
nCpuMoe,
selectedGpuIds,
});
} else {
clearRememberedLoadSettings(pid);
@ -1319,7 +1706,11 @@ export function ChatSettingsPanel({
</Button>
</div>
) : null}
<ChatTemplateFields />
{/* 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. */}
<ChatTemplateFields disabled={modelControlsDisabled} />
</div>
</CollapsibleSection>
)}
@ -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() {
<button
type="button"
onClick={openEditor}
className="cursor-pointer text-left text-[13px] font-medium tracking-nav text-nav-fg"
disabled={disabled}
className="cursor-pointer text-left text-[13px] font-medium tracking-nav text-nav-fg disabled:pointer-events-none disabled:opacity-50"
>
Chat Template
</button>
@ -2131,7 +2523,8 @@ function ChatTemplateFields() {
<button
type="button"
onClick={() => setOverride(null)}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white"
disabled={disabled}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
aria-label="Revert chat template"
>
<HugeiconsIcon
@ -2155,7 +2548,8 @@ function ChatTemplateFields() {
<button
type="button"
onClick={openEditor}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white"
disabled={disabled}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
aria-label="Edit chat template"
>
<HugeiconsIcon
@ -2218,7 +2612,13 @@ function ChatTemplateFields() {
>
Cancel
</Button>
<Button type="button" onClick={saveEditor} disabled={!draftDirty}>
{/* Also locked mid-load: an autoLoad can start with this dialog
already open, and a save then would be silently clobbered. */}
<Button
type="button"
onClick={saveEditor}
disabled={!draftDirty || disabled}
>
Save
</Button>
</div>

217
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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 {

46
studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
View file

@ -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,
]);

117
studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
View file

@ -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<ReasoningEffort, "low" | "medium" | "high">;
function sameArray<T>(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<boolean> {
if (previousCheckpoint) {
return true;
}
const previousGgufVariant = useChatRuntimeStore.getState().activeGgufVariant;
store.setCheckpoint(checkpointId, status.gguf_variant);
applyActiveModelStatusToStore(status, { previousCheckpoint });
applyActiveModelStatusToStore(status, {
previousCheckpoint,
previousGgufVariant,
});
return true;
}

31
studio/frontend/src/features/chat/presets/preset-policy.ts
View file

@ -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;
}

109
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -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),
});

369
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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<T extends object>(
resp: Parameters<typeof loadedGpuMemoryFields>[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<K extends ScalarSettingKey>(
/** 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<ChatRuntimeStore>((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<ChatRuntimeStore>((set, get) => ({
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
loadedCustomContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@ -1455,9 +1736,23 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((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,
};
});
},

38
studio/frontend/src/features/chat/types/api.ts
View file

@ -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

46
studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts Normal file
View file

@ -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<string> {
const inventoryVersion = useInventoryVersion();
const [hiddenIds, setHiddenIds] = useState<ReadonlySet<string>>(
() => 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;
}

53
studio/frontend/src/features/hub/hub-page.tsx
View file

@ -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<ModelFormatFilter>("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

1
studio/frontend/src/features/hub/index.ts
View file

@ -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,

11
studio/frontend/src/features/hub/inventory/inventory-hints.ts
View file

@ -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)];

1
studio/frontend/src/features/hub/inventory/types.ts
View file

@ -54,6 +54,7 @@ export interface CachedInventoryRow {
libraryName?: string | null;
quantMethod?: string | null;
liveDownload?: boolean;
optimistic?: boolean;
}
export interface LocalInventoryRow {

1
studio/frontend/src/features/hub/inventory/use-hub-inventory.ts
View file

@ -204,6 +204,7 @@ function liveDownloadInventoryRows(
size_bytes: job.displayBytes,
partial: true,
partial_transport: null,
optimistic: true,
},
modelFormat,
),

18
studio/frontend/src/features/hub/inventory/view-models.ts
View file

@ -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,
};
}

26
studio/frontend/src/features/hub/lib/hidden-models.ts
View file

@ -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<string>,
...values: (string | null | undefined)[]
): boolean {
return values.some(
(value) => value != null && configuredIds.has(value.trim().toLowerCase()),
);
}

17
studio/frontend/src/features/settings/api/embedding-model.ts
View file

@ -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<EmbeddingModelSettings> {
@ -87,5 +98,7 @@ export async function resetEmbeddingModelSettings(): Promise<EmbeddingModelSetti
await readFastApiError(res, "Failed to reset embedding model"),
);
}
return fromApi(await res.json());
const settings = fromApi(await res.json());
bumpInventoryVersion();
return settings;
}

1
studio/frontend/src/features/settings/index.ts
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { SettingsDialog } from "./settings-dialog";
export { loadEmbeddingModelSettings } from "./api/embedding-model";
export {
loadPersonalization,
savePersonalization,

21
studio/frontend/src/features/studio/historical-training-view.tsx
View file

@ -8,6 +8,7 @@ import { parseBackendTrainingMethod } from "@/features/training/lib/training-met
import { type ReactElement, useEffect, useState } from "react";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
import { mapRunConfigToOverride } from "./sections/run-config-override";
import { translate, useT } from "@/i18n";
type StudioT = ReturnType<typeof useT>;
@ -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 (
<div className="flex flex-col gap-6">

91
studio/frontend/src/features/studio/live-training-view.tsx
View file

@ -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<typeof setTimeout> | 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 {
)}
>
<div data-tour="studio-training-progress">
<ProgressSection key={runtime.jobId ?? "no-job"} data={viewData} />
<ProgressSection
key={runtime.jobId ?? "no-job"}
data={viewData}
configOverride={runConfigOverride}
/>
</div>
<ChartsSection
currentStep={viewData.currentStep}

40
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -28,6 +28,7 @@ import {
} from "@/features/training";
import { getTrainingMethodLabel } from "@/features/training/lib/training-methods";
import type { TrainingViewData } from "@/features/training";
import type { RunConfigOverride } from "./run-config-override";
import { useGpuUtilization } from "@/hooks";
import type { GpuUtilization } from "@/hooks/use-gpu-utilization";
import { cn } from "@/lib/utils";
@ -81,19 +82,7 @@ function configRow(
interface ProgressSectionProps {
data: TrainingViewData;
isHistorical?: boolean;
configOverride?: {
epochs?: number;
batchSize?: number;
learningRate?: string;
maxSteps?: number;
contextLength?: number;
warmupSteps?: number;
optimizerType?: string;
loraRank?: number;
loraAlpha?: number;
loraDropout?: number;
loraVariant?: string;
};
configOverride?: RunConfigOverride;
}
export function ProgressSection({
@ -183,17 +172,20 @@ export function ProgressSection({
? data.currentGradNorm
: (lastValue(data.gradNormHistory) ?? data.currentGradNorm);
const cfgEpochs = isHistorical ? configOverride?.epochs : config.epochs;
const cfgBatchSize = isHistorical ? configOverride?.batchSize : config.batchSize;
const cfgLearningRate = isHistorical ? configOverride?.learningRate : config.learningRate;
const cfgMaxSteps = isHistorical ? configOverride?.maxSteps : config.maxSteps;
const cfgContextLength = isHistorical ? configOverride?.contextLength : config.contextLength;
const cfgWarmupSteps = isHistorical ? configOverride?.warmupSteps : config.warmupSteps;
const cfgOptimizerType = isHistorical ? configOverride?.optimizerType : config.optimizerType;
const cfgLoraRank = isHistorical ? configOverride?.loraRank : config.loraRank;
const cfgLoraAlpha = isHistorical ? configOverride?.loraAlpha : config.loraAlpha;
const cfgLoraDropout = isHistorical ? configOverride?.loraDropout : config.loraDropout;
const cfgLoraVariant = isHistorical ? configOverride?.loraVariant : config.loraVariant;
// Prefer the run's saved snapshot when present (#6853). Live falls back to the
// editable form store until it loads; History shows blanks, never live form values.
const cfg = configOverride ?? (isHistorical ? undefined : config);
const cfgEpochs = cfg?.epochs;
const cfgBatchSize = cfg?.batchSize;
const cfgLearningRate = cfg?.learningRate;
const cfgMaxSteps = cfg?.maxSteps;
const cfgContextLength = cfg?.contextLength;
const cfgWarmupSteps = cfg?.warmupSteps;
const cfgOptimizerType = cfg?.optimizerType;
const cfgLoraRank = cfg?.loraRank;
const cfgLoraAlpha = cfg?.loraAlpha;
const cfgLoraDropout = cfg?.loraDropout;
const cfgLoraVariant = cfg?.loraVariant;
const optimizerLabel =
OPTIMIZER_OPTIONS.find((o) => o.value === cfgOptimizerType)?.label ??

54
studio/frontend/src/features/studio/sections/run-config-override.ts Normal file
View file

@ -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<string, unknown> | 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",
};
}

1
studio/frontend/src/features/training/index.ts
View file

@ -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";

186
studio/frontend/src/hooks/use-gpu-info.ts
View file

@ -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<GpuInfo> | null = null;
// One module-level cache so every GPU hook shares a single /api/system fetch.
let cachedSystem: SystemInfoResponse | null = null;
let systemPromise: Promise<SystemInfoResponse | null> | null = null;
async function fetchGpuOnce(): Promise<GpuInfo> {
if (cachedGpu) return cachedGpu;
if (fetchPromise) return fetchPromise;
fetchPromise = (async () => {
async function fetchSystemOnce(): Promise<SystemInfoResponse | null> {
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<GpuInfo>(
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<SystemGpuDevice[]>(
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<GpuInfo>(cachedGpu ?? DEFAULT_GPU);
export async function ensureGpuDeviceCache(): Promise<void> {
await fetchSystemOnce();
}
useEffect(() => {
if (cachedGpu) return;
let cancelled = false;
fetchGpuOnce().then((info) => {
if (!cancelled) setGpu(info);
});
return () => { cancelled = true; };
}, []);
return gpu;
}
/**
* 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) : [];
}

3
studio/frontend/src/hooks/use-system.ts
View file

@ -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;

90
tests/python/test_construct_chat_template_validation.py
View file

@ -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}</s>"
"### User: {INPUT}\n### Assistant: {OUTPUT}</s>"
)
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 = "<s>",
eos_token = "</s>",
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 = ["</s>"],
)
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 = ["</s>"],
)
rendered = _render(jinja_template, [{"role": "user", "content": "Hi"}])
assert "Hi" in rendered
_NO_SYSTEM_CHAT_TEMPLATE = (
"PREAMBLE\n"
"### User: {INPUT}\n### Assistant: {OUTPUT}</s>"
"### User: {INPUT}\n### Assistant: {OUTPUT}</s>"
)
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 = ["</s>"],
)
with pytest.raises(RuntimeError, match = "Only user and assistant roles are supported!"):
_render(
jinja_template,
[
{"role": "system", "content": "Be terse."},
{"role": "user", "content": "Hi"},
],
)

80
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -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

101
tests/sh/test_previous_torch_pin.sh Normal file
View file

@ -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"

14
tests/sh/test_torch_constraint.sh
View file

@ -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 ==="

131
tests/sh/test_unsloth_torch_override.sh Normal file
View file

@ -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

85
tests/test_model_registry.py
View file

@ -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

36
tests/version_compat/test_trl_fake_train_cpu.py
View file

@ -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"),

6
unsloth/chat_templates.py
View file

@ -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 %}"

2
unsloth/registry/_deepseek.py
View file

@ -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