🔍 Batch Finder

Find the maximum value for any dimension your PyTorch models can handle without stopping the code.

Batch Finder detects your model’s inputs (types and shapes), fixes the sizes you choose, and searches for the largest value your run can sustain without stopping the code.

✨ Features

• 🎯 One main function – You call find_max_minibatch for pretty much everything you need here.
• 🔍 Tell it what goes in – You pass names, data types, and rough shapes (input_shapes / forward_params). If your model has a config, that can help too.
• 📐 Shapes your way – Use a simple tuple or list, or a list of tuples when you have multiple tensors (for example [(-1, 128, 512), (-1, 128, 512)]). For several inputs, you can use a text string or a dict. Or skip that and use axis_to_maximize plus fixed_axis (handy for Hugging Face–style models).
• 🚀 Runs forward or full training – Turn backward on or off depending on what you want to test.
• ⚙️ Knobs you can turn – Change how fast it steps up or down (factor_up, factor_down), how many tries it gets (n_attempts), and where it starts (initial_value).
• ⏱️ Time cap (optional) – Set time_limit_seconds if you only want the search to run for so long. When time is up, you get the best size that worked so far (or None if nothing worked yet). Leave it unset to ignore the clock and rely on n_attempts and the normal stop rules.
• 🛡️ Fails without trashing your session – When things blow up, it cleans up. If even batch size 1 fails, you get None (honest “no”).
• 📊 See what it’s doing – tqdm shows a progress bar with useful status text.

📦 Installation

pip install batch-finder

Or from source:

git clone https://github.com/yourusername/batch-finder.git
cd batch-finder
pip install -e .

🚀 Quick Start

One tensor: tuple or list

Use negative integers -1 on each axis you want to maximize (the search tries a single trial size each step). In an all-integer tuple, every -1 position shares that same trial size. Use positive integers for fixed dimensions.

Other negative integers d < -1 (e.g. -2, -3) size that dimension as round(|d| × trial)—the same scaling role as negative floats below (e.g. -2 → 2× the trial size tied to -1).

If the tuple mixes integers and negative floats, you are in compact numeric mode: there must be at least one integer -1 (the searched axis). Any other dimension given as a negative float -x is sized as round(|x| × trial), where trial is the current value on the -1 axis—so |x| is the proportion you want between that axis and the searched axis (e.g. -1.5 keeps that dim about 1.5× the trial size).

from batch_finder import find_max_minibatch

def get_model():
    return MyModel()

# Maximize axis 0; other dims fixed
max_val = find_max_minibatch(get_model, input_shapes=(-1, 64, 256))

# Maximize axis 2
max_val = find_max_minibatch(get_model, input_shapes=(4, 8, -1))

# One integer -1 (searched axis) + negative float: other dims scale by |float| vs. that trial
max_val = find_max_minibatch(get_model, input_shapes=(-1, 4, -1.5, 16))

Several tensors: list of shape tuples

Pass a list or tuple of per-tensor shapes (same order as forward / forward_params). Each entry follows the same rules as a single-tensor tuple (integer -1, integer multipliers -2, -3, …, or compact floats). One trial size is searched; every -1 and every d < -1 uses that trial value as above.

# Example: two tensors, same free shape pattern
max_val = find_max_minibatch(
    get_model,
    input_shapes=[(-1, 128, 512), (-1, 128, 512)],
    forward_params=["x", "y"],
)

For symbolic names and constraints across tensors, use a string or dict below.

HuggingFace-style: axis_to_maximize + fixed_axis

from transformers import AutoModelForCausalLM
from batch_finder import find_max_minibatch

def get_model():
    return AutoModelForCausalLM.from_pretrained("distilgpt2")

# ``forward_params`` defaults to GPT2-style causal LM names; pass your own list for other architectures.
max_batch = find_max_minibatch(
    get_model,
    axis_to_maximize="batch_size",
    fixed_axis={"seq_len": 32},
)
print(f"Max batch size: {max_batch}")

Several tensors: one string

When forward takes multiple tensors, pass input_shapes as text: one (…) group per argument (same order as forward), short names for axes that must match, optional equations between names, and at least one name=-1 for the size you search.

Pattern: (dims…), (dims…), … then optional name=value rules.

• Dimensions: numbers, or names like b, t that repeat across tensors.
• One name must be set to -1 (that is the searched size).
• Rules like t=1.5b tie sizes together (non-integers are rounded for tensor shapes).

import torch
from batch_finder import find_max_minibatch

class MyModel(torch.nn.Module):
    def forward(self, x, y):
        # x: (23, b, t, 45), y: (b, t, 12)

        ...

def get_model():
    return MyModel()

max_b = find_max_minibatch(
    get_model,
    input_shapes="(23, b, t, 45),(b, t, 12), t=1.5b, b=-1",
)

Do not combine this with tuple/list single-tensor mode or with axis_to_maximize. Pass input_shapes= as a keyword.

Several tensors: dict (named arguments)

Same idea as the string form, but keys match forward parameters. Put shared rules in "#constraints" (must include exactly one symbol=-1). Values are shape text, optionally with , int or , float for dtype.

def get_model():
    return MyModel()

