mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename (#7031)
* Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename The Gemma 4 QAT GGUF repos renamed the higher-precision MTP/ subdir copies from gemma-4-...-<quant>-MTP.gguf to mtp-gemma-4-...-<quant>.gguf, so their basenames now start with the same mtp- prefix as the small repo-root drafter (mtp-gemma-4-E4B-it.gguf). The drafter selectors filtered candidates by a mtp- basename prefix and took the first in sort order. With the new names the MTP/ copies also match, and because MTP/ (uppercase) sorts before the lowercase root file, selection flipped to the large BF16 copy under MTP/ instead of the root drafter both functions document they should pick. Restrict both selectors, and the companion byte estimate, to root-level mtp-*.gguf so the MTP/ copies stay explicit-selection only: - core/inference/llama_cpp.py _pick_mtp (loader auto-download) - hub/utils/gguf_plan.py preferred_mtp_sibling (Hub variant plans) - routes/inference.py _remote_gguf_companion_bytes (VRAM headroom) Also reuse a drafter already in the local cache before downloading, so a device that already holds a copy on disk does not re-fetch it. Old-scheme names keep working (they have no root-level mtp- sibling to mis-select). Adds regression tests for the new naming, both selection paths, and the on-disk reuse. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate MTP drafter cache reuse to offline mode Reuse the cached drafter only when HF is offline. Online, route back through _download_companion_gguf/hf_hub_download so the current revision is checked (etag) and a changed drafter is refetched, matching the offline-only cross-snapshot reuse already used for the main GGUF. This avoids pairing freshly downloaded weights with a stale cached draft. Make the reuse tests offline and add an online-skips-reuse test. * Studio: prefer a root MTP drafter across all cached snapshots Offline reuse scanned snapshots one at a time and returned the first snapshot that held any drafter, only preferring root within it. A newer partial snapshot with just the MTP/ copy could shadow the small root drafter in an older snapshot. Collect drafters across all snapshots and prefer any repo-root file before an MTP/ copy. * Studio: keep newest-first snapshot order when reusing cached drafters Collecting root candidates and sorting by absolute snapshot path could pick a drafter from an older snapshot. _iter_hf_cache_snapshots yields newest first and the main GGUF is resolved in that order, so preserve it (root still preferred over MTP/ copies) to avoid pairing a fresh main weight with a stale drafter revision. --------- 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
d4fbc81d3a
commit
b5aef63c03
4 changed files with 233 additions and 10 deletions
|
|
@ -4673,6 +4673,29 @@ class LlamaCppBackend:
|
|||
cancel_event = cancel_event,
|
||||
)
|
||||
|
||||
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
|
||||
"""A drafter already in this repo's local HF cache, reused offline when a
|
||||
fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
|
||||
cached snapshots; else an existing ``MTP/`` copy (any precision -- the
|
||||
target verifies every drafted token). None if none is cached."""
|
||||
try:
|
||||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||||
|
||||
roots: list[Path] = []
|
||||
subdirs: list[Path] = []
|
||||
for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
|
||||
for f in sorted(_gguf_snapshot_files(snap)):
|
||||
if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
|
||||
(roots if "/" not in f else subdirs).append(snap / f)
|
||||
# Keep snapshot order (newest first), root before any MTP/ copy, so a
|
||||
# newer main GGUF pairs with the newest cached drafter, not a stale one.
|
||||
for cand in roots + subdirs:
|
||||
if cand.is_file():
|
||||
return str(cand)
|
||||
except Exception as e:
|
||||
logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e)
|
||||
return None
|
||||
|
||||
def _download_mtp(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -4689,11 +4712,25 @@ class LlamaCppBackend:
|
|||
are intentionally skipped. Returns the local path, or None.
|
||||
"""
|
||||
|
||||
# Offline, reuse any drafter already on disk (a fresh copy can't be
|
||||
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
|
||||
# current cached file and refetch a changed one, so skip the probe here
|
||||
# rather than pair new weights with a stale draft.
|
||||
if _hf_env_offline():
|
||||
cached = self._cached_repo_mtp_drafter(hf_repo)
|
||||
if cached:
|
||||
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
|
||||
return cached
|
||||
|
||||
def _pick_mtp(candidates: list[str]) -> Optional[str]:
|
||||
# Root-level only: MTP/ subdir copies now share the mtp- prefix but
|
||||
# are explicit-selection, not auto-fetch (they'd sort ahead of root).
|
||||
mtp_files = sorted(
|
||||
f
|
||||
for f in candidates
|
||||
if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
|
||||
if f.lower().endswith(".gguf")
|
||||
and "/" not in f
|
||||
and Path(f).name.lower().startswith("mtp-")
|
||||
)
|
||||
return mtp_files[0] if mtp_files else None
|
||||
|
||||
|
|
|
|||
|
|
@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
|
|||
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
|
||||
"""The separate MTP drafter to fetch with every variant: the repo-root
|
||||
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
|
||||
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
|
||||
prefix, first in sort order) so download and load resolve the same file;
|
||||
the higher-precision ``MTP/`` subdir copies are for explicit selection and
|
||||
are not auto-fetched. None for repos with the head baked into the main
|
||||
GGUF (Qwen)."""
|
||||
(Gemma 4). Same pick as the loader's drafter resolution (root-level
|
||||
``mtp-`` prefix, first in sort order) so download and load resolve the same
|
||||
file; the higher-precision ``MTP/`` subdir copies are for explicit
|
||||
selection and are not auto-fetched. None for repos with the head baked into
|
||||
the main GGUF (Qwen)."""
|
||||
# Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
|
||||
candidates = sorted(
|
||||
(
|
||||
s
|
||||
for s in siblings
|
||||
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
|
||||
if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
|
||||
),
|
||||
key = lambda s: getattr(s, "rfilename"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3075,10 +3075,14 @@ def _remote_gguf_companion_bytes(
|
|||
info = model_info(repo, token = hf_token, files_metadata = True)
|
||||
total = 0
|
||||
for sibling in info.siblings or []:
|
||||
base = Path(sibling.rfilename or "").name.lower()
|
||||
name = sibling.rfilename or ""
|
||||
base = Path(name).name.lower()
|
||||
if not base.endswith(".gguf"):
|
||||
continue
|
||||
if base.startswith("mtp-") or (include_mmproj and "mmproj" in base):
|
||||
# Root-level mtp- only: -hf auto-fetches the repo-root drafter, not
|
||||
# the MTP/ subdir copies (which now share the mtp- prefix too).
|
||||
is_root_mtp = "/" not in name and base.startswith("mtp-")
|
||||
if is_root_mtp or (include_mmproj and "mmproj" in base):
|
||||
total += getattr(sibling, "size", 0) or 0
|
||||
return total
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path:
|
|||
|
||||
from hub.utils.download_manifest import ExpectedFile
|
||||
from hub.utils.gguf import is_mtp_drafter_path
|
||||
from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files
|
||||
from hub.utils.gguf_plan import (
|
||||
build_gguf_variant_plans,
|
||||
plan_from_expected_files,
|
||||
preferred_mtp_sibling,
|
||||
)
|
||||
from utils.models.model_config import (
|
||||
_is_mtp_drafter,
|
||||
detect_gguf_model,
|
||||
|
|
@ -37,6 +41,8 @@ from utils.models.model_config import (
|
|||
DRAFTER_CASES = [
|
||||
("mtp-gemma-4-12b-it.gguf", True),
|
||||
("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True),
|
||||
# New-scheme MTP/ copies carry the mtp- basename prefix too.
|
||||
("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True),
|
||||
("foo/MTP/bar.gguf", True),
|
||||
("gemma-4-12b-it-Q8_0.gguf", False),
|
||||
# Baked-in Qwen MTP repos: the head is inside the main GGUF, the file
|
||||
|
|
@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path):
|
|||
assert detect_gguf_model(str(copy)) is None
|
||||
# Selecting the MTP dir itself must not surface the copies as models.
|
||||
assert detect_gguf_model(str(sub)) is None
|
||||
|
||||
|
||||
# ── Root drafter wins over new-scheme MTP/ copies ────────────────────
|
||||
# The MTP/ copies were renamed to share the mtp- basename prefix (e.g.
|
||||
# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the
|
||||
# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes
|
||||
# lowercase, so the subdir path would otherwise win).
|
||||
|
||||
NEW_SCHEME_SIBLINGS = [
|
||||
_sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"),
|
||||
_sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"),
|
||||
_sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"),
|
||||
_sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"),
|
||||
_sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"),
|
||||
_sib("mmproj-F16.gguf", 500, "mmproj"),
|
||||
]
|
||||
|
||||
|
||||
def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies():
|
||||
picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS)
|
||||
assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf"
|
||||
|
||||
|
||||
def test_variant_plans_new_scheme_uses_root_drafter():
|
||||
plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS)
|
||||
assert set(plans) == {"q4_k_m", "q8_0"}
|
||||
for plan in plans.values():
|
||||
assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames
|
||||
assert not any("MTP/" in name for name in plan.target_filenames)
|
||||
assert "drafter" in plan.companion_hashes
|
||||
# Download size = main + mmproj + root drafter (not the 200-byte BF16 copy).
|
||||
assert plans["q4_k_m"].download_size_bytes == 4_600
|
||||
|
||||
|
||||
def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
|
||||
# _pick_mtp is nested; capture it via the companion-download seam.
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe
|
||||
captured = {}
|
||||
|
||||
def _fake_companion(
|
||||
*,
|
||||
hf_repo,
|
||||
hf_token,
|
||||
pick,
|
||||
label,
|
||||
cancel_event = None,
|
||||
):
|
||||
captured["pick"] = pick
|
||||
return None
|
||||
|
||||
b = LlamaCppBackend()
|
||||
b._download_companion_gguf = _fake_companion
|
||||
b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
|
||||
|
||||
repo_files = [
|
||||
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
|
||||
"MTP/mtp-gemma-4-E4B-it-Q4_0.gguf",
|
||||
"MTP/mtp-gemma-4-E4B-it-Q8_0.gguf",
|
||||
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
|
||||
"mmproj-F16.gguf",
|
||||
"mtp-gemma-4-E4B-it.gguf",
|
||||
]
|
||||
assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf"
|
||||
|
||||
|
||||
# ── Reuse an on-disk drafter offline; fetch fresh online ─────────────
|
||||
|
||||
|
||||
def _seed_snapshot(tmp_path, names):
|
||||
snap = tmp_path / "snap"
|
||||
for rel in names:
|
||||
f = snap / rel
|
||||
f.parent.mkdir(parents = True, exist_ok = True)
|
||||
f.write_bytes(b"x")
|
||||
return snap
|
||||
|
||||
|
||||
def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch):
|
||||
import utils.models.model_config as mc
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
snap = _seed_snapshot(
|
||||
tmp_path,
|
||||
[
|
||||
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
|
||||
"mtp-gemma-4-E4B-it.gguf",
|
||||
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
|
||||
"mmproj-F16.gguf",
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
|
||||
|
||||
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
|
||||
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
|
||||
|
||||
|
||||
def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch):
|
||||
# Pre-fix build may have fetched only the MTP/ copy; reuse it offline.
|
||||
import utils.models.model_config as mc
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
snap = _seed_snapshot(
|
||||
tmp_path,
|
||||
[
|
||||
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
|
||||
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
|
||||
|
||||
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
|
||||
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf"
|
||||
|
||||
|
||||
def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch):
|
||||
# A newer partial snapshot holds only the MTP/ copy; an older one has the
|
||||
# root. Must still return the small root, not the large subdir copy.
|
||||
import utils.models.model_config as mc
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"])
|
||||
snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"])
|
||||
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full])
|
||||
|
||||
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
|
||||
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
|
||||
|
||||
|
||||
def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch):
|
||||
# Two snapshots both hold a root drafter; newest-first order must win so a
|
||||
# fresh main GGUF is not paired with a stale drafter revision.
|
||||
import utils.models.model_config as mc
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"])
|
||||
oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"])
|
||||
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest])
|
||||
|
||||
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
|
||||
assert got is not None and Path(got).parent.parent.name == "newest"
|
||||
|
||||
|
||||
def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
|
||||
# Online, do not reuse a cached copy: go to the download path so a changed
|
||||
# drafter is refetched (hf_hub_download checks the current revision).
|
||||
import utils.models.model_config as mc
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
|
||||
snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"])
|
||||
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
|
||||
|
||||
reached = {}
|
||||
|
||||
def _fake_companion(
|
||||
*,
|
||||
hf_repo,
|
||||
hf_token,
|
||||
pick,
|
||||
label,
|
||||
cancel_event = None,
|
||||
):
|
||||
reached["hit"] = True
|
||||
return None
|
||||
|
||||
b = LlamaCppBackend()
|
||||
b._download_companion_gguf = _fake_companion
|
||||
assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None
|
||||
assert reached.get("hit") is True
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue