diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index b5bffd46b..e57e1e065 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -311,7 +311,11 @@ def __init__( # Context Params self.context_params = llama_cpp.llama_context_default_params() self.context_params.n_ctx = n_ctx + # Fix for issue #1762: ensure n_ubatch is at least n_batch for embedding + # When n_batch > 512, n_ubatch must also be increased to avoid crashes self.context_params.n_batch = self.n_batch + if self.context_params.n_ubatch < n_batch: + self.set_n_ubatch(n_batch) self.context_params.n_ubatch = min(self.n_batch, n_ubatch) self.context_params.n_threads = self.n_threads self.context_params.n_threads_batch = self.n_threads_batch diff --git a/tests/test_embed_n_ubatch_fix.py b/tests/test_embed_n_ubatch_fix.py new file mode 100644 index 000000000..567b388c5 --- /dev/null +++ b/tests/test_embed_n_ubatch_fix.py @@ -0,0 +1,69 @@ +"""Test for issue #1762 fix - auto-adjust n_ubatch for embedding""" +import pytest +from unittest.mock import Mock, MagicMock, patch +from llama_cpp import Llama + + +def test_embed_auto_adjusts_n_ubatch(): + """Test that embed() automatically adjusts n_ubatch when n_batch > 512""" + # Mock the Llama class + with patch('llama_cpp.Llama.__init__', return_value=None): + llama = Llama.__new__(Llama) + + # Setup mock attributes + llama.n_batch = 4096 + llama.context_params = Mock() + llama.context_params.n_ubatch = 512 # Default value + llama.context_params.embeddings = True + llama.context_params.n_seq_max = 1 + llama._model = Mock() + llama._model.n_embd = Mock(return_value=768) + llama._ctx = Mock() + llama.verbose = False + + # Mock set_n_ubatch method + llama.set_n_ubatch = Mock() + + # Call embed - should auto-adjust n_ubatch + try: + llama.embed("test input") + except Exception: + pass # We expect it to fail later, we just want to check n_ubatch adjustment + + # Verify set_n_ubatch was called with n_batch value + llama.set_n_ubatch.assert_called_once_with(4096) + + +def test_embed_no_adjust_when_n_ubatch_sufficient(): + """Test that embed() doesn't adjust n_ubatch when it's already >= n_batch""" + with patch('llama_cpp.Llama.__init__', return_value=None): + llama = Llama.__new__(Llama) + + # Setup mock attributes with sufficient n_ubatch + llama.n_batch = 4096 + llama.context_params = Mock() + llama.context_params.n_ubatch = 4096 # Already sufficient + llama.context_params.embeddings = True + llama.context_params.n_seq_max = 1 + llama._model = Mock() + llama._model.n_embd = Mock(return_value=768) + llama._ctx = Mock() + llama.verbose = False + + # Mock set_n_ubatch method + llama.set_n_ubatch = Mock() + + # Call embed - should NOT adjust n_ubatch + try: + llama.embed("test input") + except Exception: + pass + + # Verify set_n_ubatch was NOT called + llama.set_n_ubatch.assert_not_called() + + +if __name__ == "__main__": + test_embed_auto_adjusts_n_ubatch() + test_embed_no_adjust_when_n_ubatch_sufficient() + print("All tests passed!")