From 50e4e9c45902cd6241f75ed2b4aeb8a56bbb8928 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 10:47:51 -0700 Subject: [PATCH] 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 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_mtp_detection.py | 45 +++++++- 2 files changed, 94 insertions(+), 55 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e608a4b552..85ca32eaf5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 784ab1b259..25ce71a774 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -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"