From 5205bc0ed679d64ba8f5daf5684c14ec8e6a0671 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 13 May 2026 04:48:15 -0700 Subject: [PATCH] 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> --- studio/backend/core/inference/llama_cpp.py | 104 ++++++++++--- .../tests/test_llama_cpp_context_fit.py | 138 +++++++++++++++++- 2 files changed, 224 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6c97ef8cb2..35933e6685 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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() diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index caa6397901..1ea76edd15 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -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