batch-probe 0.1.0__tar.gz

batch_probe-0.1.0/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/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,165 @@
+# 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/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "batch-probe"
+version = "0.1.0"
+description = "Find the maximum batch size that fits in GPU memory. Binary search with OOM recovery."
+readme = "README.md"
+requires-python = ">=3.9"
+license = "MIT"
+authors = [
+    {name = "Andrew H. Bond"},
+]
+keywords = ["pytorch", "gpu", "memory", "batch-size", "oom", "cuda"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+    "torch>=1.13.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+    "ruff>=0.4.0",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/ahb-sjsu/batch-probe"
+"Bug Tracker" = "https://github.com/ahb-sjsu/batch-probe/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py39"

batch_probe-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

batch_probe-0.1.0/src/batch_probe.egg-info/PKG-INFO

@@ -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/src/batch_probe.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+LICENSE
+README.md
+pyproject.toml
+src/batch_probe.egg-info/PKG-INFO
+src/batch_probe.egg-info/SOURCES.txt
+src/batch_probe.egg-info/dependency_links.txt
+src/batch_probe.egg-info/requires.txt
+src/batch_probe.egg-info/top_level.txt
+src/torch_probe/__init__.py
+src/torch_probe/_cache.py
+src/torch_probe/_cleanup.py
+src/torch_probe/_probe.py
+src/torch_probe/py.typed
+tests/test_cache.py
+tests/test_probe.py

batch_probe-0.1.0/src/batch_probe.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

batch_probe-0.1.0/src/batch_probe.egg-info/requires.txt

@@ -0,0 +1,6 @@
+torch>=1.13.0
+
+[dev]
+pytest>=7.0.0
+pytest-cov>=4.0.0
+ruff>=0.4.0

batch_probe-0.1.0/src/batch_probe.egg-info/top_level.txt

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

batch_probe-0.1.0/src/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"

batch_probe-0.1.0/src/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()

batch_probe-0.1.0/src/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()

batch_probe-0.1.0/src/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

batch_probe-0.1.0/tests/test_cache.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""Tests for the caching layer."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from torch_probe import cached_probe, clear_cache
+
+
+class CountingModel(nn.Module):
+    """Model that counts how many times forward() is called."""
+
+    def __init__(self, oom_threshold: int = 8):
+        super().__init__()
+        self.linear = nn.Linear(10, 2)
+        self.oom_threshold = oom_threshold
+        self.call_count = 0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.call_count += 1
+        if x.shape[0] > self.oom_threshold:
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory.")
+        return self.linear(x)
+
+
+class TestCachedProbe:
+    def setup_method(self):
+        clear_cache()
+
+    def test_cache_hit(self):
+        model = CountingModel(oom_threshold=8)
+        input_fn = lambda bs: {"x": torch.randn(bs, 10)}
+
+        r1 = cached_probe(model, input_fn, mode="infer", high=32, verbose=False)
+        calls_after_first = model.call_count
+
+        r2 = cached_probe(model, input_fn, mode="infer", high=32, verbose=False)
+        calls_after_second = model.call_count
+
+        assert r1 == r2
+        assert calls_after_second == calls_after_first  # No new forward calls
+
+    def test_different_modes_separate_cache(self):
+        model = CountingModel(oom_threshold=8)
+        input_fn = lambda bs: {"x": torch.randn(bs, 10)}
+
+        r_train = cached_probe(model, input_fn, mode="train", high=32, verbose=False)
+        r_infer = cached_probe(model, input_fn, mode="infer", high=32, verbose=False)
+
+        # Both should probe (different modes = different cache keys)
+        # Results may differ since train mode uses backward pass
+        assert isinstance(r_train, int)
+        assert isinstance(r_infer, int)
+
+    def test_clear_cache(self):
+        model = CountingModel(oom_threshold=8)
+        input_fn = lambda bs: {"x": torch.randn(bs, 10)}
+
+        cached_probe(model, input_fn, mode="infer", high=32, verbose=False)
+        calls_first = model.call_count
+
+        clear_cache()
+        cached_probe(model, input_fn, mode="infer", high=32, verbose=False)
+        calls_second = model.call_count
+
+        assert calls_second > calls_first  # Had to re-probe after clearing

batch_probe-0.1.0/tests/test_probe.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2026 Andrew H. Bond
+# Licensed under the MIT License.
+
+"""Tests for the core probe_batch_size function.
+
+These tests work without a GPU by using models that simulate OOM
+via a Python-side threshold check.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from torch_probe import probe_batch_size
+
+
+class FakeOOMModel(nn.Module):
+    """Model that raises OOM above a configurable batch size threshold."""
+
+    def __init__(self, oom_threshold: int = 16):
+        super().__init__()
+        self.linear = nn.Linear(10, 2)
+        self.oom_threshold = oom_threshold
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.shape[0] > self.oom_threshold:
+            raise torch.cuda.OutOfMemoryError(
+                "CUDA out of memory. Tried to allocate 2.00 GiB"
+            )
+        return self.linear(x)
+
+
+class FakeDictModel(nn.Module):
+    """Model that returns a dict with 'loss' key (like HuggingFace)."""
+
+    def __init__(self, oom_threshold: int = 8):
+        super().__init__()
+        self.linear = nn.Linear(10, 2)
+        self.oom_threshold = oom_threshold
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> dict:
+        if input_ids.shape[0] > self.oom_threshold:
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory.")
+        out = self.linear(input_ids.float())
+        return {"loss": out.mean(), "logits": out}
+
+
+def _make_input_fn(feature_dim: int = 10):
+    """Create an input_fn for FakeOOMModel."""
+    return lambda bs: {"x": torch.randn(bs, feature_dim)}
+
+
+def _make_dict_input_fn(feature_dim: int = 10):
+    """Create an input_fn for FakeDictModel."""
+    return lambda bs: {
+        "input_ids": torch.zeros(bs, feature_dim, dtype=torch.long),
+        "attention_mask": torch.ones(bs, feature_dim, dtype=torch.long),
+    }
+
+
+class TestProbeBatchSize:
+    def test_finds_max_batch_infer(self):
+        model = FakeOOMModel(oom_threshold=32)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=128,
+            headroom=0.2,
+            verbose=False,
+        )
+        # max successful = 32, safe = int(32 * 0.8) = 25
+        assert result == 25
+
+    def test_finds_max_batch_train(self):
+        model = FakeOOMModel(oom_threshold=32)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="train",
+            high=128,
+            headroom=0.2,
+            verbose=False,
+        )
+        assert result == 25
+
+    def test_headroom_zero(self):
+        model = FakeOOMModel(oom_threshold=32)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=128,
+            headroom=0.0,
+            verbose=False,
+        )
+        assert result == 32
+
+    def test_headroom_fifty_percent(self):
+        model = FakeOOMModel(oom_threshold=32)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=128,
+            headroom=0.5,
+            verbose=False,
+        )
+        assert result == 16
+
+    def test_returns_at_least_one(self):
+        # Even with high headroom and low threshold, never returns 0
+        model = FakeOOMModel(oom_threshold=2)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=64,
+            headroom=0.9,
+            verbose=False,
+        )
+        assert result >= 1
+
+    def test_all_oom(self):
+        # Model OOMs even at batch=1
+        model = FakeOOMModel(oom_threshold=0)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            low=1,
+            high=64,
+            headroom=0.2,
+            verbose=False,
+        )
+        # best stays at low=1 but it OOMs, so best stays at initial=1
+        # This is an edge case — we return 1 as the floor
+        assert result >= 1
+
+    def test_dict_output_model(self):
+        model = FakeDictModel(oom_threshold=8)
+        result = probe_batch_size(
+            model,
+            _make_dict_input_fn(),
+            mode="train",
+            high=64,
+            headroom=0.2,
+            verbose=False,
+        )
+        assert result == max(1, int(8 * 0.8))  # 6
+
+    def test_cpu_probes_normally(self):
+        # On CPU, probe still runs (no CUDA cleanup, but OOM still detected)
+        model = FakeOOMModel(oom_threshold=16)
+        result = probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=64,
+            device="cpu",
+            headroom=0.2,
+            verbose=False,
+        )
+        assert result == int(16 * 0.8)  # 12
+
+    def test_non_oom_error_propagates(self):
+        class BadModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 2)
+
+            def forward(self, x):
+                raise ValueError("Something went wrong")
+
+        model = BadModel()
+        with pytest.raises(ValueError, match="Something went wrong"):
+            probe_batch_size(
+                model,
+                _make_input_fn(),
+                mode="infer",
+                high=16,
+                verbose=False,
+            )
+
+    def test_restores_training_mode(self):
+        model = FakeOOMModel(oom_threshold=16)
+
+        # Start in eval mode
+        model.eval()
+        assert not model.training
+
+        probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="train",
+            high=32,
+            verbose=False,
+        )
+        # Should restore eval mode
+        assert not model.training
+
+        # Start in train mode
+        model.train()
+        assert model.training
+
+        probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=32,
+            verbose=False,
+        )
+        # Should restore train mode
+        assert model.training
+
+    def test_verbose_output(self, capsys):
+        model = FakeOOMModel(oom_threshold=8)
+        probe_batch_size(
+            model,
+            _make_input_fn(),
+            mode="infer",
+            high=32,
+            verbose=True,
+        )
+        captured = capsys.readouterr()
+        assert "torch-probe" in captured.out
+        assert "max=8" in captured.out
+        assert "safe=" in captured.out