ml-mlip-nvalchemi

Files (7)

• SKILL.md
• README.md · examples/cu_batch_static_accuracy/
• results_mace.json · examples/cu_batch_static_accuracy/
• results_matgl.json · examples/cu_batch_static_accuracy/
• benchmark_results.md · resources/
• run_md_benchmark.py · scripts/
• run_nvalchemi_benchmark.py · scripts/

SKILL.md (first lines)

# 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:

| MLIP | NValchemi wrapper | Location |
|------|------------------|----------|
| MACE | `nvalchemi.models.mace.MACEWrapper` | upstream (nvalchemi-toolkit) |
| MatGL TensorNet | `matgl.ext._alchmtk.TensorNetWrapper` | matgl package |
| MatGL CHGNet | `CHGNetWrapper` | `src/utils/mlips/nvalchemi/matgl_wrappers.py` |
| MatGL QET | `QETWrapper` | matgl package |
| FairChem UMA | `FairChemWrapper` | `src/utils/mlips/nvalchemi/fairchem_nv.py` |

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`, `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), 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 | Behaviour |
|-----|-----------|
| `max_batch_atoms=None` (default) | Auto-estimated: `free_VRAM × 0.5 / bytes_per_atom` using per-architecture calibration (FairChem 0.15 B/param/atom, MACE 0.5, TensorNet 1.5, CHGNet 3.0, fallback 5 MB/atom) |
| Explicit integer (e.g. `500`) | Forces inflight for almost any real dataset; recommended on shared GPUs |
| Very large integer | Forces fixed-batch (all structures in one GPU call) |

### How to tell which backend ran

Every batch result dict carries a `"backend"` key:

| `"backend"` value | Meaning |
|---|---|
| `"nvalchemi_inflight"` | Inflight rolling-window (large datasets / shared GPU) |
| `"nvalchemi"` | Fixed-batch NValchemi (entire set fits in one GPU pass) |
| `"sequential"` | Plain ASE FIRE, one structure at a time |

```python
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](../../docs/verification/nvalchemi-neighbor-cache.md).

The separate upstream inactive-output defect can still return live energies

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills