From a5adce7f5aa588a2d3d795b27c786ad024127dc0 Mon Sep 17 00:00:00 2001 From: Noah Date: Sun, 7 Jun 2026 13:38:39 +0100 Subject: [PATCH 1/4] Update main.py Refine auto batch size with binary search after OOM. After the exponential probe hits CUDA OOM, binary-search between the last successful and first failed size to find higher-throughput batch sizes, still picking the size with the best measured tokens/s. --- src/heretic/main.py | 143 +++++++++++++++++++++++++++++++------------- 1 file changed, 100 insertions(+), 43 deletions(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index b99b1acc..48dd6067 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -68,6 +68,7 @@ def _is_help_invocation() -> bool: from .reproduce import collect_reproducibles from .system import empty_cache, get_accelerator_info from .utils import ( + Prompt, format_duration, get_readme_intro, get_trial_parameters, @@ -83,6 +84,102 @@ def _is_help_invocation() -> bool: upload_reproduce_folder, ) +BATCH_SIZE_REFINEMENT_THRESHOLD = 8 + + +def _is_oom_error(error: Exception) -> bool: + if isinstance(error, torch.cuda.OutOfMemoryError): + return True + if isinstance(error, RuntimeError) and "out of memory" in str(error).lower(): + return True + return False + + +def _benchmark_batch_size( + model: Model, + good_prompts: list[Prompt], + batch_size: int, +) -> float: + print(f"* Trying batch size [bold]{batch_size}[/]... ", end="") + + prompts = good_prompts * math.ceil(batch_size / len(good_prompts)) + prompts = prompts[:batch_size] + + # 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() + + 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)") + + return performance + + +def _determine_batch_size( + model: Model, + good_prompts: list[Prompt], + max_batch_size: int, +) -> int: + print() + print("Determining optimal batch size...") + + batch_size = 1 + best_batch_size = -1 + best_performance = -1.0 + last_successful_batch_size = 0 + failed_batch_size: int | None = None + + while batch_size <= max_batch_size: + try: + performance = _benchmark_batch_size(model, good_prompts, batch_size) + except Exception as error: + if batch_size == 1: + # Even a batch size of 1 already fails. + # We cannot recover from this. + raise + + print(f"[red]Failed[/] ({error})") + if _is_oom_error(error): + failed_batch_size = batch_size + empty_cache() + break + + if performance > best_performance: + best_batch_size = batch_size + best_performance = performance + + last_successful_batch_size = batch_size + batch_size *= 2 + + if failed_batch_size is not None: + low = last_successful_batch_size + high = failed_batch_size + + while high - low > BATCH_SIZE_REFINEMENT_THRESHOLD: + mid = (low + high) // 2 + try: + performance = _benchmark_batch_size(model, good_prompts, mid) + except Exception as error: + if not _is_oom_error(error): + raise + + print(f"[red]Failed[/] ({error})") + empty_cache() + high = mid + continue + + if performance > best_performance: + best_batch_size = mid + best_performance = performance + low = mid + + return best_batch_size + def obtain_merge_strategy(settings: Settings, model: Model) -> str | None: """ @@ -338,49 +435,9 @@ def run(): 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 - - print(f"[red]Failed[/] ({error})") - 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 + settings.batch_size = _determine_batch_size( + model, good_prompts, settings.max_batch_size + ) print(f"* Chosen batch size: [bold]{settings.batch_size}[/]") if settings.response_prefix is None: From 6630ddeec3ce0714e02ed7208cf97a00bab499ea Mon Sep 17 00:00:00 2001 From: noah Date: Sun, 7 Jun 2026 14:46:58 +0100 Subject: [PATCH 2/4] change max batch --- config.default.toml | 2 +- src/heretic/config.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/config.default.toml b/config.default.toml index 9424bfb5..00a1cd4b 100644 --- a/config.default.toml +++ b/config.default.toml @@ -37,7 +37,7 @@ offload_outputs_to_cpu = true batch_size = 0 # auto # Maximum batch size to try when automatically determining the optimal batch size. -max_batch_size = 128 +max_batch_size = 1024 # 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 ada5792b..0e31c6ea 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -180,7 +180,7 @@ class Settings(BaseSettings): ) max_batch_size: int = Field( - default=128, + default=1024, 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. From e96ca40c4d031e0da8fc83de39b5a40499160546 Mon Sep 17 00:00:00 2001 From: Noah <99681487+NoahOksuz@users.noreply.github.com> Date: Sun, 7 Jun 2026 16:53:20 +0100 Subject: [PATCH 3/4] Update src/heretic/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/heretic/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index 48dd6067..02b44137 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -88,7 +88,7 @@ def _is_help_invocation() -> bool: def _is_oom_error(error: Exception) -> bool: - if isinstance(error, torch.cuda.OutOfMemoryError): + if isinstance(error, torch.OutOfMemoryError): return True if isinstance(error, RuntimeError) and "out of memory" in str(error).lower(): return True From 1e7f1e1439b12dc9f55cc0869a2f77bac134fac8 Mon Sep 17 00:00:00 2001 From: Noah <99681487+NoahOksuz@users.noreply.github.com> Date: Sun, 7 Jun 2026 16:53:29 +0100 Subject: [PATCH 4/4] Update src/heretic/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/heretic/main.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/heretic/main.py b/src/heretic/main.py index 02b44137..1ed6b034 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -125,6 +125,11 @@ def _determine_batch_size( good_prompts: list[Prompt], max_batch_size: int, ) -> int: + if not good_prompts: + raise ValueError("The list of good prompts must not be empty.") + if max_batch_size < 1: + raise ValueError("max_batch_size must be at least 1.") + print() print("Determining optimal batch size...")