From 8f3720f01c27965b81b47a50e06d71064bfb5199 Mon Sep 17 00:00:00 2001 From: Ard van Noordenne Date: Mon, 5 Oct 2026 10:32:47 +0200 Subject: [PATCH] optimizer: keep the probe's margin when the whole job fits The whole-job shortcut ran a job as one batch whenever its forward-pass peak fit in max_memory_padding of free memory. Optimization needs more than one forward pass (optimizer state, fragmentation), so jobs at 0.78-0.80 of the budget passed the check and then ran out of memory mid-run (UMA-s, 878 structures, peak 18.1-18.8 of 23.4 GB). Require peak * memory_scaling_factor to fit instead. The probe returns a batch one growth step below a size that fit, so this gives the shortcut the same margin. Jobs over the threshold fall back to the probe. --- docs/config.md | 2 +- src/gpuma/optimizer.py | 6 ++++-- tests/optimization/test_batch.py | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/docs/config.md b/docs/config.md index b3081fb..6276af5 100644 --- a/docs/config.md +++ b/docs/config.md @@ -223,7 +223,7 @@ technical: |---|---|---| | `device` | `"cuda"` | `"cpu"`, `"cuda"`, or `"cuda:N"` (e.g. `"cuda:0"`). Falls back to `cuda:0` if the requested index doesn't exist | | `max_memory_padding` | `0.95` | Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches | -| `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. Must be > 1 | +| `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. A job runs as one batch without probing only if its forward-pass peak times this factor fits in `max_memory_padding` of free memory. Must be > 1 | | `max_atoms_to_try` | `100000` | Upper bound on the autobatcher's calibration probe size (atoms) | | `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower | | `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"` | diff --git a/src/gpuma/optimizer.py b/src/gpuma/optimizer.py index 6758437..12226e1 100644 --- a/src/gpuma/optimizer.py +++ b/src/gpuma/optimizer.py @@ -377,9 +377,11 @@ def _whole_job_memory_scaler(state: Any, model: Any, batcher: Any) -> float | No free, _ = torch.cuda.mem_get_info() budget_gb = (free + torch.cuda.memory_reserved()) / 1024**3 - if peak_gb > batcher.max_memory_padding * budget_gb: + # Keep the probe's margin: its batch sits one growth step below a size that fit. + if peak_gb * batcher.memory_scaling_factor > batcher.max_memory_padding * budget_gb: logger.info( - "Job peaks at %.2f of %.2f GB, over max_memory_padding; probing GPU memory", + "Job peaks at %.2f of %.2f GB, too close to the limit for one batch; " + "probing GPU memory", peak_gb, budget_gb, ) diff --git a/tests/optimization/test_batch.py b/tests/optimization/test_batch.py index c19824c..f47702d 100644 --- a/tests/optimization/test_batch.py +++ b/tests/optimization/test_batch.py @@ -390,6 +390,24 @@ def oom(state, model, real): assert all(r is not None for r in results) assert probes + def test_peak_within_one_growth_step_falls_back_to_probe( + self, monkeypatch, methane, ethanol + ): + """Used to fail: a job at 0.8 of the budget ran as one batch and OOMed mid-run.""" + import torch + + def near_limit(state, model, real): + real(state, model) + free, _ = torch.cuda.mem_get_info() + return 0.8 * (free + torch.cuda.memory_reserved()) / 1024**3 + + jobs = [methane, ethanol, self._largest()] + self._patch_whole_job_forward(monkeypatch, jobs, near_limit) + probes = self._count_probes(monkeypatch) + results = optimize_structure_batch(jobs, self._config()) + assert all(r is not None for r in results) + assert probes + def test_insufficient_headroom_falls_back_to_probe(self, monkeypatch, methane, ethanol): def huge_peak(state, model, real): real(state, model)