From 427829c7061c34775345efd2f9422383becc8301 Mon Sep 17 00:00:00 2001 From: Ard van Noordenne Date: Mon, 5 Oct 2026 10:08:57 +0200 Subject: [PATCH] optimizer: size UMA batches by n_atoms, configurable via memory_scales_with Add technical.memory_scales_with (auto|n_atoms|n_edges). auto picks n_atoms for UMA, whose large fixed per-structure memory makes edge counts overcharge small molecules, and keeps n_edges for ORB and SevenNet. --- docs/config.md | 3 +++ src/gpuma/config.py | 12 +++++++++ src/gpuma/optimizer.py | 28 +++++++++++++++++--- tests/optimization/test_memory_metric.py | 33 ++++++++++++++++++++++++ tests/test_config.py | 23 +++++++++++++++++ 5 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 tests/optimization/test_memory_metric.py diff --git a/docs/config.md b/docs/config.md index b3081fb..f8f0d1f 100644 --- a/docs/config.md +++ b/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"` | diff --git a/src/gpuma/config.py b/src/gpuma/config.py index 9ae1a44..817b70d 100644 --- a/src/gpuma/config.py +++ b/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. diff --git a/src/gpuma/optimizer.py b/src/gpuma/optimizer.py index 6758437..da7b763 100644 --- a/src/gpuma/optimizer.py +++ b/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, diff --git a/tests/optimization/test_memory_metric.py b/tests/optimization/test_memory_metric.py new file mode 100644 index 0000000..22d0b62 --- /dev/null +++ b/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) diff --git a/tests/test_config.py b/tests/test_config.py index b774c39..3828e45 100644 --- a/tests/test_config.py +++ b/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 ---