FastMD

FastMD is a single-device simulation framework for machine-learned interatomic potentials (MLIPs). It provides a common interface for ASE calculators, structure relaxation, ASE molecular dynamics, and GPU-resident molecular dynamics, with eager, model_graph, and whole_step execution modes.

Supported models

Checkpoint names link to official weight files or model pages.

Model	Name or entry point	Model implementation
MatRIS	matris_10m_oam, matris_10m_mp, or model_path=...	Bundled single-device implementation; weights loaded separately
MACE	mace_mp_0, mace_mpa_0, mace + model_path, or a .model path	Optional mace-torch package
SevenNet	sevennet_0, 7net-0, or sevennet + model_path	Optional sevenn package
ORB	orb_v3, with optional model_path	Optional orb-models package
CHGNet	chgnet_0_3_0	Pretrained weights supplied by the optional chgnet package
eSEN	esen + model_path	Legacy fairchem-core==1.10.0; isolated environment and local weights required
NequIP	nequip-oam-l, nequip-oam-xl, nequip-mp-l, or nequip, each with model_path	Official energy network; CPU/GPU eager, with opt-in validated padding and CUDA Graph capture

S01 benchmark results

The table contains 41 checkpoints across five model families, with a complete official/FastMD pair for each checkpoint.

Setting	Value
System	S01: CsSnI3_500K_Ivor_VASP, 135 atoms
Hardware	NVIDIA H100 80 GB HBM3; one GPU per trajectory
Ensemble	NVT Nose–Hoover chain at 500 K
Integration	0.25 fs × 80,000 steps = 20 ps
Thermostat	tdamp=25 fs, tchain=3, tloop=1
Recording	One frame every 10 steps; 8,001 frames including the initial state
Precision	FP32 IEEE
Initial state	Shared coordinates, cell, and momenta; seed 0; no center-of-mass momentum removal
Timing	GPU-synchronized wall time for the complete MD run, including trajectory output; model loading, compilation, capture, and post-run validation excluded

Speedup = official full-run time / FastMD full-run time. The displayed step times are the complete MD wall times divided by 80,000.

Model family	Completed checkpoint pairs	Recorded speedup range
MatRIS	2	5.13–5.52×
MACE	15	12.88–24.73×
SevenNet	8	14.87–24.47×
ORB	13	3.23–7.19×
CHGNet	3	18.45–21.05×

MatRIS

Checkpoint	Official (ms/step)	FastMD (ms/step)	Speedup
matris_10m_oam	154.5975	28.0180	5.5178×
matris_10m_mp	164.6329	32.0987	5.1290×

MACE

Checkpoint	Official (ms/step)	FastMD (ms/step)	Speedup
mace-small	31.8842	1.5295	20.8460×
mace-medium	43.0477	1.7405	24.7327×
mace-large	40.2385	2.1759	18.4924×
mace-small-0b	31.5469	2.0938	15.0665×
mace-medium-0b	39.0663	2.3195	16.8429×
mace-small-0b2	28.2565	2.1946	12.8752×
mace-medium-0b2	44.8428	2.5230	17.7738×
mace-large-0b2	38.3883	2.8898	13.2842×
mace-medium-0b3	52.3621	2.5283	20.7101×
mace-medium-mpa-0	43.0448	2.3935	17.9840×
mace-small-omat-0	41.8795	2.1667	19.3284×
mace-medium-omat-0	41.1945	2.4880	16.5574×
mace-mace-matpes-pbe-0	51.4085	2.5274	20.3405×
mace-mace-matpes-r2scan-0	36.5004	2.4031	15.1891×
mace-large-2024-01-07	36.9483	2.1359	17.2991×

SevenNet

Checkpoint	Official (ms/step)	FastMD (ms/step)	Speedup
7net-0	53.6565	2.5952	20.6754×
7net-0_22may2024	62.0417	2.5770	24.0756×
7net-l3i5	80.1659	3.2757	24.4730×
7net-omat	59.0921	3.8197	15.4705×
7net-nano-4.5	34.0357	1.4519	23.4430×
7net-nano-5.0	20.3308	1.3669	14.8732×
7net-nano-5.5	32.0208	1.4654	21.8514×
7net-nano-6.0	23.2114	1.3616	17.0477×

ORB

Checkpoint	Official (ms/step)	FastMD (ms/step)	Speedup
orb-v3-conservative-20-omat	34.3849	5.2446	6.5562×
orb-v3-conservative-inf-omat	36.7809	5.1188	7.1855×
orb-v3-direct-20-omat	8.8192	2.5000	3.5277×
orb-v3-direct-inf-omat	9.6175	2.7153	3.5420×
orb-v3-conservative-20-mpa	28.9682	4.7892	6.0487×
orb-v3-conservative-inf-mpa	33.8396	5.0916	6.6461×
orb-v3-direct-20-mpa	9.0028	2.7347	3.2921×
orb-v3-direct-inf-mpa	8.9940	2.7543	3.2654×
orb-v2	11.1864	3.4686	3.2251×
orb-d3-v2	11.5781	3.4350	3.3706×
orb-d3-sm-v2	8.6669	2.5232	3.4349×
orb-d3-xs-v2	6.7959	1.7014	3.9943×
orb-mptraj-only-v2	10.9139	3.2807	3.3267×

