mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-24 15:31:01 +00:00
Studio: a kept load short on its own GPU no longer unloads models on other GPUs or on the CPU (#14234)
## Summary
With **Keep multiple models loaded** set to **When they fit**, a model
pinned to one GPU that does not fully fit there unloads models that
cannot make room for it:
1. Idle CPU-only kept models (a GGUF started with no GPU layers, for
example a small LLM or TTS held in RAM) are picked as eviction victims,
one per retry, though unloading them frees no VRAM on a discrete card.
2. When no kept model on the short GPUs can go, the active model is
replaced even when it is pinned to other GPUs, or runs on the CPU. That
frees nothing on the GPUs the new model is pinned to, so the new model
lands on the same GPUs with the same `--fit on` offload it would have
had anyway, and the user just loses the other models. The Model memory
hint already promises the opposite: "When they fit and Always never
unload a model on another GPU."
This is the setup in #9257 (a big GPU model plus small CPU-only models
in one Studio) and the "8 GPUs, two models split 4 + 4, sometimes it
unloads the other model" report on #8488.
## Changes
- `GpuMemoryShortError` gains `discrete`: true only when none of the
short GPUs shares host RAM. It is false for a probe total of 0 (APU,
iGPU), for the probe's own ROCm-unified and integrated-CUDA classifiers
(DGX Spark, Jetson), for the Metal shortfall, and for a non-GGUF
out-of-memory error.
- `model_slots.eviction_victims(..., discrete)`: on a discrete
shortfall, a kept llama-server with no VRAM plan and no GPU layers is
not a victim. A GPU server that fell back to the CPU still has its plan
and is evicted as before.
- `load_model_gated`: when nothing kept can be evicted, the active model
is kept and the load retries beside it with what fits
(`force_alongside`) only when the active model's placement is one its
child cannot leave and is disjoint from the short GPUs: an explicit
`gpu_ids` pin (the recorded device mask, already widened for companion
devices) or a CPU-only server. An unpinned, safetensors or otherwise
unknown placement is replaced as before.
Always and Off are unchanged. Without `gpu_ids`, the short GPUs include
the active model's, so single-pool users keep today's replace.
## Testing
Unit tests in `studio/backend/tests/test_multi_model_serving.py`:
- `test_a_cpu_only_kept_model_is_never_evicted_for_gpu_room`, which also
covers a shared pool and a GPU server that fell back to CPU.
- `test_a_model_short_on_its_own_gpu_keeps_the_models_on_others`, which
also covers a pinned `gpu_layers=0` active model.
- `test_a_model_short_where_the_active_model_runs_still_replaces_it`.
- `test_a_short_load_still_replaces_an_active_model_it_may_share_with`,
parametrized over a shared pool, an unpinned active model, and a
companion-widened pin.
The Metal guard test now also asserts `discrete` is false. The
multi-model, GPU arbiter, residency, Metal, training VRAM and shortfall
suites pass: 1339 passed. The 3 failures in
`test_ctx_checkpoint_host_budget.py` fail identically on the merge base.
End to end on a Kaggle T4x2: one `install.sh --local` at merge base
`3697c7d62`, then the same install with this patch applied, same models
and steps:
- Qwen3-0.6B Q4_K_M with `gpu_ids=[0]`, gemma-3-270m with `gpu_ids=[1]`
kept alongside, and SmolLM2-135M kept alongside with
`gpu_memory_mode="manual", gpu_layers=0`.
- Each llama-server ran where asked (`CUDA_VISIBLE_DEVICES=0`, `1`,
`-1`) and answered the requests naming it. Unloading gemma left the
other two serving. This was the same on both sides.
- GPU 1 was then filled to about 1 GB free, and Qwen3-1.7B was loaded
alongside with `gpu_ids=[1]` under When they fit:
| | merge base | this PR |
|---|---|---|
| serving before | Qwen3-0.6B (GPU 0), SmolLM2 (CPU) | Qwen3-0.6B (GPU
0), SmolLM2 (CPU) |
| serving after | Qwen3-1.7B only | Qwen3-0.6B, SmolLM2, Qwen3-1.7B |
| Qwen3-1.7B placement | GPU 1, `--fit on`, 518 MiB | GPU 1, `--fit on`,
518 MiB |
Under Always, both sides refuse the same load with a 409 and unload
nothing.
Fixes #9257
This commit is contained in:
parent
40e20017af
commit
c708b1d727
5 changed files with 142 additions and 5 deletions
|
|
@ -603,9 +603,10 @@ def _gpu_short_message(
|
|||
|
||||
class GpuMemoryShortError(RuntimeError):
|
||||
"""A load that must fit next to the loaded models does not. ``capped``: it would, at a
|
||||
smaller context than asked for. ``gpu_indices``: the GPUs ``short_mib`` was measured on."""
|
||||
smaller context than asked for. ``gpu_indices``: the GPUs ``short_mib`` was measured on.
|
||||
``discrete``: none of them shares host RAM, so a server off them or on the CPU frees nothing there."""
|
||||
|
||||
__slots__ = ("capped", "short_mib", "gpu_indices")
|
||||
__slots__ = ("capped", "short_mib", "gpu_indices", "discrete")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -613,11 +614,13 @@ class GpuMemoryShortError(RuntimeError):
|
|||
capped: bool = False,
|
||||
short_mib: int = 0,
|
||||
gpu_indices = None,
|
||||
discrete: bool = False,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.capped = capped
|
||||
self.short_mib = short_mib
|
||||
self.gpu_indices = gpu_indices
|
||||
self.discrete = discrete
|
||||
|
||||
|
||||
def _apple_wired_headroom_bytes() -> Optional[int]:
|
||||
|
|
@ -26642,6 +26645,8 @@ class LlamaCppBackend:
|
|||
# Bound before the try for the same reason as _detected_gpus.
|
||||
_placement_verdict_partial = False
|
||||
_ctx_capped_for_vram = False
|
||||
# Probed GPUs reporting total 0: a pool shared with host RAM (APU, iGPU, SoC).
|
||||
_shared_pool_gpus: set[int] = set()
|
||||
_metal_alongside_shortfall: Optional[tuple[int, int]] = None
|
||||
# A path that never prices the launch must not commit the previous one's plan.
|
||||
self._pending_plan_mib = {}
|
||||
|
|
@ -26834,6 +26839,7 @@ class LlamaCppBackend:
|
|||
_sel_gpus = [g for g in gpus if g[0] in _wanted_ids]
|
||||
gpus = _sel_gpus if _sel_gpus else gpus
|
||||
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
|
||||
_shared_pool_gpus = {idx for idx, total in total_by_idx.items() if total <= 0}
|
||||
# GPU picker: restrict every mode to the chosen devices, so
|
||||
# auto selection only considers them and manual mask to
|
||||
# them (the env block below pins CUDA/HIP_VISIBLE_DEVICES).
|
||||
|
|
@ -29918,11 +29924,21 @@ class LlamaCppBackend:
|
|||
need_mib = (gguf_size + mmproj_size + _kv_bytes(asked_ctx)) // (1024 * 1024)
|
||||
free_mib = sum(free for _, free in gpus)
|
||||
need_gb, free_gb = need_mib / 1024, free_mib / 1024
|
||||
# The probe's own classifiers: an integrated CUDA SoC can report a real total.
|
||||
try:
|
||||
_shared = (
|
||||
_shared_pool_gpus
|
||||
| self._rocm_unified_memory_gpu_ids()
|
||||
| self._integrated_cuda_gpu_ids()
|
||||
)
|
||||
except Exception:
|
||||
_shared = None
|
||||
raise GpuMemoryShortError(
|
||||
_gpu_short_message(need_gb, free_gb, asked_ctx, effective_ctx, capped_only),
|
||||
capped = capped_only,
|
||||
short_mib = max(0, need_mib - free_mib),
|
||||
gpu_indices = tuple(idx for idx, _ in gpus),
|
||||
discrete = _shared is not None and not _shared & {idx for idx, _ in gpus},
|
||||
)
|
||||
|
||||
# An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly
|
||||
|
|
|
|||
|
|
@ -76,6 +76,11 @@ def _worker_alive(orchestrator) -> bool:
|
|||
return bool(callable(worker_alive) and worker_alive())
|
||||
|
||||
|
||||
def cpu_only(llama) -> bool:
|
||||
"""A running llama-server started with no GPU layers: unloading it frees no VRAM."""
|
||||
return bool(llama.is_active) and getattr(llama, "_gpu_offload_active", None) is False
|
||||
|
||||
|
||||
def holds_vram() -> bool:
|
||||
"""A kept model on the GPU or one loading; a llama-server started with no GPU layers is not."""
|
||||
return any_loading() or any(
|
||||
|
|
@ -143,6 +148,7 @@ def eviction_victims(
|
|||
exclude,
|
||||
short_mib: int,
|
||||
gpu_indices = None,
|
||||
discrete: bool = False,
|
||||
) -> list[ExtraSlot]:
|
||||
"""Least recently used idle slots, as many as it takes to free ``short_mib`` on ``gpu_indices``;
|
||||
one at a time when a footprint is unknown, since the retry prices the rest. A slot still
|
||||
|
|
@ -154,6 +160,8 @@ def eviction_victims(
|
|||
held = sum(mib for idx, mib in planned.items() if gpu_indices is None or idx in gpu_indices)
|
||||
if slot.generations or slot.refs or (planned and not held):
|
||||
continue
|
||||
if discrete and not planned and cpu_only(slot.llama):
|
||||
continue
|
||||
victims.append(slot)
|
||||
freed += held
|
||||
if not held or freed >= short_mib:
|
||||
|
|
|
|||
|
|
@ -18910,7 +18910,11 @@ async def load_model_gated(
|
|||
# No chat may start on a victim between its pick and teardown.
|
||||
async with inference_lifecycle_gate() if new_slot else nullcontext():
|
||||
victims = await asyncio.to_thread(
|
||||
model_slots.eviction_victims, extra, exc.short_mib, exc.gpu_indices
|
||||
model_slots.eviction_victims,
|
||||
extra,
|
||||
exc.short_mib,
|
||||
exc.gpu_indices,
|
||||
exc.discrete,
|
||||
)
|
||||
# Named before the drop, while their public ids are still live.
|
||||
reports = await asyncio.to_thread(
|
||||
|
|
@ -18919,7 +18923,15 @@ async def load_model_gated(
|
|||
dropped = await asyncio.to_thread(model_slots.evict_these, victims)
|
||||
note_keyed([r for slot in dropped for r in reports.get(slot, ())])
|
||||
if not dropped:
|
||||
if exc.capped:
|
||||
# An active model off the short GPUs frees nothing there: load beside it.
|
||||
if exc.capped or (
|
||||
not request.force_alongside
|
||||
and await asyncio.to_thread(
|
||||
model_slots.in_slot,
|
||||
None,
|
||||
lambda: _active_model_holds_none_of(exc),
|
||||
)
|
||||
):
|
||||
request = request.model_copy(update = {"force_alongside": True})
|
||||
continue
|
||||
# Fits nowhere beside the others: replace the active model, as before,
|
||||
|
|
@ -19076,6 +19088,23 @@ def _unload_primary() -> None:
|
|||
backend.unload_model(name)
|
||||
|
||||
|
||||
def _active_model_holds_none_of(exc) -> bool:
|
||||
"""The active llama-server holds nothing on the discrete GPUs ``exc`` was short on; call via
|
||||
``in_slot(None, ...)``. Only a placement the child cannot leave counts: an explicit pin (its
|
||||
device mask, widened for companion devices) or no GPU at all. Anything else answers False."""
|
||||
llama = get_llama_cpp_backend()
|
||||
if not (exc.discrete and exc.gpu_indices and llama.is_active):
|
||||
return False
|
||||
if getattr(_peek_inference_backend(), "active_model_name", None):
|
||||
return False
|
||||
if not getattr(llama, "_planned_vram_mib", None) and model_slots.cpu_only(llama):
|
||||
return True
|
||||
pin = getattr(llama, "_gpu_ids", None)
|
||||
return bool(getattr(llama, "_requested_gpu_ids", None) and pin) and not set(pin) & set(
|
||||
exc.gpu_indices
|
||||
)
|
||||
|
||||
|
||||
def _primary_model_label() -> Optional[str]:
|
||||
llama = get_llama_cpp_backend()
|
||||
if getattr(llama, "is_loaded", False):
|
||||
|
|
|
|||
|
|
@ -1985,7 +1985,7 @@ class TestMetalAlongsideAdmission:
|
|||
assert caught.value.short_mib > 0
|
||||
assert not caught.value.capped
|
||||
# Metal and MLX residents share one pool, not discrete GPU ordinals.
|
||||
assert caught.value.gpu_indices is None
|
||||
assert caught.value.gpu_indices is None and not caught.value.discrete
|
||||
|
||||
def test_minimum_context_must_fit_too(self, tmp_path, monkeypatch):
|
||||
from core.inference.llama_cpp import GpuMemoryShortError
|
||||
|
|
|
|||
|
|
@ -1526,6 +1526,71 @@ def test_the_replace_fallback_frees_the_active_model_before_placing_the_new_one(
|
|||
assert seen == [False] and primary.model_identifier == "org/C-GGUF"
|
||||
|
||||
|
||||
def test_a_cpu_only_kept_model_is_never_evicted_for_gpu_room(backends):
|
||||
_, extra = backends
|
||||
extra.llama._gpu_offload_active = False
|
||||
on_gpu = _slot("org/D-GGUF", last_used = 5.0)
|
||||
model_slots.slots.append(on_gpu)
|
||||
assert model_slots.eviction_victims(None, 5000, (0,), discrete = True) == [on_gpu]
|
||||
model_slots.slots.remove(on_gpu)
|
||||
assert model_slots.eviction_victims(None, 5000, (0,), discrete = True) == []
|
||||
# A pool shared with host RAM (APU, Metal): the CPU-only server holds it too.
|
||||
assert model_slots.eviction_victims(None, 5000, (0,)) == [extra]
|
||||
# A GPU server that fell back to the CPU keeps the plan it is still charged for.
|
||||
extra.llama._planned_vram_mib = {0: 6000}
|
||||
assert model_slots.eviction_victims(None, 5000, (0,), discrete = True) == [extra]
|
||||
|
||||
|
||||
def _short_on_gpu_one(monkeypatch, discrete = True):
|
||||
monkeypatch.setattr(inf, "release_chat_gpu_claim", lambda: None)
|
||||
calls = []
|
||||
|
||||
async def load(request, *args, **kwargs):
|
||||
calls.append((request.alongside, request.force_alongside))
|
||||
if request.alongside and not request.force_alongside:
|
||||
raise GpuMemoryShortError(
|
||||
"needs 13 GB, 8 GB free", short_mib = 5000, gpu_indices = (1,), discrete = discrete
|
||||
)
|
||||
llama = inf.get_llama_cpp_backend()
|
||||
llama.model_identifier, llama.hf_variant = request.model_path, request.gguf_variant
|
||||
llama.is_loaded = llama.is_active = True
|
||||
return "loaded"
|
||||
|
||||
monkeypatch.setattr(inf, "_run_tracked_load_model_impl", load)
|
||||
return calls
|
||||
|
||||
|
||||
def test_a_model_short_on_its_own_gpu_keeps_the_models_on_others(backends, monkeypatch):
|
||||
primary, extra = backends
|
||||
primary._requested_gpu_ids = primary._gpu_ids = [0]
|
||||
extra.request = LoadRequest(model_path = "org/B-GGUF", gguf_variant = "Q8_0")
|
||||
extra.llama._gpu_offload_active = False
|
||||
calls = _short_on_gpu_one(monkeypatch)
|
||||
request = LoadRequest(model_path = "org/C-GGUF", alongside = True, gpu_ids = [1])
|
||||
assert asyncio.run(inf.load_model_gated(request, None, "s")) == "loaded"
|
||||
assert calls == [(True, False), (True, True)]
|
||||
assert primary.model_identifier == "org/A-GGUF" and primary.is_active
|
||||
# A pin to the short GPU with no GPU layers (manual gpu_layers=0) holds nothing there either.
|
||||
primary._requested_gpu_ids = primary._gpu_ids = [1]
|
||||
primary._gpu_offload_active = False
|
||||
short = GpuMemoryShortError("short", gpu_indices = (1,), discrete = True)
|
||||
assert model_slots.in_slot(None, lambda: inf._active_model_holds_none_of(short))
|
||||
assert [s.llama.model_identifier for s in model_slots.slots] == ["org/B-GGUF", "org/C-GGUF"]
|
||||
|
||||
|
||||
def test_a_model_short_where_the_active_model_runs_still_replaces_it(backends, monkeypatch):
|
||||
primary, extra = backends
|
||||
primary._requested_gpu_ids = primary._gpu_ids = [1]
|
||||
extra.request = LoadRequest(model_path = "org/B-GGUF", gguf_variant = "Q8_0")
|
||||
extra.llama._gpu_offload_active = False
|
||||
calls = _short_on_gpu_one(monkeypatch)
|
||||
request = LoadRequest(model_path = "org/C-GGUF", alongside = True, gpu_ids = [1])
|
||||
assert asyncio.run(inf.load_model_gated(request, None, "s")) == "loaded"
|
||||
assert calls == [(True, False), (False, False)]
|
||||
assert primary.model_identifier == "org/C-GGUF"
|
||||
assert [s.llama.model_identifier for s in model_slots.slots] == ["org/B-GGUF"]
|
||||
|
||||
|
||||
def test_metal_shortfall_evicts_then_rechecks_shared_memory(backends, monkeypatch, tmp_path):
|
||||
from .test_metal_explicit_context_guard import _launch
|
||||
|
||||
|
|
@ -1560,3 +1625,22 @@ def test_metal_shortfall_evicts_then_rechecks_shared_memory(backends, monkeypatc
|
|||
assert len(launches) == 1
|
||||
assert response.evicted == ["org/B-GGUF:Q8_0"]
|
||||
assert primary.is_loaded
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", ["shared_pool", "unpinned", "companion_widened"])
|
||||
def test_a_short_load_still_replaces_an_active_model_it_may_share_with(backends, monkeypatch, case):
|
||||
primary, extra = backends
|
||||
primary._requested_gpu_ids = primary._gpu_ids = [0]
|
||||
if case == "unpinned":
|
||||
# Auto placement: a user --device or an inherited LLAMA_ARG_DEVICE may outrank the plan.
|
||||
primary._requested_gpu_ids, primary._gpu_ids = None, None
|
||||
primary._planned_vram_mib = {0: 9000}
|
||||
elif case == "companion_widened":
|
||||
primary._gpu_ids = [0, 1]
|
||||
extra.request = LoadRequest(model_path = "org/B-GGUF", gguf_variant = "Q8_0")
|
||||
extra.llama._gpu_offload_active = False
|
||||
calls = _short_on_gpu_one(monkeypatch, discrete = case != "shared_pool")
|
||||
request = LoadRequest(model_path = "org/C-GGUF", alongside = True, gpu_ids = [1])
|
||||
assert asyncio.run(inf.load_model_gated(request, None, "s")) == "loaded"
|
||||
assert calls[-1] == (False, False)
|
||||
assert primary.model_identifier == "org/C-GGUF"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue