[avnoordenne]

torch-sim's autobatchers compute the edge count for every system so they can bin-pack batches, and they do it in one neighbor-list pass over the whole state. That pass allocates fixed-width buffers (192 neighbor slots per atom), about 4.4 KiB per atom, so large jobs run out of GPU memory before the first optimization step. A cheetah acids run with all conformers kept (276k conformers, 10.7M atoms) needs about 45 GiB for this pass alone and fails on both a 24 GB card and an A40.

This PR splits that pass into chunks along system boundaries. Each chunk holds as many atoms as fit in half the device's free memory, and the per-system results are identical to the single pass. There is no new config option.

Changes

• src/gpuma/autobatching.py (new): n_edges_scalers_chunked, which replaces torch-sim's _n_edges_scalers while _optimize_batch runs. A system too large for one chunk runs alone. On CPU, or when the state fits in one chunk, the call goes straight to torch-sim's single pass.
• src/gpuma/optimizer.py: installs the chunked pass before the batcher is built.
• docs/config.md: one paragraph on the behavior.
• tests/optimization/test_autobatching.py: chunk boundaries, oversized systems, the CPU fallback, results equal to the single pass, and the routing through torch-sim.

Testing

• CPU test suite passes (256 tests).
• GPU, Quadro RTX 6000 (24 GB), on the cheetah acids batch with all conformers kept (275,988 conformers, 10.74M atoms, orb_v3_direct_omol + D3, LBFGS):
  • without this PR: CUDA OOM 14 s into the run, requesting 23.04 GiB in batch_naive_neighbor_list
  • with this PR: the pass runs in 10 chunks (3.4 s, 5.5 GiB peak), the in-flight batcher starts, and the first optimization step completes
  • per-system edge counts match torch-sim's single pass exactly on the largest part of the batch that fits in one pass (63,666 systems, 3.0M atoms)

Notes for review

• This patches a private torch-sim function (torch_sim.autobatching._n_edges_scalers), so a torch-sim upgrade could break it. Tested against torch-sim 0.6.2. The cleaner long-term fix is to chunk this pass inside torch-sim itself. This is something I will work on. However, I do need this fix so I thought this was faster to release.
• The run above stopped after the first optimization step; a full end-to-end run on this dataset was not repeated with this exact commit.