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

[niklashoelter]

Problem

Some models return graph-carrying outputs — notably UMA, whose energy keeps
requires_grad=True even though its forces are already detached. Any state that is
retained after a forward pass therefore pins that pass's entire forward autograd graph.
torch_sim.optimize retains such states in two places, and both leak GPU memory until the
device fills. Because the leak is allocated memory (not cached), torch.cuda.empty_cache()
cannot reclaim it, and cache-clearing / allocator-config workarounds don't help.

1. InFlightAutoBatcher swap loop

optimize accumulates every converged system (via InFlightAutoBatcher.next_batch → the
caller's completed-states list) for the entire run so it can restore the original order at
the end. A state popped out of the running batch preserves grad_fn, so each retained
completed state pins the full forward graph of the swap it finished in — one live graph per
finished system. torch.cuda.memory_allocated climbs monotonically (tens of MB per
structure) until OOM.

Reproduced with uma-s-1p1 over a ~4300-structure molecular library (FIRE, D3 on): live
memory climbs from ~13 GB to 24 GB over ~356 swaps and OOMs; with this fix it stays flat
at ~13 GB across the entire run (validated to full completion, twice).

2. _chunked_apply optimizer-state init pass

Before the swap loop, optimize initializes the optimizer state over all systems via
_chunked_apply, which runs the init fn per bin through a BinningAutoBatcher and
accumulates every result in a GPU-resident list. The same grad-carrying-energy states are
held here too, so live memory grows by ~one graph per bin. Unlike the swap loop,
BinningAutoBatcher has no runtime OOM recovery, so a bin that tips over the device is
fatal — and this bites larger models before the swap loop is ever reached.

Reproduced with uma-m-1p1 over the same ~4300-structure library (FIRE, D3 on): OOMs
~halfway through the init pass (bin 73/154), completely unaffected by max_memory_padding
(0.95 → bin 73, 0.90 → bin 78 — padding only shifts the crash a few bins, as expected for
an accumulation leak). With this fix the init pass turns from a monotonic climb into a
bounded sawtooth and the full run completes.

This is distinct from the LBFGS/BFGS history-buffer accumulation fixed in #559/#568 — it
affects the plain optimization loop (e.g. FIRE) via the model's own energy tensor.

Fix

Add a small detach_state_graph() helper in torch_sim/autobatching.py and call it on
states as they are retained, at both sites:

• InFlightAutoBatcher.next_batch, on completed states before they are accumulated;
• _chunked_apply (in torch_sim/runners.py), on each initialized state before it is
  accumulated.

Retained states are only ever read for their values (never differentiated), so dropping the
graph is safe and changes no numerical results. Because the helper is now used across
modules, it is promoted from a private _detach_state_graph to a public
torch_sim.detach_state_graph (exported in __all__).

Tests / validation

tests/test_autobatching.py: test_detach_state_graph_drops_grad_but_keeps_values —
asserts grad-carrying tensors are detached in place, non-grad tensors are left untouched,
and values are unchanged. Full autobatching suite green; ruff + ty clean.

Full 4272-structure benchmark now completes for all fire configs with the default
max_memory_padding=0.95: uma-s-1p1, uma-s-1p2, uma-m-1p1, orb-v3-direct,
orb-v3-conservative — all 4272/4272, no regression. uma-m-1p1 was fatal at every padding
without the _chunked_apply fix.

[CompRhys]

This should be fixed upstream in fairchem rather than here?

[niklashoelter]

This should be fixed upstream in fairchem rather than here?

Fair point, and worth raising upstream — though FWIW it's currently unfiled and unfixed in fairchem main: conservative UMA returns a graph-carrying energy (forces are detached, energy isn't).

However, I still think this is a valuable defence layer in torch-sim:

• The leak comes from retaining states across the run — torch-sim's own pattern. ModelInterface doesn't forbid grad-carrying outputs, so any such model leaks here, not just UMA.
• These sites are exactly where a state is finished and kept only for its values — so detaching is safe and near-free.
• Notably, fairchem's own torch-sim wrapper already detaches at the boundary (torchsim_interface.py does {k: v.detach() for k, v in results.items()}) — so fairchem itself treats this as torch-sim's side to handle. This just extends the same guard to the native adapter path, which bypasses that wrapper.

Happy to open the fairchem issue in parallel too.

[niklashoelter]

Opened the upstream issue: facebookresearch/fairchem#2072

[CompRhys]

• The leak comes from retaining states across the run — torch-sim's own pattern. ModelInterface doesn't forbid grad-carrying outputs, so any such model leaks here, not just UMA.

I did toy around with trying to add the 'retain_graph' arg to models to try and have an opinion on whether the graph should be detatched or not but didn't think that the pattern was clean enough to insist it's part of the interface, mostly because the use cases for fully differentiable simulation appear to be niche.