diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 82f7a8a7d..6b877ca29 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -16,7 +16,10 @@ env: HYBRID_MODEL_FILE: Falcon-H1-Tiny-90M-Instruct-Q2_K.gguf EMBEDDING_REPO_ID: CompendiumLabs/bge-small-en-v1.5-gguf EMBEDDING_MODEL_FILE: bge-small-en-v1.5-q4_k_m.gguf - MODEL_CACHE_KEY: qwen35-q8-mamba130m-q2-falconh1tiny-q2-bge-small-q4 + LONG_EMBEDDING_REPO_ID: lm-kit/bge-m3-gguf + LONG_EMBEDDING_MODEL_FILE: bge-m3-F16.gguf + LONG_EMBEDDING_REVISION: 9379ce497e8814b200f2dc0d18eb4045426dcb8c + MODEL_CACHE_KEY: qwen35-q8-mamba130m-q2-falconh1tiny-q2-bge-small-q4-bge-m3-f16 jobs: download-model: @@ -34,6 +37,7 @@ jobs: path: ~/.cache/huggingface/hub key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} restore-keys: | + ${{ runner.os }}-model-qwen35-q8-mamba130m-q2-falconh1tiny-q2-bge-small-q4 ${{ runner.os }}-model-qwen35-q8-mamba130m-q2-falconh1tiny-q2 - name: Download model run: | @@ -41,6 +45,7 @@ jobs: hf download ${{ env.RECURRENT_REPO_ID }} ${{ env.RECURRENT_MODEL_FILE }} hf download ${{ env.HYBRID_REPO_ID }} ${{ env.HYBRID_MODEL_FILE }} hf download ${{ env.EMBEDDING_REPO_ID }} ${{ env.EMBEDDING_MODEL_FILE }} + hf download ${{ env.LONG_EMBEDDING_REPO_ID }} ${{ env.LONG_EMBEDDING_MODEL_FILE }} --revision ${{ env.LONG_EMBEDDING_REVISION }} build-linux: needs: download-model diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..d01daeef8 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(embed): respect physical batch limits for non-causal embeddings by @original4422 in #2387 + ## [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..12087d937 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -74,7 +74,7 @@ def __init__( seed: int = llama_cpp.LLAMA_DEFAULT_SEED, n_ctx: int = 512, n_batch: int = 512, - n_ubatch: int = 512, + n_ubatch: Optional[int] = None, n_threads: Optional[int] = None, n_threads_batch: Optional[int] = None, rope_scaling_type: Optional[ @@ -159,7 +159,7 @@ def __init__( seed: RNG seed, -1 for random n_ctx: Text context, 0 = from model n_batch: Prompt processing maximum batch size - n_ubatch: Physical batch size + n_ubatch: Physical batch size. Defaults to n_batch in embedding mode and 512 otherwise. n_threads: Number of threads to use for generation n_threads_batch: Number of threads to use for batch processing rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggerganov/llama.cpp/pull/2054 @@ -308,6 +308,8 @@ def __init__( ].key = b"\0" # ensure sentinel element is zeroed self.model_params.kv_overrides = self._kv_overrides_array + if n_ubatch is None: + n_ubatch = n_batch if embedding else 512 self.n_batch = min(n_ctx, n_batch) # ??? self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1) self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count() @@ -1098,6 +1100,14 @@ def embed( "Llama model must be created with embedding=True to call this method" ) + batch_capacity = n_batch + if not llama_cpp.llama_get_causal_attn( + self._ctx.ctx + ) or not llama_cpp.llama_get_memory(self._ctx.ctx): + # Non-causal sequences must fit in one physical batch. Respect an + # explicit microbatch limit without truncating a sequence to fit it. + batch_capacity = min(n_batch, llama_cpp.llama_n_ubatch(self._ctx.ctx)) + if self.verbose: llama_cpp.llama_perf_context_reset(self._ctx.ctx) @@ -1160,9 +1170,14 @@ def decode_batch(seq_sizes: List[int]): raise ValueError( f"Requested tokens ({n_tokens}) exceed batch size of {n_batch}" ) + if n_tokens > batch_capacity: + raise ValueError( + f"Requested tokens ({n_tokens}) exceed physical batch size of " + f"{batch_capacity} for non-causal embeddings; increase n_ubatch" + ) # time to eval batch - if t_batch + n_tokens > n_batch or p_batch >= n_seq_max: + if t_batch + n_tokens > batch_capacity or p_batch >= n_seq_max: decode_batch(s_batch) s_batch = [] t_batch = 0 diff --git a/tests/test_embedding_batches.py b/tests/test_embedding_batches.py new file mode 100644 index 000000000..4e4a7ee0a --- /dev/null +++ b/tests/test_embedding_batches.py @@ -0,0 +1,139 @@ +"""Exercise native embedding batch limits in subprocesses (regressions can abort).""" + +import contextlib +import json +from pathlib import Path +import subprocess +import sys + +from huggingface_hub import hf_hub_download +import numpy as np +import pytest + +import llama_cpp + + +@pytest.fixture(scope="module") +def long_embedding_model_path(): + return hf_hub_download( + "lm-kit/bge-m3-gguf", + "bge-m3-F16.gguf", + revision="9379ce497e8814b200f2dc0d18eb4045426dcb8c", + ) + + +def _check_embedding_case(model_path, case, params): + with contextlib.closing( + llama_cpp.Llama( + model_path, + n_gpu_layers=0, + n_threads=2, + n_threads_batch=2, + embedding=True, + pooling_type=llama_cpp.LLAMA_POOLING_TYPE_NONE, + verbose=False, + **params, + ) + ) as model: + assert not llama_cpp.llama_get_causal_attn(model._ctx.ctx) + native_ubatch = llama_cpp.llama_n_ubatch(model._ctx.ctx) + if case == "explicit-small": + assert native_ubatch == 256 + text = "a " * 255 + assert len(model.tokenize(text.encode())) == 257 + with pytest.raises(ValueError, match="physical batch size"): + model.embed(text) + return + if case == "explicit-small-multi": + assert native_ubatch == 256 + inputs = ["a " * 126, "a " * 127] + counts = [128, 129] + elif case == "multi": + assert native_ubatch == 513 + inputs = ["a " * 254, "a " * 255] + counts = [256, 257] + elif case == "short": + assert native_ubatch == min(params["n_ctx"], params["n_batch"]) + inputs = ["a " * 3] + counts = [5] + else: + assert native_ubatch == 513 + inputs = ["a " * 511] + counts = [513] + assert [len(model.tokenize(text.encode())) for text in inputs] == counts + embeddings, count = model.embed(inputs, truncate=False, return_count=True) + assert count == sum(counts) + assert len(embeddings) == len(inputs) + for embedding, tokens in zip(embeddings, counts): + array = np.asarray(embedding) + assert array.shape == (tokens, model.n_embd()) + assert np.isfinite(array).all() + + +@pytest.mark.parametrize( + "case,params", + [ + ("long", {"n_ctx": 1024, "n_batch": 513}), + ("long", {"n_ctx": 0, "n_batch": 513}), + ("long", {"n_ctx": 1024, "n_batch": 513, "n_ubatch": 513}), + ("long", {"n_ctx": 1024, "n_batch": 513, "n_ubatch": 0}), + ("multi", {"n_ctx": 1024, "n_batch": 513}), + ("explicit-small", {"n_ctx": 1024, "n_batch": 513, "n_ubatch": 256}), + ("explicit-small-multi", {"n_ctx": 1024, "n_batch": 513, "n_ubatch": 256}), + ("short", {"n_ctx": 256, "n_batch": 513}), + ("short", {"n_ctx": 512, "n_batch": 512}), + ], +) +def test_native_embedding_batch_limits(long_embedding_model_path, case, params): + code = ( + "import json, runpy, sys\n" + "if sys.platform != 'win32':\n" + " import resource\n" + " resource.setrlimit(resource.RLIMIT_CORE, (0, 0))\n" + "runpy.run_path(sys.argv[1])['_check_embedding_case'](" + "sys.argv[2], sys.argv[3], json.loads(sys.argv[4]))" + ) + completed = subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + long_embedding_model_path, + case, + json.dumps(params), + ], + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.parametrize("n_ubatch,expected", [(None, 512), (128, 128), (0, 513)]) +@pytest.mark.parametrize("embedding", [False, True]) +def test_causal_microbatch_defaults(n_ubatch, expected, embedding): + model_path = hf_hub_download("ggml-org/models", "tinyllamas/stories15M-q4_0.gguf") + params = {} if n_ubatch is None else {"n_ubatch": n_ubatch} + if embedding and n_ubatch is None: + expected = 513 + with contextlib.closing( + llama_cpp.Llama( + model_path, + n_ctx=1024, + n_batch=513, + n_gpu_layers=0, + n_threads=2, + n_threads_batch=2, + verbose=False, + embedding=embedding, + **params, + ) + ) as model: + assert llama_cpp.llama_get_causal_attn(model._ctx.ctx) + assert llama_cpp.llama_n_ubatch(model._ctx.ctx) == expected + result = model.create_completion( + "Once upon a time", max_tokens=2, temperature=0 + ) + assert result["usage"]["completion_tokens"] == 2 + assert result["choices"][0]["text"]