Compute n_edges memory scalers in atom-bounded chunks

[avnoordenne]

Summary

• Fix: _n_edges_scalers no longer builds one neighbor list over the whole state. It runs the neighbor list over contiguous runs of whole systems, at most N_EDGES_MAX_ATOMS_PER_PASS (250,000) atoms per pass. A system larger than that runs alone.

Problem

With memory_scales_with="n_edges", the autobatchers get each system's edge count by building one neighbor list over the whole submitted state. With the Alchemiops backend, that list uses fixed-width buffers (192 neighbor slots per atom at a 6 Å cutoff), about 4.4 KiB per atom. Memory therefore grows with the size of the job, not of a batch, and large jobs run out of GPU memory before the first optimization step.

Example: a molecular conformer job of 275,988 systems (10.74M atoms, ORB v3 + D3, LBFGS) fails in BinningAutoBatcher.load_states → calculate_memory_scalers → _n_edges_scalers with a 23 GiB allocation for the shift matrix. A linear fit puts the single pass at ~45 GiB, more than an A40.

Fix

A system's edge count does not depend on the other systems in the state, so splitting the pass along system boundaries returns exactly the same scalers. The batchers and calculate_memory_scalers are unchanged, and so is the public API. max_atoms_per_pass is a keyword argument on the private function, used by the tests.

250k atoms keeps the pass at about 1 GiB with the default 6 Å width. That is small enough to sit next to a loaded model on common GPUs and large enough that the number of passes doesn't matter: the 10.74M-atom job above takes 43 passes and 0.8 s in total. This pass only produces the scalers; batch size still comes from the memory probe.

Results (Quadro RTX 6000, 24 GB, the job above)

	single pass (main)	chunked (this PR)
full state, 10.74M atoms	OOM (tried to allocate 23.05 GiB)	1.3 GiB peak, 0.8 s
per-system scalers vs single pass	—	identical on the largest prefix that fits in one pass (63,666 systems, 3.0M atoms)

Checklist

• Doc strings have been added in the Google docstring format.
• Run ruff on your code. (ruff format clean; ruff check adds no new findings.)
• Tests have been added for any new functionality or bug fixes.
  • test_n_edges_scalers_chunked_matches_single_pass: chunked scalers equal the single pass for limits of 1, 16, 40 and 10,000 atoms on a mixed periodic state (32-, 8-, 8- and 32-atom systems).
  • test_n_edges_scalers_chunk_bounds: every pass holds whole systems, covers each system and atom exactly once, and exceeds the limit only for a lone system.
  • tests/test_autobatching.py (49 passed on GPU) and tests/test_runners.py (38 passed) run locally.

Notes

• The 250k limit is set in atoms, so the memory per pass grows with the neighbor-matrix width. Alchemiops' default width grows roughly with cutoff³ (~4.6 GiB per pass at 10 Å). If that matters for long-cutoff periodic models, a follow-up could derive the limit from a byte budget and estimate_max_neighbors(cutoff).
• For non-periodic states, the same width cap as in the D3 model (Size the D3 neighbor matrix to the largest system for non-periodic states #610) could also be applied here: max_neighbors = largest system in the pass − 1. For the job above, that cuts each pass from 192 to at most 94 neighbor slots per atom. We left it out because chunking alone bounds this pass at ~1 GiB, the pass runs once per load_states, and the cap does not change the scalers or batch sizes. In D3, the neighbor list is rebuilt on every forward pass, so the cap there frees memory for larger batches.

[avnoordenne]

The CI run on 732e8c5 was hit by a GitHub runner outage: 22 jobs were cancelled with "The job was not acquired by Runner of type hosted even after multiple attempts" and never started. Could someone with access use "Re-run failed jobs"? The full run on 4112b71 passed all test-core jobs and prek; the second commit only renames a parameter. The metatomic failures also occur on main since 33caef4 and are not required checks.