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)