Pin aimnet 0.2.6 and drop the AIMNet2 L-BFGS warning

AGENTS.md

@@ -43,8 +43,4 @@ Guidance for AI coding agents working on the GPUMA repository. (Agents that
 - The `--json` summary is a contract for agents. Only add keys; bump
   `SUMMARY_SCHEMA_VERSION` in `cli_report.py` for anything else.
 - No backward-compatibility shims unless asked for.
-- The AIMNet2 L-BFGS warning (`_warn_batch_optimizer_memory` in `optimizer.py`)
-  is temporary: drop it, and its docs and skill lines, once an `aimnet` release
-  contains isayevlab/aimnetcentral#150. Bump the pin, then rerun a large mixed
-  L-BFGS job to confirm.
 - Don't add AI co-author trailers or "generated with" footers to commits or PRs.

README.md

@@ -196,8 +196,6 @@ When a run is started from SMILES, an RDKit/MMFF force field (via nvMolKit on GP
 - AIMNet2: models are downloaded automatically on first use. They cover 14
   elements, and only `aimnet2-nse` handles open-shell structures; other inputs
   are rejected before the model loads (see [docs/config.md](docs/config.md)).
-  Use the `fire` batch optimizer: `lbfgs` can run out of GPU memory on large
-  AIMNet2 jobs.
 - BFGS: `"batch_optimizer": "bfgs"` keeps a dense Hessian per structure and can
   run out of GPU memory on large batch jobs with any backend; prefer `lbfgs` or `fire`.
 - D3 dispersion correction: off by default. Set `"d3_correction": true` to

docs/config.md

@@ -200,7 +200,7 @@ technical:
 | Parameter | Default | Description |
 |---|---|---|
 | `batch_optimization_mode` | `"batch"` | `"sequential"` (ASE per structure) or `"batch"` (torch-sim GPU-accelerated). If a batch fails, its structures are retried one at a time, so one bad input fails only itself |
-| `batch_optimizer` | `"fire"` | Optimizer for both single and batch modes: `"fire"`, `"gradient_descent"`, `"lbfgs"`, or `"bfgs"`. In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS (`"gradient_descent"` falls back to FIRE). Defaults to `"fire"` if unset; an unknown value is an error. In batch mode, `"bfgs"` keeps a dense Hessian per structure that the memory estimate leaves out and can run out of GPU memory on large jobs, and `"lbfgs"` can do the same with AIMNet2; gpuma warns for both |
+| `batch_optimizer` | `"fire"` | Optimizer for both single and batch modes: `"fire"`, `"gradient_descent"`, `"lbfgs"`, or `"bfgs"`. In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS (`"gradient_descent"` falls back to FIRE). Defaults to `"fire"` if unset; an unknown value is an error. In batch mode, `"bfgs"` keeps a dense Hessian per structure that the memory estimate leaves out and can run out of GPU memory on large jobs; gpuma warns for it |
 | `charge` | `0` | Total charge of the system (inferred from SMILES, not overridden). Must be an integer |
 | `multiplicity` | `1` | Spin multiplicity of the system. A multiplicity the electron count can't have (e.g. a singlet radical) is an error |
 | `force_convergence_criterion` | `5e-2` | Force convergence threshold (eV/A). Used for both single and batch modes |
@@ -237,7 +237,7 @@ technical:
 | `max_memory_padding` | `0.95` | Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches. Must be between 0.05 and 0.95 |
 | `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. A job runs as one batch without probing only if its forward-pass peak times this factor fits in `max_memory_padding` of free memory. 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 for AIMNet2, whose Coulomb term grows with the square of the batch's atom count; `n_edges` for ORB and SevenNet |
+| `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 for AIMNet2, whose all-pairs Coulomb within each molecule edge counts don't track; `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. Must be at least 1 |
 | `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"`, `"CRITICAL"` |
 

docs/index.md

@@ -163,8 +163,6 @@ When a run is started from SMILES, an RDKit/MMFF force field (via nvMolKit on GP
 - AIMNet2: set `"model_type": "aimnet"`. Structures with other elements, an
   open-shell structure on a closed-shell model, or an H–H distance below 0.8 Å
   (for `aimnet2` and `aimnet2-2025`) are rejected before the model loads.
-  Use the `fire` batch optimizer: `lbfgs` can run out of GPU memory on large
-  AIMNet2 jobs.
 - BFGS: `"batch_optimizer": "bfgs"` keeps a dense Hessian per structure and can
   run out of GPU memory on large batch jobs with any backend; prefer `lbfgs` or `fire`.
 - D3 dispersion correction: off by default. Enable it with

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
   "warp-lang==1.17.0",
   "sevenn[torchsim]==0.13.0",
   "nvmolkit==0.5.1",
-  "aimnet[ase,torchsim]==0.2.5"
+  "aimnet[ase,torchsim]==0.2.6"
 ]
 
 [project.optional-dependencies]

src/gpuma/optimizer.py

@@ -489,7 +489,7 @@ def _memory_scales_with(config: Config) -> str:
     if choice != "auto":
         return choice
     # UMA's large per-system memory makes edge counts overcharge small molecules; AIMNet2's
-    # all-pairs Coulomb matrix grows with the square of the batch's atoms, whatever its edges.
+    # all-pairs Coulomb within each molecule isn't tracked by edge counts either.
     return "n_atoms" if resolve_model_type(config) in ("fairchem", "aimnet") else "n_edges"
 
 
@@ -501,12 +501,6 @@ def _warn_batch_optimizer_memory(config: Config) -> None:
             "BFGS keeps a dense Hessian per structure that the batch memory estimate leaves out, "
             "so large jobs can run out of GPU memory; use lbfgs or fire"
         )
-    elif name == "lbfgs" and resolve_model_type(config) == "aimnet":
-        # aimnet <= 0.2.5 sizes its Coulomb neighbor matrix by the batch's total atoms.
-        logger.warning(
-            "L-BFGS with AIMNet2 can run out of GPU memory on large jobs, as AIMNet2 batches "
-            "fill the GPU and leave no room for its history; use fire"
-        )
 
 
 @_reporting._instrument("mlip_optimize", "torchsim")

src/gpuma/skills/gpuma/SKILL.md

@@ -87,8 +87,7 @@ For any code other than 0, read
   Se, Br and I; use `aimnet2-nse` for radicals and triplets. For conformer
   searches on organic molecules the model choice changes the results little.
 - **Optimizer.** Set `optimization.batch_optimizer` to `lbfgs`. It is faster
-  than the default `fire` and at least as accurate. With AIMNet2 use `fire`:
-  `lbfgs` can run out of GPU memory on large jobs. Avoid `bfgs`, which is
+  than the default `fire` and at least as accurate. Avoid `bfgs`, which is
   slower than `fire` and can run out of GPU memory on large jobs.
 - **Dispersion.** Leave `model.d3_correction` off for UMA and the ORB OMol
   models. Their training data (ωB97M-D4rev or ωB97M-V) already includes dispersion.

src/gpuma/skills/gpuma/references/config.md

@@ -23,7 +23,7 @@ them for molecules.
 
 | Key | Default | Guidance |
 | --- | --- | --- |
-| `optimization.batch_optimizer` | `fire` | Use `lbfgs`, but `fire` with AIMNet2, where `lbfgs` can run out of GPU memory on large jobs. Never `bfgs`, which is slower than `fire` and can run out of GPU memory. |
+| `optimization.batch_optimizer` | `fire` | Use `lbfgs`. Never `bfgs`, which is slower than `fire` and can run out of GPU memory. |
 | `optimization.batch_optimization_mode` | `batch` | Keep `batch`. `sequential` optimizes one structure at a time and is only for debugging. |
 | `optimization.force_convergence_criterion` | `0.05` eV/Å | Tighten (e.g. `0.01`) only when the geometry feeds something force-sensitive, such as frequencies. |
 | `optimization.energy_convergence_criterion` | `null` | Leave unset; the force criterion is the reliable one. |

tests/optimization/test_memory_metric.py

@@ -45,13 +45,13 @@ def test_value_set_after_construction_is_checked():
         ("orb", "bfgs", "BFGS"),
         ("fairchem", "bfgs", "BFGS"),
         ("aimnet", "bfgs", "BFGS"),
-        ("aimnet", "lbfgs", "L-BFGS"),
+        ("aimnet", "lbfgs", None),
         ("aimnet", "fire", None),
         ("orb", "lbfgs", None),
     ],
 )
 def test_batch_optimizer_memory_warning(caplog, model_type, optimizer, warned):
-    """The memory estimate misses BFGS's Hessian, and AIMNet2 batches leave no room for L-BFGS."""
+    """The memory estimate misses BFGS's Hessian; other optimizers fit."""
     cfg = Config(
         {"model": {"model_type": model_type}, "optimization": {"batch_optimizer": optimizer}}
     )