max_b = find_max_minibatch(
    get_model,
    input_shapes={
        "input_ids": "(b, t), int",
        "attention_mask": "(b, t), int",
        "input_ids_encoder": "(d, b, t), int",
        "attention_mask_encoder": "(d, b, t), int",
        "labels": "(b, t)",
        "#constraints": "t=2b, b=-1",
    },
)

Use the literal "#constraints" for the rules entry, or import FINDER_CONSTRAINTS_KEY from batch_finder to avoid typos (CONSTRAINTS_KEY is the same string, kept for compatibility).

Custom search parameters

def get_model():
    return MyModel()

max_val = find_max_minibatch(
    get_model,
    input_shapes=(-1, 128, 512),
    initial_value=8,
    n_attempts=30,
    factor_down=3.0,   # divide by 3 on failure
    factor_up=2.0,     # multiply by 2 on success
    time_limit_seconds=120.0,  # optional: return best trial so far before timeout
)

📖 API Reference

find_max_minibatch(...)

Find the largest workable size for the free axis without OOM.

Parameters:

Returns: Shape tuple, tuple of shape tuples (multi-tensor list input_shapes), or int (string/dict input_shapes or axis_to_maximize), or None if nothing worked. With time_limit_seconds, you may get the best value found before the deadline even if n_attempts was not reached.

Modes:

• Tuple/list: first forward argument, or a list/tuple of shape tuples for several arguments; -1 marks the free axis, or add negative floats to scale other axes off that trial size.
• String: one shape group per tensor; names, one =-1, optional equations between names.
• Dict: like the string form, keys = parameter names, "#constraints" for the rules.
• axis_to_maximize + fixed_axis: when you skip input_shapes.

Example output (axis_to_maximize + fixed_axis):

--- Detected inputs (type, estimated shape) ---
  input_ids: integer, (32, 64)
  attention_mask: integer, (32, 64)
---

batch_size fixed={'seq_len': 32}: 100%|████████████████████| 22/50 [01:26<00:00,  3.9s/it, gpus=1, i=22/50, max_ok=1919, min_fail=1920, status=✅, value=1919

✅ Max value that passed: 1919

🔧 How It Works

1. Inputs – Uses input_shapes dict keys or forward_params (no live module in the parent during the search).
2. Types – Integers for *ids, *mask, labels; floats otherwise.
3. Shapes – From optional config, or from get_model().config (one probe if you omit config) plus argument names.
4. Search – On success, grow (memory-guided or factor_up); on failure without a bracket, shrink by factor_down. Once a success and a failure bracket the limit, the next trial is always the midpoint (max_ok + min_fail) // 2 until the bracket tightens to one step. Stops at failure at size 1, when n_attempts is reached, or when time_limit_seconds elapses (then returns the best successful trial so far, if any).
5. Loss – Uses output.loss if present, otherwise sums output tensors.

⚠️ Important Notes

• Memory: Lower initial_value on small GPUs.
• Speed: inference_only=True is faster.
• Training: inference_only=False runs backward too.
• Size 1: If size 1 fails, the function returns None.
• OOM: CUDA OOM and similar errors are caught and the search continues. On Linux/macOS, attempts default to a subprocess so the parent keeps running if the host kills the worker (SIGKILL). Set BATCH_FINDER_SUBPROCESS=0 for in-process attempts if repeated spawn reloads are worse on your host. Each worker run ends with del model, gc.collect(), and CUDA/MPS cache clears where applicable. After failures, the same cleanup runs.
• Defaults: Omitting initial_value and delay picks smaller starts and shorter pauses on CPU/MPS than on CUDA.
• Login nodes: Prefer a real GPU job for meaningful limits. On CPU-only hosts, use a smaller initial_value and/or inference_only=True if RAM is tight.
• Memory-guided steps: After each good run, peak GPU memory vs total guides the next step (capped by max_growth_multiplier). Install psutil for a rough CPU hint. Set memory_guided=False to use only factor_up / factor_down.
• Multi-GPU (one OS process): Use cuda_mem_devices="all" or a list so every visible GPU is measured; memory-guided growth uses the tightest GPU (minimum headroom).
• DDP / torchrun / Accelerate (several processes): Each rank runs the search on its assigned GPU (cuda_mem_devices default is that GPU only). When WORLD_SIZE > 1, find_max_minibatch takes the minimum successful trial size across ranks (JSON sync) before returning. Use BATCH_FINDER_SYNC_DIR or the default $WORK/.cache / $HOME/.cache — do not pass a per-process random output_dir (e.g. each worker re-importing a script with a new random experiment id); every rank must resolve the same sync directory. Layout: <base>/find_max_minibatch_sync/<job_tag>/. The first rank done waits (poll every 4s, status log about every 30s) until others finish searching.
• Time limit: Pass time_limit_seconds=… to cap wall-clock time for the search loop (after the one-time probe). You get the best batch found so far when the limit hits. Subprocess workers are stopped if they would run past the remaining budget; an in-process forward/backward still runs to completion once started.

🤝 Contributing

Contributions are welcome. Please open a Pull Request.

📝 License

MIT — see the LICENSE file.

Made with ❤️ for the PyTorch community