optimizer: measure n_edges scalers in memory-bounded chunks

docs/config.md

@@ -230,6 +230,8 @@ technical:
 
 > You can also control the GPU with the `CUDA_VISIBLE_DEVICES` environment variable.
 
+Before calibration, torch-sim measures every system's neighbor-list edge count so it can bin-pack batches. torch-sim runs that measurement over the whole submitted batch in one pass, which holds roughly ten kilobytes per atom and overflows a 44 GB card at a few million atoms. GPUMA routes the measurement through chunks sized from the device's free memory, so batch size is bounded by the optimization itself, not by this bookkeeping. The per-system numbers are identical either way, and there is no knob to set.
+
 ---
 
 ## Available Models

src/gpuma/autobatching.py

@@ -0,0 +1,121 @@
+"""Chunked ``n_edges`` memory scalers for torch-sim's autobatchers.
+
+torch-sim sizes every system for bin-packing by running the neighbor list over
+the whole batched state in one pass. That pass holds fixed-width buffers of
+roughly ten kilobytes per atom, so a few million atoms fill a 44 GB card before
+the first optimization step. Per-system edge counts do not depend on other
+systems, so the same numbers come out of passes over atom-bounded chunks.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch_sim.autobatching as _autobatching
+from torch_sim.neighbors import torchsim_nl
+from torch_sim.state import SimState
+
+logger = logging.getLogger(__name__)
+
+# Worst-case bytes per neighbor slot: int32 matrix, int32x3 shifts, bool mask,
+# and the int64/int32 pair-list temporaries built from them.
+_BYTES_PER_NEIGHBOR_SLOT = 53
+_DEVICE_FRACTION = 0.5
+_FALLBACK_MAX_NEIGHBORS = 192
+
+_torch_sim_n_edges_scalers = _autobatching._n_edges_scalers
+
+
+def _max_neighbors(cutoff: float) -> int:
+    try:
+        from nvalchemiops.neighbors.neighbor_utils import estimate_max_neighbors
+    except ImportError:
+        return _FALLBACK_MAX_NEIGHBORS
+    return int(estimate_max_neighbors(cutoff))
+
+
+def chunk_atom_budget(device: torch.device, cutoff: float) -> int | None:
+    """Atoms one neighbor-list pass may hold on ``device``.
+
+    Parameters
+    ----------
+    device : torch.device
+        Device the state lives on; only CUDA devices are budgeted.
+    cutoff : float
+        Neighbor-list cutoff in Å, which sets the matrix width.
+
+    Returns
+    -------
+    int or None
+        Atom budget, or ``None`` when the device is not budgeted.
+    """
+    if device.type != "cuda":
+        return None
+    free, _total = torch.cuda.mem_get_info(device)
+    per_atom = _BYTES_PER_NEIGHBOR_SLOT * _max_neighbors(cutoff)
+    return max(1, int(free * _DEVICE_FRACTION / per_atom))
+
+
+def _chunk_bounds(n_atoms_per_system: list[int], max_atoms: int) -> list[tuple[int, int, int, int]]:
+    """Split systems into runs of at most ``max_atoms`` atoms; an oversized system runs alone."""
+    bounds = []
+    s0 = a0 = 0
+    atoms = 0
+    for s, n in enumerate(n_atoms_per_system):
+        if atoms and atoms + n > max_atoms:
+            bounds.append((s0, s, a0, a0 + atoms))
+            s0, a0, atoms = s, a0 + atoms, 0
+        atoms += n
+    bounds.append((s0, len(n_atoms_per_system), a0, a0 + atoms))
+    return bounds
+
+
+def n_edges_scalers_chunked(
+    state: SimState, cutoff: float, max_atoms: int | None = None
+) -> list[float]:
+    """Per-system neighbor-list edge counts, computed in atom-bounded chunks.
+
+    Parameters
+    ----------
+    state : SimState
+        Batched state whose atoms are grouped contiguously by system.
+    cutoff : float
+        Neighbor-list cutoff in Å.
+    max_atoms : int, optional
+        Atoms per pass; defaults to :func:`chunk_atom_budget` for the state's device.
+
+    Returns
+    -------
+    list of float
+        One edge count per system, identical to torch-sim's single-pass result.
+    """
+    if max_atoms is None:
+        max_atoms = chunk_atom_budget(state.device, cutoff)
+    if max_atoms is None or state.n_atoms <= max_atoms:
+        return _torch_sim_n_edges_scalers(state, cutoff)
+
+    bounds = _chunk_bounds(state.n_atoms_per_system.tolist(), max_atoms)
+    logger.debug(
+        "n_edges scalers: %d atoms in %d chunks of <= %d atoms",
+        state.n_atoms,
+        len(bounds),
+        max_atoms,
+    )
+    scalers: list[float] = []
+    for s0, s1, a0, a1 in bounds:
+        _, system_mapping, _ = torchsim_nl(
+            positions=state.positions[a0:a1],
+            cell=state.cell[s0:s1],
+            pbc=state.pbc,
+            cutoff=cutoff,
+            system_idx=state.system_idx[a0:a1] - s0,
+        )
+        scalers.extend(system_mapping.bincount(minlength=s1 - s0).float().tolist())
+    return scalers
+
+
+def install_chunked_edge_scalers() -> None:
+    """Route torch-sim's ``n_edges`` scalers through the chunked pass. Idempotent."""
+    if _autobatching._n_edges_scalers is not n_edges_scalers_chunked:
+        _autobatching._n_edges_scalers = n_edges_scalers_chunked

