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

Studio: reserve fallback MTP RAM, bound NUMA on smallest node, capture abort early

Browse source 
Follow-up review on the CPU-only / NUMA hardening:

- Reserve a flat MTP fraction in the CPU context fit when MTP engages but the draft
  KV cannot be byte-sized (mtp_overhead_fn is None): budget_frac otherwise skips the
  flat reserve, so the fit could pick a context that OOMs once the draft allocates.
- Bound the NUMA local-placement decision on the smallest node's free RAM, not the
  largest. The loader is not pinned, so first-touch may land on any node; interleave
  a footprint that fits only the larger node instead of gambling on placement.
- Capture the first launch's scheduler abort before the no-spec MTP fallback resets
  the stdout tail, then fold it into the terminal memo decision, so a fallback that
  fails for another reason does not drop the memo and let the UI replay the load.

Unit tests and the offline simulations updated and passing.
This commit is contained in:
Daniel Han 2026-06-29 08:32:31 +00:00
commit 1813035d50
5 changed files with 103 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -5754,6 +5754,12 @@ class LlamaCppBackend:
# the minimum so the tightest fit gets the smallest context.
_cpu_cap = 4096
else:
# MTP engages but the draft KV can't be byte-sized:
# _mtp_bytes returns 0 and budget_frac skips the flat
# reserve, so trim the budget to still hold back MTP RAM.
_cpu_budget = _CPU_RAM_BUDGET_FRAC
if _mtp_will_engage_cpu and mtp_overhead_fn is None:
_cpu_budget -= _MTP_VRAM_RESERVE_FRAC
_fit = self._fit_context_to_vram(
requested_ctx = _ctx_ceiling,
available_mib = _avail_mib,
@ -5768,7 +5774,7 @@ class LlamaCppBackend:
mtp_overhead_fn = (
_mtp_bytes if _mtp_will_engage_cpu else None
),
budget_frac = _CPU_RAM_BUDGET_FRAC,
budget_frac = _cpu_budget,
)
_cpu_cap = max(4096, min(_ctx_ceiling, _fit))
except Exception as _cap_exc: # best-effort; fall back to ceiling
@ -6308,6 +6314,9 @@ class LlamaCppBackend:
)
healthy = _spawn_and_wait(cmd)
# A scheduler abort on the first launch must still be memoed even if a later
# no-spec/text-only fallback overwrites the stdout tail; track it across them.
_pre_fallback_sched_abort = False
# #6415 split-mode tensor warmup abort. Latch it on THIS first spawn:
# the flash-attn-off retry below can't run tensor (needs flash_attn),
# so its output drops the marker and recording later would miss it,
@ -6407,6 +6416,12 @@ class LlamaCppBackend:
# cancel check stops an /unload-killed attempt respawning. A
# decode-probe failure above also routes here.
if not healthy and _spec_requested_mtp and not self._cancel_event.is_set():
# The no-spec retry below resets the stdout tail; capture a scheduler
# abort from this first launch now so it can still be memoed if the
# retries then fail for another reason (else the UI replays the load).
_pre_fallback_sched_abort = _pre_fallback_sched_abort or (
self._is_sched_reserve_abort("\n".join(self._stdout_lines[-200:]))
)
# Blame the binary only when the output shows MTP itself
# failing (unknown arch / draft or context build); an
# unrelated crash (e.g. OOM) gets a neutral message.
@ -6471,10 +6486,14 @@ class LlamaCppBackend:
# mmproj fallback is ruled out, so a VLM that recovers text-only is not
# blocked by the fail-fast guard on its next load. Wider slice as the
# GGML_ASSERT line can scroll past the [New LWP] dump.
_was_sched_abort = (
not self._cancel_event.is_set()
and (self._is_signal_crash(_crash_rc) or self._is_abort_exit(_crash_rc))
and self._is_sched_reserve_abort("\n".join(self._stdout_lines[-200:]))
_was_sched_abort = not self._cancel_event.is_set() and (
# A scheduler abort captured before the no-spec fallback reset stdout,
_pre_fallback_sched_abort
# or one still visible in the current tail.
or (
(self._is_signal_crash(_crash_rc) or self._is_abort_exit(_crash_rc))
and self._is_sched_reserve_abort("\n".join(self._stdout_lines[-200:]))
)
)
self._kill_process()
# The #6415 split-axis abort is latched earlier (first spawn).

32
studio/backend/core/inference/numa.py
View file

@ -36,6 +36,10 @@ class NumaTopology:
def largest_node_free_mib(self) -> int:
return max(self.node_free_mib.values(), default = 0)
@property
def smallest_node_free_mib(self) -> int:
return min(self.node_free_mib.values(), default = 0)
def _parse_online(spec: str) -> list[int]:
"""Parse a sysfs cpulist-style range, e.g. '0-1' or '0,2-3', into node ids."""
@ -110,10 +114,11 @@ def decide_interleave(
topology: NumaTopology | None = None,
has_numactl: bool | None = None,
) -> InterleaveDecision:
"""Interleave only when CPU-only, multi-node, and the footprint exceeds the largest
node's free RAM but fits across all nodes; otherwise leave placement local.
model_size_bytes should be the resident footprint (weights + KV), not weights
alone, so a model whose weights fit a node but whose footprint does not still
"""Interleave only when CPU-only, multi-node, and the footprint exceeds the smallest
node's free RAM but fits across all nodes; otherwise leave placement local. The
smallest node is the bound because the loader is not pinned, so first-touch may land
on any node. model_size_bytes should be the resident footprint (weights + KV), not
weights alone, so a model whose weights fit a node but whose footprint does not still
interleaves."""
if not cpu_only:
return InterleaveDecision(False, "not cpu-only; leaving NUMA placement to the OS")
@ -125,14 +130,17 @@ def decide_interleave(
return InterleaveDecision(False, "single NUMA node; interleave not needed")
model_mib = model_size_bytes // (1024 * 1024)
largest = topo.largest_node_free_mib
smallest = topo.smallest_node_free_mib
total = topo.total_free_mib
if model_mib <= largest:
# Keep local only when the footprint fits EVERY node (the smallest). We don't bind
# the loader, so first-touch may land on any node; local placement is safe only when
# any node can hold it. A footprint that fits only the larger node still interleaves.
if model_mib <= smallest:
return InterleaveDecision(
False,
f"model ~{model_mib} MiB fits the largest node's free RAM "
f"(~{largest} MiB); keeping local placement",
f"model ~{model_mib} MiB fits every node's free RAM "
f"(smallest ~{smallest} MiB); keeping local placement",
)
# Impossible across all nodes regardless of numactl: surface the smaller-quant /
@ -150,16 +158,16 @@ def decide_interleave(
# Needed but unavailable: surface it; caller decides whether to block.
return InterleaveDecision(
False,
f"model ~{model_mib} MiB exceeds the largest NUMA node's free RAM "
f"(~{largest} MiB) and needs interleaving across {topo.node_count} nodes, "
f"model ~{model_mib} MiB exceeds the smallest NUMA node's free RAM "
f"(~{smallest} MiB) and needs interleaving across {topo.node_count} nodes, "
f"but `numactl` is not installed. Install numactl (e.g. `apt install "
f"numactl`) or the model may fail to fit a single node.",
)
return InterleaveDecision(
True,
f"model ~{model_mib} MiB exceeds the largest NUMA node's free RAM "
f"(~{largest} MiB) but fits across {topo.node_count} nodes (~{total} MiB total "
f"model ~{model_mib} MiB exceeds the smallest NUMA node's free RAM "
f"(~{smallest} MiB) but fits across {topo.node_count} nodes (~{total} MiB total "
f"free); wrapping with numactl --interleave=all",
prefix = ("numactl", "--interleave=all"),
)

13
studio/backend/tests/test_cpu_only_defaults.py
View file

@ -174,7 +174,8 @@ def test_cpu_context_cap_reuses_fit_helper_against_ram():
src = _load_model_src()
assert "_available_system_memory_mib()" in src
assert "self._fit_context_to_vram(" in src
assert "budget_frac = _CPU_RAM_BUDGET_FRAC" in src
assert "_cpu_budget = _CPU_RAM_BUDGET_FRAC" in src
assert "budget_frac = _cpu_budget" in src
def test_cpu_ram_preflight_warns_when_weights_exceed_ram():
@ -275,6 +276,16 @@ def test_cpu_fit_skips_mtp_reserve_when_mla_auto_drops():
assert _nows("_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0") in src
def test_cpu_fit_reserves_flat_mtp_when_draft_unsized():
"""When MTP engages but the draft KV can't be byte-sized (mtp_overhead_fn is None),
budget_frac skips the flat reserve, so the CPU budget is trimmed to still hold MTP
RAM back instead of fitting a context that OOMs once the draft allocates (PR review)."""
src = _load_model_src()
assert "if _mtp_will_engage_cpu and mtp_overhead_fn is None:" in src
assert "_cpu_budget -= _MTP_VRAM_RESERVE_FRAC" in src
assert "budget_frac = _cpu_budget" in src
def test_zero_offload_folds_into_cpu_only():
"""A visible GPU plus a user -ngl 0 must be treated as CPU-only: the GPU list is
dropped before _cpu_only is computed so the CPU safe defaults apply (PR review fix)."""

15
studio/backend/tests/test_numa_interleave.py
View file

@ -87,6 +87,7 @@ def test_parse_online_ranges():
def test_topology_aggregates():
assert _USER_TOPO.node_count == 2
assert _USER_TOPO.largest_node_free_mib == 465594
assert _USER_TOPO.smallest_node_free_mib == 223814
assert _USER_TOPO.total_free_mib == 465594 + 223814
@ -98,12 +99,20 @@ def test_interleaves_when_model_exceeds_largest_node_but_fits_across():
assert "interleave=all" in d.reason
def test_no_interleave_when_model_fits_largest_node():
# A 200 GB model fits node 0's 465 GB free -> keep local placement.
def test_no_interleave_when_model_fits_every_node():
# A 200 GB model fits even the smaller node (223 GB) -> safe on any node, keep local.
d = decide_interleave(200 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True)
assert d.interleave is False
assert d.prefix == ()
assert "fits" in d.reason
assert "fits every node" in d.reason
def test_interleaves_when_fits_larger_node_but_not_smaller():
# 300 GB fits node 0 (465) but not node 1 (223); the loader is not bound, so first-
# touch could land on node 1. Interleave instead of gambling on placement (PR review).
d = decide_interleave(300 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True)
assert d.interleave is True
assert d.prefix == ("numactl", "--interleave=all")
def test_no_interleave_on_gpu_host():

36
studio/backend/tests/test_sched_reserve_abort.py
View file

@ -281,7 +281,7 @@ def test_abort_memo_deferred_until_mmproj_fallback_ruled_out():
raise, after the text-only mmproj fallback is ruled out, so a VLM that recovers
text-only is not blocked by the fail-fast guard next time (PR review fix)."""
src = _load_model_src()
assert "_was_sched_abort = (" in src
assert "_was_sched_abort = " in src
assert "if _was_sched_abort:" in src
fn = ast.parse(src).body[0]
strip_line = _call_line(fn, "_strip_mmproj_args")
@ -289,3 +289,37 @@ def test_abort_memo_deferred_until_mmproj_fallback_ruled_out():
# Recording happens after the projector strip, i.e. only once the fallback is tried.
assert strip_line is not None and record_line is not None
assert record_line > strip_line
def test_sched_abort_captured_before_mtp_fallback():
"""The no-spec MTP fallback resets the stdout tail, so the first launch's scheduler
abort must be captured before it runs and folded into the terminal decision; else a
differently-failing fallback drops the memo and the UI replays the load (PR review fix)."""
src = _load_model_src()
assert "_pre_fallback_sched_abort = False" in src
assert "_pre_fallback_sched_abort = _pre_fallback_sched_abort or (" in src
# The capture is set before the no-spec fallback spawns.
fn = ast.parse(src).body[0]
capture_line = next(
(
n.lineno
for n in ast.walk(fn)
if isinstance(n, ast.Assign)
and any(
isinstance(t, ast.Name) and t.id == "_pre_fallback_sched_abort" for t in n.targets
)
and isinstance(n.value, ast.BoolOp)
),
None,
)
fallback_line = next(
(
n.lineno
for n in ast.walk(fn)
if isinstance(n, ast.Call)
and any(isinstance(a, ast.Name) and a.id == "fallback_cmd" for a in n.args)
),
None,
)
assert capture_line is not None and fallback_line is not None
assert capture_line < fallback_line