187144d4e7Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison.
559 lines
19 KiB
Python
559 lines
19 KiB
Python
# 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 GGUF load-time context auto-fit decision.
|
|
|
|
Guards two regressions in ``LlamaCppBackend.load_model``:
|
|
|
|
1. Auto mode (``n_ctx == 0``) when weights exceed every GPU subset's free
|
|
memory: auto-pick should fall back to 4096 (a usable slider value) rather
|
|
than leaving native ctx. User can still drag higher onto ``--fit on``.
|
|
2. Explicit ctx must never be silently shrunk: when KV overflows fittable
|
|
weights, honor the explicit ctx with ``--fit on`` flexing ``-ngl``.
|
|
|
|
Drives the post-metadata decision block against a stubbed instance: no GPU,
|
|
network, subprocess, or GGUF I/O. Cross-platform.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stub heavy/unavailable deps before importing the module under test.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
_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")
|
|
sys.modules.setdefault("structlog", _structlog_stub)
|
|
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc_name in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
):
|
|
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
|
|
|
|
|
|
class _FakeTimeout:
|
|
def __init__(self, *a, **kw):
|
|
pass
|
|
|
|
|
|
_httpx_stub.Timeout = _FakeTimeout
|
|
_httpx_stub.Client = type(
|
|
"Client",
|
|
(),
|
|
{
|
|
"__init__": lambda self, **kw: None,
|
|
"__enter__": lambda self: self,
|
|
"__exit__": lambda self, *a: None,
|
|
},
|
|
)
|
|
sys.modules.setdefault("httpx", _httpx_stub)
|
|
|
|
from core.inference.llama_cpp import LlamaCppBackend
|
|
from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
GIB = 1024**3
|
|
FALLBACK_CTX = 4096
|
|
|
|
|
|
def _make_backend(
|
|
native_ctx = 131072,
|
|
n_layers = 80,
|
|
n_kv_heads = 8,
|
|
n_heads = 64,
|
|
kv_key_length = 128,
|
|
kv_value_length = 128,
|
|
):
|
|
"""LlamaCppBackend with GGUF metadata set and decision helpers stubbed."""
|
|
inst = LlamaCppBackend.__new__(LlamaCppBackend)
|
|
inst._context_length = native_ctx
|
|
inst._n_layers = n_layers
|
|
inst._n_kv_heads = n_kv_heads
|
|
inst._n_heads = n_heads
|
|
inst._embedding_length = 8192
|
|
inst._kv_key_length = kv_key_length
|
|
inst._kv_value_length = kv_value_length
|
|
inst._kv_lora_rank = None
|
|
inst._sliding_window = None
|
|
inst._sliding_window_pattern = None
|
|
inst._ssm_inner_size = None
|
|
inst._full_attention_interval = None
|
|
inst._key_length_mla = None
|
|
inst._n_kv_heads_by_layer = None
|
|
inst._kv_key_length_swa = None
|
|
inst._kv_value_length_swa = None
|
|
return inst
|
|
|
|
|
|
def _drive(
|
|
n_ctx,
|
|
model_gib,
|
|
gpus,
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 325_000,
|
|
can_estimate_kv = True,
|
|
extra_args = None,
|
|
):
|
|
"""Drive the post-metadata portion of load_model with stubbed inputs.
|
|
|
|
Mirrors llama_cpp.py:1137-1296 to assert the built command, without
|
|
subprocesses or GPU probes.
|
|
"""
|
|
inst = _make_backend(native_ctx = native_ctx)
|
|
model_size = int(model_gib * GIB)
|
|
cache_type_kv = None
|
|
|
|
def fake_estimate(
|
|
n_ctx_,
|
|
_type = None,
|
|
**_kwargs,
|
|
):
|
|
return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes
|
|
|
|
inst._estimate_kv_cache_bytes = fake_estimate
|
|
inst._can_estimate_kv = lambda: can_estimate_kv
|
|
|
|
context_length = inst._context_length
|
|
# Use the production helper, not a reimplementation, to avoid testing our own logic.
|
|
ctx_override = parse_ctx_override(extra_args)
|
|
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
|
|
|
|
effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0)
|
|
max_available_ctx = context_length or effective_ctx
|
|
if requested_ctx > 0:
|
|
effective_ctx = requested_ctx
|
|
elif context_length is not None:
|
|
effective_ctx = context_length
|
|
else:
|
|
effective_ctx = 0
|
|
original_ctx = effective_ctx
|
|
max_available_ctx = context_length or effective_ctx
|
|
|
|
gpu_indices, use_fit = None, True
|
|
explicit_ctx = requested_ctx > 0
|
|
|
|
if gpus and inst._can_estimate_kv() and effective_ctx > 0:
|
|
native_ctx_for_cap = context_length or effective_ctx
|
|
if native_ctx_for_cap > 0:
|
|
ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
|
|
best_cap = 0
|
|
for n_gpus in range(1, len(ranked_for_cap) + 1):
|
|
subset = ranked_for_cap[:n_gpus]
|
|
pool_mib = sum(free for _, free in subset)
|
|
capped = inst._fit_context_to_vram(
|
|
native_ctx_for_cap,
|
|
pool_mib,
|
|
model_size,
|
|
cache_type_kv,
|
|
)
|
|
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
|
|
total_mib = (model_size + kv) / (1024 * 1024)
|
|
if total_mib <= pool_mib * 0.90:
|
|
best_cap = max(best_cap, capped)
|
|
if best_cap > 0:
|
|
max_available_ctx = best_cap
|
|
|
|
if explicit_ctx:
|
|
requested_total = model_size + inst._estimate_kv_cache_bytes(
|
|
effective_ctx, cache_type_kv
|
|
)
|
|
gpu_indices, use_fit = inst._select_gpus(requested_total, gpus)
|
|
else:
|
|
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
|
matched = False
|
|
pin_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
|
for n_gpus in range(1, len(ranked) + 1):
|
|
subset = ranked[:n_gpus]
|
|
pool_mib = sum(free for _, free in subset)
|
|
capped = inst._fit_context_to_vram(
|
|
effective_ctx,
|
|
pool_mib,
|
|
model_size,
|
|
cache_type_kv,
|
|
)
|
|
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
|
|
total_mib = (model_size + kv) / (1024 * 1024)
|
|
if total_mib <= pool_mib * pin_fraction:
|
|
effective_ctx = capped
|
|
gpu_indices = sorted(idx for idx, _ in subset)
|
|
use_fit = False
|
|
matched = True
|
|
break
|
|
if not matched:
|
|
effective_ctx = min(FALLBACK_CTX, effective_ctx)
|
|
# Mirror llama_cpp.py: re-check fit at FALLBACK_CTX.
|
|
if effective_ctx > 0:
|
|
for n_gpus in range(1, len(ranked) + 1):
|
|
subset = ranked[:n_gpus]
|
|
pool_mib = sum(free for _, free in subset)
|
|
kv = inst._estimate_kv_cache_bytes(effective_ctx, cache_type_kv)
|
|
total_mib = (model_size + kv) / (1024 * 1024)
|
|
if total_mib <= pool_mib * pin_fraction:
|
|
gpu_indices = sorted(idx for idx, _ in subset)
|
|
use_fit = False
|
|
break
|
|
elif gpus:
|
|
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
|
|
if use_fit and not explicit_ctx:
|
|
effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
|
|
|
|
return {
|
|
"c_arg": effective_ctx if effective_ctx > 0 else 0,
|
|
"use_fit": use_fit,
|
|
"gpu_indices": gpu_indices,
|
|
"max_available_ctx": max_available_ctx,
|
|
"original_ctx": original_ctx,
|
|
"ctx_override": ctx_override,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Auto mode, model weights exceed VRAM (Bug A guard)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAutoModeWeightsExceedVRAM:
|
|
"""``n_ctx == 0`` on a model whose weights don't fit anywhere."""
|
|
|
|
def test_minimax_like_single_gpu(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
)
|
|
assert plan["c_arg"] == FALLBACK_CTX
|
|
assert plan["use_fit"] is True
|
|
assert plan["gpu_indices"] is None
|
|
# UI slider ceiling stays at native: user can drag higher and get
|
|
# the "might be slower" path.
|
|
assert plan["max_available_ctx"] == 196608
|
|
|
|
def test_multi_gpu_all_subsets_fail(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 400,
|
|
gpus = [(0, 80_000), (1, 80_000), (2, 80_000), (3, 80_000)],
|
|
native_ctx = 131072,
|
|
)
|
|
assert plan["c_arg"] == FALLBACK_CTX
|
|
assert plan["use_fit"] is True
|
|
assert plan["gpu_indices"] is None
|
|
|
|
def test_no_kv_metadata_auto(self):
|
|
"""File-size-only fallback path also defaults to 4096."""
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
can_estimate_kv = False,
|
|
)
|
|
assert plan["c_arg"] == FALLBACK_CTX
|
|
assert plan["use_fit"] is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Explicit ctx, KV overflows fittable weights (Bug B guard)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExplicitCtxRespectsUser:
|
|
"""``n_ctx > 0`` must never be silently shrunk."""
|
|
|
|
def test_fittable_weights_oversized_kv(self):
|
|
# 8 GB weights + 131k ctx KV on 24 GB VRAM. Budget = 21.6 GB, KV
|
|
# at 131k >> 13.6 GB remaining, so _select_gpus flips use_fit=True.
|
|
plan = _drive(
|
|
n_ctx = 131072,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
native_ctx = 131072,
|
|
)
|
|
assert plan["c_arg"] == 131072
|
|
assert plan["use_fit"] is True
|
|
assert plan["gpu_indices"] is None
|
|
|
|
def test_explicit_that_fits_uses_ngl(self):
|
|
plan = _drive(
|
|
n_ctx = 8192,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
native_ctx = 131072,
|
|
)
|
|
assert plan["c_arg"] == 8192
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0]
|
|
|
|
def test_explicit_on_weights_exceed_vram(self):
|
|
# User drags the slider to 32k on a too-big model: honored.
|
|
plan = _drive(
|
|
n_ctx = 32768,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
)
|
|
assert plan["c_arg"] == 32768
|
|
assert plan["use_fit"] is True
|
|
|
|
def test_explicit_at_fallback_on_too_big(self):
|
|
plan = _drive(
|
|
n_ctx = FALLBACK_CTX,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
)
|
|
assert plan["c_arg"] == FALLBACK_CTX
|
|
assert plan["use_fit"] is True
|
|
|
|
def test_explicit_below_floor_honored(self):
|
|
# 2048 is below --fit-ctx default; honored since user set it.
|
|
plan = _drive(
|
|
n_ctx = 2048,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
)
|
|
assert plan["c_arg"] == 2048
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pass-through --ctx-size participates in context fit (#5676).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtraArgsCtxOverride:
|
|
def test_ctx_size_extra_honored_over_auto(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
extra_args = ["--ctx-size", "128000"],
|
|
)
|
|
assert plan["ctx_override"] == 128000
|
|
assert plan["original_ctx"] == 128000
|
|
assert plan["c_arg"] == 128000
|
|
assert plan["use_fit"] is True
|
|
|
|
def test_ctx_size_short_alias_honored_over_auto(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 131,
|
|
gpus = [(0, 97_000)],
|
|
native_ctx = 196608,
|
|
extra_args = ["-c", "128000"],
|
|
)
|
|
assert plan["c_arg"] == 128000
|
|
assert plan["use_fit"] is True
|
|
|
|
def test_ctx_size_extra_wins_over_first_class_field(self):
|
|
plan = _drive(
|
|
n_ctx = 4096,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
native_ctx = 131072,
|
|
extra_args = ["--ctx-size", "128000"],
|
|
)
|
|
assert plan["original_ctx"] == 128000
|
|
assert plan["c_arg"] == 128000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Non-regression: fittable + auto still auto-picks largest fitting ctx
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFittableAutoPickRegressions:
|
|
def test_small_model_one_gpu(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 8192,
|
|
)
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0]
|
|
assert plan["c_arg"] > FALLBACK_CTX
|
|
|
|
def test_medium_model_needs_multi_gpu(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 60,
|
|
gpus = [(0, 40_000), (1, 40_000)],
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 8192,
|
|
)
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0, 1]
|
|
|
|
def test_no_kv_metadata_fittable_auto(self):
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 8,
|
|
gpus = [(0, 24_000)],
|
|
native_ctx = 131072,
|
|
can_estimate_kv = False,
|
|
)
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #5106 regression: 91-95% utilization must still pin GPU.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTightFitPinsToGPU:
|
|
"""Models that fit at 91-95% of free VRAM must use the GPU."""
|
|
|
|
def test_rtx_4090_qwen_24gb_class(self):
|
|
# noahterbest's #5106 log: 20.8 GB model on 22805 MiB free GPU,
|
|
# ctx=4096 -> ~94% utilization, ~1.4 GiB headroom.
|
|
plan = _drive(
|
|
n_ctx = 0,
|
|
model_gib = 20.8,
|
|
gpus = [(0, 22_805)],
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 25_000,
|
|
)
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0]
|
|
|
|
def test_explicit_ctx_at_94_pct_pins_to_gpu(self):
|
|
# Explicit-ctx branch must agree with auto-ctx on headroom.
|
|
plan = _drive(
|
|
n_ctx = 4096,
|
|
model_gib = 20.8,
|
|
gpus = [(0, 22_805)],
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 25_000,
|
|
)
|
|
assert plan["use_fit"] is False
|
|
assert plan["gpu_indices"] == [0]
|
|
|
|
def test_genuine_overflow_still_uses_fit(self):
|
|
# Beyond 95% must still defer to --fit on.
|
|
plan = _drive(
|
|
n_ctx = 4096,
|
|
model_gib = 23,
|
|
gpus = [(0, 22_000)],
|
|
native_ctx = 131072,
|
|
kv_per_token_bytes = 25_000,
|
|
)
|
|
assert plan["use_fit"] is True
|
|
assert plan["gpu_indices"] is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Platform-agnostic input shape
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"])
|
|
def test_identical_decision_across_platforms(platform_tag):
|
|
"""Decision takes ``[(gpu_idx, free_mib), ...]`` regardless of source;
|
|
identical inputs must yield identical plans."""
|
|
plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
|
|
plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
|
|
assert plan_a == plan_b, platform_tag
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _classify_gpu_offload: detect silent CPU fallback (#5106).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestClassifyGpuOffload:
|
|
def _backend(self, stdout_lines):
|
|
inst = LlamaCppBackend.__new__(LlamaCppBackend)
|
|
inst._stdout_lines = list(stdout_lines)
|
|
return inst
|
|
|
|
def test_cuda_buffer_present_returns_true(self):
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: offloaded 33/33 layers to GPU",
|
|
"load_tensors: CUDA0 model buffer size = 21000.0 MiB",
|
|
"load_tensors: CPU_Mapped model buffer size = 0.6 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|
|
|
|
def test_cpu_only_buffer_returns_false(self):
|
|
# Buffer lines printed but only CPU buffers -- the silent CPU
|
|
# fallback symptom we want to catch.
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
|
"load_tensors: CPU model buffer size = 0.6 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
|
|
|
|
def test_no_buffer_lines_returns_none(self):
|
|
# If we can't see buffer-allocation lines at all, don't guess.
|
|
inst = self._backend(
|
|
[
|
|
"INFO [main] starting server",
|
|
"load_tensors: file format = GGUF V3",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(True, [(0, 22805)]) is None
|
|
|
|
def test_no_gpus_detected_returns_none(self):
|
|
# CPU-only systems are valid; suppress the warning entirely.
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(False, []) is None
|
|
|
|
def test_user_did_not_intend_gpu_returns_none(self):
|
|
# Studio called start_llama_server without expecting GPU; don't warn.
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(False, [(0, 22805)]) is None
|
|
|
|
def test_rocm_buffer_marker_returns_true(self):
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: ROCm0 model buffer size = 21000.0 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|
|
|
|
def test_metal_buffer_marker_returns_true(self):
|
|
inst = self._backend(
|
|
[
|
|
"load_tensors: Metal model buffer size = 8000.0 MiB",
|
|
]
|
|
)
|
|
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|