From fc3552ec63742979388fab342dec3f0c4daa0608 Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Tue, 10 Nov 2026 11:36:32 +0100 Subject: [PATCH 1/6] Reduce number of memory probes --- torch_sim/autobatching.py | 30 ++++++++++++++++++++++++++++-- torch_sim/runners.py | 13 +++++++++---- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 530e685a..08c54d59 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -480,8 +480,11 @@ 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 " @@ -489,9 +492,12 @@ def estimate_max_memory_scaler( 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() ) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 4278dce1..0a3bb986 100644 --- a/torch_sim/runners.py +++ b/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,16 @@ 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, max_atoms_to_try=autobatcher.max_atoms_to_try, oom_error_message=autobatcher.oom_error_message, max_memory_padding=autobatcher.max_memory_padding, From ed9ab5583e8c92f0ff9ea8e53066b6f0fadcaad7 Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Tue, 10 Nov 2026 11:38:49 +0100 Subject: [PATCH 2/6] test the new number of probes per autobatcher call --- tests/test_autobatching.py | 160 +++++++++++++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index feeff7a3..499fe251 100644 --- a/tests/test_autobatching.py +++ b/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,162 @@ def test_binning_auto_batcher_with_lbfgs( all_finished_states.extend(batch.split()) assert len(all_finished_states) == len(lbfgs_states) + + +def test_estimate_max_memory_scaler_probes_identical_systems_once( + si_sim_state: ts.SimState, + lj_model: LennardJonesModel, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A single system, or identical ones, is probed once, not twice.""" + probed_systems = [] + + def mock_measure(state: ts.SimState, *_args: Any) -> float: + probed_systems.append(state.n_systems) + if state.n_atoms > 1_000: + raise RuntimeError("CUDA out of memory") + return 0.1 + + monkeypatch.setattr( + "torch_sim.autobatching.measure_model_memory_forward", mock_measure + ) + states = ts.concatenate_states([si_sim_state, si_sim_state]) + metrics = calculate_memory_scalers(states, "n_atoms") + + estimate_max_memory_scaler(states, lj_model, metrics) + + # Each probe starts with a single system + assert probed_systems.count(1) == 1 + + +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, + monkeypatch: pytest.MonkeyPatch, +) -> None: + probed_systems = [] + + def mock_measure(state: ts.SimState, *_args: Any) -> float: + probed_systems.append(state.n_systems) + if state.n_atoms > 1_000: + raise RuntimeError("CUDA out of memory") + return 0.1 + + monkeypatch.setattr( + "torch_sim.autobatching.measure_model_memory_forward", mock_measure + ) + states = ts.concatenate_states([si_sim_state, fe_supercell_sim_state]) + metrics = calculate_memory_scalers(states, "n_atoms") + + estimate_max_memory_scaler(states, lj_model, metrics) + + assert probed_systems.count(1) == 2 + + +@pytest.mark.parametrize( + ("make_states", "n_probes"), + [ + (list, 2), # all the states are known: the smallest and the largest + (ts.concatenate_states, 2), + (iter, 3), # lazy: a bootstrap probe of the first state, then those two + ], +) +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, + monkeypatch: pytest.MonkeyPatch, + make_states: Any, + n_probes: int, +) -> None: + probed_systems = [] + + def mock_measure(state: ts.SimState, *_args: Any) -> float: + probed_systems.append(state.n_systems) + if state.n_atoms > 1_000: + raise RuntimeError("CUDA out of memory") + return 0.1 + + monkeypatch.setattr( + "torch_sim.autobatching.measure_model_memory_forward", mock_measure + ) + batcher = InFlightAutoBatcher(model=lj_model, memory_scales_with="n_atoms") + + batcher.load_states(make_states([si_sim_state, fe_supercell_sim_state])) + + assert probed_systems.count(1) == n_probes + + +@pytest.mark.parametrize(("identical", "n_probes"), [(True, 1), (False, 2)]) +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, + monkeypatch: pytest.MonkeyPatch, + *, + identical: bool, + n_probes: int, +) -> None: + """The in-flight autobatcher takes the limit estimated for the init.""" + probed_systems = [] + + def mock_measure(state: ts.SimState, *_args: Any) -> float: + probed_systems.append(state.n_systems) + if state.n_atoms > 1_000: + raise RuntimeError("CUDA out of memory") + return 0.1 + + monkeypatch.setattr( + "torch_sim.autobatching.measure_model_memory_forward", mock_measure + ) + 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, + ) + + # Each probe starts with a single system: the smallest and the largest ones + assert probed_systems.count(1) == n_probes + + +def test_optimize_estimates_the_memory_at_the_cutoff_of_the_autobatcher( + ar_supercell_sim_state: ts.SimState, + lj_model: LennardJonesModel, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The limit estimated for the init counts edges at the autobatcher's cutoff.""" + + def mock_measure(state: ts.SimState, *_args: Any) -> float: + if state.n_atoms > 1_000: + raise RuntimeError("CUDA out of memory") + return 0.1 + + monkeypatch.setattr( + "torch_sim.autobatching.measure_model_memory_forward", mock_measure + ) + # 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 + ) + + 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( + estimate_max_memory_scaler(ar_supercell_sim_state, lj_model, metrics) + ) From 1cb99f2b1d44b6a8f23ddbe2169da95bf469223b Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Thu, 12 Nov 2026 15:52:41 +0100 Subject: [PATCH 3/6] Pass memory_scaling_factor to avoid silencing custom scaling factors --- torch_sim/runners.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 0a3bb986..005c7305 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -737,6 +737,7 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 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, From 16753c37618b811a4b54e796326f1636305f818a Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Thu, 12 Nov 2026 15:55:13 +0100 Subject: [PATCH 4/6] fix: lazy iterators were not using custom cutoffs to compute the memory metric --- torch_sim/autobatching.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 08c54d59..93fb57a3 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -1051,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 From c640c670ae099fc1e393dfd7f4ddf86f06f1a3ee Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Thu, 12 Nov 2026 15:56:29 +0100 Subject: [PATCH 5/6] Simplify the tests by using a shared fixture --- tests/test_autobatching.py | 125 +++++++++++++------------------------ 1 file changed, 42 insertions(+), 83 deletions(-) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 499fe251..1db9ebc3 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -1075,113 +1075,73 @@ def test_binning_auto_batcher_with_lbfgs( assert len(all_finished_states) == len(lbfgs_states) -def test_estimate_max_memory_scaler_probes_identical_systems_once( - si_sim_state: ts.SimState, - lj_model: LennardJonesModel, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A single system, or identical ones, is probed once, not twice.""" - probed_systems = [] - - def mock_measure(state: ts.SimState, *_args: Any) -> float: - probed_systems.append(state.n_systems) - if state.n_atoms > 1_000: - raise RuntimeError("CUDA out of memory") - return 0.1 +@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] = [] - monkeypatch.setattr( - "torch_sim.autobatching.measure_model_memory_forward", mock_measure - ) - states = ts.concatenate_states([si_sim_state, si_sim_state]) - metrics = calculate_memory_scalers(states, "n_atoms") + def recording_probe(state: ts.SimState, *_args: Any, **_kwargs: Any) -> int: + probed.append(state.n_atoms) + return 200 - estimate_max_memory_scaler(states, lj_model, metrics) - - # Each probe starts with a single system - assert probed_systems.count(1) == 1 + 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, - monkeypatch: pytest.MonkeyPatch, + probed_n_atoms: list[int], + *, + identical: bool, + expected: list[int], ) -> None: - probed_systems = [] - - def mock_measure(state: ts.SimState, *_args: Any) -> float: - probed_systems.append(state.n_systems) - if state.n_atoms > 1_000: - raise RuntimeError("CUDA out of memory") - return 0.1 - - monkeypatch.setattr( - "torch_sim.autobatching.measure_model_memory_forward", mock_measure - ) - states = ts.concatenate_states([si_sim_state, fe_supercell_sim_state]) + """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_systems.count(1) == 2 + assert probed_n_atoms == expected @pytest.mark.parametrize( - ("make_states", "n_probes"), + ("make_states", "expected"), [ - (list, 2), # all the states are known: the smallest and the largest - (ts.concatenate_states, 2), - (iter, 3), # lazy: a bootstrap probe of the first state, then those two + (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, - monkeypatch: pytest.MonkeyPatch, + probed_n_atoms: list[int], make_states: Any, - n_probes: int, + expected: list[int], ) -> None: - probed_systems = [] - - def mock_measure(state: ts.SimState, *_args: Any) -> float: - probed_systems.append(state.n_systems) - if state.n_atoms > 1_000: - raise RuntimeError("CUDA out of memory") - return 0.1 - - monkeypatch.setattr( - "torch_sim.autobatching.measure_model_memory_forward", mock_measure - ) batcher = InFlightAutoBatcher(model=lj_model, memory_scales_with="n_atoms") batcher.load_states(make_states([si_sim_state, fe_supercell_sim_state])) - assert probed_systems.count(1) == n_probes + assert probed_n_atoms == expected -@pytest.mark.parametrize(("identical", "n_probes"), [(True, 1), (False, 2)]) +@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, - monkeypatch: pytest.MonkeyPatch, + probed_n_atoms: list[int], *, identical: bool, - n_probes: int, + expected: list[int], ) -> None: """The in-flight autobatcher takes the limit estimated for the init.""" - probed_systems = [] - - def mock_measure(state: ts.SimState, *_args: Any) -> float: - probed_systems.append(state.n_systems) - if state.n_atoms > 1_000: - raise RuntimeError("CUDA out of memory") - return 0.1 - - monkeypatch.setattr( - "torch_sim.autobatching.measure_model_memory_forward", mock_measure - ) 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 @@ -1195,29 +1155,29 @@ def mock_measure(state: ts.SimState, *_args: Any) -> float: autobatcher=True, ) - # Each probe starts with a single system: the smallest and the largest ones - assert probed_systems.count(1) == n_probes + assert probed_n_atoms == expected -def test_optimize_estimates_the_memory_at_the_cutoff_of_the_autobatcher( +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 counts edges at the autobatcher's cutoff.""" + """The limit estimated for the init uses the autobatcher's cutoff and growth.""" + scale_factors: list[float] = [] - def mock_measure(state: ts.SimState, *_args: Any) -> float: - if state.n_atoms > 1_000: - raise RuntimeError("CUDA out of memory") - return 0.1 + def recording_probe(*_args: Any, scale_factor: float, **_kwargs: Any) -> int: + scale_factors.append(scale_factor) + return 200 - monkeypatch.setattr( - "torch_sim.autobatching.measure_model_memory_forward", mock_measure - ) + 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 + model=lj_model, + memory_scales_with="n_edges", + cutoff=cutoff, + memory_scaling_factor=2.0, ) ts.optimize( @@ -1229,6 +1189,5 @@ def mock_measure(state: ts.SimState, *_args: Any) -> float: ) metrics = calculate_memory_scalers(ar_supercell_sim_state, "n_edges", cutoff) - assert autobatcher.max_memory_scaler == pytest.approx( - estimate_max_memory_scaler(ar_supercell_sim_state, lj_model, metrics) - ) + assert autobatcher.max_memory_scaler == pytest.approx(200 * metrics[0]) + assert scale_factors == [2.0] From ec18c58b266f996bfbebea4167aa7aa353adb304 Mon Sep 17 00:00:00 2001 From: Gael Huynh Date: Thu, 12 Nov 2026 15:57:10 +0100 Subject: [PATCH 6/6] Add test to check if lazy iterators are correctly using the cutoff --- tests/test_autobatching.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 1db9ebc3..4ddb9fc3 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -1191,3 +1191,22 @@ def recording_probe(*_args: Any, scale_factor: float, **_kwargs: Any) -> int: 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])