Studio: pin GPU at 95% headroom and warn on silent CPU fallback (#5323)
* Studio: pin GPU at 95% headroom and warn on silent CPU fallback
Two related runtime-side fixes for unslothai/unsloth#5106 ("model
loaded fully on RAM instead of VRAM"):
1. GPU pin threshold bump 0.90 -> 0.95
-------------------------------------
``_select_gpus`` and the auto-ctx pin loop in ``start_llama_server``
used a ``pool * 0.90`` threshold to decide whether the model fits on
GPU. Models that needed 91-94% of free VRAM were classified as "does
not fit", so Studio set ``gpu_indices = None`` and shipped
``--fit on`` to llama-server without ``-ngl``. The unsloth
llama.cpp fork's ``--fit on`` then ran with its default
``--fit-target 1024`` (1 GiB margin per device, an upstream default
inherited from ggml-org#18679). On a tight fit where compute
buffers + CUDA context push the projected free below the 1 GiB
target, the fork's fit logic shaves layer weights off the GPU --
slow inference for users whose models would have loaded comfortably
with ``-ngl -1``.
The classic reproducer from #5106 (noahterbest's log):
GGUF size: 20.8 GB, est. KV cache: 0.1 GB, context: 4096,
GPUs free: [(0, 22805)], selected: None, fit: True
20.8 GiB on a 22.27 GiB free RTX 4090 is 94% utilization. The model
fits (1.4 GiB headroom), but the 0.90 threshold kicks it to fit
mode. Bumping to 0.95 keeps these in the fits-on-GPU branch and
emits ``-ngl -1`` directly. The fork's ``--fit on`` still serves as
the safety net for the genuinely-too-large case.
The auto-ctx fallback also re-checks fit at 4096 before handing off
to ``--fit on``: a 20.8 GiB model with a 131072 native context fails
the auto loop at native ctx, falls back to ``min(4096, ctx)``, but
its weights + 4096 KV pin to the GPU comfortably. Without the
re-check we still emitted ``--fit on``.
``_fit_context_to_vram``'s 0.90 budget for context binary search is
intentionally left tighter than the pin fraction. That routine
chooses the slider value, where over-promising would OOM at runtime.
``_select_gpus`` decides whether to pin at all, where being
conservative pushes layers to CPU.
2. Belt-and-suspenders: warn on silent CPU fallback
---------------------------------------------------
After ``_wait_for_health`` succeeds, scan llama-server's stdout for
``model buffer size`` lines. If Studio detected GPUs and intended
GPU use but only CPU buffers were allocated, log a structured
warning citing #5106. Markers cover CUDA / ROCm / Metal / Vulkan /
OpenCL / SYCL backends. New ``_gpu_offload_active: Optional[bool]``
field surfaces the result for any future API consumer.
This catches runtime-load failures the install-time fix cannot
cover (cudart bundle pairing PR #5322 is the install-side
companion): user overriding ``--fit-target``, uncommon driver +
toolkit configurations, future regressions in the install path.
Tests: 10 new cases in studio/backend/tests/test_llama_cpp_context_fit.py:
* TestTightFitPinsToGPU x3: noahterbest's exact reproducer (auto and
explicit ctx pins to GPU at 94%); guard against threshold over-
broadening (genuine overflow still falls back to ``--fit on``).
* TestClassifyGpuOffload x7: CUDA / ROCm / Metal buffer markers
return True; CPU-only buffer lines return False; absent buffer
lines or no GPUs detected return None (no warning).
25 context-fit tests pass (15 baseline + 10 new). 511 tests total
across the affected test files. No regressions.
Refs #5106
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments to be more succinct
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
213666a23f
commit
5205bc0ed6
2 changed files with 224 additions and 18 deletions
|
|
@ -433,6 +433,8 @@ class LlamaCppBackend:
|
|||
self._hf_variant: Optional[str] = None
|
||||
self._is_vision: bool = False
|
||||
self._healthy = False
|
||||
# Set by _classify_gpu_offload after _wait_for_health.
|
||||
self._gpu_offload_active: Optional[bool] = None
|
||||
self._context_length: Optional[int] = None
|
||||
self._effective_context_length: Optional[int] = None
|
||||
self._max_context_length: Optional[int] = None
|
||||
|
|
@ -956,6 +958,13 @@ class LlamaCppBackend:
|
|||
logger.debug(f"torch GPU probe failed: {e}")
|
||||
return []
|
||||
|
||||
# Free-VRAM fraction at which Studio pins the GPU directly instead
|
||||
# of deferring to ``--fit on``. 5% headroom covers CUDA context +
|
||||
# compute buffers; 0.90 was too conservative and dropped 91-94%
|
||||
# fits to CPU offload (#5106). The fork's --fit on still catches
|
||||
# the truly-too-large case.
|
||||
_GPU_PIN_VRAM_FRACTION = 0.95
|
||||
|
||||
@staticmethod
|
||||
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
|
||||
"""Return DLL dirs from pip-installed CUDA wheels under
|
||||
|
|
@ -1024,11 +1033,11 @@ class LlamaCppBackend:
|
|||
"""Pick GPU(s) for a model based on estimated VRAM and free memory.
|
||||
|
||||
``model_size_bytes`` should include both model weights and estimated
|
||||
KV cache. The 90% threshold provides headroom for compute buffers,
|
||||
CUDA context, and other runtime overhead.
|
||||
KV cache. The ``_GPU_PIN_VRAM_FRACTION`` threshold provides headroom
|
||||
for compute buffers, CUDA context, and other runtime overhead.
|
||||
|
||||
Returns (gpu_indices, use_fit):
|
||||
- ([1], False) model fits on 1 GPU at 90% of free
|
||||
- ([1], False) model fits on 1 GPU at the headroom threshold
|
||||
- ([1, 2], False) model needs 2 GPUs
|
||||
- (None, True) model too large, let --fit handle it
|
||||
"""
|
||||
|
|
@ -1036,12 +1045,13 @@ class LlamaCppBackend:
|
|||
return None, True
|
||||
|
||||
model_size_mib = model_size_bytes / (1024 * 1024)
|
||||
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
||||
|
||||
# Sort GPUs by free memory descending
|
||||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||||
|
||||
# Try fitting on 1 GPU (90% of free memory threshold)
|
||||
if ranked[0][1] * 0.90 >= model_size_mib:
|
||||
# Try fitting on 1 GPU at the usable-VRAM threshold.
|
||||
if ranked[0][1] * usable_fraction >= model_size_mib:
|
||||
return [ranked[0][0]], False
|
||||
|
||||
# Try fitting on N GPUs (accumulate free memory from most-free)
|
||||
|
|
@ -1049,7 +1059,7 @@ class LlamaCppBackend:
|
|||
selected = []
|
||||
for idx, free_mib in ranked:
|
||||
selected.append(idx)
|
||||
cumulative += free_mib * 0.90
|
||||
cumulative += free_mib * usable_fraction
|
||||
if cumulative >= model_size_mib:
|
||||
return sorted(selected), False
|
||||
|
||||
|
|
@ -1282,10 +1292,11 @@ class LlamaCppBackend:
|
|||
) -> int:
|
||||
"""Return the largest context length that fits in GPU VRAM.
|
||||
|
||||
Uses 90% of available VRAM as the budget (matching _select_gpus
|
||||
threshold -- 10% reserved for compute buffers, CUDA context,
|
||||
scratch space, flash-attn workspace, etc.).
|
||||
If the model weights alone don't fit, returns min_ctx unchanged.
|
||||
Uses 90% of available VRAM as the ctx-fit budget. Tighter than
|
||||
``_GPU_PIN_VRAM_FRACTION`` on purpose: over-promising context
|
||||
OOMs at runtime, while pinning conservatively just defers to
|
||||
--fit on. If the weights alone don't fit, returns
|
||||
``requested_ctx`` unchanged.
|
||||
|
||||
``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False
|
||||
the KV cache lives in CPU RAM and doesn't compete with weights
|
||||
|
|
@ -2031,6 +2042,7 @@ class LlamaCppBackend:
|
|||
# still has valid state to publish.
|
||||
effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
|
||||
max_available_ctx = self._context_length or effective_ctx
|
||||
gpus: list[tuple[int, int]] = []
|
||||
try:
|
||||
model_size = self._get_gguf_size_bytes(model_path)
|
||||
gpus = self._get_gpu_free_memory()
|
||||
|
|
@ -2114,8 +2126,11 @@ class LlamaCppBackend:
|
|||
gpu_indices, use_fit = self._select_gpus(requested_total, gpus)
|
||||
# No silent shrink: effective_ctx stays == n_ctx.
|
||||
else:
|
||||
# Auto context: prefer fewer GPUs, cap context to fit.
|
||||
# Auto context: prefer fewer GPUs, cap context
|
||||
# to fit. Same headroom threshold as
|
||||
# _select_gpus (#5106).
|
||||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||||
pin_fraction = self._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)
|
||||
|
|
@ -2130,18 +2145,31 @@ class LlamaCppBackend:
|
|||
capped, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * 0.90:
|
||||
if total_mib <= pool_mib * pin_fraction:
|
||||
effective_ctx = capped
|
||||
gpu_indices = sorted(idx for idx, _ in subset)
|
||||
use_fit = False
|
||||
break
|
||||
else:
|
||||
# No subset can host the weights (weights alone
|
||||
# exceed 90% of every pool). Per spec, default
|
||||
# the UI-visible context to 4096 and let
|
||||
# --fit on flex -ngl so llama-server offloads
|
||||
# layers to CPU RAM.
|
||||
# Native ctx doesn't fit. Drop to 4096 and
|
||||
# re-check before deferring to --fit on:
|
||||
# a model that overflows at 131k may pin
|
||||
# comfortably with a 4096 KV cache (#5106).
|
||||
effective_ctx = min(4096, effective_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 = self._estimate_kv_cache_bytes(
|
||||
effective_ctx,
|
||||
cache_type_kv,
|
||||
n_parallel = n_parallel,
|
||||
)
|
||||
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:
|
||||
# Can't estimate KV -- fall back to file-size-only check.
|
||||
|
|
@ -2570,12 +2598,54 @@ class LlamaCppBackend:
|
|||
|
||||
self._healthy = True
|
||||
|
||||
# Catch silent CPU fallback when GPU was intended (#5106).
|
||||
self._gpu_offload_active = self._classify_gpu_offload(
|
||||
gpu_indices is not None or use_fit, gpus or []
|
||||
)
|
||||
if self._gpu_offload_active is False:
|
||||
logger.warning(
|
||||
"llama-server appears to have loaded the model entirely "
|
||||
"on CPU even though Studio detected at least one GPU. "
|
||||
"This usually means the prebuilt binary's GPU backend "
|
||||
"failed to load -- on Windows, cudart64_X.dll / "
|
||||
"cublas64_X.dll could not be resolved. Reinstall the "
|
||||
"Studio llama.cpp prebuilt or install a matching CUDA "
|
||||
"toolkit (issue unslothai/unsloth#5106).",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"llama-server ready on port {self._port} "
|
||||
f"for model '{model_identifier}'"
|
||||
)
|
||||
return True
|
||||
|
||||
def _classify_gpu_offload(
|
||||
self,
|
||||
expected_gpu: bool,
|
||||
detected_gpus: list[tuple[int, int]],
|
||||
) -> Optional[bool]:
|
||||
"""True if a GPU model buffer was allocated, False if only CPU
|
||||
buffers landed despite GPU intent, None when there's no signal
|
||||
(no GPU detected, no buffer-size lines, etc.)."""
|
||||
if not detected_gpus or not expected_gpu:
|
||||
return None
|
||||
# llama-server logs one ``... model buffer size = N MiB`` line
|
||||
# per backend buffer; CUDA0 / ROCm0 / Metal / Vulkan0 /
|
||||
# OpenCL0 / SYCL0 are GPU, CPU / CPU_Mapped are not.
|
||||
gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
|
||||
saw_buffer_line = False
|
||||
saw_gpu_buffer = False
|
||||
for line in self._stdout_lines:
|
||||
if "model buffer size" not in line:
|
||||
continue
|
||||
saw_buffer_line = True
|
||||
if any(marker in line for marker in gpu_markers):
|
||||
saw_gpu_buffer = True
|
||||
break
|
||||
if not saw_buffer_line:
|
||||
return None
|
||||
return saw_gpu_buffer
|
||||
|
||||
def unload_model(self) -> bool:
|
||||
"""Terminate the llama-server subprocess and cancel any in-flight download."""
|
||||
self._cancel_event.set()
|
||||
|
|
|
|||
|
|
@ -192,6 +192,7 @@ def _drive(
|
|||
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)
|
||||
|
|
@ -203,7 +204,7 @@ def _drive(
|
|||
)
|
||||
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * 0.90:
|
||||
if total_mib <= pool_mib * pin_fraction:
|
||||
effective_ctx = capped
|
||||
gpu_indices = sorted(idx for idx, _ in subset)
|
||||
use_fit = False
|
||||
|
|
@ -211,6 +212,17 @@ def _drive(
|
|||
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:
|
||||
|
|
@ -378,6 +390,52 @@ class TestFittableAutoPickRegressions:
|
|||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -391,3 +449,81 @@ def test_identical_decision_across_platforms(platform_tag):
|
|||
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):
|
||||
# llama-server printed buffer lines but only CPU buffers --
|
||||
# this is 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 use;
|
||||
# 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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue