Studio: gracefully disable MTP when the model has no head or drafter (#6159)
* Studio: gracefully disable MTP when the model has no head or drafter
Selecting MTP or MTP+Ngram in Speculative Decoding on a GGUF with no nextn
head and no separate drafter aborted the whole load. llama-server does not
no-op an empty draft-mtp request: it exits with 'failed to measure MTP
context memory: failed to create llama_context', surfaced to the user as a
generic 'llama-server failed to start. Check that the GGUF file is valid
and you have enough memory.'
Build-time fix in _build_speculative_flags: when a forced mtp / mtp+ngram
mode targets a model with no MTP head and no drafter (is_mtp_model is
False), default back instead of emitting draft-mtp. mtp falls back to
--spec-default; mtp+ngram keeps the ngram-mod half, which needs no head.
Real MTP models (embedded head or separate drafter), sub-3B MTP overrides,
and the auto path are unchanged.
Runtime hardening: the existing post-launch MTP retry only fired for
separate-file drafters (--model-draft in spec_flags), so an embedded-head
model that the binary cannot build still hard-failed. Gate the retry on the
spec block requesting MTP, recognise the embedded-head abort strings
('failed to measure MTP context memory', 'failed to create llama_context'),
and make the drafter name None-safe in the warning.
Tests: extend the resolver matrix (forced mtp / mtp+ngram on a non-MTP
model) and add two cases asserting the default-back emission.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
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
dab0b77673
commit
50e4e9c459
2 changed files with 94 additions and 55 deletions
|
|
@ -3572,53 +3572,43 @@ class LlamaCppBackend:
|
|||
)
|
||||
|
||||
healthy = _spawn_and_wait(cmd)
|
||||
# A separate MTP drafter (e.g. Gemma's gemma4-assistant head)
|
||||
# can fail to load on a llama-server that advertises the
|
||||
# --spec-type draft-mtp flag but predates the drafter's
|
||||
# architecture -- and that aborts the whole server. Retry once
|
||||
# with the whole spec block replaced by --spec-default so the
|
||||
# main model still loads (without speculation). Replacing the
|
||||
# slice -- rather than re-resolving with mtp_draft_path=None --
|
||||
# is what guarantees MTP is off even for a forced mtp / mtp+ngram
|
||||
# request (which would otherwise re-emit --spec-type draft-mtp).
|
||||
# _requested_spec_mode (the user's choice) is left intact so a
|
||||
# duplicate /load doesn't thrash a reload.
|
||||
# Gate on the flag actually being in the command, not on the
|
||||
# drafter merely existing on disk: local loads pass the path
|
||||
# even in off/ngram modes (and auto drops MTP sub-3B), where a
|
||||
# retry would blame the drafter for an unrelated failure and
|
||||
# override the user's spec choice. The cancel check keeps an
|
||||
# /unload that killed the first attempt from respawning.
|
||||
if (
|
||||
not healthy
|
||||
and "--model-draft" in spec_flags
|
||||
and not self._cancel_event.is_set()
|
||||
):
|
||||
# Only blame the binary's age when the output shows the
|
||||
# drafter actually failing (unknown arch / draft load);
|
||||
# an unrelated crash (e.g. OOM) gets a neutral message.
|
||||
# Substrings are upstream llama.cpp messages
|
||||
# (llama_model_load / srv load_model [spec]); if they
|
||||
# drift, only the wording degrades -- the retry fires
|
||||
# either way.
|
||||
_attempt_output = "\n".join(self._stdout_lines)
|
||||
# Any MTP request can abort the server: a separate drafter
|
||||
# (Gemma) on a binary that predates its arch, or an embedded
|
||||
# head (Qwen) the binary cannot build. Retry once with the
|
||||
# spec slice replaced by --spec-default so the main model still
|
||||
# loads. Gate on the spec block (not the drafter path, which
|
||||
# off/ngram local loads also carry) and keep
|
||||
# _requested_spec_mode so a duplicate /load doesn't thrash. The
|
||||
# cancel check stops an /unload-killed attempt respawning.
|
||||
_spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags)
|
||||
if not healthy and _spec_requested_mtp and not self._cancel_event.is_set():
|
||||
# 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.
|
||||
_lo = "\n".join(self._stdout_lines).lower()
|
||||
if (
|
||||
"unknown model architecture" in _attempt_output
|
||||
or "failed to measure draft model memory" in _attempt_output
|
||||
"unknown model architecture" in _lo
|
||||
or "failed to measure draft model memory" in _lo
|
||||
or "failed to measure mtp context memory" in _lo
|
||||
or "failed to create llama_context" in _lo
|
||||
):
|
||||
_retry_reason = (
|
||||
"the prebuilt may predate its architecture; retrying "
|
||||
"without speculative decoding -- run "
|
||||
"`unsloth studio update` for MTP"
|
||||
"the prebuilt may predate it; retrying without "
|
||||
"speculative decoding -- run `unsloth studio "
|
||||
"update` for MTP"
|
||||
)
|
||||
else:
|
||||
_retry_reason = (
|
||||
"retrying without speculative decoding in case the "
|
||||
"drafter is the cause"
|
||||
"retrying without speculative decoding in case MTP is the cause"
|
||||
)
|
||||
_drafter = (
|
||||
Path(launch_mtp_draft_path).name
|
||||
if launch_mtp_draft_path
|
||||
else "embedded head"
|
||||
)
|
||||
logger.warning(
|
||||
"llama-server failed to start with MTP drafter %s; %s.",
|
||||
Path(launch_mtp_draft_path).name,
|
||||
"llama-server failed to start with MTP (%s); %s.",
|
||||
_drafter,
|
||||
_retry_reason,
|
||||
)
|
||||
self._kill_process()
|
||||
|
|
@ -3790,8 +3780,9 @@ class LlamaCppBackend:
|
|||
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
|
||||
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
|
||||
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
|
||||
diverse content); forced MTP variants engage anyway and just log a
|
||||
warning per the user's choice.
|
||||
diverse content); forced MTP on a model with no head/drafter defaults
|
||||
back (mtp -> spec-default, mtp+ngram -> ngram-mod) since llama-server
|
||||
aborts otherwise; sub-3B real-MTP engages with a warning.
|
||||
"""
|
||||
flags: List[str] = []
|
||||
# Reset; emit branches re-set on the resolved emission.
|
||||
|
|
@ -3909,33 +3900,40 @@ class LlamaCppBackend:
|
|||
_emit_ngram_mod()
|
||||
return flags
|
||||
if effective_mode == "mtp":
|
||||
if not is_mtp_model:
|
||||
# No head and no drafter: llama-server aborts on draft-mtp
|
||||
# instead of no-op'ing, so default back.
|
||||
logger.warning(
|
||||
"MTP requested but this GGUF has no MTP head or drafter; "
|
||||
"loading without speculative decoding."
|
||||
)
|
||||
flags.append("--spec-default")
|
||||
self._speculative_type = "default"
|
||||
return flags
|
||||
if _mtp_too_small:
|
||||
logger.warning(
|
||||
f"Forcing MTP on a {_mtp_size_b:.1f}B model; "
|
||||
"the bench shows draft-mtp regresses below 3B. "
|
||||
"Engaging anyway (user override)."
|
||||
)
|
||||
elif not is_mtp_model:
|
||||
logger.warning(
|
||||
"Forcing MTP on a non-MTP GGUF; llama-server may "
|
||||
"fall back to spec-off if no nextn head is present. "
|
||||
"Engaging anyway (user override)."
|
||||
)
|
||||
_emit_mtp(chain_ngram = False)
|
||||
return flags
|
||||
if effective_mode == "mtp+ngram":
|
||||
if not is_mtp_model:
|
||||
# No head/drafter: keep the ngram half (needs no head),
|
||||
# drop the draft-mtp chain that would abort the server.
|
||||
logger.warning(
|
||||
"MTP+Ngram requested but this GGUF has no MTP head or "
|
||||
"drafter; loading ngram-mod only."
|
||||
)
|
||||
_emit_ngram_mod()
|
||||
return flags
|
||||
if _mtp_too_small:
|
||||
logger.warning(
|
||||
f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; "
|
||||
"the bench shows the chain regresses below 3B. "
|
||||
"Engaging anyway (user override)."
|
||||
)
|
||||
elif not is_mtp_model:
|
||||
logger.warning(
|
||||
"Forcing MTP+Ngram on a non-MTP GGUF; llama-server "
|
||||
"may fall back to ngram-only if no nextn head is "
|
||||
"present. Engaging anyway (user override)."
|
||||
)
|
||||
_emit_mtp(chain_ngram = True)
|
||||
return flags
|
||||
|
||||
|
|
|
|||
|
|
@ -1056,8 +1056,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
|
|||
("mtp", False, _MTP_MODEL, "draft-mtp", "3", False),
|
||||
# ── mtp forced on sub-3B: engage anyway ──
|
||||
("mtp", True, _SUB_3B_MTP_MODEL, "draft-mtp", "2", False),
|
||||
# ── mtp forced on non-MTP: engage anyway ──
|
||||
("mtp", True, _NON_MTP_MODEL, "draft-mtp", "2", False),
|
||||
# ── mtp forced on non-MTP: default back (no head/drafter) ──
|
||||
("mtp", True, _NON_MTP_MODEL, None, None, False),
|
||||
# ── ngram forced: ngram-mod alone on BOTH platforms ──
|
||||
("ngram", True, _MTP_MODEL, "ngram-mod", None, True),
|
||||
("ngram", False, _MTP_MODEL, "ngram-mod", None, True),
|
||||
|
|
@ -1066,6 +1066,8 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
|
|||
("mtp+ngram", True, _MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
|
||||
("mtp+ngram", False, _MTP_MODEL, "ngram-mod,draft-mtp", "3", True),
|
||||
("mtp+ngram", True, _SUB_3B_MTP_MODEL, "ngram-mod,draft-mtp", "2", True),
|
||||
# ── mtp+ngram forced on non-MTP: keep ngram, drop draft-mtp ──
|
||||
("mtp+ngram", True, _NON_MTP_MODEL, "ngram-mod", None, True),
|
||||
# ── off: nothing emitted ──
|
||||
("off", True, _MTP_MODEL, None, None, False),
|
||||
("off", False, _MTP_MODEL, None, None, False),
|
||||
|
|
@ -1178,3 +1180,42 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
|
|||
# choice is still reflected in _requested_spec_mode.
|
||||
assert backend.requested_spec_mode == "mtp"
|
||||
assert backend.speculative_type is None
|
||||
|
||||
|
||||
def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch):
|
||||
# Forcing MTP on a model with no head/drafter must NOT emit draft-mtp:
|
||||
# llama-server aborts on it ("failed to measure MTP context memory")
|
||||
# rather than no-op'ing. Default back to --spec-default instead.
|
||||
backend = _resolver_backend(monkeypatch)
|
||||
flags = backend._build_speculative_flags(
|
||||
speculative_type = "mtp",
|
||||
spec_draft_n_max = None,
|
||||
extra_args = None,
|
||||
model_identifier = _NON_MTP_MODEL,
|
||||
model_path = None,
|
||||
gpus = True,
|
||||
binary = "/fake/llama-server",
|
||||
)
|
||||
assert "--spec-type" not in flags
|
||||
assert "--spec-default" in flags
|
||||
assert backend.speculative_type == "default"
|
||||
assert backend.requested_spec_mode == "mtp"
|
||||
|
||||
|
||||
def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch):
|
||||
# mtp+ngram on a non-MTP model drops the doomed draft-mtp chain but keeps
|
||||
# the ngram half, which needs no head.
|
||||
backend = _resolver_backend(monkeypatch)
|
||||
flags = backend._build_speculative_flags(
|
||||
speculative_type = "mtp+ngram",
|
||||
spec_draft_n_max = None,
|
||||
extra_args = None,
|
||||
model_identifier = _NON_MTP_MODEL,
|
||||
model_path = None,
|
||||
gpus = True,
|
||||
binary = "/fake/llama-server",
|
||||
)
|
||||
parsed = _flags_dict(flags)
|
||||
assert parsed.get("--spec-type") == "ngram-mod"
|
||||
assert backend.speculative_type == "ngram-mod"
|
||||
assert backend.requested_spec_mode == "mtp+ngram"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue