perf(autobatcher): reduce the number of memory probes

tests/test_autobatching.py

@@ -10,6 +10,7 @@
     _n_edges_scalers,
     calculate_memory_scalers,
     determine_max_batch_size,
+    estimate_max_memory_scaler,
     to_constant_volume_bins,
 )
 from torch_sim.models.lennard_jones import LennardJonesModel
@@ -1072,3 +1073,140 @@ def test_binning_auto_batcher_with_lbfgs(
         all_finished_states.extend(batch.split())
 
     assert len(all_finished_states) == len(lbfgs_states)
+
+
+@pytest.fixture
+def probed_n_atoms(monkeypatch: pytest.MonkeyPatch) -> list[int]:
+    """Record the number of atoms of each system whose batch size is probed."""
+    probed: list[int] = []
+
+    def recording_probe(state: ts.SimState, *_args: Any, **_kwargs: Any) -> int:
+        probed.append(state.n_atoms)
+        return 200
+
+    monkeypatch.setattr(ts.autobatching, "determine_max_batch_size", recording_probe)
+    return probed
+
+
+@pytest.mark.parametrize(("identical", "expected"), [(True, [8]), (False, [8, 256])])
+def test_estimate_max_memory_scaler_probes_the_smallest_and_largest_systems(
+    si_sim_state: ts.SimState,
+    fe_supercell_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    probed_n_atoms: list[int],
+    *,
+    identical: bool,
+    expected: list[int],
+) -> None:
+    """Identical systems are probed once, different ones the smallest and largest."""
+    second = si_sim_state if identical else fe_supercell_sim_state
+    states = ts.concatenate_states([si_sim_state, second])
+    metrics = calculate_memory_scalers(states, "n_atoms")
+
+    estimate_max_memory_scaler(states, lj_model, metrics)
+
+    assert probed_n_atoms == expected
+
+
+@pytest.mark.parametrize(
+    ("make_states", "expected"),
+    [
+        (list, [8, 256]),  # all the states are known: the smallest and the largest
+        (ts.concatenate_states, [8, 256]),
+        (iter, [8, 8, 256]),  # lazy: a bootstrap probe of the first state first
+    ],
+)
+def test_in_flight_auto_batcher_probes_known_states_without_bootstrap(
+    si_sim_state: ts.SimState,
+    fe_supercell_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    probed_n_atoms: list[int],
+    make_states: Any,
+    expected: list[int],
+) -> None:
+    batcher = InFlightAutoBatcher(model=lj_model, memory_scales_with="n_atoms")
+
+    batcher.load_states(make_states([si_sim_state, fe_supercell_sim_state]))
+
+    assert probed_n_atoms == expected
+
+
+@pytest.mark.parametrize(("identical", "expected"), [(True, [32]), (False, [32, 256])])
+def test_optimize_reuses_the_memory_estimate_of_the_init(
+    ar_supercell_sim_state: ts.SimState,
+    fe_supercell_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    probed_n_atoms: list[int],
+    *,
+    identical: bool,
+    expected: list[int],
+) -> None:
+    """The in-flight autobatcher takes the limit estimated for the init."""
+    second = ar_supercell_sim_state if identical else fe_supercell_sim_state
+    state = ts.initialize_state(
+        [ar_supercell_sim_state, second], lj_model.device, lj_model.dtype
+    )
+
+    ts.optimize(
+        system=state,
+        model=lj_model,
+        optimizer=ts.Optimizer.fire,
+        max_steps=5,
+        autobatcher=True,
+    )
+
+    assert probed_n_atoms == expected
+
+
+def test_optimize_estimates_the_init_with_the_settings_of_the_autobatcher(
+    ar_supercell_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The limit estimated for the init uses the autobatcher's cutoff and growth."""
+    scale_factors: list[float] = []
+
+    def recording_probe(*_args: Any, scale_factor: float, **_kwargs: Any) -> int:
+        scale_factors.append(scale_factor)
+        return 200
+
+    monkeypatch.setattr(ts.autobatching, "determine_max_batch_size", recording_probe)
+    # First neighbors only (3.7 A), where the default of 6.0 also counts the second
+    cutoff = 4.0
+    autobatcher = InFlightAutoBatcher(
+        model=lj_model,
+        memory_scales_with="n_edges",
+        cutoff=cutoff,
+        memory_scaling_factor=2.0,
+    )
+
+    ts.optimize(
+        system=ar_supercell_sim_state,
+        model=lj_model,
+        optimizer=ts.Optimizer.fire,
+        max_steps=5,
+        autobatcher=autobatcher,
+    )
+
+    metrics = calculate_memory_scalers(ar_supercell_sim_state, "n_edges", cutoff)
+    assert autobatcher.max_memory_scaler == pytest.approx(200 * metrics[0])
+    assert scale_factors == [2.0]
+
+
+def test_in_flight_auto_batcher_counts_lazy_states_at_its_cutoff(
+    ar_supercell_sim_state: ts.SimState,
+    lj_model: LennardJonesModel,
+    probed_n_atoms: list[int],
+) -> None:
+    """The first state of a lazy iterator counts edges at the autobatcher's cutoff."""
+    # First neighbors only (3.7 A), where the default of 6.0 also counts the second
+    cutoff = 4.0
+    batcher = InFlightAutoBatcher(
+        model=lj_model, memory_scales_with="n_edges", cutoff=cutoff
+    )
+
+    batcher.load_states(iter([ar_supercell_sim_state]))
+
+    metrics = calculate_memory_scalers(ar_supercell_sim_state, "n_edges", cutoff)
+    assert probed_n_atoms == [32, 32]
+    assert batcher.max_memory_scaler == pytest.approx(200 * metrics[0])

torch_sim/autobatching.py

@@ -480,18 +480,24 @@ def estimate_max_memory_scaler(
     min_metric = metric_values.min()
     max_metric = metric_values.max()
 
-    min_state = states[int(metric_values.argmin())]
-    max_state = states[int(metric_values.argmax())]
+    min_idx = int(metric_values.argmin())
+    max_idx = int(metric_values.argmax())
+
+    min_state = states[min_idx]
+    max_state = states[max_idx]
 
     print(  # noqa: T201
         "Model Memory Estimation: Estimating memory from worst case of "
         f"largest and smallest system. Largest system has {max_state.n_atoms} atoms "
         f"and {max_state.n_systems} batches, and smallest system has "
         f"{min_state.n_atoms} atoms and {min_state.n_systems} batches.",
     )
+
     # determine_max_batch_size releases the caching allocator on its way out, so
     # the device is already handed back before the real optimization run starts.
     min_state_max_batches = determine_max_batch_size(min_state, model, **kwargs)
+    if max_idx == min_idx:
+        return min_state_max_batches * min_metric.item()
     max_state_max_batches = determine_max_batch_size(max_state, model, **kwargs)
 
     return min(
@@ -939,6 +945,26 @@ def load_states(self, states: Sequence[T] | Iterator[T] | T) -> float | None:
             This method resets the current state indices and completed state tracking,
             so any ongoing processing will be restarted when this method is called.
         """
+        # When all the states are known, estimate the limit over them as
+        # BinningAutoBatcher does, instead of probing the first state to pick the
+        # first batch. Lazy iterators keep that bootstrap probe.
+        if not self.max_memory_scaler and not isinstance(states, Iterator):
+            batched = (
+                states if isinstance(states, SimState) else ts.concatenate_states(states)
+            )
+            metrics = calculate_memory_scalers(
+                batched, self.memory_scales_with, self.cutoff
+            )
+            _max_memory_scaler = estimate_max_memory_scaler(
+                batched,
+                self.model,
+                metrics,
+                max_atoms=self.max_atoms_to_try,
+                scale_factor=self.memory_scaling_factor,
+                oom_error_message=self.oom_error_message,
+            )
+            self.max_memory_scaler = _max_memory_scaler * self.max_memory_padding
+
         self.states_iterator = iter(
             states if isinstance(states, Iterable) else states.split()
         )
@@ -1025,7 +1051,9 @@ def _get_first_batch(self) -> T:
         # we need to sample a state and use it to estimate the max metric
         # for the first batch
         first_state = next(self.states_iterator)
-        first_metric = calculate_memory_scalers(first_state, self.memory_scales_with)[0]
+        first_metric = calculate_memory_scalers(
+            first_state, self.memory_scales_with, self.cutoff
+        )[0]
         self.current_scalers += [first_metric]
         self.current_idx += [0]
         self.iteration_count.append(0)  # Initialize attempt counter for first state

torch_sim/runners.py

@@ -534,7 +534,7 @@ def _chunked_apply[T: SimState](
     model: ModelInterface,
     init_kwargs: Any,
     **batcher_kwargs: Any,
-) -> T:
+) -> tuple[T, float | None]:
     """Apply a function to a state in chunks.
 
     This prevents us from running out of memory when applying a function to a large
@@ -548,7 +548,9 @@ def _chunked_apply[T: SimState](
         **batcher_kwargs: Additional keyword arguments for the autobatcher
 
     Returns:
-        A state with the function applied
+        tuple[T, float | None]: A state with the function applied, and the
+            max_memory_scaler of the autobatcher, given or estimated, so that
+            callers can reuse it instead of estimating it again.
     """
     autobatcher = BinningAutoBatcher(model=model, **batcher_kwargs)
     autobatcher.load_states(states)
@@ -566,7 +568,7 @@ def _chunked_apply[T: SimState](
     ]
 
     ordered_states = autobatcher.restore_original_order(initialized_states)
-    return ts.concatenate_states(ordered_states)
+    return ts.concatenate_states(ordered_states), autobatcher.max_memory_scaler
 
 
 def generate_force_convergence_fn[T: MDState | FireState](
@@ -725,13 +727,17 @@ def optimize[T: OptimState](  # noqa: C901, PLR0915
             "optimize: initializing optimizer state via chunked apply "
             "(BinningAutoBatcher); InFlightAutoBatcher will be used for optimization"
         )
-        state = _chunked_apply(
+        # The in-flight batcher reuses the limit estimated for the init, instead of
+        # probing the same systems again
+        state, autobatcher.max_memory_scaler = _chunked_apply(
             init_fn,
             initial_state,
             model,
             init_kwargs=dict(**init_kwargs or {}),
             max_memory_scaler=autobatcher.max_memory_scaler,
             memory_scales_with=autobatcher.memory_scales_with,
+            cutoff=autobatcher.cutoff,
+            memory_scaling_factor=autobatcher.memory_scaling_factor,
             max_atoms_to_try=autobatcher.max_atoms_to_try,
             oom_error_message=autobatcher.oom_error_message,
             max_memory_padding=autobatcher.max_memory_padding,