fix(autobatching): detach states to stop UMA autograd-graph leak (swap loop + chunked init)

tests/test_autobatching.py

@@ -13,6 +13,7 @@
     to_constant_volume_bins,
 )
 from torch_sim.models.lennard_jones import LennardJonesModel
+from torch_sim.state import detach_state_graph
 
 
 def test_exact_fit():
@@ -488,6 +489,32 @@ 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.
+    """
+    # 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(
     "oom_message",
     [

torch_sim/__init__.py

@@ -93,7 +93,12 @@
     optimize,
     static,
 )
-from torch_sim.state import SimState, concatenate_states, initialize_state
+from torch_sim.state import (
+    SimState,
+    concatenate_states,
+    detach_state_graph,
+    initialize_state,
+)
 from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter
 
 
@@ -145,6 +150,7 @@
     "calc_temperature",
     "concatenate_states",
     "constraints",
+    "detach_state_graph",
     "elastic",
     "fire_init",
     "fire_step",

torch_sim/autobatching.py

@@ -30,7 +30,7 @@
 import torch_sim as ts
 from torch_sim.models.interface import ModelInterface
 from torch_sim.neighbors import torchsim_nl
-from torch_sim.state import SimState
+from torch_sim.state import SimState, detach_state_graph
 from torch_sim.typing import MemoryScaling
 
 
@@ -1110,6 +1110,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)

torch_sim/runners.py

@@ -22,7 +22,7 @@
 from torch_sim.integrators.md import MDState
 from torch_sim.models.interface import ModelInterface
 from torch_sim.optimizers import OPTIM_REGISTRY, FireState, Optimizer, OptimState
-from torch_sim.state import _CANONICAL_MODEL_KEYS, SimState
+from torch_sim.state import _CANONICAL_MODEL_KEYS, SimState, detach_state_graph
 from torch_sim.trajectory import TrajectoryReporter
 from torch_sim.typing import StateLike
 from torch_sim.units import UnitSystem
@@ -472,10 +472,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)

torch_sim/state.py

@@ -1003,6 +1003,32 @@ def _state_to_device[T: SimState](  # noqa: C901
     return type(state)(**attrs)
 
 
+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 get_attrs_for_scope(
     state: SimState, scope: Literal["per-atom", "per-system", "global"]
 ) -> Generator[tuple[str, Any], None, None]: