fix(embed): respect physical batch limits for non-causal embeddings

.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,13 +37,15 @@ 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: |
           hf download ${{ env.REPO_ID }} ${{ env.MODEL_FILE }}
           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

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

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

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"]