CHGNet

Checkpoint	Official (ms/step)	FastMD (ms/step)	Speedup
chgnet-0.3.0	53.3781	2.7588	19.3484×
chgnet-0.2.0	52.0884	2.4750	21.0455×
chgnet-r2scan	51.5925	2.7956	18.4548×

Installation

Use Python 3.10 or newer and run these commands from the repository root. Base and optional dependencies are managed in pyproject.toml:

python -m pip install -e .
# GPU neighbor-graph construction:
python -m pip install -e '.[gpu-graph]'
# CHGNet with GPU graph construction:
python -m pip install -e '.[gpu-graph,chgnet]'
# Install other model backends as needed:
python -m pip install -e '.[gpu-graph,mace]'
# Experimental M3GNet integration:
python -m pip install -e '.[m3gnet]'

The base installation requires Torch ≥ 2.6. The gpu-graph extra requires Torch ≥ 2.8, nvalchemi-toolkit-ops[torch]>=0.4.1,<0.5, and warp-lang>=1.13.0. Installing this extra can upgrade an older Torch installation; prepare a Torch build matching your CUDA environment first. These dependencies provide the GPU neighbor-list backend. The former graph/setup.py and fastmd_op build are no longer required. Use -e . and -e '.[m3gnet]' in place of the old requirements.txt and requirements-m3gnet.txt installation paths.

The MACE and SevenNet topology profiles use cuEquivariance. The upstream CUDA 13 dependency set is available through python -m pip install -e '.[gpu-graph,mace,cueq-cu13]'; substitute sevennet for mace as needed. Use optimization_profile="generic" without that acceleration environment. MACE and SevenNet require different e3nn versions and should use separate environments.

requirements-{mace,sevennet,orb,chgnet}.txt preserve upstream test-environment versions and hashes. They do not replace the base installation, contain Python/platform-specific wheel hashes, and are not complete dependency lockfiles. Use the extras above for normal installation. The upstream CI scripts mentioned in those records are not included here.

eSEN environment

Legacy eSEN uses an isolated environment because fairchem 1.10.0 conflicts with the base package's dependency constraints. The recorded environment uses Python 3.11, fairchem-core 1.10.0, Torch 2.8.0+cu128, NumPy 1.26.4, SciPy 1.14.1, e3nn 0.5.9, ASE 3.26.0, nvalchemi-toolkit-ops 0.2.0, and warp-lang 1.10.0, with PyG extensions matching Torch/CUDA. Torch 2.8 is an explicit override of fairchem's declared Torch 2.4 requirement. There is no esen extra: use the prepared environment and load this checkout through PYTHONPATH, for example PYTHONPATH="$PWD" python your_script.py from the repository root. Supply an eSEN-30M-OAM checkpoint explicitly; fairchem 2.x/UMA and direct-force heads are not supported by this adapter.

NequIP CPU usage

Install python -m pip install -e '.[nequip]' in a separate environment. NequIP 0.14 requires e3nn < 0.6. Download the official local .nequip.zip package from the model page. Current .ckpt training checkpoints are also accepted; legacy deployed and compiled artifacts are not. The adapter preserves the saved weight precision, double-precision geometry, per-type energy shifts, and ZBL potential. FastMD builds the radius graph and differentiates the upstream atomic energies for forces and stress.

import torch
from ase.io import read
from fastmd import InferenceConfig, MLIPCalculator

torch.set_num_threads(4)
atoms = read("S02.traj")
with MLIPCalculator(
    model="nequip-oam-l",
    model_path="NequIP-OAM-L-0.1.nequip.zip",
    device="cpu",
    task="efs",
    config=InferenceConfig(execution="eager", optimization_profile="generic"),
) as calculator:
    atoms.calc = calculator
    print(atoms.get_potential_energy(), atoms.get_forces(), atoms.get_stress())

NequIP defaults to eager execution with dynamic edge filtering. To opt into CUDA Graph capture, load an adapter with fastmd.models.nequip.load_nequip(..., fixed_shape=True) and pass it as model to the calculator or GPU MD driver. This checks the supported upstream energy pipeline and verifies zero energy/gradient contributions from far self-loop padding before permitting capture; unsupported pipelines raise an error. CPU and GPU calculators preserve double-precision geometry independently of the FP32 neural weights. GPU-resident MD retains that precision with ensemble="nvt_nhc", including whole-step force, energy, and stress buffers.

