diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..40adc4148 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix: size `Llama.scores` from `self._logits_all` instead of the raw `logits_all` argument, so passing `draft_model` without also explicitly passing `logits_all=True` no longer crashes generation past `n_batch` tokens with `ValueError: could not broadcast input array ... into shape (0,)` by @yadukpb + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..ff86ed9a6 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -482,7 +482,7 @@ def free_lora_adapter(): self._requires_eval = True self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc) self.scores: npt.NDArray[np.single] = np.ndarray( - (n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single + (n_ctx if self._logits_all else n_batch, self._n_vocab), dtype=np.single ) self._mirostat_mu = ctypes.c_float( diff --git a/tests/test_llama.py b/tests/test_llama.py index c1c16e30c..c8adff920 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -267,6 +267,35 @@ def logit_processor_func(input_ids, logits): assert number_1 == number_3 +def test_draft_model_scores_buffer_sized_for_n_ctx(llama_cpp_transformer_model_path): + # Llama.__init__ sets self._logits_all = True whenever draft_model is + # given (see the `logits_all if draft_model is None else True` line), + # regardless of whether the caller also passed logits_all=True -- but + # self.scores used to be allocated from the raw logits_all argument + # instead of self._logits_all, so it came out n_batch-sized instead of + # n_ctx-sized. eval() writes into self.scores using self._logits_all, + # so generation crashed with "could not broadcast input array ... + # into shape (0,)" as soon as n_past exceeded n_batch. + from llama_cpp.llama_speculative import LlamaPromptLookupDecoding + + llm = llama_cpp.Llama( + model_path=llama_cpp_transformer_model_path, + n_ctx=128, + n_batch=16, + draft_model=LlamaPromptLookupDecoding(max_ngram_size=2, num_pred_tokens=4), + verbose=False, + # Deliberately not passing logits_all=True: draft_model alone + # should be enough, per self._logits_all's own logic above. + ) + + assert llm.scores.shape[0] >= llm.n_batch + + # More tokens than n_batch so eval() must cross that boundary at least + # once; this is what used to raise ValueError before the fix. + output = llm.create_completion("Once upon a time", max_tokens=30) + assert output["choices"][0]["text"] + + def test_real_llama_repeated_prompt_cache(llama_cpp_model_path): model = llama_cpp.Llama( llama_cpp_model_path,