diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 5169a53a..2e99e775 100644 --- a/tests/test_autobatching.py +++ b/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", [ diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py index b24fab12..9a848f9f 100644 --- a/torch_sim/__init__.py +++ b/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", diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 07ee9525..fa8a9209 100644 --- a/torch_sim/autobatching.py +++ b/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) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index cd46cc7c..571016e7 100644 --- a/torch_sim/runners.py +++ b/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) diff --git a/torch_sim/state.py b/torch_sim/state.py index 2c3e28b8..768e0f9b 100644 --- a/torch_sim/state.py +++ b/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]: