From 24d7a13292cb8fec03ff4c9c729a719e57137135 Mon Sep 17 00:00:00 2001 From: niklashoelter Date: Sat, 4 Jul 2026 13:59:49 +0200 Subject: [PATCH 1/4] fix(autobatching): detach completed states to stop UMA forward-graph memory leak InFlightAutoBatcher accumulates every converged system in the caller's completed-states list for the whole run. A popped state preserves grad_fn, and some models return graph-carrying outputs - UMA's energy keeps requires_grad=True while its forces are already detached - so each completed state pins its swap's entire forward autograd graph. Across hundreds of in-flight swaps that is one retained graph per finished system (~tens of MB each), growing live (allocated, not cached) GPU memory monotonically until the device OOMs deep into a long optimization. empty_cache cannot reclaim it because the memory is allocated, not cached, which is why it surfaces as a late 'reserved but unallocated' fragmentation OOM. Detach any grad-carrying tensors on completed states as they leave next_batch, before they are accumulated. Completed states are only read for their values, never differentiated, so dropping the graph is safe. Live memory then stays flat across the whole run. --- torch_sim/autobatching.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index d4324166..182ea846 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -186,6 +186,32 @@ def measure_model_memory_forward(state: SimState, model: ModelInterface) -> floa return torch.cuda.max_memory_allocated() / 1024**3 # Convert to GB +def _detach_state_graph[T: SimState](state: T) -> T: + """Detach any autograd-graph-carrying tensors on a state, in place. + + Some models return graph-carrying outputs - notably UMA, whose ``energy`` + keeps ``requires_grad=True`` even though its forces are already detached. + When a converged system is popped out of the running batch and accumulated + for the remainder of the run, such a tensor pins that swap's *entire* + forward autograd graph, so live GPU memory grows by roughly one graph per + finished system until the device fills - independent of the batch size and + unreclaimable by ``empty_cache`` (it is allocated, not cached). Completed + states are only ever read for their values, never differentiated, so + dropping the graph is safe. + + Args: + state (SimState): State whose grad-carrying tensor attributes are + detached in place. + + Returns: + SimState: The same state instance, with grad-carrying tensors detached. + """ + for attr_name, attr_value in state.attributes.items(): + if torch.is_tensor(attr_value) and attr_value.requires_grad: + setattr(state, attr_name, attr_value.detach()) + return state + + def determine_max_batch_size( state: SimState, model: ModelInterface, @@ -1093,6 +1119,14 @@ def next_batch( # noqa: C901 completed_states = updated_state.pop(completed_idx) + # Drop any retained autograd graph before these states are accumulated + # for the rest of the run. A popped state preserves grad_fn, and models + # such as UMA return a graph-carrying energy, so without this each + # completed state would pin its swap's full forward graph - leaking GPU + # memory (one graph per finished system) until the device fills. + for completed_state in completed_states: + _detach_state_graph(completed_state) + # necessary to ensure states that finish at the same time are ordered properly completed_states.reverse() completed_idx.sort(reverse=True) From de80139689ef645c21141c88fb0d6776be9794ac Mon Sep 17 00:00:00 2001 From: niklashoelter Date: Sat, 4 Jul 2026 15:11:26 +0200 Subject: [PATCH 2/4] test(autobatching): cover _detach_state_graph grad-drop for the UMA leak fix --- tests/test_autobatching.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index a05e0608..3dd6c84f 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -488,6 +488,34 @@ def mock_measure(*_args: Any, **_kwargs: Any) -> float: assert max_size == 2 +def test_detach_state_graph_drops_grad_but_keeps_values( + si_sim_state: ts.SimState, +) -> None: + """`_detach_state_graph` strips grad graphs (the UMA leak) but preserves data. + + Models such as UMA return a graph-carrying ``energy`` (``requires_grad=True``) + while their forces are detached; accumulating those graph-carrying states for + the whole run is the memory leak. The helper must detach grad-carrying tensors + in place, leave non-grad tensors untouched, and not change any values. + """ + from torch_sim.autobatching import _detach_state_graph + + # Give one tensor attribute an autograd graph, as UMA's energy would carry. + grad_positions = (si_sim_state.positions.detach().clone().requires_grad_()) * 2 + values_before = grad_positions.detach().clone() + si_sim_state.positions = grad_positions + masses_before = si_sim_state.masses # a plain, non-grad tensor + assert si_sim_state.positions.requires_grad + + returned = _detach_state_graph(si_sim_state) + + assert returned is si_sim_state # detaches in place + assert not si_sim_state.positions.requires_grad + assert si_sim_state.positions.grad_fn is None + assert torch.allclose(si_sim_state.positions, values_before) # values unchanged + assert si_sim_state.masses is masses_before # non-grad tensors left as-is + + @pytest.mark.parametrize("scale_factor", [1.1, 1.4]) def test_determine_max_batch_size_small_scale_factor_no_infinite_loop( si_sim_state: ts.SimState, From 65c4af412b49acab915d348308078695130027be Mon Sep 17 00:00:00 2001 From: niklashoelter Date: Sun, 5 Jul 2026 11:25:06 +0200 Subject: [PATCH 3/4] fix(autobatching): detach states in _chunked_apply init pass The same graph-carrying-energy leak fixed for InFlightAutoBatcher also occurs in _chunked_apply, which optimize() uses to initialize optimizer state over all systems before the swap loop. It runs the init fn per bin through a BinningAutoBatcher and accumulates every result in a GPU-resident list. Models such as UMA return an energy that keeps requires_grad=True, so each accumulated state pins its bin's full forward autograd graph and live memory grows by ~one graph per bin until the device fills mid-pass. Unlike the swap loop, BinningAutoBatcher has no runtime OOM recovery, so this is fatal for larger models (reproduced with uma-m-1p1 on a 4272-structure run: OOM ~halfway through the init pass, unaffected by max_memory_padding). Detach each initialized state before it is accumulated, mirroring the InFlightAutoBatcher fix. The initialized states are only read for their values downstream, so dropping the graph is safe. Verified: full 4272-structure benchmark now completes for uma-m-1p1/uma-s/orb fire runs with no regression. --- torch_sim/runners.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index cd46cc7c..23c6bc49 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -17,7 +17,11 @@ from tqdm import tqdm import torch_sim as ts -from torch_sim.autobatching import BinningAutoBatcher, InFlightAutoBatcher +from torch_sim.autobatching import ( + BinningAutoBatcher, + InFlightAutoBatcher, + _detach_state_graph, +) from torch_sim.integrators import INTEGRATOR_REGISTRY, Integrator from torch_sim.integrators.md import MDState from torch_sim.models.interface import ModelInterface @@ -472,10 +476,17 @@ def _chunked_apply[T: SimState]( """ autobatcher = BinningAutoBatcher(model=model, **batcher_kwargs) autobatcher.load_states(states) - initialized_states = [] + # Each initialized bin is accumulated and held until every bin is done, then + # concatenated. Models such as UMA return a graph-carrying energy, so without + # detaching, every accumulated state would pin its bin's full forward autograd + # graph - live GPU memory then grows by roughly one graph per bin until the + # device fills mid-pass (fatal here: BinningAutoBatcher has no OOM recovery). + # The initialized states are only read for their values downstream, so + # dropping the graph is safe. Mirrors the InFlightAutoBatcher fix. initialized_states = [ - fn(model=model, state=system, **init_kwargs) for system, _indices in autobatcher + _detach_state_graph(fn(model=model, state=system, **init_kwargs)) + for system, _indices in autobatcher ] ordered_states = autobatcher.restore_original_order(initialized_states) From 677ccef5e1be8bb9232a70853fc3c51918e4ccb5 Mon Sep 17 00:00:00 2001 From: niklashoelter Date: Sun, 5 Jul 2026 11:29:19 +0200 Subject: [PATCH 4/4] refactor(autobatching): promote detach_state_graph to public API detach_state_graph is now used across modules (autobatching's InFlightAutoBatcher and runners' _chunked_apply) and is a small, reusable utility for dropping retained autograd graphs on completed states. Drop the leading underscore and export it from the top-level torch_sim namespace so the cross-module use is a supported import rather than reaching into a private helper. --- tests/test_autobatching.py | 6 +++--- torch_sim/__init__.py | 7 ++++++- torch_sim/autobatching.py | 4 ++-- torch_sim/runners.py | 4 ++-- 4 files changed, 13 insertions(+), 8 deletions(-) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 3dd6c84f..781e5816 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -491,14 +491,14 @@ def mock_measure(*_args: Any, **_kwargs: Any) -> float: def test_detach_state_graph_drops_grad_but_keeps_values( si_sim_state: ts.SimState, ) -> None: - """`_detach_state_graph` strips grad graphs (the UMA leak) but preserves data. + """`detach_state_graph` strips grad graphs (the UMA leak) but preserves data. Models such as UMA return a graph-carrying ``energy`` (``requires_grad=True``) while their forces are detached; accumulating those graph-carrying states for the whole run is the memory leak. The helper must detach grad-carrying tensors in place, leave non-grad tensors untouched, and not change any values. """ - from torch_sim.autobatching import _detach_state_graph + from torch_sim.autobatching import detach_state_graph # Give one tensor attribute an autograd graph, as UMA's energy would carry. grad_positions = (si_sim_state.positions.detach().clone().requires_grad_()) * 2 @@ -507,7 +507,7 @@ def test_detach_state_graph_drops_grad_but_keeps_values( masses_before = si_sim_state.masses # a plain, non-grad tensor assert si_sim_state.positions.requires_grad - returned = _detach_state_graph(si_sim_state) + returned = detach_state_graph(si_sim_state) assert returned is si_sim_state # detaches in place assert not si_sim_state.positions.requires_grad diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py index b24fab12..3be40775 100644 --- a/torch_sim/__init__.py +++ b/torch_sim/__init__.py @@ -22,7 +22,11 @@ transforms, units, ) -from torch_sim.autobatching import BinningAutoBatcher, InFlightAutoBatcher +from torch_sim.autobatching import ( + BinningAutoBatcher, + InFlightAutoBatcher, + detach_state_graph, +) from torch_sim.integrators import ( INTEGRATOR_REGISTRY, Integrator, @@ -145,6 +149,7 @@ "calc_temperature", "concatenate_states", "constraints", + "detach_state_graph", "elastic", "fire_init", "fire_step", diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 182ea846..d0db549d 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -186,7 +186,7 @@ def measure_model_memory_forward(state: SimState, model: ModelInterface) -> floa return torch.cuda.max_memory_allocated() / 1024**3 # Convert to GB -def _detach_state_graph[T: SimState](state: T) -> T: +def detach_state_graph[T: SimState](state: T) -> T: """Detach any autograd-graph-carrying tensors on a state, in place. Some models return graph-carrying outputs - notably UMA, whose ``energy`` @@ -1125,7 +1125,7 @@ def next_batch( # noqa: C901 # completed state would pin its swap's full forward graph - leaking GPU # memory (one graph per finished system) until the device fills. for completed_state in completed_states: - _detach_state_graph(completed_state) + detach_state_graph(completed_state) # necessary to ensure states that finish at the same time are ordered properly completed_states.reverse() diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 23c6bc49..a76a7468 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -20,7 +20,7 @@ from torch_sim.autobatching import ( BinningAutoBatcher, InFlightAutoBatcher, - _detach_state_graph, + detach_state_graph, ) from torch_sim.integrators import INTEGRATOR_REGISTRY, Integrator from torch_sim.integrators.md import MDState @@ -485,7 +485,7 @@ def _chunked_apply[T: SimState]( # The initialized states are only read for their values downstream, so # dropping the graph is safe. Mirrors the InFlightAutoBatcher fix. initialized_states = [ - _detach_state_graph(fn(model=model, state=system, **init_kwargs)) + detach_state_graph(fn(model=model, state=system, **init_kwargs)) for system, _indices in autobatcher ]