Molecular dynamics with nvalchemi-toolkit

11.7. Molecular dynamics with nvalchemi-toolkit#

nvalchemi-toolkit is NVIDIA’s GPU-accelerated framework for batched molecular dynamics and structure optimization with machine-learning interatomic potentials. DeePMD-kit ships a thin adapter, DPA4Wrapper, that exposes a trained DPA-4 / SeZM model to any nvalchemi dynamics engine (NVE, NVT, NPT, FIRE, …). The model itself runs unmodified; the wrapper only translates between the nvalchemi graph batch and the model’s internal interface.

Note

Supported backends: PyTorch PyTorch, for DPA-4 / SeZM energy models. nvalchemi-toolkit is an optional dependency and must be installed separately. A CUDA device is recommended, since nvalchemi’s neighbour-list and integrator kernels are GPU-accelerated.

11.7.1. Installation#

Install the optional toolkit through the DeePMD-kit extra:

pip install deepmd-kit[nvalchemi]

This pulls in the nvalchemi-toolkit package; equivalently, install it directly with pip install nvalchemi-toolkit. Refer to the nvalchemi-toolkit documentation for the build that matches your Python, platform, and CUDA environment.

The DeePMD-kit adapter lives in deepmd.pt.nvalchemi; importing it without nvalchemi-toolkit present raises an actionable error.

11.7.2. Loading a model#

A trained DeePMD-kit checkpoint (.pt) is loaded and wrapped in one call:

import torch
from deepmd.pt.nvalchemi import DPA4Wrapper

model = DPA4Wrapper.from_checkpoint(
    "model.ckpt.pt",
    device=torch.device("cuda"),
    compute_stress=True,  # enable the Cauchy stress output (needs a periodic cell)
)

For a multi-task checkpoint, pass the branch name with head="...". An already-instantiated model can be wrapped directly with DPA4Wrapper(model).

from_checkpoint also accepts a frozen .pt2 (AOTInductor) package produced by dp --pt freeze; it is loaded through its precompiled callable (float64 I/O, and device-locked to the host it was frozen on).

11.7.2.1. Performance#

The model runs eagerly by default. To use DeePMD-kit’s compiled inference path, set the environment variables before loading the model:

  • DP_COMPILE_INFER=1 — compile the model. The first call pays a one-time compile cost (~1–2 min); subsequent steps are roughly 3x faster, and the dynamic-shape graph handles the changing neighbour count during MD without recompiling.

  • DP_TRITON_INFER=1 — additionally enable the Triton inference kernels for a further speedup on larger cells.

A frozen .pt2 package bakes the compilation in — and, when you run dp --pt freeze with DP_TRITON_INFER=1 set, the Triton kernels too — so it skips the warm-up at the cost of being device-locked.

11.7.3. Single-point evaluation#

Build an AtomicData object, batch it, compute a neighbour list, and call the model. The wrapper returns a dictionary with energy (shape (B, 1)), forces (shape (N, 3)), and, when enabled, stress (shape (B, 3, 3)):

from nvalchemi.data import AtomicData, Batch
from nvalchemi.neighbors import compute_neighbors

data = AtomicData(
    atomic_numbers=atomic_numbers,  # (N,) integer atomic numbers
    positions=positions,  # (N, 3) in Angstrom
    cell=cell,  # (1, 3, 3) lattice vectors, or omit for a cluster
    pbc=pbc,  # (1, 3) booleans, or omit for a cluster
)
batch = Batch.from_data_list([data], device="cuda")
compute_neighbors(batch, config=model.model_config.neighbor_config)

out = model(batch)
energy = out["energy"]  # (B, 1) eV
forces = out["forces"]  # (N, 3) eV/A
stress = out["stress"]  # (B, 3, 3) eV/A^3 (Cauchy stress = virial / volume)

Forces and stress are computed conservatively inside the model and returned directly, so no gradient bookkeeping is required on the caller side.

11.7.4. Molecular dynamics#

For dynamics, register a neighbour-list hook so the list is rebuilt before each force evaluation, then drive the batch with an integrator. The following snippet runs canonical (NVT) dynamics with a Langevin thermostat:

