llama: let an MTP draft borrow the target's embeddings and lm head - #142
llama: let an MTP draft borrow the target's embeddings and lm head#142danielhanchen wants to merge 1 commit into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4e71418da0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (src == nullptr) { | ||
| throw std::runtime_error(format("%s: draft needs tensor '%s' from the target, which does not have it", | ||
| __func__, name.c_str())); |
There was a problem hiding this comment.
When the target has tied embeddings and omits output.weight, its loader falls back to a duplicate named token_embd.weight; consequently this name-based lookup cannot find output.weight. A stripped draft still requests the optional output tensor before its own tie fallback, reaches this exception, and cannot load against an otherwise compatible tied target. Resolve the target model's effective output tensor, including its tied fallback, rather than requiring an entry with the requested GGUF name.
Useful? React with 👍 / 👎.
| if args.mtp_shared_embd: | ||
| if not args.mtp: | ||
| logger.error("--mtp-shared-embd only applies together with --mtp") | ||
| sys.exit(1) | ||
| model_class.mtp_shared_embd = True |
There was a problem hiding this comment.
This enables mtp_shared_embd for every architecture supporting --mtp, but several such exporters still unconditionally retain their shared tensors: DeepSeek V3.2/V4, HY V3, Nemotron-H, and Step3.5 have unchanged mtp_only keep lists that do not inspect this flag. For those inputs, --mtp-shared-embd succeeds and writes the shared-target metadata but produces essentially the ordinary large sidecar, contrary to the CLI promise. Either update each supported filter or reject this option for exporters that do not implement stripping.
Useful? React with 👍 / 👎.
|
I don't understand, what do you need |
|
@CISC I'll make it auto detect |
A NextN/MTP draft exported with --mtp carries the token embeddings, output norm and lm head so it can be loaded as a standalone model. For every current sidecar those three tensors are most of the file: ggml-org/Qwen3.8-27B-GGUF mtp-Qwen3.8-27B-Q4_0.gguf is 1.565 GiB, of which 1.332 GiB (85%) is the copy, against 0.223 GiB for the MTP block itself. Add an opt-in --mtp-shared-embd that leaves them out. The loader resolves those names against the already loaded target model instead. The graph side needs no change: the nextn blocks of seventeen archs already fall back to model.tok_embd and model.output. No new metadata: a file missing token_embd is what marks a draft that shares the target's tensors, and such a file does not load at all today, so no existing sidecar can reach the new path. The missing token_embd is also what gates borrowing output and output_norm, so an arch that ties its head to its own token_embd when output.weight is absent (qwen3.5, qwen3-next) is not affected. Shapes are checked against the target and a mismatch is refused, as is loading such a file on its own.
4e71418 to
f36eda8
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
💡 Codex Reviewllama.cpp/src/llama-model-loader.h Lines 245 to 247 in f36eda8 Condense the hard-wrapped comment
Rewrite this new three-line comment as one or two concise lines rather than splitting its sentences at a fixed width; the repository explicitly requires concise comments and prohibits splitting a sentence across lines. AGENTS.md reference: AGENTS.md:L74-L81 Line 2732 in f36eda8 Preserve the existing lazy-read CLI aliases
Keep ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
* unsloth: pin ggml-org#27941, #152 and #154 ggml-org#27941 at 02eb201, no longer a draft and MERGEABLE upstream. It fixes four qwen4exp correctness defects and is the only account of the Flash-Next reporter who is on Vulkan only, where the allocator variable is never set. The previous pin PR #148 carried the stale 8161d11 and is closed. #152 at 258345e and #154 at 31e432e both merge onto b10708 on their own. #137 is already pinned at 4e1865e, which is still its head, so no change. #142 and #144 are NOT pinned here: both predate an upstream refactor that folded lazy_mode and model_shared into a lazy struct, so they conflict on src/llama.cpp and src/llama-model-loader.cpp against b10708. b10708 has 'ml.lazy.mode = params.lazy_mode' where those branches still write 'ml.lazy_mode' and 'ml.model_shared'. They need rebasing onto current upstream before they can be pinned; pinning them now would fail the resolve. * unsloth: repin ggml-org#27941 to 6b2b85c The PR moved on at 10:44Z, after 02eb201 was pinned. Verified to merge onto the current base tag: 6 files changed, 480 insertions, 138 deletions. * unsloth: pin #144, rebased, and leave #142 out as contained in it #144 rebased onto b10709 at 6fc8df1. Two of its nine commits were dropped rather than carried: 'qwen4exp: fix QSA correctness defects and harden metadata loading' and the comment tidy on top of it. That work is what ggml-org#27941 supersedes, and keeping both copies is what made the two pins conflict in llama-kv-cells.h and llama-memory-hybrid-idx.cpp. Listed after ggml-org#27941 so the upstream version of that work lands first. #142 is deliberately NOT pinned: its single commit is byte-identical to #144's c7bd6f2 apart from the lazy API spelling, so #144 already contains it and pinning both would apply the same change twice. Verified: b10709 + ggml-org#27941 + #144 merges clean, 24 files, +612/-94. * unsloth: repin ggml-org#27754 to 949f7ef The PR moved on at 10:39Z; 5796547 was pinned by #159 earlier today. Verified to merge onto b10709: 44 files changed, 2673 insertions, 38 deletions. Every entry in the set is now at its PR's current head. --------- Co-authored-by: danielhanchen <elliegouldingstuff@gmail.com>
2 participants
What this does
A NextN/MTP draft exported with
--mtpcarriestoken_embd,output_normandlm_headso that it can be loaded as a standalone model. For every sidecar published so far those three tensors are most of the file.Taking ggml-org's own,
ggml-org/Qwen3.8-27B-GGUFmtp-Qwen3.8-27B-Q4_0.gguf, archqwen35, 18 tensors:token_embd+output(both[5120, 248320]q4_0) +output_normblk.64.*, 15 tensorsTheir BF16 (5.538 GiB) and Q8_0 (2.947 GiB) sidecars sit at the same ratio.
This adds an opt-in
--mtp-shared-embdthat leaves those tensors out and marks the file withnextn_shared_target_tensors. The loader resolves the missing names against the already loaded target model instead.The graph side needs no change at all. Twelve archs already fall back to
model.tok_embd/model.outputfor their nextn blocks (qwen3next.cpp:660,815and the same in deepseek32, deepseek4, glm4-moe, glm-dsa, nemotron-h-moe, mimo2, cohere2moe, hy-v3, qwen35, qwen35moe, step35). Only the loader half was missing.Result
Stripping the three tensors from ggml-org's sidecar with
scripts/mtp_strip_shared.pyand running it against the unmodifiedQwen3.8-27B-Q4_K_M.gguftarget:2873913f642b14fb6c6ebc46aa64a47ellama-speculative-simple, greedy, fixed seed, one B200.Existing sidecars are not affected
The borrow is gated on the new metadata key, so a file published before this change cannot reach the new code path. There are four independent checks, in this order, and a normal model load exits on the first:
token_embd/output/output_normnextn_shared_target_tensorsVerified against the two published sidecars I could find, on the same binary, before and after:
unsloth/gemma-4-E2B-it-GGUFmtp-gemma-4-E2B-it-Q8_0c9af508b9a1d596f5bb29bc2eb8592d6, unchangedggml-org/Qwen3.8-27B-GGUFmtp-Qwen3.8-27B-Q4_02873913f642b14fb6c6ebc46aa64a47e, unchangedOnly decode speed moves, and in both directions, which is run to run noise.
The gemma-4 file is the stricter of the two and is worth describing, because it is what shaped the design. It is arch
gemma4-assistant, a self contained 4 block draft running at hidden 256 against the target's 1536, bridging withnextn.pre_projection [3072, 256]andnextn.post_projection [256, 1536]. It has its owntoken_embdat[256, 262144]and nooutput.weight, becausegemma4-assistant.cpp:37-38ties the head to its own embeddings unconditionally. A draft may legitimately run at a different width from its target, so "the tensor is missing" is not a safe trigger on its own.That also decides where the hook sits. It runs before
check_tensor_dims, which is both before the throw for a required tensor and before theTENSOR_NOT_REQUIREDtoTENSOR_DUPLICATEDtie fallback. Without that, an arch whose lm head falls back totoken_embdwould alias the borrowed embedding rather than borrow the target's realoutput.weight, which is silently wrong for any model with untied embeddings.Errors
Loading such a file on its own:
Pointing it at the wrong target:
Shapes are compared across all four dimensions, because the draft uses the tensor directly rather than only needing a compatible vocabulary.
Scope
126 insertions, 13 deletions, 17 files.
git diff --statis empty for bothsrc/models/andggml/.Borrowed tensors are deliberately not counted in
n_created, whichdone_getting_tensorsrequires to equal the tensor count in the file; they are not in the file and are neither allocated nor freed by the draft.--mtp-shared-embdis opt in and is rejected unless--mtpis also given, so files produced by--mtptoday are unchanged and keep loading on any build. It applies to the shared_QwenMtpMixinplus the parallel keep lists inglm.py,dots3.py,bailingmoe3.pyandcommand_r.py.Tests
test-llama-archs -a qwen35: OK, 8.39e-08 CUDA, 0.00e+00 CPU, roundtrip OKtest-llama-archs -a glm4moe: OK, 8.52e-08 CUDA, 0.00e+00 CPU, roundtrip OKNotes
Based on
base/upstream-6fe749801, a pushed copy of ggml-org6fe749801, because this fork's master does not carry the NextN converter machinery this builds on.The borrow log line only appears under
-v, the same aslazy read enabled. A user cannot see the mechanism engage at default verbosity.scripts/mtp_strip_shared.py, used to produce the stripped sidecar from a published one, is in the workspace rather than in this diff. It is a test fixture, not part of the change.