batch-probe 0.1.0__py3-none-any.whl

batch_probe-0.1.0.dist-info/METADATA

@@ -0,0 +1,194 @@
+Metadata-Version: 2.4
+Name: batch-probe
+Version: 0.1.0
+Summary: Find the maximum batch size that fits in GPU memory. Binary search with OOM recovery.
+Author: Andrew H. Bond
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/ahb-sjsu/batch-probe
+Project-URL: Bug Tracker, https://github.com/ahb-sjsu/batch-probe/issues
+Keywords: pytorch,gpu,memory,batch-size,oom,cuda
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.13.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: ruff>=0.4.0; extra == "dev"
+Dynamic: license-file
+
+# batch-probe
+
+Find the maximum batch size that fits in GPU memory.
+
+Binary search with OOM recovery, configurable safety headroom, no framework required.
+
+## The Problem
+
+Every ML practitioner has done this:
+
+```
+batch_size = 64   # OOM
+batch_size = 32   # OOM
+batch_size = 16   # OOM
+batch_size = 8    # works... but am I leaving GPU memory on the table?
+```
+
+`batch-probe` automates this. It binary-searches for the largest batch size your model can handle, with a safety margin so you don't OOM during real training.
+
+## Install
+
+```bash
+pip install batch-probe
+```
+
+## Quick Start
+
+```python
+from torch_probe import probe_batch_size
+
+batch_size = probe_batch_size(
+    model,
+    lambda bs: {
+        "input_ids": torch.zeros(bs, 512, dtype=torch.long, device="cuda"),
+        "attention_mask": torch.ones(bs, 512, dtype=torch.long, device="cuda"),
+    },
+)
+# torch-probe: probing batch size (mode=train, range=[1, 4096], headroom=20%)... max=6, safe=4
+```
+
+That's it. Three lines. Works with any `nn.Module`.
+
+## Usage
+
+### Encoder models (BERT, RoBERTa, etc.)
+
+```python
+batch_size = probe_batch_size(
+    model,
+    lambda bs: {
+        "input_ids": torch.zeros(bs, 128, dtype=torch.long, device="cuda"),
+        "attention_mask": torch.ones(bs, 128, dtype=torch.long, device="cuda"),
+    },
+    mode="train",
+)
+```
+
+### Seq2seq models (T5, BART, etc.)
+
+```python
+batch_size = probe_batch_size(
+    model,
+    lambda bs: {
+        "input_ids": torch.zeros(bs, 512, dtype=torch.long, device="cuda"),
+        "attention_mask": torch.ones(bs, 512, dtype=torch.long, device="cuda"),
+        "labels": torch.zeros(bs, 512, dtype=torch.long, device="cuda"),
+    },
+    mode="train",
+)
+```
+
+### Vision models
+
+```python
+batch_size = probe_batch_size(
+    model,
+    lambda bs: {"x": torch.randn(bs, 3, 224, 224, device="cuda")},
+    mode="infer",
+)
+```
+
+### Inference-only probing
+
+Inference uses ~2-4x less memory than training (no gradients stored):
+
+```python
+infer_batch = probe_batch_size(model, input_fn, mode="infer")
+train_batch = probe_batch_size(model, input_fn, mode="train")
+# infer_batch >> train_batch
+```
+
+### Custom headroom
+
+Default is 20% safety margin. Adjust for your risk tolerance:
+
+```python
+# Conservative (40% headroom) — for long training runs
+batch_size = probe_batch_size(model, input_fn, headroom=0.4)
+
+# Aggressive (5% headroom) — squeeze every last sample
+batch_size = probe_batch_size(model, input_fn, headroom=0.05)
+```
+
+### Caching
+
+Use `cached_probe` to avoid re-probing the same model:
+
+```python
+from torch_probe import cached_probe, clear_cache
+
+batch_size = cached_probe(model, input_fn, mode="train")  # probes
+batch_size = cached_probe(model, input_fn, mode="train")  # cache hit
+
+clear_cache()  # reset if model changed
+```
+
+## How It Works
+
+1. Binary search between `low` (default 1) and `high` (default 4096)
+2. At each midpoint, create dummy tensors via your `input_fn`
+3. Run a forward pass (+ backward pass in train mode)
+4. If OOM: upper bound ← midpoint − 1, clean GPU memory
+5. If success: lower bound ← midpoint + 1
+6. Return `int(max_successful × (1 − headroom))`
+
+The OOM recovery uses `gc.collect()` + `torch.cuda.empty_cache()` + `torch.cuda.synchronize()` to fully reclaim memory between iterations.
+
+## vs. Alternatives
+
+| Feature | batch-probe | Lightning BatchSizeFinder | HF `auto_find_batch_size` |
+|---|---|---|---|
+| Works with raw PyTorch | Yes | No (needs LightningModule) | No (needs HF Trainer) |
+| Algorithm | Binary search | Power-of-2 scaling | Halve on OOM |
+| Configurable headroom | Yes | No | No |
+| Train + infer modes | Yes | Train only | Train only |
+| Dependencies | torch only | pytorch-lightning | accelerate |
+
+## API Reference
+
+### `probe_batch_size(model, input_fn, *, mode, low, high, headroom, device, verbose)`
+
+Find the maximum safe batch size.
+
+- **model** (`nn.Module`): Your model, already on the target device.
+- **input_fn** (`Callable[[int], dict[str, Tensor]]`): Takes batch size, returns dict of tensors for `model(**inputs)`.
+- **mode** (`"train"` | `"infer"`): Train mode runs forward + backward. Default: `"train"`.
+- **low** (`int`): Minimum batch size. Default: `1`.
+- **high** (`int`): Upper bound for search. Default: `4096`.
+- **headroom** (`float`): Safety margin. Default: `0.2` (20%).
+- **device** (`str | torch.device | None`): Override device. Default: model's device.
+- **verbose** (`bool`): Print progress. Default: `True`.
+
+Returns: `int` — safe batch size.
+
+### `cached_probe(model, input_fn, *, mode, **kwargs)`
+
+Same as `probe_batch_size` but caches results keyed on model class, param count, input shapes, and mode.
+
+### `clear_cache()`
+
+Clear all cached probe results.
+
+## License
+
+MIT