src/gpuma/optimizer.py

@@ -404,6 +404,7 @@ def _optimize_batch(
     from torch_sim.autobatching import InFlightAutoBatcher
     from torch_sim.typing import SystemExtras
 
+    from .autobatching import install_chunked_edge_scalers
     from .models import _device_for_torch
 
     logger.info("Starting batch optimization of %d structures", len(structures))
@@ -455,6 +456,10 @@ def _optimize_batch(
     max_atoms_to_try = int(config.technical.max_atoms_to_try)
     steps_between_swaps = int(config.technical.steps_between_swaps)
 
+    # torch-sim measures every system's edge count over the whole state at once;
+    # the chunked pass keeps that under the device's free memory.
+    install_chunked_edge_scalers()
+
     # max_atoms_to_try bounds the device probe, not the job: capping it at the
     # job size put the ceiling below the job's own largest structure.
     batcher = InFlightAutoBatcher(

tests/optimization/test_autobatching.py

@@ -0,0 +1,66 @@
+"""Tests for the chunked n_edges memory scalers. Runs on CPU."""
+
+import pytest
+import torch
+import torch_sim
+import torch_sim.autobatching as ts_autobatching
+from ase.build import molecule
+
+from gpuma.autobatching import (
+    _chunk_bounds,
+    _torch_sim_n_edges_scalers,
+    install_chunked_edge_scalers,
+    n_edges_scalers_chunked,
+)
+
+CUTOFF = 6.0
+
+
+@pytest.fixture
+def mixed_state():
+    """Seven molecules of unequal size, so chunk edges never fall on equal counts."""
+    names = ["H2O", "CH4", "C6H6", "C2H6", "NH3", "CH3CH2OH", "C60"]
+    atoms = [molecule(n) for n in names]
+    for a in atoms:
+        a.info["charge"] = 0
+        a.info["spin"] = 1
+    return torch_sim.io.atoms_to_state(atoms, device=torch.device("cpu"), dtype=torch.float64)
+
+
+class TestChunkBounds:
+    def test_single_chunk_when_under_budget(self):
+        assert _chunk_bounds([3, 4, 5], 100) == [(0, 3, 0, 12)]
+
+    def test_splits_at_system_boundaries(self):
+        assert _chunk_bounds([3, 4, 5, 2], 7) == [(0, 2, 0, 7), (2, 4, 7, 14)]
+        assert _chunk_bounds([3, 4, 5, 2], 6) == [
+            (0, 1, 0, 3),
+            (1, 2, 3, 7),
+            (2, 3, 7, 12),
+            (3, 4, 12, 14),
+        ]
+
+    def test_oversized_system_runs_alone(self):
+        assert _chunk_bounds([2, 50, 2], 10) == [(0, 1, 0, 2), (1, 2, 2, 52), (2, 3, 52, 54)]
+
+
+class TestChunkedScalers:
+    def test_matches_single_pass(self, mixed_state):
+        reference = _torch_sim_n_edges_scalers(mixed_state, CUTOFF)
+        assert len(reference) == mixed_state.n_systems
+        for max_atoms in (1, 10, 25, 1000):
+            chunked = n_edges_scalers_chunked(mixed_state, CUTOFF, max_atoms=max_atoms)
+            assert chunked == reference, f"max_atoms={max_atoms}"
+
+    def test_no_budget_on_cpu_uses_single_pass(self, mixed_state):
+        reference = _torch_sim_n_edges_scalers(mixed_state, CUTOFF)
+        assert n_edges_scalers_chunked(mixed_state, CUTOFF) == reference
+
+
+class TestInstall:
+    def test_routes_torch_sim_scalers(self, mixed_state):
+        install_chunked_edge_scalers()
+        install_chunked_edge_scalers()
+        assert ts_autobatching._n_edges_scalers is n_edges_scalers_chunked
+        via_torch_sim = ts_autobatching.calculate_memory_scalers(mixed_state, "n_edges", CUTOFF)
+        assert via_torch_sim == _torch_sim_n_edges_scalers(mixed_state, CUTOFF)