The loader also accepts kernel="cueq" to apply NequIP's official cuEquivariance modifier. This optional path requires cuequivariance, cuequivariance-torch and matching CUDA ops (tested with version 0.9.1); the default is kernel="upstream". Conversion verifies that every original saved state tensor is unchanged. The execution plan reports the actual backend, and CUDA loading rejects a fallback to naive kernels. Combine kernel="cueq" with fixed_shape=True to request both optimizations.

The S01–S05 GPU validation and full MD campaign is recorded in the benchmark manifest and live results. The separate S01 optimization experiment distinguishes short correctness/timing checks from validated full trajectories.

Unified calculator

from ase.build import bulk
from fastmd import MLIPCalculator, InferenceConfig

atoms = bulk("Si", "diamond", a=5.43, cubic=True)
atoms.calc = MLIPCalculator(
    model="chgnet_0_3_0", device="cuda", task="efs",
    config=InferenceConfig(execution="eager", optimization_profile="generic"),
)
print(atoms.get_potential_energy())  # eV, total energy
print(atoms.get_forces())           # eV/Å
print(atoms.get_stress())           # eV/ų, ASE Voigt order

For a CPU example, use model="matris_10m_oam" and device="cpu". The adapters and shared derivative layer handle model energy normalization, gradients, and stress units while preserving ASE conventions at the interface.

Molecular dynamics

python examples/simulate.py --structure input.extxyz \
  --model chgnet_0_3_0 --driver gpu --device cuda \
  --ensemble nvt_nhc --temperature 300 --timestep 0.5 \
  --execution auto --steps 100

You can also instantiate GPUResidentMolecularDynamics directly. Use it as a context manager to finish writing logs and trajectories on exit. MolecularDynamics provides ASE-driven MD, and StructOptimizer provides structure relaxation. execution="auto" selects a path using model capabilities, available memory, and applicable constraints; it does not guarantee whole-step capture for every model.

fastmd-explain --structure input.extxyz --model chgnet_0_3_0 --device cuda --json

Structure relaxation

The relaxation API ports the task-level behavior of hsy_relax onto the current MLIPCalculator. Relaxation and MD share model adapters, energy derivatives, operator optimizations, neighbor construction, and InferenceConfig. There is no separate relaxation model backend. The older branch's FastMDCalculator and CUDAGraphConfig are replaced by MLIPCalculator and InferenceConfig here.

from ase.io import read, write
from fastmd import InferenceConfig, RelaxationConfig, make_relaxation_calculator, relax

atoms = read("input.extxyz")
settings = RelaxationConfig(relax_cell=True, optimizer="FIRE", fmax=0.05, steps=500)
with make_relaxation_calculator(
    model="chgnet_0_3_0", device="cuda", relax_cell=settings.relax_cell,
    config=InferenceConfig(execution="auto", optimization_profile="topology"),
) as calculator:
    result = relax(atoms, calculator, settings)
    print(result.converged, result.steps, result.final_energy, result.max_filter_force)
    print(calculator.config_summary()["effective"])
    # Reuse this same calculator for further structures; each gets a fresh optimizer.
write("relaxed.extxyz", result.atoms)

Use relax_cell=False in both settings and the calculator factory for fixed-cell relaxation. The factory fixes the prediction task to ef for a fixed cell and efs for a variable cell, so ASE's repeated property requests reuse each geometry's evaluated results. The optimizer choices are FIRE, BFGS, LBFGS, BFGSLineSearch, LBFGSLineSearch, and MDMin. Pass optimizer parameters through optimizer_kwargs, and FrechetCellFilter parameters such as mask, hydrostatic_strain, or scalar_pressure through filter_kwargs. Pressure uses ASE units (eV/ų).

The relaxation factory and StructOptimizer default to auto_policy="latency". With execution="auto", a short task skips expensive online memory calibration when there is no compatible cached measurement. Set auto_policy="throughput" to permit calibration for a longer workload. Ordinary MLIPCalculator instances retain the throughput default, and a supplied calculator retains its policy. calculator.memory_probe_summary() reports measurements, cache hits, and skips. Measured bounds are reused within the process for the same adapter/model, configuration and exact geometry, subject to model mutation and capacity checks; they are not extrapolated to new cells or different structures.

An experimental auto_policy="latency_trial" permits a normal model-graph capture when no fitted memory bound is cached. It uses the first eager peak and current device headroom for admission, measures every actual capture (including later growth), and falls back to eager on budget rejection or capture failure. The estimate is a heuristic, not a guaranteed memory bound. Inspect calculator.capture_trial_summary() for capture time and outcomes. This policy skips multi-point calibration but still pays capture cost; it may be slower for short tasks and does not retain graphs across atom-count changes. The default remains latency until the workload has been measured.