batch_probe-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+batch_probe-0.1.0.dist-info/licenses/LICENSE,sha256=26X3SJ_z1LyCWoldsXSVmDu5ttk5d4SAxrgnxmlBlMo,1092
+torch_probe/__init__.py,sha256=JYZ4O7Psx9WemtWsp_0tMS3DWB5Wjo7qWgUV6QgF3ys,344
+torch_probe/_cache.py,sha256=yEoiSleY1KSrcsdfVs72aEUw7jTDFDKc4GkXoBNFF4w,1736
+torch_probe/_cleanup.py,sha256=pPy84HVpV8LSuAHxIBjktnNGBZSrVFtyqP0RaJ06TeA,413
+torch_probe/_probe.py,sha256=ADUz1cBPyP7G_Ud9FG0KWk2SYxi-FVsbcQ8WNo5GdlY,5589
+torch_probe/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+batch_probe-0.1.0.dist-info/METADATA,sha256=bFfYvcTbVFofV4KwpNrSanE0rOG8TgFDCjO8yRFkmpc,6098
+batch_probe-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+batch_probe-0.1.0.dist-info/top_level.txt,sha256=65Vtl91J5NNyFLwI8ywIHO_2_GKHxsPFUMkNUa1mOJs,12
+batch_probe-0.1.0.dist-info/RECORD,,

batch_probe-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

batch_probe-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Andrew H. Bond
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

batch_probe-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_probe