from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import initialize_velocities
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.integrators import NVTLangevin
from nvalchemi.hooks import NeighborListHook
from nvalchemi.neighbors import compute_neighbors

# ``forces`` and ``energy`` are pre-allocated so the integrator can read forces
# and the engine can write results back into the batch in place.
data = AtomicData(
    atomic_numbers=atomic_numbers,
    positions=positions,
    cell=cell,
    pbc=pbc,
    forces=torch.zeros_like(positions),
    energy=torch.zeros((1, 1), dtype=positions.dtype, device=positions.device),
)
batch = Batch.from_data_list([data], device="cuda")

# Seed Maxwell-Boltzmann velocities at the target temperature.
temperature = torch.full(
    (batch.num_graphs,), 330.0, dtype=positions.dtype, device="cuda"
)
initialize_velocities(
    batch.velocities, batch.atomic_masses, temperature, batch.batch_idx.int()
)

nl_hook = NeighborListHook(
    model.model_config.neighbor_config, stage=DynamicsStage.BEFORE_COMPUTE
)
nvt = NVTLangevin(model, dt=0.5, temperature=330.0, friction=0.01, hooks=[nl_hook])

# Prime the neighbour list and forces, then integrate.
compute_neighbors(batch, config=model.model_config.neighbor_config)
nvt.compute(batch)
batch = nvt.run(batch, n_steps=1000)

Switching ensemble is a one-line change: use NVE(model, dt=...) for the microcanonical ensemble or NPT(...) for constant pressure (which consumes the stress output). nvalchemi also provides logging and monitoring hooks (e.g. LoggingHook, EnergyDriftMonitorHook) that attach to the same engine.

11.7.5. Geometry optimization#

The same model drives the FIRE optimizer. Convergence is controlled by a maximum-force criterion:

from nvalchemi.dynamics.base import ConvergenceHook
from nvalchemi.dynamics.optimizers import FIRE2

opt = FIRE2(
    model,
    dt=1.0,
    hooks=[nl_hook],
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.05),
)
compute_neighbors(batch, config=model.model_config.neighbor_config)
opt.compute(batch)
batch = opt.run(batch, n_steps=200)  # stops early once fmax <= 0.05 eV/A

11.7.6. Outputs and configuration#

The wrapper advertises its capabilities through model.model_config (an nvalchemi ModelConfig):

  • outputs — energy, forces, and stress.

  • active_outputs — the subset computed on each call. energy and forces are active by default; stress is added when compute_stress=True (or via model.set_config("active_outputs", {"energy", "forces", "stress"})).

  • neighbor_config — the cutoff and neighbour-list format the model requires. Pass it to compute_neighbors or NeighborListHook as shown above.

11.7.7. Heterogeneous batches#

A single Batch may contain several structures of different sizes and cells. The wrapper evaluates the whole batch in one pass and returns per-structure energy and stress ((B, 1) and (B, 3, 3)) together with the concatenated per-atom forces ((N, 3)), making it straightforward to evaluate many systems at once.

11.7.8. Units and conventions#

  • Lengths are in Angstrom, energies in eV, masses in amu, and time in femtoseconds.

  • Atomic numbers are mapped to model types using the checkpoint type_map. Pass atomic_number_to_type={Z: type_index, ...} to DPA4Wrapper to override this for non-standard type maps.

  • The reported stress is the Cauchy stress, equal to the virial divided by the cell volume; it requires a periodic cell.

11.7.9. Limitations#

  • Only DPA-4 / SeZM energy models are supported.

  • Acceleration uses DeePMD-kit’s own compiled inference (DP_COMPILE_INFER or a frozen .pt2); nvalchemi’s FusedStage torch.compile is not used.

  • Embeddings (compute_embeddings) require the .pt backend, not .pt2.

  • Charge / spin conditioning is applied as a single global value per batch.

11.7.10. Examples#

Complete, runnable scripts for single-point evaluation, NVE, NVT, and geometry optimization are provided in examples/water/dpa4/nvalchemi/.