Use adaptive per-method batch sizing, drop startup benchmark

config.default.toml

@@ -33,11 +33,9 @@ device_map = "auto"
 # but may slightly reduce performance due to host/device transfers.
 offload_outputs_to_cpu = true
 
-# Number of input sequences to process in parallel (0 = auto).
-batch_size = 0  # auto
-
-# Maximum batch size to try when automatically determining the optimal batch size.
-max_batch_size = 128
+# Number of input sequences to process in parallel.
+# If an out-of-memory error occurs, the batch size is halved automatically.
+batch_size = 128
 
 # Maximum number of tokens to generate for each response.
 max_response_length = 100

src/heretic/config.py

@@ -182,16 +182,9 @@ class Settings(BaseSettings):
     )
 
     batch_size: int = Field(
-        default=0,  # auto
-        description="Number of input sequences to process in parallel (0 = auto).",
-    )
-
-    max_batch_size: int = Field(
         default=128,
-        description="Maximum batch size to try when automatically determining the optimal batch size.",
-        # When storing a settings object, the batch size is already fixed,
-        # either determined by the automatic mechanism or by explicit user choice.
-        exclude=True,
+        description="Number of input sequences to process in parallel. "
+        "If an out-of-memory error occurs, the batch size is halved automatically.",
     )
 
     max_response_length: int = Field(

src/heretic/main.py

@@ -386,57 +386,6 @@ def run():
     bad_prompts = load_prompts(settings, settings.bad_prompts)
     print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
 
-    if settings.batch_size == 0:
-        print()
-        print("Determining optimal batch size...")
-
-        batch_size = 1
-        best_batch_size = -1
-        best_performance = -1
-
-        while batch_size <= settings.max_batch_size:
-            print(f"* Trying batch size [bold]{batch_size}[/]... ", end="")
-
-            prompts = good_prompts * math.ceil(batch_size / len(good_prompts))
-            prompts = prompts[:batch_size]
-
-            try:
-                # Warmup run to build the computation graph so that part isn't benchmarked.
-                model.get_responses(prompts)
-
-                start_time = time.perf_counter()
-                responses = model.get_responses(prompts)
-                end_time = time.perf_counter()
-            except Exception as error:
-                if batch_size == 1:
-                    # Even a batch size of 1 already fails.
-                    # We cannot recover from this.
-                    raise
-
-                formatted = format_exception(error)
-                if "\n" in formatted:
-                    print(f"[red]Failed:\n{formatted}[/]")
-                else:
-                    print(f"[red]Failed ({formatted})[/]")
-
-                break
-
-            response_lengths = [
-                len(model.tokenizer.encode(response)) for response in responses
-            ]
-            performance = sum(response_lengths) / (end_time - start_time)
-
-            print(f"[green]Ok[/] ([bold]{performance:.0f}[/] tokens/s)")
-
-            if performance > best_performance:
-                best_batch_size = batch_size
-                best_performance = performance
-
-            batch_size *= 2
-
-        settings.batch_size = best_batch_size
-        print(f"* Chosen batch size: [bold]{settings.batch_size}[/]")
-
     if settings.response_prefix is None:
         print()
         print("Checking for common response prefix...")

src/heretic/model.py

@@ -2,6 +2,7 @@
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 import math
+from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass
 from typing import Any, Type, cast
@@ -33,7 +34,7 @@
 
 from .config import QuantizationMethod, RowNormalization, Settings
 from .system import empty_cache
-from .utils import Prompt, batchify, format_exception, print
+from .utils import Prompt, format_exception, print
 
 
 def get_model_class(
@@ -97,6 +98,10 @@ def __init__(self, settings: Settings):
         self.tokenizer.padding_side = "left"
 
         self.model = None  # ty:ignore[invalid-assignment]
+        self._batch_sizes: dict[str, int] = {
+            key: settings.batch_size if settings.batch_size > 0 else 128
+            for key in ("responses", "residuals", "logprobs")
+        }
         self.max_memory = (
             {int(k) if k.isdigit() else k: v for k, v in settings.max_memory.items()}
             if settings.max_memory
@@ -672,20 +677,43 @@ def get_responses(
             skip_special_tokens=skip_special_tokens,
         )
 
+    def _batched(
+        self, prompts: list[Prompt], key: str, fn: Callable[[list[Prompt]], Any]
+    ) -> list[Any]:
+        """Run fn on prompt batches, halving the batch size on OOM."""
+        results = []
+        i = 0
+        while i < len(prompts):
+            n = self._batch_sizes[key]
+            try:
+                results.append(fn(prompts[i : i + n]))
+                i += n
+            except (torch.cuda.OutOfMemoryError, RuntimeError) as error:
+                # On non-CUDA platforms (MPS, XPU), OOM raises RuntimeError
+                # with an "out of memory" message rather than the typed exception.
+                if (
+                    isinstance(error, RuntimeError)
+                    and "out of memory" not in str(error).lower()
+                ):
+                    raise
+                if n == 1:
+                    raise
+                self._batch_sizes[key] = n // 2
+                empty_cache()
+        return results
+
     def get_responses_batched(
         self,
         prompts: list[Prompt],
         skip_special_tokens: bool = False,
     ) -> list[str]:
         responses = []
-
-        for batch in batchify(prompts, self.settings.batch_size):
-            for response in self.get_responses(
-                batch,
-                skip_special_tokens=skip_special_tokens,
-            ):
-                responses.append(response)
-
+        for batch in self._batched(
+            prompts,
+            "responses",
+            lambda b: self.get_responses(b, skip_special_tokens=skip_special_tokens),
+        ):
+            responses.extend(batch)
         return responses
 
     def get_residuals(self, prompts: list[Prompt]) -> Tensor:
@@ -741,12 +769,10 @@ def get_residuals(self, prompts: list[Prompt]) -> Tensor:
         return residuals
 
     def get_residuals_batched(self, prompts: list[Prompt]) -> Tensor:
-        residuals = []
-
-        for batch in batchify(prompts, self.settings.batch_size):
-            residuals.append(self.get_residuals(batch))
-
-        return torch.cat(residuals, dim=0)
+        return torch.cat(
+            self._batched(prompts, "residuals", self.get_residuals),
+            dim=0,
+        )
 
     def get_residuals_mean(self, prompts: list[Prompt]) -> Tensor:
         if not prompts:
@@ -755,18 +781,14 @@ def get_residuals_mean(self, prompts: list[Prompt]) -> Tensor:
         running_sum = None
         total_count = 0
 
-        for batch in batchify(prompts, self.settings.batch_size):
-            batch_residuals = self.get_residuals(batch)
-
+        for batch in self._batched(prompts, "residuals", self.get_residuals):
             # Accumulate in high precision on CPU to reduce peak VRAM usage.
-            batch_sum = batch_residuals.sum(dim=0, dtype=torch.float64).cpu()
-
+            batch_sum = batch.sum(dim=0, dtype=torch.float64).cpu()
             if running_sum is None:
                 running_sum = batch_sum
             else:
                 running_sum += batch_sum
-
-            total_count += batch_residuals.shape[0]
+            total_count += batch.shape[0]
 
         assert running_sum is not None
 
@@ -806,12 +828,10 @@ def get_logprobs(self, prompts: list[Prompt]) -> Tensor:
         return logprobs
 
     def get_logprobs_batched(self, prompts: list[Prompt]) -> Tensor:
-        logprobs = []
-
-        for batch in batchify(prompts, self.settings.batch_size):
-            logprobs.append(self.get_logprobs(batch))
-
-        return torch.cat(logprobs, dim=0)
+        return torch.cat(
+            self._batched(prompts, "logprobs", self.get_logprobs),
+            dim=0,
+        )
 
     def stream_chat_response(self, chat: list[dict[str, str]]) -> str:
         # This cast is valid because str is the return type