diff --git a/config.default.toml b/config.default.toml index 7ce6a5a1..43f25fdd 100644 --- a/config.default.toml +++ b/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 diff --git a/src/heretic/config.py b/src/heretic/config.py index 7bc8a4d6..3702f326 100644 --- a/src/heretic/config.py +++ b/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( diff --git a/src/heretic/main.py b/src/heretic/main.py index c232ada3..e5bbada7 100644 --- a/src/heretic/main.py +++ b/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...") diff --git a/src/heretic/model.py b/src/heretic/model.py index 3ea72fc9..98428e8f 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -2,6 +2,7 @@ # Copyright (C) 2025-2026 Philipp Emanuel Weidmann + 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