diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index a05e0608..781e5816 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, 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 d4324166..d0db549d 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) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index cd46cc7c..a76a7468 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)