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

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,

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",

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)

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)