optimizer: size UMA batches by n_atoms, configurable via memory_scales_with

docs/config.md

@@ -55,6 +55,7 @@ Unknown fields are preserved. **Always use a config file for CLI and API calls.*
     "memory_scaling_factor": 1.75,
 
     "max_atoms_to_try": 100000,
+    "memory_scales_with": "auto",
     "steps_between_swaps": 1,
 
     "logging_level": "INFO"
@@ -175,6 +176,7 @@ technical:
   memory_scaling_factor: 1.75
 
   max_atoms_to_try: 100000
+  memory_scales_with: auto
   steps_between_swaps: 1
 
   logging_level: INFO
@@ -225,6 +227,7 @@ technical:
 | `max_memory_padding` | `0.95` | Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches |
 | `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. Must be > 1 |
 | `max_atoms_to_try` | `100000` | Upper bound on the autobatcher's calibration probe size (atoms) |
+| `memory_scales_with` | `"auto"` | Per-structure memory metric for batch sizing: `"n_atoms"`, `"n_edges"` or `"auto"`. `"auto"` uses `n_atoms` for UMA, whose large fixed memory per structure makes edge counts overcharge small molecules, and `n_edges` for ORB and SevenNet |
 | `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower |
 | `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"` |
 

src/gpuma/config.py

@@ -58,6 +58,7 @@
         "max_memory_padding": 0.95,
         "memory_scaling_factor": 1.75,
         "max_atoms_to_try": 100_000,
+        "memory_scales_with": "auto",
         "steps_between_swaps": 1,
         "logging_level": "INFO",
     },
@@ -93,6 +94,9 @@ def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any
     {"fire", "gradient_descent", "lbfgs", "bfgs"}
 )
 
