From b32ccbc6f4c38cf579db73bfb60eff0d07196715 Mon Sep 17 00:00:00 2001 From: Ard van Noordenne Date: Wed, 26 Aug 2026 16:56:59 +0200 Subject: [PATCH] optimizer: size the autobatcher by the GPU, not by the input _optimize_batch capped the autobatcher's memory probe with effective_max_atoms = min(batched_state.n_atoms, max_atoms_to_try) but max_atoms_to_try bounds how much InFlightAutoBatcher is allowed to try fitting on the device. Tying it to the total atom count of the submitted job made the memory ceiling a function of how much work was handed in, so a small job probed a small ceiling. InFlightAutoBatcher then raises outright when any state's metric exceeds that ceiling. Checked against the real _optimize_batch path on a Quadro RTX 6000, with a Lennard-Jones potential standing in for the MLIP so only the batching logic varies. Before, three of four shapes fail: single structure ValueError: Max metric ... 18.0 > 17.1 small then large (10, 200) ValueError: State metric=398.0 > 302.4 3 equal conformers ok mixed library (12/40/8/150) ValueError: State metric=298.0 > 211.2 After, all four complete. A single structure in batch mode failed deterministically: the probe is capped at exactly one structure's worth of atoms, so max_memory_padding=0.95 scales the ceiling just below that same structure's own metric. Equal-sized conformers were the shape that happened to survive, which is why this reads as intermittent. Also drop the batcher.load_states() call. torch_sim.optimize calls it itself, on the OptimState from _chunked_apply -- which carries the optimizer's own tensors and so reflects the real footprint. Pre-loading was not just redundant: _get_first_batch skips the estimate when max_memory_scaler is already set, so the pre-load pinned the ceiling to a measurement taken on the bare SimState and torch_sim never re-measured, packing batches against an underestimate. Worst with LBFGS, which stores history per system. The "Memory estimation" phase disappears with it, since the probe now happens inside torch_sim.optimize where gpuma cannot time it. TimingCapture keeps the attribute so examples/full_benchmark.py still runs; it reads 0.0 and the cost lands in "Optimization". Documented on the property. --- src/gpuma/optimizer.py | 47 +++++++++++++++++++++-------------- src/gpuma/utils/decorators.py | 14 +++++++++++ 2 files changed, 43 insertions(+), 18 deletions(-) diff --git a/src/gpuma/optimizer.py b/src/gpuma/optimizer.py index a3b63b7..0cef299 100644 --- a/src/gpuma/optimizer.py +++ b/src/gpuma/optimizer.py @@ -456,25 +456,36 @@ def _optimize_batch( max_atoms_to_try = int(config.technical.max_atoms_to_try) steps_between_swaps = int(config.technical.steps_between_swaps) - effective_max_atoms = min(batched_state.n_atoms, max_atoms_to_try) - with timed_block("Memory estimation"): - batcher = InFlightAutoBatcher( - model, - memory_scales_with="n_edges", - memory_scaling_factor=memory_scaling_factor, - max_memory_padding=max_memory_padding, - max_atoms_to_try=effective_max_atoms, - ) + # max_atoms_to_try bounds the autobatcher's probe for how much fits on the + # GPU, so it must describe the *device*, not the job. Passing + # min(batched_state.n_atoms, ...) tied the ceiling to the size of whatever + # was submitted: a two-structure run probed at most two structures' worth + # of memory, and the resulting max_memory_scaler was small enough that + # InFlightAutoBatcher rejected any later structure larger than the first + # ("State metric=N is greater than max_metric M") or, short of that, packed + # one system per batch and threw away the batching entirely. + batcher = InFlightAutoBatcher( + model, + memory_scales_with="n_edges", + memory_scaling_factor=memory_scaling_factor, + max_memory_padding=max_memory_padding, + max_atoms_to_try=max_atoms_to_try, + ) + logger.debug( + "Autobatcher params: memory_scales_with=n_edges, " + "max_memory_padding=%.2f, max_atoms_to_try=%d, steps_between_swaps=%d", + max_memory_padding, + max_atoms_to_try, + steps_between_swaps, + ) - batcher.load_states(batched_state) - logger.debug( - "Autobatcher params: memory_scales_with=n_edges, max_memory_scaler=%.0f, " - "max_memory_padding=%.2f, max_atoms=%d, steps_between_swaps=%d", - batcher.max_memory_scaler, - max_memory_padding, - effective_max_atoms, - steps_between_swaps, - ) + # No batcher.load_states() here. torch_sim.optimize calls it itself, on the + # OptimState it builds via _chunked_apply -- which carries the optimizer's + # own tensors (FIRE velocities, LBFGS history) and so reflects the real + # footprint. Loading first is not merely redundant: _get_first_batch skips + # the estimate when max_memory_scaler is already set, so a pre-load pinned + # the ceiling to a measurement taken on the bare SimState and torch_sim + # never re-measured, leaving the batch packed against an underestimate. with timed_block("Optimization"): final_state = torch_sim.optimize( diff --git a/src/gpuma/utils/decorators.py b/src/gpuma/utils/decorators.py index 8b28ed4..1d5feb3 100644 --- a/src/gpuma/utils/decorators.py +++ b/src/gpuma/utils/decorators.py @@ -92,6 +92,12 @@ class TimingCapture: - ``optimization`` — ``"Optimization"`` - ``total`` — ``"Total optimization"`` + ``memory_estimation`` reads ``0.0`` on the current batch path. The + autobatcher probe now runs inside ``torch_sim.optimize``, where gpuma + cannot time it separately, so its cost is included in ``optimization``. + The attribute is kept so existing benchmark scripts keep working, and + ``overhead`` is correspondingly just ``model_loading`` in practice. + Other ``timed_block`` events are still recorded; access them via :meth:`get` or :attr:`raw`. Missing phases default to ``0.0``. """ @@ -118,6 +124,14 @@ def model_loading(self) -> float: @property def memory_estimation(self) -> float: + """Autobatcher probe time; ``0.0`` now that it runs inside torch-sim. + + gpuma used to call ``batcher.load_states()`` itself, which triggered + the probe somewhere it could be timed. That pre-load also pinned the + memory ceiling to a measurement taken before the optimizer's own + tensors existed, so it was removed; the probe moved inside + ``torch_sim.optimize`` and its cost now lands in ``optimization``. + """ return self.raw.get("Memory estimation", 0.0) @property