ML Mlip Nvalchemi

ml-mlip-nvalchemi

Goal

Run optional GPU-parallel static predictions, relaxation and MD across multiple structures with NValchemi. This backend is experimental and disabled by default in AtomisticSkills 2.0.0. List and directory inputs normally run each structure through the model's native calculator and ASE/MatCalc. Installing nvalchemi-toolkit or selecting a GPU does not enable the batch engine.

[!WARNING] The locked toolkit 0.2.0 has confirmed batch-dynamics correctness defects. Same-size refills and variable-cell changes can omit periodic neighbors; AtomisticSkills guards this cache defect for 0.2.x, with measurable overhead. A separate defect can leave live energies inconsistent with frozen geometry after convergence; preserving the first converged snapshot reduces exposure but does not repair the upstream live batch. MatGL inflight remains disabled. Native NPT currently falls back to ASE because initial stress is missing. Static predictions do not reuse the defective dynamics cache, but this does not establish correctness for every model or structure. Historical speedups are workload-specific and do not certify the release's dynamics paths.

Explicit selection: pass use_nvalchemi=True on each batch call to static_calculation, relax_structure or run_md. MCP tools expose the same flag on predict_structure, relax_structure and run_md. The default is False; the choice does not persist into later calls. Runtime logs identify experimental use, and batch results report the actual backend.

When offering this option, explain these limitations. Do not enable it merely because batching could be faster. Validate a small representative case against the default path before a larger run, and check backend for any fallback. Single-structure calls continue using their normal calculator.

Background

NValchemi provides batched dynamics integrators (FIRE, NVT Nose-Hoover, NPT, etc.) and a BaseModelMixin interface. AtomisticSkills wraps each MLIP in a BaseModelMixin-compatible class:

With use_nvalchemi=True, src/utils/mlips/base.py dispatches as follows:

• static_calculation(list, use_nvalchemi=True) → _batch_static_nvalchemi() → single batched forward
• relax_structure(list, use_nvalchemi=True) → _batch_relax_nvalchemi() → batched FIRE
• run_md(list, use_nvalchemi=True) → _batch_md_nvalchemi() → batched NVT/NVE/NPT integrator

Inflight batching (relaxation)

After explicit opt-in, relaxation can choose fixed-batch or inflight execution:

_batch_relax()
 ├─ use_nvalchemi=True AND nvalchemi available AND model loads?
 │    YES → _batch_relax_nvalchemi()
 │              └─ sum(atoms) > max_batch_atoms AND model._nvalchemi_supports_inflight?
 │                   YES → _batch_relax_nvalchemi_inflight()   ← rolling GPU window
 │                   NO  → fixed-batch NValchemi               ← all structures at once
 │    NO  → _batch_relax_sequential()                          ← plain ASE FIRE, one by one

Note: All MatGL wrappers (TensorNetWrapper, M3GNetWrapper, CHGNetWrapper) set _nvalchemi_supports_inflight=False and use fixed-batch NValchemi regardless of structure count, because after graduation, energies are wrong (TensorNet Cu −83.70 vs −86.57 eV fixed-batch; CHGNet 0.26 eV; M3GNet 28 meV), MACE inflight is available after opt-in with the cache guard.

Inflight batching keeps only max_batch_atoms atoms on the GPU at once. As each structure converges or exhausts its step budget it is evicted and a new one is loaded. This is necessary when the full set of structures would exceed GPU memory.

max_batch_atoms — how it is set:

How to tell which backend ran

Every batch result dict carries a "backend" key:

result = wrapper.relax_structure(structures, fmax=0.05, steps=500, use_nvalchemi=True)
print(result["backend"])   # "nvalchemi_inflight" / "nvalchemi" / "sequential"

Logger messages (written to stderr / MCP server log) also signal transitions:

• "Total atoms (N) exceeds batch limit (M); switching to inflight batching."
• "NValchemi inflight relax: N structures, live batch ≤M atoms, ≤S steps/structure."

Variable-cell relaxation validation

Variable-cell relaxation uses upstream FIRE2VariableCell with dt=0.05, tmax=0.5, delaystep=5, and maxstep=0.2. Its cell force is already normalized by system size; AtomisticSkills no longer applies an additional atom-count scaling. Standard-form cell preparation is retained.

