Repository navigation
Conversation
Llama.__init__ already sets self._logits_all = True whenever draft_model
is given, regardless of the raw logits_all constructor argument -- but
self.scores was allocated from that raw argument directly instead of
self._logits_all. That mismatch leaves self.scores sized n_batch instead
of n_ctx whenever a caller passes draft_model=... without also explicitly
passing logits_all=True.
eval() writes into self.scores based on self._logits_all (not the raw
argument), so the undersized buffer crashes generation as soon as n_past
exceeds n_batch:
ValueError: could not broadcast input array from shape (32000,) into shape (0,)
Fix: size self.scores from self._logits_all, consistent with how eval()
already decides whether to write into it at all.
Found while building a small CPU-only voice agent
(https://github.com/yadukpb/potato-agent) that uses
LlamaPromptLookupDecoding for speculative decoding -- hit this on every
turn past ~512 tokens.
abetlen#657 reports a superficially similar "could not broadcast ... into shape
(0,)" message, but that one is input_ids overflowing past n_ctx itself
(genuine context-window exhaustion), a different root cause, so this
isn't folded into that thread. abetlen#2110 is also draft_model-related but a
separate (performance) complaint.
Tested:
- git submodule update --init --recursive, make deps (builds the C++
extension from this checkout with the fix applied), make lint -- clean.
- Added test_draft_model_scores_buffer_sized_for_n_ctx in tests/test_llama.py;
verified it fails with the exact ValueError above against the pre-fix
code and passes against the fix.
- Ran tests/test_llama.py, tests/test_llama_speculative.py,
tests/test_llama_chat_format.py, and tests/test_llama_grammar.py (27
tests total) against the from-source build: all pass, no regressions.
- CPU only, Linux x86_64, no special build flags.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
Llama.__init__already setsself._logits_all = Truewheneverdraft_modelis given, regardless of the rawlogits_allconstructor argument — butself.scoreswas allocated from that raw argument directly instead ofself._logits_all. One-line fix: sizeself.scoresfromself._logits_all, consistent with howeval()already usesself._logits_allto decide whether to write into it at all.Why it matters
Passing
draft_model=...without also explicitly passinglogits_all=Trueleavesself.scoressizedn_batchinstead ofn_ctx.eval()writes into it based onself._logits_all, so generation crashes as soon asn_pastexceedsn_batch:Found while building a small CPU-only voice agent (potato-agent) that uses
LlamaPromptLookupDecodingfor speculative decoding — hit this on every turn past ~512 tokens.#657 reports a superficially similar
"could not broadcast ... into shape (0,)"message, but on closer read that one isinput_idsoverflowing pastn_ctxitself (genuine context-window exhaustion) — a different root cause, so I didn't fold this into that thread. #2110 is alsodraft_model-related but a separate (performance) complaint.How tested
git submodule update --init --recursive,make deps(builds the C++ extension from this checkout with the fix applied),make lint— clean.test_draft_model_scores_buffer_sized_for_n_ctxintests/test_llama.py: constructs aLlamawithdraft_model=LlamaPromptLookupDecoding(...)and no explicitlogits_all, generates pastn_batchtokens. Verified it fails with the exactValueErrorabove against the pre-fix code and passes against the fix.tests/test_llama.py,tests/test_llama_speculative.py,tests/test_llama_chat_format.py, andtests/test_llama_grammar.py(27 tests total) against the from-source build: all pass, no regressions.