torch_probe/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""torch-probe: Find the maximum batch size that fits in GPU memory."""
+
+from torch_probe._cache import cached_probe, clear_cache
+from torch_probe._probe import probe_batch_size
+
+__all__ = ["probe_batch_size", "cached_probe", "clear_cache"]
+__version__ = "0.1.0"

torch_probe/_cache.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""In-memory cache for probe results."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Literal, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from torch_probe._probe import probe_batch_size
+
+_cache: dict[str, int] = {}
+
+
+def _make_key(
+    model: nn.Module,
+    input_fn: Callable[[int], Dict[str, torch.Tensor]],
+    mode: str,
+) -> str:
+    """Build a cache key from model identity and input shape."""
+    # Model class + param count gives a stable identity
+    model_id = f"{model.__class__.__name__}_{sum(p.numel() for p in model.parameters())}"
+
+    # Probe input shapes at batch=1
+    try:
+        sample = input_fn(1)
+        shapes = "_".join(
+            f"{k}:{tuple(v.shape)}:{v.dtype}" for k, v in sorted(sample.items())
+        )
+        # Clean up sample tensors
+        del sample
+    except Exception:
+        shapes = "unknown"
+
+    return f"{model_id}__{mode}__{shapes}"
+
+
+def cached_probe(
+    model: nn.Module,
+    input_fn: Callable[[int], Dict[str, torch.Tensor]],
+    *,
+    mode: Literal["train", "infer"] = "train",
+    **kwargs: Any,
+) -> int:
+    """Like :func:`probe_batch_size` but caches results.
+
+    Same arguments as :func:`probe_batch_size`. Returns a cached result
+    if the same model class, parameter count, input shapes, and mode
+    have been probed before.
+    """
+    key = _make_key(model, input_fn, mode)
+    if key not in _cache:
+        _cache[key] = probe_batch_size(model, input_fn, mode=mode, **kwargs)
+    return _cache[key]
+
+
+def clear_cache() -> None:
+    """Clear all cached probe results."""
+    _cache.clear()

torch_probe/_cleanup.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""GPU memory cleanup utilities."""
+
+from __future__ import annotations
+
+import gc
+
+import torch
+
+
+def gpu_cleanup() -> None:
+    """Aggressively free GPU memory after an OOM or between probe iterations."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()

torch_probe/_probe.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""Core binary-search GPU memory probe."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Literal, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from torch_probe._cleanup import gpu_cleanup
+
+
+def _extract_loss(outputs: Any) -> torch.Tensor:
+    """Extract a scalar loss from model outputs for the backward pass.
+
+    Handles:
+      - HuggingFace ModelOutput / dataclass with .loss attribute
+      - dict with "loss" key
+      - plain Tensor
+      - tuple (uses first element)
+      - dict without "loss" (uses first value)
+    """
+    # .loss attribute (HuggingFace ModelOutput, dataclasses)
+    if hasattr(outputs, "loss") and outputs.loss is not None:
+        return outputs.loss
+
+    # dict with "loss" key
+    if isinstance(outputs, dict):
+        if "loss" in outputs:
+            return outputs["loss"]
+        # Fall back to first tensor value
+        for v in outputs.values():
+            if isinstance(v, torch.Tensor):
+                return v.mean()
+
+    # plain Tensor
+    if isinstance(outputs, torch.Tensor):
+        return outputs.mean()
+
+    # tuple / list
+    if isinstance(outputs, (tuple, list)) and len(outputs) > 0:
+        first = outputs[0]
+        if isinstance(first, torch.Tensor):
+            return first.mean()
+
+    raise TypeError(
+        f"Cannot extract a loss from model output of type {type(outputs)}. "
+        "Ensure your model returns a tensor, a dict with a 'loss' key, "
+        "or an object with a .loss attribute."
+    )
+
+
+def probe_batch_size(
+    model: nn.Module,
+    input_fn: Callable[[int], Dict[str, torch.Tensor]],
+    *,
+    mode: Literal["train", "infer"] = "train",
+    low: int = 1,
+    high: int = 4096,
+    headroom: float = 0.2,
+    device: Optional[Union[torch.device, str]] = None,
+    verbose: bool = True,
+) -> int:
+    """Find the maximum batch size that fits in GPU memory.
+
+    Uses binary search with OOM recovery. Tries a forward pass (and backward
+    pass in train mode) at each candidate batch size. Returns the largest
+    successful size minus a safety margin.
+
+    Args:
+        model: Any ``nn.Module``, already on the target device.
+        input_fn: A callable that takes a batch size ``int`` and returns a dict
+            of tensors to pass as ``**kwargs`` to ``model()``. Tensors must
+            already be on the correct device.
+        mode: ``"train"`` runs forward + backward (2-4x more memory).
+            ``"infer"`` runs forward only under ``torch.no_grad()``.
+        low: Minimum batch size to try (and the floor for the return value).
+        high: Starting upper bound for binary search.
+        headroom: Fraction of headroom to subtract. ``0.2`` (default) means
+            the returned batch size is ``int(max_successful * 0.8)``.
+        device: Device to check. Defaults to the device of the model's first
+            parameter. On CPU the probe still runs but skips CUDA-specific cleanup.
+        verbose: Print probe progress.
+
+    Returns:
+        Safe batch size (``int``), guaranteed ``>= low``.
+
+    Example::
+
+        from torch_probe import probe_batch_size
+
+        batch_size = probe_batch_size(
+            model,
+            lambda bs: {
+                "input_ids": torch.zeros(bs, 512, dtype=torch.long, device="cuda"),
+                "attention_mask": torch.ones(bs, 512, dtype=torch.long, device="cuda"),
+            },
+        )
+    """
+    # Resolve device
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+    device = torch.device(device) if isinstance(device, str) else device
+
+    is_cuda = device.type == "cuda"
+
+    # Save and restore model state
+    was_training = model.training
+    best = low
+
+    if verbose:
+        print(f"torch-probe: probing batch size (mode={mode}, range=[{low}, {high}], "
+              f"headroom={headroom:.0%})...", end="", flush=True)
+
+    while low <= high:
+        mid = (low + high) // 2
+        success = False
+        inputs = None
+
+        try:
+            if is_cuda:
+                gpu_cleanup()
+            inputs = input_fn(mid)
+
+            if mode == "train":
+                model.train()
+                outputs = model(**inputs)
+                loss = _extract_loss(outputs)
+                loss.backward()
+                model.zero_grad(set_to_none=True)
+            else:
+                model.eval()
+                with torch.no_grad():
+                    model(**inputs)
+
+            success = True
+
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+            err_msg = str(e).lower()
+            if "out of memory" not in err_msg and "cuda" not in err_msg:
+                # Not an OOM — re-raise
+                model.train(was_training)
+                raise
+        finally:
+            # Always clean up tensors
+            if inputs is not None:
+                del inputs
+            if is_cuda:
+                gpu_cleanup()
+
+        if success:
+            best = mid
+            low = mid + 1
+        else:
+            high = mid - 1
+
+    # Restore model state
+    model.train(was_training)
+
+    safe = max(1, int(best * (1.0 - headroom)))
+    # Never go below the user's requested minimum
+    safe = max(safe, 1)
+
+    if verbose:
+        print(f" max={best}, safe={safe}")
+
+    return safe