+# n_atoms_x_density is excluded: it divides by a bounding box, meaningless for molecules.
+VALID_MEMORY_SCALES_WITH: frozenset[str] = frozenset({"auto", "n_atoms", "n_edges"})
+
 
 def resolve_model_type(config: Config | dict[str, Any]) -> str:
     """Normalize a ``model_type`` value to its canonical form.
@@ -404,6 +408,14 @@ def validate_config(config: Config) -> None:
     if dev != "cpu" and not dev.startswith("cuda"):
         raise ValueError(f"Device must be 'cpu', 'cuda' or 'cuda:N' (e.g. 'cuda:0'), got {dev!r}")
 
+    scales_with = str(getattr(tech, "memory_scales_with", "auto") or "auto").strip().lower()
+    if scales_with not in VALID_MEMORY_SCALES_WITH:
+        raise ValueError(
+            f"memory_scales_with must be one of {sorted(VALID_MEMORY_SCALES_WITH)}, "
+            f"got {scales_with!r}"
+        )
+    tech.memory_scales_with = scales_with
+
     # If CUDA is requested but not available, we don't fail here; the
     # runtime will transparently fall back to CPU via model helpers.
 

src/gpuma/optimizer.py

@@ -21,7 +21,13 @@
 from ase import Atoms
 from ase.optimize import BFGS, FIRE, LBFGS
 
-from .config import DEFAULT_CONFIG, Config, load_config_from_file, resolve_model_type
+from .config import (
+    DEFAULT_CONFIG,
+    VALID_MEMORY_SCALES_WITH,
+    Config,
+    load_config_from_file,
+    resolve_model_type,
+)
 from .utils.decorators import timed_block
 from .utils.logging_utils import log_optimization_summary
 from .models import _parse_device_string, load_calculator, load_torchsim_model
@@ -394,6 +400,20 @@ def _whole_job_memory_scaler(state: Any, model: Any, batcher: Any) -> float | No
     )))
 
 
+def _memory_scales_with(config: Config) -> str:
+    """Return the autobatcher memory metric: ``n_atoms`` for UMA, else ``n_edges``, unless set."""
+    choice = str(config.technical.get("memory_scales_with", "auto") or "auto").strip().lower()
+    if choice not in VALID_MEMORY_SCALES_WITH:
+        raise ValueError(
+            f"memory_scales_with must be one of {sorted(VALID_MEMORY_SCALES_WITH)}, "
+            f"got {choice!r}"
+        )
+    if choice != "auto":
+        return choice
+    # UMA's large per-system memory makes edge counts overcharge small molecules.
+    return "n_atoms" if resolve_model_type(config) == "fairchem" else "n_edges"
+
+
 def _optimize_batch(
     structures: list[Structure],
     config: Config,
@@ -457,16 +477,18 @@ def _optimize_batch(
 
     # 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.
+    memory_scales_with = _memory_scales_with(config)
     batcher = InFlightAutoBatcher(
         model,
-        memory_scales_with="n_edges",
+        memory_scales_with=memory_scales_with,
         memory_scaling_factor=memory_scaling_factor,
         max_memory_padding=max_memory_padding,
         max_atoms_to_try=max_atoms_to_try,
     )
     logger.debug(
-        "Autobatcher params: memory_scales_with=n_edges, "
+        "Autobatcher params: memory_scales_with=%s, "
         "max_memory_padding=%.2f, max_atoms_to_try=%d, steps_between_swaps=%d",
+        memory_scales_with,
         max_memory_padding,
         max_atoms_to_try,
         steps_between_swaps,

tests/optimization/test_memory_metric.py

@@ -0,0 +1,33 @@
+"""Tests for choosing the autobatcher memory metric per backend."""
+
+import pytest
+
+from gpuma.config import Config
+from gpuma.optimizer import _memory_scales_with
+
+
+@pytest.mark.parametrize(
+    ("model_type", "expected"),
+    [("fairchem", "n_atoms"), ("uma", "n_atoms"), ("orb", "n_edges"), ("sevennet", "n_edges")],
+)
+def test_auto_picks_metric_by_backend(model_type, expected):
+    cfg = Config({"model": {"model_type": model_type}})
+    assert _memory_scales_with(cfg) == expected
+
+
+@pytest.mark.parametrize("explicit", ["n_atoms", "n_edges"])
+def test_explicit_metric_overrides_backend(explicit):
+    cfg = Config({
+        "model": {"model_type": "fairchem"},
+        "technical": {"memory_scales_with": explicit},
+    })
+    assert _memory_scales_with(cfg) == explicit
+
+
+def test_value_set_after_construction_is_checked():
+    cfg = Config({"model": {"model_type": "orb"}})
+    cfg.technical.memory_scales_with = "N_ATOMS"
+    assert _memory_scales_with(cfg) == "n_atoms"
+    cfg.technical.memory_scales_with = "n_atoms_x_density"
+    with pytest.raises(ValueError, match="memory_scales_with"):
+        _memory_scales_with(cfg)

tests/test_config.py

@@ -244,6 +244,29 @@ def test_technical_to_dict():
     assert d["technical"]["logging_level"] == "INFO"
 
 
+# --- memory_scales_with validation ---
+
+
+def test_memory_scales_with_defaults_to_auto():
+    assert Config().technical.memory_scales_with == "auto"
+
+
+def test_memory_scales_with_is_normalized():
+    cfg = Config({"technical": {"memory_scales_with": " N_Atoms "}})
+    assert cfg.technical.memory_scales_with == "n_atoms"
+
+
+@pytest.mark.parametrize("bad", ["n_atoms_x_density", "edges"])
+def test_memory_scales_with_rejects_unknown(bad):
+    with pytest.raises(ValueError, match="memory_scales_with"):
+        Config({"technical": {"memory_scales_with": bad}})
+
+
+def test_memory_scales_with_empty_falls_back_to_auto():
+    cfg = Config({"technical": {"memory_scales_with": ""}})
+    assert cfg.technical.memory_scales_with == "auto"
+
+
 # --- batch_optimizer validation ---