relax copies the input atoms, retains constraints, and returns an independent final geometry with a SinglePointCalculator. Reusing or closing the inference calculator cannot invalidate previous results. It never closes a calculator supplied by the caller. steps=0 evaluates the initial structure. max_force is the maximum constrained atomic force; for variable-cell relaxation, convergence uses max_filter_force, including scaled cell forces, rather than atomic forces alone. fmax is therefore not a separate stress tolerance. Non-finite results raise an error. Timing includes initial evaluation and graph capture, but excludes model loading.

The existing StructOptimizer interface also accepts calculator=calculator, keeps its final_structure / trajectory return keys, and reports convergence and optimizer steps. For an active ASE MolecularDynamics instance, relax(md.atoms, md.calculator, settings) reuses its calculator and cached inference state; a cell relaxation requires its task to include stress. Use a shared calculator sequentially. To transfer a model from GPU-resident MD, pass model=md.adapter to the factory: the weights and adapter are reused, while the ASE calculator manages its own execution state.

Path	Optimization reuse	Requirements
relax / StructOptimizer	Common calculator, operator kernels, compilation where supported, eager or model_graph	Model's calculator support; stress and a full periodic cell for variable-cell relaxation
BatchedMLIPModel / relax_batch	Borrowed model adapters and installed kernels; one heterogeneous GPU batch per forward with TorchSim FIRE	Optional fastmd[torchsim]; small, fully periodic, unconstrained structures; eager tensor batch
MultiSystemRelaxation	Existing multi-system engine, slot batching, GPU FIRE and whole-step CUDA Graphs	CUDA, union/per-edge-shift adapter capabilities, full periodic cells, no ASE constraints

See TorchSim batch relaxation for batching examples, timing boundaries and supported model families.

execution="auto" uses the calculator's existing capability and memory decisions. Explicit model_graph requests must satisfy those checks. ASE optimization runs on the host, so whole_step is not an execution mode for the functional API. Variable cells continue through the shared calculator's cell/topology invalidation and graph-capacity logic. Graph capture is not guaranteed to improve short relaxations; inspect the effective configuration and measure the workload. The third-party CUDA/PBC requirements in Installation also apply to relaxation.

The CLI can process one structure or a multi-frame file with one loaded model:

fastmd-relax --structure input.extxyz --model chgnet_0_3_0 --device cuda \
  --execution auto --fmax 0.05 --steps 500 --output relaxed.extxyz

# Fixed cell; all frames, fresh optimizers, shared model and inference caches:
python -m fastmd.relax_cli --structure inputs.extxyz --index ':' --fixed-cell \
  --model mace --model-path /path/to/model.model --device cuda \
  --trajectory-dir relax-trajectories --output relaxed.extxyz --summary relaxation.json

python examples/relax.py exposes the same options. Final structures are written as .extxyz or .traj, and the JSON summary includes per-structure energies, convergence, step counts, timing, and the final calculator configuration. Exit code 0 means every structure converged; exit code 1 means at least one reached the step limit, with its results still written. This CLI processes frames sequentially. For supported GPU batched workloads:

from ase.io import read
from fastmd import InferenceConfig, MultiSystemRelaxation

with MultiSystemRelaxation(
    model="mace", model_path="/path/to/model.model",
    structures=read("inputs.extxyz", ":"), device="cuda",
    cell_filter="frechet", fmax=0.05, max_steps=500,
    config=InferenceConfig(execution="whole_step", optimization_profile="topology"),
) as runner:
    results = runner.run()

Use cell_filter=None for fixed cells. Batched execution retains its existing capability checks; it does not accept every calculator backend. The old branch's Discovery dataset workflow is separate from the relaxation API and is not included in this port.

Layout and extension points

fastmd/
  models/            Model adapters, capabilities, patches, and the MatRIS implementation
  engine/            Configuration, model protocols, derivatives, capacity, and caches
  execution/         Integrators, thermostats, CUDA Graphs, logging, and batched execution
  interfaces/        ASE calculators, MD, relaxation, and multi-system interfaces
  graph/             Neighbor lists, atom/line graphs, fixed-capacity layouts, and topology
  ops/               Shared PyTorch, Triton, and CUDA operators
  integrations/      Experimental M3GNet topology integration
  adapters/          Compatibility forwarding for legacy adapter imports
  applications/      Compatibility forwarding for legacy application imports
  config.py          Configuration compatibility entry point
  precision.py       Precision configuration and records

To add an MLIP, implement ModelAdapter (available through adapters/protocol.py) and expose an EnergyTap, SystemEnergyTap, or ModelOwnedDerivative interface. Pass the adapter as model= to reuse the simulation runtime. Declare fixed-shape capture, padding, and cutoff-continuity capabilities accurately so the runtime can select a supported execution path.