Note
Go to the end to download the full example code.
MACE NPT: domain-decomposed constant-pressure MD with a barostat#
The constant-pressure sibling of example 03. Load a MACE checkpoint, run
a short NPT trajectory across multiple ranks
under DomainParallel, and record the
trajectory — including the evolving cell — to an xyz file from rank 0.
Why NPT-under-DD needs more than NVT: a barostat and a thermostat both
couple to global thermodynamic quantities — the total kinetic energy,
the total degrees of freedom, and the full-system pressure tensor — none
of which any single rank can see from its owned subdomain alone. The
framework handles this transparently: NPT
declares __dd_thermo_kind__ = "npt", and on partition()
DomainParallel installs a dynamics
coordinator that
all-reduces the per-rank kinetic energy and kinetic-pressure tensor into mesh-global values (the consolidated virial is already global),
replaces the integrator’s per-shard degrees-of-freedom with the global count (and rescales the Nosé–Hoover chain masses accordingly), and
keeps the replicated barostat state and the cell byte-identical across ranks by broadcasting them from rank 0 each step.
So the wrapper and the integrator stay ensemble-correct; the only user
change from example 03 is asking the model for stress (the barostat
needs the virial) and swapping NVTLangevin for NPT.
System: alpha-quartz SiO2 supercell, isotropic barostat at zero external pressure — the box relaxes toward its equilibrium volume while the thermostat holds temperature.
Note
Run with:
torchrun --nproc_per_node=2 examples/distributed/06_mace_npt_distributed.py
For multi-GPU MACE+cuEquivariance, set the env var below to avoid a JIT-compilation race across ranks:
CUEQUIVARIANCE_OPS_PARALLEL_COMPILE=0 \
torchrun --nproc_per_node=N \
examples/distributed/06_mace_npt_distributed.py
Output xyz file at ./mace_npt_trajectory.xyz (rank 0 only); each frame
carries the cell at that step, so the volume relaxation is visible in
OVITO/VMD.
from __future__ import annotations
import argparse
import os
import sys
import warnings
from pathlib import Path
import torch
from loguru import logger
from nvalchemi.data import AtomicData, Batch
from nvalchemi.distributed import DomainConfig, DomainParallel, HookScope
from nvalchemi.dynamics import NPT, HostMemory
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.hooks import SnapshotHook
from nvalchemi.hooks import NeighborListHook
# Skip the heavy distributed launch during the Sphinx-Gallery docs build (it has
# no torchrun environment), mirroring examples 01-03.
_DOCS_BUILD = os.environ.get("NVALCHEMI_SPHINX_BUILD") == "1"
_DISTRIBUTED_ENV = "RANK" in os.environ and "WORLD_SIZE" in os.environ
# Reuse the SiO2 supercell builder from the benchmark suite (see example 03).
# Sphinx-Gallery executes examples without defining ``__file__`` and skips the
# launcher path below, so only resolve the benchmark import for a real run.
if not _DOCS_BUILD:
sys.path.insert(
0, str(Path(__file__).resolve().parents[2] / "benchmark" / "distributed")
)
from _benchmark_common import build_sio2_supercell # noqa: E402
# ----------------------------------------------------------------------
# System construction (rank 0 — DomainParallel scatters from there)
# ----------------------------------------------------------------------
def build_initial_batch(
repeats: tuple[int, int, int], dtype: torch.dtype, device: torch.device
) -> Batch:
pos, numbers, masses, cell, velocities = build_sio2_supercell(
repeats=repeats, dtype=dtype, seed=0
)
data = AtomicData(
positions=pos.to(device),
atomic_numbers=numbers.to(device),
atomic_masses=masses.to(device),
cell=cell.to(device).unsqueeze(0),
pbc=torch.tensor([[True, True, True]], device=device),
)
data.add_node_property("velocities", velocities.to(device))
# NPT reads a stress tensor each step; pre-allocate so the field exists
# before the first forward (the model overwrites it).
data["stress"] = torch.zeros(1, 3, 3, dtype=dtype, device=device)
return Batch.from_data_list([data], device=device)
def cell_volume(batch: Batch) -> float:
"""Volume (ų) of a single-system batch's (replicated) cell."""
cell = batch.cell
if cell.dim() == 3:
cell = cell[0]
return float(torch.linalg.det(cell).abs())
# ----------------------------------------------------------------------
# Trajectory persistence (rank 0 only)
# ----------------------------------------------------------------------
def write_trajectory_xyz(sink: HostMemory, path: Path) -> int:
"""Decode the :class:`HostMemory` sink into per-frame
:class:`ase.Atoms` (with the step's cell) and write an extxyz
trajectory. Returns the number of frames written.
"""
from ase import Atoms
from ase.io import write as ase_write
trajectory_batch = sink.read()
n_frames = trajectory_batch.num_graphs
if path.exists():
path.unlink()
for frame in range(n_frames):
single = trajectory_batch.index_select(torch.tensor([frame]))
cell = single.cell
if cell.dim() == 3:
cell = cell.squeeze(0)
atoms = Atoms(
numbers=single.atomic_numbers.detach().cpu().numpy(),
positions=single.positions.detach().cpu().numpy(),
cell=cell.detach().cpu().numpy(),
pbc=True,
)
atoms.info["frame"] = frame
ase_write(str(path), atoms, format="extxyz", append=True)
return n_frames
# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main() -> None:
parser = argparse.ArgumentParser(
description="MACE NPT (barostat) under DomainParallel."
)
parser.add_argument(
"--checkpoint",
default="medium-0b2",
help="MACE foundation model checkpoint name. "
"Default fetches MACE-MP-0b2 from HuggingFace.",
)
parser.add_argument(
"--repeats",
type=int,
nargs=3,
default=[3, 3, 3],
help="SiO2 unit-cell repeats along (a, b, c). 3x3x3 → 243 atoms.",
)
parser.add_argument("--n-steps", type=int, default=200)
parser.add_argument("--temperature-k", type=float, default=300.0)
parser.add_argument(
"--pressure",
type=float,
default=0.0,
help="Target external pressure in eV/ų (1 bar ≈ 6.32e-7). "
"Default 0 → relax toward the equilibrium volume.",
)
parser.add_argument(
"--dt-fs", type=float, default=1.0, help="MD timestep in femtoseconds."
)
parser.add_argument(
"--barostat-time-fs",
type=float,
default=1000.0,
help="Barostat coupling time τ_P (fs). Larger = gentler cell motion.",
)
parser.add_argument(
"--thermostat-time-fs",
type=float,
default=100.0,
help="Thermostat coupling time τ_T (fs).",
)
parser.add_argument(
"--snapshot-every",
type=int,
default=10,
help="Persist a frame to the trajectory sink every N steps.",
)
parser.add_argument(
"--output-xyz",
type=Path,
default=Path("mace_npt_trajectory.xyz"),
help="xyz file path (rank 0 only).",
)
parser.add_argument(
"--dtype",
default="float32",
choices=["float32", "float64"],
help="Model / simulation dtype.",
)
args = parser.parse_args()
dtype = torch.float64 if args.dtype == "float64" else torch.float32
# Docs build / no torchrun: no process group to join, so skip the launch.
if _DOCS_BUILD or not _DISTRIBUTED_ENV:
logger.info(
"Not running under torchrun — skipping the distributed run. "
"Launch with: torchrun --nproc_per_node=N "
"examples/distributed/06_mace_npt_distributed.py"
)
return
# ----- Distributed bootstrap via DistributedManager (see example 03) -----
from nvalchemi.distributed import DistributedManager
DistributedManager.initialize()
try:
dm = DistributedManager()
rank, world_size, device = dm.rank, dm.world_size, torch.device(dm.device)
mesh = dm.initialize_mesh(mesh_shape=(world_size,), mesh_dim_names=("domain",))
if rank == 0:
logger.info(
"MACE NPT distributed: world_size={ws} device={dev} "
"checkpoint={ckpt} repeats={r} n_steps={n} T={T}K "
"P={P} eV/ų dt={dt}fs",
ws=world_size,
dev=device,
ckpt=args.checkpoint,
r=tuple(args.repeats),
n=args.n_steps,
T=args.temperature_k,
P=args.pressure,
dt=args.dt_fs,
)
# ----- Load MACE wrapper -----
with warnings.catch_warnings():
warnings.simplefilter("ignore")
from nvalchemi.models.mace import MACEWrapper
wrapper = MACEWrapper.from_checkpoint(
args.checkpoint, dtype=dtype, device=device
).eval()
# The barostat needs the virial → ask MACE for stress (example 03 only
# needed energy + forces). Everything else is unchanged.
wrapper.set_config("active_outputs", {"energy", "forces", "stress"})
if rank == 0:
logger.info("MACE wrapper ready: cutoff={c} Å", c=wrapper.cutoff)
# ----- Domain config -----
domain_cfg = DomainConfig(cutoff=float(wrapper.cutoff), skin=0.5, mesh=mesh)
# ----- Hooks -----
nl_hook = NeighborListHook(
wrapper.model_config.neighbor_config,
skin=0.5,
stage=DynamicsStage.BEFORE_COMPUTE,
)
n_frames_expected = (args.n_steps // args.snapshot_every) + 1
trajectory_sink = HostMemory(capacity=n_frames_expected)
snapshot_hook = SnapshotHook(
sink=trajectory_sink,
frequency=args.snapshot_every,
)
# RANK_ZERO scope: gather the full system (with the current cell) onto
# rank 0 so each frame is the whole box, not one rank's shard.
snapshot_hook.scope = HookScope.RANK_ZERO
# ----- Inner integrator: NPT -----
# NPT couples a Nosé–Hoover thermostat chain to the particle velocities
# and a barostat to the cell. Under DomainParallel the dynamics
# coordinator globalises the kinetic energy / DOF / pressure tensor and
# broadcasts the replicated barostat state + cell each step (see the
# module docstring), so this stays ensemble-correct across ranks.
integrator = NPT(
model=wrapper,
dt=args.dt_fs,
temperature=args.temperature_k,
pressure=args.pressure,
barostat_time=args.barostat_time_fs,
thermostat_time=args.thermostat_time_fs,
pressure_coupling="isotropic",
chain_length=3,
hooks=[nl_hook],
n_steps=args.n_steps,
)
# ----- DomainParallel wrapping (SnapshotHook on the outer, as in 03) -----
# ``with dynamics:`` makes teardown exception-safe; the process-group
# lifecycle stays at launcher scope (``DistributedManager.cleanup()`` below).
with DomainParallel(
dynamics=integrator,
config=domain_cfg,
n_steps=args.n_steps,
hooks=[snapshot_hook],
) as dynamics:
# ----- Build the initial batch on rank 0 and partition -----
initial_batch = (
build_initial_batch(tuple(args.repeats), dtype=dtype, device=device)
if rank == 0
else None
)
v0 = cell_volume(initial_batch) if rank == 0 else None
owned_batch = dynamics.partition(initial_batch)
if rank == 0:
logger.info(
"Partitioned: n_owned (rank 0) = {n} of {tot} atoms; V0 = {v:.2f} ų",
n=int(owned_batch.positions.shape[0]),
tot=int(initial_batch.positions.shape[0]),
v=v0,
)
# ----- Run the trajectory -----
final_batch = dynamics.run(owned_batch)
# ----- Report volume relaxation + persist trajectory -----
# The cell is replicated (broadcast from rank 0 each step), so any rank's
# final cell is the global cell; report it from rank 0.
if rank == 0:
v1 = cell_volume(final_batch)
logger.info(
"Volume: {v0:.2f} → {v1:.2f} ų ({pct:+.2f}%) over {n} steps",
v0=v0,
v1=v1,
pct=100.0 * (v1 - v0) / v0,
n=args.n_steps,
)
n_frames = write_trajectory_xyz(trajectory_sink, args.output_xyz)
logger.info(
"Done. Wrote {f} xyz frames to {p}.", f=n_frames, p=args.output_xyz
)
finally:
# Process-group teardown stays at launcher scope.
DistributedManager.cleanup()
if __name__ == "__main__":
main()
Total running time of the script: (0 minutes 0.002 seconds)