fit: measure a draft head that borrows the target's embeddings - #227
danielhanchen wants to merge 1 commit into
Conversation
A draft head can ship without token_embd/output and borrow them from the model it drafts for through ctx_other, so it cannot take a context on its own. The fit path loads it independently, that fails, and the measurement falls back to zero on every device, so the draft is budgeted at nothing and the load OOMs. Lend the extra model a metadata-only context of the main model, on retry only, so nothing that measures today changes. Assisted-by: Claude
|
You have reached your Codex usage limits for security reviews. Please try again later. |
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. |
|
Heads-up: unslothai/unsloth#10149 works around this same gap from the Unsloth Studio side, by adding its own draft estimate to |
LeoBorcherding
left a comment
LeoBorcherding
left a comment
There was a problem hiding this comment.
tried to run this on an R9700 (32 GB, windows 11, vulkan) with the same Qwen3.8-Flash-Next UD-IQ1_S target and both Q8_0 sidecars from your table. short version: the idea is sound, but on the tree Unsloth Studio actually ships, this path never runs.
what i found
- this PR's base (
40a6d6d) has no qwen4exp, so Flash-Next can't load there at all and i couldn't get a before/after on the base itself. - the prebuilt Unsloth Studio ships today (
b11160-mix-a6922cc) already measures the shared head. its loader lets a borrowing draft stand in a placeholder shape when the fitter opens it alone (borrow_shared_tensor, theno_allocbranch), so the draft gets a context and gets priced:
| draft | draft measured on the GPU | fitter's projected GPU use | "failed to measure" warning |
|---|---|---|---|
| none | – | 42179 MiB | – |
mtp-...-shared-Q8_0 |
2800 MiB (2647 weights + 32 KV + 121 compute) | 45316 MiB | no |
mtp-...-Q8_0 |
3444 MiB | 45960 MiB | no |
-fit on -c 16384 -np 1, all three measured on the same box with the shipped llama-server.exe.
- since the first
llama_init_from_modelsucceeds there, the lend-and-retry added here never fires. i got that by reading the code, not from a run with this patch applied (it also doesn't apply cleanly on the mix tree, only the#includehunk conflicts).
where this does matter
it's the fix you want once qwen4exp borrows through ctx_other the way ggml-org#28243 does it, because there the context really does throw requires ctx_other to be set. so i'd pair it with that change instead of landing it alone. when both are in, there will be two mechanisms for the same thing (the loader stand-in and the context lend), and it's worth deciding whether the stand-in goes away.
the gemma4-assistant question from your "not covered" section is still open; i don't have that model either.
next: building ggml-org#28243 with and without this PR on the same R9700, so there's a before/after on a tree where the retry actually fires. will post those numbers here.
This branch has not been deployed
2 participants
Overview
A draft head can ship without
token_embdandoutputand borrow them from the model it drafts for throughctx_other(src/llama-context.cpp:156, used byqwen4exp,eagle3,dflash, and unconditionally bygemma4-assistantat:146). Such a model cannot take a context on its own.--fit onloads the draft independently (common/fit.cpp), so that context creation fails, andadd_extra_memoryfalls back to a zero-filled measurement on every device whilecommon_fit_paramsstill returns SUCCESS. The target is then fitted as if the draft's weights, KV cache and compute buffers cost nothing, and the real load runs out of memory.This lends the extra model a metadata-only (
no_alloc+LLAMA_LOAD_MODE_NONE) context of the main model asctx_other, and only on retry after the firstllama_init_from_modelhas already returnednullptr, so nothing that measures correctly today takes a different path.Measurements
Qwen3.8-Flash-Next
UD-IQ1_Starget on one B200,-fit on -fitt 92160, comparing what the fitter projects it will use. The shared sidecar borrows; the standalone sidecar carries its owntoken_embd+output.MTP/mtp-Qwen3.8-Flash-Next-shared-Q8_0.ggufMTP/mtp-Qwen3.8-Flash-Next-Q8_0.ggufBefore, the shared draft was budgeted 226 MiB over having no draft at all. The two arms that already measured are unchanged to the MiB. The remaining gap between the two sidecars is expected: the shared one does not carry the ~1.3 GB of
token_embd+outputthat the standalone one does.W failed to measure the memory of the extra model, fitting without itno longer appears.E llama_init_from_model: failed to initialize the context: ... requires ctx_other to be setstill prints once per measurement, since the lend happens after the first attempt fails.Not covered
gemma4-assistantdrafts become measurable for the first time with this change and then reach the KV-cache sharing path (src/llama-model.cpp:2688) against ano_alloccontext. I have not tested that combination and do not have such a model here; if it is unsound it would abort rather than throw, so the existingtry/catchinadd_extra_memorywould not contain it. Worth checking before this goes upstream.Builds clean with
LLAMA_FATAL_WARNINGS=ON.