All three relaxation modes report success only on convergence, not_converged at the step limit, and failed for execution errors. Results include converged and per-structure steps, with separate aggregate counts. For variable-cell runs, convergence includes the per-atom virial row norm as well as atomic forces. This is the small-strain counterpart of ASE's FrechetCellFilter criterion, not exact equality at finite strain.

Neighbor-cache handling in 2.0.0: NValchemi 0.2.x can omit neighbors after same-size refills or gradual cell changes. AtomisticSkills uses a version-gated guard for every neighbor hook in fixed relaxation, inflight relaxation and batch MD. Changes to cell, PBC or atom partition trigger a complete allocation refresh; ordinary fixed-cell steps retain their buffers. Installed packages are unchanged. MACE inflight stays available after opt-in, and FairChem builds its own graph without this hook. See the exposure and timing report.

The separate upstream inactive-output defect can still return live energies inconsistent with frozen coordinates after convergence. MatGL inflight remains disabled pending validation of that repair. Trajectory extraction preserves the first converged snapshot; turning extraction off exposes the live-batch output limitation. A corrected upstream release is still needed to retire the compatibility guard.

The former variable-cell speedup tables used force-only convergence and are withdrawn. New speedups remain pending a supported upstream release containing the neighbor-cache and inactive-output repairs. A source-checkout validation run must not be presented as performance of the committed dependency locks. Static-inference and MD tables below measure different operations.

Molecular Dynamics (MD) Benchmark: Sequential vs. Batched (20 structures, 100 steps)

Speedup comparison for a 100-step MD simulation under the nvt_nose_hoover ensemble at 300 K on 20 strained Cu FCC structures, each expanded to a fixed 108-atom cubic supercell ($\ge 10\text{ \AA}$ sides). Sequential = NValchemi disabled, structures run one at a time; Batched = all 20 driven through NValchemi integrators in a single GPU batch. Best-of-2 wall time, measured serially (one environment at a time to avoid GPU contention).

MACE-OMAT-0-small (mlip)

• Sequential MD: 54.48 s
• Batched MD (NValchemi): 11.12 s (4.90x speedup)

TensorNet-PES-MatPES-PBE-2025.2 (mlip)

• Sequential MD: baseline
• Batched MD (NValchemi): 1.8× speedup over sequential (16 structures × 200 steps on GB10; batch NVE energy drift equals sequential, max 0.02 meV/atom)

FairChem uma-s-1p2 (fairchem)

• Sequential MD: 339.74 s
• Batched MD: disabled — routed to sequential (see note below; measured ~0.64x, i.e. slower, before being disabled)

When does batched MD help? Batched MD yields substantial speedups for launch-latency-bound models at small system sizes (e.g. MACE at 4.90x, TensorNet at 1.8x). For heavy models like FairChem uma-s-1p2 whose single-system path is already compute-bound, batching provides no speedup and MD is routed to sequential. The wrappers still accept a list of structures (and batch static/relax remain available); only FairChem's MD path is gated. Measured on NVIDIA GB10 (aarch64, CUDA 13, Warp 1.14).

FairChem batched MD disabled (_nvalchemi_supports_batch_md = False): uma-s-1p2's forward scales superlinearly per atom — ≈1.57 ms/atom at batch=1 (108 atoms) rising to ≈2.43 ms/atom at batch=20 (2160 atoms), 1.55x worse — so a single large batched step is slower than running the structures one at a time through the model's optimized single-system path (batched 0.64x). Two facts pin this down: (1) the cost is intrinsic to the eSCN/MoE forward, not the neighbor list — correcting the wrapper cutoff (12 A → the model's true 6 A) cut adapt_input edges from 530 to 78 per atom but left the per-step time unchanged at ~5.25 s; (2) uma-s-1p2 runs with external_graph_gen=False, so it rebuilds its own graph internally and ignores the edges adapt_input provides (energies are identical for any cutoff we pass, including a 0-edge 2 A list). Batched MD is therefore correct (energies match sequential to 0.00 meV/atom) but never a speedup, so run_md falls back to sequential.

TensorNet batched MD enabled with stream fix: TensorNet batch MD is re-enabled. The NeighborListHook "race" it was previously disabled for was a CUDA stream mismatch, now fixed by warp_on_torch_stream in src/utils/mlips/nvalchemi/nvalchemi_utils.py. Batch NVE energy drift equals sequential (max 0.02 meV/atom), and batch MD is 1.8× faster than sequential for 16 structures × 200 steps on GB10.