Isothermal-Isobaric Molecular Dynamics with a Foundation Interatomic Potential in LAMMPS across NVIDIA and Intel GPUs

Introduction

Foundation machine-learning interatomic potentials (MLIPs) trained on hundreds of millions of structures [1] occupy a different point in the parallel design space from the short-ranged potentials that weak-scale to a full machine. A single Universal Model for Atoms (UMA) checkpoint, with task heads for materials (omat), catalysis (oc20), etc., models chemistries that lack a transferable classical interatomic potential; that transferrability costs the locality multi-node parallelism needs. UMA stacks four message-passing blocks at a 6 Å cutoff [1, 8], a ∼24 Å receptive field, and two global couplings: channel balancing [8] and the mixture-of-linear-experts (MoLE) routing [1]. A correct spatial halo is four shells deep, not one pair cutoff. At the cell sizes this ensemble can hold, even one shell is comparable to the owned subdomain (Sec. VI-B), so a spatial split leaves the halo as large as the work it surrounds. What remains correct without model surgery is graph parallelism, which keeps the configuration global and shards only the edges. Every block then gathers the full node-feature tensor, so communication is $O(N)$ per layer and per-device node memory is $O(N)$ independent of world size. Both properties are acceptable inside a node, but they tend to slow down performance when inter-node communication happens: more nodes no longer means better simulation performance. (Sec. VI-B). One node is therefore the unit of speed for this model class, and further nodes provide capacity. The systems question is how to put that evaluation inside the Large-scale Atomic/Molecular Massively Parallel Simulator (LAMMPS) [5] so that microcanonical (NVE), canonical (NVT), and especially isothermal-isobaric (NPT) simulations run on the NVIDIA and Intel data-center GPUs already installed in high-performance computing (HPC) centers.

The distinction that matters in this systems paper is between an ad hoc and a foundation potential. The million- and billion-atom MLIP results that define the current scaling frontier were all obtained with a potential trained for one system: Spectral Neighbor Analysis Potential (SNAP) [14] for carbon at 20 billion atoms [2], Deep Potential for copper and water at 100 million [3], Allegro for one solvated biomolecular assembly at 44 million [4]. Each is short-ranged and locally decomposable by construction, which is precisely why it weak-scales to a full machine; the cost is a training campaign per chemistry. We use transferable in the strict sense: a potential is transferable if changing the chemical elements or composition requires no new parameter fit or retraining. By that measure the established approaches fail: classical embedded-atom [15], polarizable-ion and fixed-charge force fields, and ad hoc per-system MLIPs alike, all refit per chemistry. UMA takes the opposite position. One checkpoint, selected at export to a task head, evaluates forces across molecules, materials and catalysis, and the released model matches or beats specialized models without fine-tuning [1]. For the demonstration workload in this paper (heating ramps of NaCl, equiatomic Pt–Ru, and equiatomic (HfNbTaTiZr)N from one omat checkpoint), a classical potential would need three separately parameterized tables for a salt, a binary alloy, and a five-cation nitride, whereas one foundation model runs all three.

Three constraints then dominate the design.

First, parallelism here pools memory rather than providing speed. Edge sharding across the devices of one node raises the largest cell that fits.

A second node does not shorten the step: per-rank node-feature work stays $O(N)$, and the node-boundary cost itself is $\sim2\%$ of the NPT step (Sec. VI-B). Capacity grows by about 23 thousand atoms per node, but the step time grows with the cell. The remaining multi-node axis is data parallelism over independent configurations.

Second, NPT runs with checkpointing off. The barostat needs a strain gradient, which lets the cell change volume through a transition, and we have not validated that gradient through our activation-checkpointing path (§ 2-C). Running with checkpointing off costs capacity: the checkpoint-on NVT ceiling on the same node is 20 times higher.

Third, portability is a requirement rather than a feature. Centers are deploying both NVIDIA and Intel accelerators, and a molecular dynamics (MD) engine whose results depend on which vendor’s hardware runs it is not suitable for scientific use. We therefore treat cross-device agreement as a correctness property to be measured.

Graph parallelism and the autograd virial are FairChem’s (§ 2-B). Our contribution is a new capability, an FP64, NPT-capable LibTorch pair style that also runs on Intel GPUs, not a speedup over an existing solution. It comprises three contributions:

1. A LibTorch interface from LAMMPS to UMA. A TorchScript export and a C++ engine behind pair_style uma, with one MPI rank per device and FP64 end to end, built from one source tree for NVIDIA GPUs (NCCL collectives) and Intel GPUs (oneCCL collectives, named XCCL in PyTorch). Thermodynamic output at steps 0 and 10 agrees across every device count that fits on each GPU type, and forces match an ASE FairChem reference to $\sim10^{-14}$ eV/Å (Sec. VI-C).

2. The strain-gradient virial through that interface, for NPT. FairChem’s predictor already derives stress from position and cell gradients; we carry it through exported modules, graph parallelism, and fix npt, and show it stays nearly free. On the 22,000-atom cell, ten NPT steps take 8.54 s against 8.52 s for NVT, because the virial rides the same backward pass as the forces (Sec. VI-D).

3. NPT molecular dynamics of up to 47,000 atoms at 0.157 ns/day, with delivered 100 ps runs. We measure the checkpoint-off NPT ceilings on $4 \times$ A100-40GB, one 12-tile Intel node, and $8 \times$ H200 (Sec. VI-D), and the Intel ceiling out to 32 nodes (Sec. VI-B). On all three node types, one omat checkpoint then runs 100 ps NPT heating ramps of NaCl, equiatomic Pt–Ru, and equiatomic (HfNbTaTiZr)N with no new fit (Sec. VII).

Scope

The headline results are node-local by design; Sec. VI-B adds one multi-node capacity measurement, to 32 nodes. A spatial-DD path exists in the source but is not validated and is not claimed; the halo comparison above is why a thin halo does not apply. We claim a capacity gain only over a single device, the regime of the ASE FairChem calculator, and a precision difference, not a capacity or speed advantage, over FairChem’s LAMMPS wrap. We also do not claim nanosecond trajectories (Sec. VI-D). The heating-ramp jumps of Sec. VII are not melting points; Sec. VII-B states why. Section III places this against custom million-atom MLIP MD, other vendors’ UMA stacks, and Python drivers.

Background

UMA in MD

UMA [1] represents a configuration as a graph with edges inside a cutoff (6 Å in the exported uma--s1p2 artifacts used here). Equivariant message-passing layers update node features; energy is a sum of per-atom terms and forces are its gradient. Task heads (omat, oc20, …) are selected at export. Mixed-element systems are legal at runtime; the neighbor graph, not a classical cross-term file, supplies the multi-element interactions. Two properties drive everything below: the model is deep, so its receptive field is $L$ times the cutoff for $L$ blocks ($4 \times6 = 24$ Å here), and it carries global couplings (channel balancing and the mixture-of-linear-experts routing) that are not expressible as a neighborhood reduction.

Why graph parallel, not spatial domain decomposition

LAMMPS spatial decomposition assigns atoms by box region and ghosts a halo of width equal to the pair cutoff, and this is what lets custom, short-ranged MLIPs weak-scale to whole machines [2, 3, 4]. For a deep graph neural network (GNN) the physical cutoff is not the receptive field: a correct halo must be $L$ shells deep, and at $L=4$ with a 6 Å cutoff the ghost region can exceed the owned partition, which we observe directly ($n_{\mathrm{all}} \approx62\mathrm{k}$ ghosts per rank for a 74,000-atom cell at $k=1$). The global couplings then remain unhandled even with a deep halo.

Graph (edge) parallelism instead keeps the configuration global and shards edge work across devices, using collectives to gather node features. This is the multi-GPU mode of FairChem’s own predictor [8]: edges are split by index, node features are all-gathered before each block, and forces and the cell gradient are all-reduced. We reimplement it in C++ over exported TorchScript modules behind a LAMMPS pair style, with NCCL and oneCCL collectives. The design is FairChem’s; the engine, the Intel GPU port, and the measurements are ours. It is correct for this architecture without model surgery, and it is the mode in which we have both a working NPT barostat and an implementation that runs on both CUDA and XPU devices. Its two costs are structural and set the scope of this paper: each block gathers the full $[N,C]$ node tensor onto every rank, so communication is $O(N)$ per layer and per-device node memory is $O(N)$ independent of world size. Only edge work is divided. Inside a node, over NVLink or Xe Link, this is a good trade; Sec. VI-B shows a second node provides memory but not speed.

Why checkpoint-off NPT

The NPT virial is a derivative through cell strain. Activation checkpointing [22] is compatible with it in principle: FairChem’s eager model checkpoints its edge-chunked convolutions with torch.utils.checkpoint and still differentiates through the cell. A traced module cannot use that API, so our engine recomputes chunks and blocks in custom C++ autograd functions, and we have not validated the cell gradient through them. The pair style therefore refuses the combination, and every NPT result below runs with checkpointing off (UMA_AC=off, UMA_CKPT unset); if either switch is on, the pair style aborts fix npt. The cost below is a limit of our implementation, not of the physics. The price is capacity: on one 12-tile Intel node the checkpoint-on NVT ceiling is $N=38$ (439,000 atoms, 28.8 GiB peak), 20 times the 22,000 atoms of the checkpoint-off NPT ceiling.

Related Work

We do not claim the largest MLIP MD. We claim that the foundation-MLIP regime is a different systems problem from the custom-MLIP regime that defines the scaling frontier, and that in this regime a correct barostat and cross-device agreement, not atom count, are the binding constraints.

Scale frontier and foundation models

The record-setting MLIP MD results are all single-chemistry: Nguyen-Cong et al. [2] run 20 billion carbon atoms at 50.0 peta floating-point operations per second (PFLOPS) on Summit; Jia et al. [3] reach 100 million atoms at 91 PFLOPS with Deep Potential; Kozinsky et al. [4] run Allegro [13] on a 44-million-atom solvated HIV capsid and report 70% weak scaling to 5120 A100s. Each uses a potential trained for its target system, with a short receptive field that decomposes spatially with a thin halo. UMA’s $\sim24$ Å receptive field and global couplings force graph parallelism, which pools memory within a node and caps the cell at tens of thousands of atoms per node (Sec. VI-B). In exchange, one checkpoint spans molecules, materials, and catalysis without fine-tuning [1]: NaCl, equiatomic Pt–Ru, and an equiatomic five-cation nitride from one binary with no retraining, while no single transferable classical potential spans that set. Papers at the scale frontier report NVE/NVT; we know of no published strain-gradient virial NPT result at foundation-MLIP cost.

Integration and baseline context

Deep Potential [3, 10], SNAP [2, 14], Allegro/Neural Equivariant Interatomic Potentials (NequIP) [4, 12, 13], and MACE [11] are delivered as LAMMPS pair styles; the integration plumbing is established prior art. What changes at foundation scale is the parallel decomposition, memory envelope, and virial path. NVIDIA ALCHEMI [38] runs UMA across several NVIDIA GPUs from a Python driver on the Warp/CUDA stack; to our knowledge it has no Intel XPU path, and this paper does not compare against it. Other MLIP families reach LAMMPS through spatial decomposition: SevenNet [39] exchanges node features between message-passing layers for NequIP-based models, and chemtrain-deploy [40] deploys semilocal MACE, Allegro, and PaiNN potentials through LAMMPS domain decomposition. Both build on message passing with a short per-layer reach and are established options for large cells. A UMA user trades speed per atom for one checkpoint spanning molecules, materials, and catalysis; this paper makes that choice available as an FP64, NPT-capable pair style on NVIDIA and Intel GPUs, and does not claim it is faster than those engines. For context, SevenNet reports over 80% weak-scaling efficiency on 32 GPUs [39] and chemtrain-deploy scales to millions of atoms [40]; the models and hardware differ, so these are not comparisons. To our knowledge no systematic measurement of multi-device UMA capacity ceilings and NPT step cost has been published. The UMA paper reports that UMA-S fits 100,000 or more atoms on one 80 GB GPU in its default inference settings [1]. Our NPT ceilings are lower because they are FP64, differentiate through the cell, and retain activations (checkpointing off); the checkpoint-on NVT ceiling of 439,000 atoms on one Intel node (§ 2-C) is the closer comparison. The Atomic Simulation Environment (ASE) [7] and FairChem [8] Python drivers are the quantitative baseline throughout: identical geometry, tiles, and FP64 model. FairChem also ships a LAMMPS path (fairchem-lammps) that drives LAMMPS through fix external rather than a pair style [5, 9]; that wrap is a second baseline in Sec. VI-C, not the engine under test. The Python driver on that geometry recomputes the stress under a different integrator (Nosé–Hoover chain) [18, 19, 20], Berendsen-style in the driver [21], so the timing in Sec. VI-D is cost per step. FairChem’s LAMMPS wrap is the closer baseline, and it is not slower than pair_style uma. Classical embedded-atom method (EAM) [15, 17] and modified embedded-atom method (MEAM) [16] are fit per element set; one umat checkpoint evaluates NaCl, Pt–Ru, and (HfNbTaTiZr)N with no new fit, and a classical step on any one table is faster (this paper carries no speed comparison).

Design

Pair style

pair_style uma builds the neighbor graph from LAMMPS neighbor lists, runs a traced PyTorch C++ (LibTorch) module, and returns per-atom energies and forces. newton off, units metal, timestep 0.001 ps are the production conventions. Export is offline: a geometry (structure.xyz) is traced once per $(N,W,\mathrm{task})$ triple because edge_pad_cap is baked from that graph. NaCl artifacts are not reused for the alloy cell (different coordination at the same cutoff), nor a smoke export for a 20k-atom cell.

Launch contract

The launch contract is enforced by the pair style, not merely followed by convention:

• $W$ Message Passing Interface (MPI) ranks [37] on the $W$ devices of one node, or across nodes in the multi-node runs, rank $r$ bound to device $r$; pair_style uma ...devices 1 on every rank, so the graph-parallel world size equals the MPI world size.

• Combining multiple ranks with devices > 1 (one rank forking several device workers) is a hard error: the two are alternative realizations of the same partition and mixing them double-shards the graph.

• Multi-rank operation additionally requires precision double. The multi-device runtime is FP64-only, so capacity cannot be traded for precision on this path, a limitation we return to in § 11.

• Collectives use NCCL [33] on NVIDIA GPUs (the engine is compiled for CUDA [36]) and XCCL, native oneCCL [34], on Intel GPUs (the engine is compiled with icpx for the SYCL C++ heterogeneous programming model [35]), behind one internal peer interface offering all_gather_nodes and all_reduce_sum.

Throughout, $W$ counts devices as the runtime sees them. An Intel node exposes 12 devices (GPU tiles) to the runtime, so $W=12$ per node and all Intel-node efficiencies are per tile. On NVIDIA nodes, $W$ equals the number of GPUs. Atom counts are rounded to the nearest thousand; ratios, per-atom costs, and the cube edge $N$ (in unit cells) use the exact cells.

NPT virial path

Isotropic fix npt at 300 K / 1 bar is the bulk default. The pair style enables position+cell autograd when a barostat is present and checkpointing is off. NVE/NVT do not need the strain gradient and are used for vacuum nanoparticles, where padded-box NPT is unphysical.

Portability

A thin device layer selects CUDA versus Intel XPU at build and run time, and the LAMMPS input, task head, and checkpoint-off NPT contract are identical on both GPU types. The claim we defend is stronger than input compatibility: on one GPU type, results agree across device counts. Step-0 energies are identical to the last printed digit across every $W$ on NVIDIA GPUs (NCCL) and across every $W$ on Intel GPUs (XCCL). This does not claim that NVIDIA and Intel GPUs produce the same bits: the NVIDIA cell is a different random realization. The ASE energy and force deviations are in Sec. VI-C. For a deep equivariant network evaluated in FP64 through a reordered all-reduce this is not automatic, and it is what makes a result from either GPU type suitable for scientific use. Intel GPU numbers come from 12-tile Intel nodes, one to 32 nodes. NVIDIA GPU numbers come from 4×A100-SXM4-40GB nodes and one 8×H200 node.

Implementation

The engine is a C++ LAMMPS package plus a Python export of per-block TorchScript modules (uma-s-1p2) [6]. Export writes model_traced.pt and per-block/chunk artifacts; LAMMPS loads them at pair_coeff. Environment contract:

• UMA_CHECKPOINT: path to the FairChem .pt (export only).

• UMA_AC=off, UMA_CKPT unset for NPT.

• UMA_MP_NATOMS: atom count for the graph-parallel gather.

Activation checkpointing is off by default and is enabled only by opting in (UMA_AC=chunk|block|full or UMA_CKPT=1); when either is set and a barostat is present, the pair style aborts with a diagnostic rather than integrating against an incorrect pressure. All NPT results below therefore run the default. UMA_CHECKPOINT names the source checkpoint at export only, and UMA_MP_NATOMS selects the artifact because partition offsets are baked into the TorchScript at trace time, so an artifact is valid only for the $(N,W)$ it was traced at, which is a real constraint of tracing.

The graph is rebuilt inside every force evaluation. The LAMMPS neighbor list is requested only to register the cutoff and is never read; the engine constructs its own UMA graph from positions each step. This is why the LAMMPS timing breakdown attributes 99.99% of loop time to Pair and ≈0.00% to Neigh: graph construction, network evaluation and the backward pass are all inside the pair call. Double precision is required for the multi-device and NPT virial paths.

Evaluation

We report the multi-node capacity and cost sweep (Sec. VI-B), correctness, node-local capacity, and where the step goes. The demonstration that one checkpoint covers three chemistries, and the accuracy of those heating-ramp jumps, is Sec. VII.

Platforms

All runs use uma-s-1p2 in FP64. NaCl rocksalt at $8N^3$ atoms is the capacity and parity workload; the heating ramps of Sec. VII are the demonstration workload, not the capacity cells.

• NVIDIA, tight memory: one 4×A100-SXM4-40GB node. Checkpoint-off NPT ceiling and the node-boundary comparison.

• NVIDIA, high memory: one 8×H200 node (141 GB/device). Largest checkpoint-off NPT cell, 47,000 atoms; the ten-step loop rate is 0.157 ns/day.

• Intel: 12-tile Intel nodes, one to 32 nodes. Checkpoint-off NaCl NPT ceiling 22,000 atoms on one node and 779,000 on 32; the ten-step loop rate is 0.10 ns/day on one node.

Beyond one node: capacity and cost

Graph parallelism gathers the full node-feature tensor every block, so it communicates $O(N)$ per layer and its per-device memory is $O(N)$ regardless of world size.

We tested whether that is survivable across nodes on the identical NaCl 4,000-atom cell, checkpoint-off NPT, over each machine’s own network fabric.

Two comparisons isolate two effects. At equal world size (four A100s on one node versus two plus two, and twelve Intel tiles on one node versus six plus six), crossing the node boundary adds 2.5% and 2.1% per NPT step. Energy, forces, and step-0 thermo agree across the split ($|\Delta E| = 2 \times10^{-11}$ eV, max $|\Delta F| = 9 \times10^{-16}$ eV/Å; the A100 pressure is 10,978.496 bar on one node and on two). Doubling the device count (four A100s versus eight, twelve Intel tiles versus twenty-four) does not shorten the step: 646 ms becomes 788 ms (1.22× longer) on A100, and 1,009 ms becomes 1,063 ms (1.05× longer) on the Intel node. On A100, about 80% of that gap is per-rank GPU compute that grows with world size; collectives are 14%. The step is compute on both machines: the collective is 2.0–4.5% of the A100 step and 0.5–2.1% of the Intel step. On Intel we separate communication from waiting by draining each rank’s queue and running a barrier before every timed collective. A host timer around the all-reduce reads 665 ms on one node, but 646 ms of that is the rank’s own queued compute and 20 ms is rank skew; the collective itself is 5.5 ms. Crossing the node boundary at $W=12$ adds 16.5 ms of collective time (1.6% of the step). That time is host-observed: device-side collective kernels take 3.5 ms at $W=12$ on one node and 3.8 ms at $W=24$ on two, so most of the added time is not kernel time. TCP-socket two-node runs inflate the Intel step 2.8× and are not the numbers above. The strain-gradient virial adds 0.2 ms (A100) to 0.3 ms (Intel), so NPT and NVT differ by 1% or less at every layout. A second node is therefore a memory increment, and NPT is valid across it.

Capacity and speed across nodes. We ran cubic NaCl for ten NPT steps at 300 K and 1 bar on 1 to 32 12-tile Intel nodes ($W=12\times\mathrm{nodes}$, FP64, checkpointing off, each machine’s own network fabric), taking on each node count the largest cell that finished (Figure 1). Capacity grows linearly with nodes, from 22,000 atoms to 779,000, 22–24 thousand per node throughout; the next size runs out of device memory at every node count. The replicated $[N,C]$ tensor costs 9.2 kB per atom, 7.2 GB (11% of one tile’s memory) at 779,000 atoms, so capacity must eventually lag the node count, but it has not at 32 nodes. Speed does not follow. The step grows from 0.82 s (the single-node baseline of Sec. VI-D) to 11.6 s, so the ten-step loop falls from 0.102 to 0.007 ns/day and weak-scaling efficiency from 74% at two nodes to 7% at 32. That is not the ∼2% boundary cost: the 2% is at equal world size on a fixed 4,000-atom cell, whereas weak scaling roughly doubles the cell with the devices. Forward and backward grow 8.6× and 10.5×. Graph build, a full-cell neighbor list that every rank repeats, grows 38×, from 127 ms to 4.8 s: 15% of the step on one node and 41% on 32. That is the lowest-effort remedy, since a shared or distributed neighbor list needs no model surgery, though it leaves the forward and backward growth. Wall time per atom falls from 37 μs on one node to 14–15 μs from eight nodes on only because the cell grows 35× while the step grows 14×; in device time, the cost per atom rises about 13× from one node to 32 (37 μs on one node against 14.8 μs on 32). More devices therefore provide cell size, not a shorter step. Graph build alone is 5.8–6.2 μs per atom at every size, as expected of work every rank repeats for the whole cell. At 32 nodes the cell is 259 Å across, so equal cubes for 384 ranks are 36 Å wide against a 24 Å receptive field: a 24 Å halo would be ∼12 times the owned volume ($\left(84/36\right)^3 \approx12.7$ for the padded region). One OpenMP thread per rank instead of the default three count leaves capacity unchanged; the default is up to 12% slower (16 nodes: ten-step loop 59.9 against 53.3 s).

For a foundation MLIP with a ∼24 Å receptive field, one node is the unit of speed, and further nodes pool memory. The remaining multi-node axis is data parallelism over independent configurations, which needs no cross-node collectives and scales trivially. Recovering multi-node speed for a single configuration would take a shared neighbor build and, for the forward and backward growth, replacing the per-block all-gather with a depth-$L$ halo exchange, which is model surgery on a released checkpoint that we do not attempt. That is a route to shorter steps. It does not answer a measured communication bottleneck: the all-gather is ∼1% of the 4,000-atom A100 step, the only size at which we split it, and we did not measure the share at multi-node cell sizes. A per-block exchange of one 6 Å shell around 36 Å cubes would carry ∼1.4 times the owned features per rank, against 384 times for the all-gather at 32 nodes (an estimate, not measured). Custom MLIPs of Sec. III remain the appropriate comparison: their potentials are local, so their halo is thin and their scaling is real.

Correctness

Against the ASE FairChem oracle on NaCl, Table I lists the energy and per-atom force deviations of pair_style uma (all atoms compared, not sampled). An autograd (AG) versus finite-difference (FD) check on the energy at the checkpoint-off $N=14$, $W=12$ cell gives $\max\lvert\mathrm{AG}-\mathrm{FD}\rvert= 5.18 \times10^{-7}$. The virial path is gated behind the barostat, so NVT graph-parallel forces are bit-identical to the non-virial build.

Node	Path	$N$	$\lvert\Delta E\rvert$	$\max\lvert\Delta F\rvert$
1$\times$A100	uma	6	$1.3\times10^{-10}$	$1.9\times10^{-14}$
1$\times$A100	wrap	6	$3.0\times10^{-6}$	$6.4\times10^{-6}$
1$\times$H200	uma	6	$1.3\times10^{-10}$	$2.0\times10^{-14}$
1$\times$H200	wrap	6	$3.0\times10^{-6}$	$6.4\times10^{-6}$
4$\times$A100	uma	10	$4.0\times10^{-9}$	$2.9\times10^{-14}$
4$\times$A100	wrap	10	$2.7\times10^{-5}$	$1.1\times10^{-5}$
4$\times$H200	uma	10	$4.0\times10^{-9}$	$2.9\times10^{-14}$
4$\times$H200	wrap	10	$2.7\times10^{-5}$	$1.1\times10^{-5}$
12-tile Intel	uma	14	$3.3\times10^{-8}$	$8.8\times10^{-14}$
12-tile Intel	wrap	14	$1.5\times10^{-5}$	$1.9\times10^{-5}$

Parity is against the reference implementation, not against density-functional theory (DFT). The numbers below are the deviation of this engine from the released UMA model. They say nothing about how well UMA reproduces first-principles forces, which is the model’s question and is addressed in [1].

Device-count consistency, checked on each GPU type. The stronger correctness result is that the answer does not depend on how many devices produced it. On the checkpoint-off NPT ladders, step-0 and step-10 thermo agree across every $W$ that fits: $W=1$–$12$ on one 12-tile Intel node, $W=1, 2, 4$ on $4\times$A100-40GB, and $W=1, 2, 4, 8$ on H200. Table I reports the ASE deviations for pair_style uma at $W=1$ and $W=4$ on A100 and on H200, and at $W=12$ on one 12-tile Intel node. The Intel ladder uses oneCCL; the CUDA ladders use NCCL. Trajectories, not just single points, agree: multi-device step-10 energies match the single-device trajectory. We obtain this by keeping the reduction order fixed with respect to $W$: the partition is a deterministic tensor_split of a fixed atom ordering and the per-shard reduction tree is built from that ordering, not from arrival order, so the all-reduce result does not depend on message timing.

Two LAMMPS paths versus ASE. The numbers above compare pair_style uma to the FairChem graph-parallel Python driver. FairChem also ships a LAMMPS path (fairchem-lammps) that injects forces through fix external rather than a pair style [5], [9]. Table I holds both paths to the same ASE FAIRChemCalculator baseline on rattled NaCl rocksalt (FP64, task omat; energy and every per-atom force, not a subsample) across $N=6$ ($W=1$), $N=10$ ($W=4$), and $N=14$ ($W=12$) on NVIDIA and Intel GPUs. The result is a clear separation that holds at every cell and on every device: pair_style uma tracks the reference at the FP64 floor, $\lvert\Delta E\rvert\sim10^{-10}$–$10^{-8}\,\mathrm{eV}$ and $\max\lvert\Delta F\rvert\sim10^{-14}\,\mathrm{eV/\mathring{A}}$, while the wrap’s deviations lie four orders of magnitude higher in energy and eight in force, $10^{-6}$–$10^{-5}\,\mathrm{eV}$ and $10^{-6}$–$10^{-5}\,\mathrm{eV/\mathring{A}}$ (Table I). The gap is the input precision, not the model: the wrap’s fix external callback builds positions and the cell as FP32 tensors and upcasts them for the FP64 model, so the coordinates, not the UMA weights, set that residual. FP32 unit roundoff is $\epsilon_{\mathrm{FP32}} \approx1.2 \times10^{-7}$; the measured wrap forces lie $\sim$50–100$\times$ above that unit roundoff. This is a numerical-precision distinction, and it is the one GPU MD engines are built around. GPU codes expose precision modes: all-double (DPDP), single-precision math with double-precision or 64-bit fixed-point accumulation (mixed, SPDP/SPFP), and all-single (SPSP) [41], [42]. The field’s consensus is that all-single is unsafe for production, because its rounding error accumulates over a long trajectory. AMBER reports SPSP producing unphysical trajectories in generalized-Born simulations and recommends mixed precision as the floor [41]; GROMACS runs its GPU path in mixed precision [44]; and the HALMD single-precision GPU integrator drifts about 0.1% in total energy per $10^7$ steps [43]. The mixed-precision recipe is not available for UMA. Its forces are the autograd gradient of the energy and inherit the precision of the network, so the released model runs in single or double precision, not in between. pair_style uma takes the double-precision side of that choice, the whole force and energy path in FP64 (DPDP), a mode usually reserved for CPUs because GPU FP64 throughput is lower; graph parallelism is what makes it affordable here. Against the FP64 ASE baseline the pair_style uma force residual is $10^{-14}\,\mathrm{eV/\mathring{A}}$ and the wrap’s is $10^{-6}$–$10^{-5}\,\mathrm{eV/\mathring{A}}$, with the energy gap four orders wide. The wrap’s model also runs in FP64, so that residual comes from its single-precision coordinate input; it is identical on A100 and H200. We did not run the wrap’s default single-precision model, whose residual would be larger. The comparison is fidelity to the released model, not accuracy against first-principles forces, and we do not compare long-run trajectories between the two.

NPT step cost and capacity

Step cost. At the largest node-local NPT size on one 12-tile Intel node (22,000 atoms), ten MD steps cost 8.52 s under NVT and 8.54 s under NPT: the strain gradient and its all-reduce add 0.2%. A device-synchronized phase profile of these runs puts the virial at 0.2–1.2 ms and the NPT–NVT step difference at 0.6–1.9 ms. Backward does not grow. The virial is already available from the same backward pass that produces forces, and reducing six additional scalars across the world is not measurable against an $O(N)$ node-feature gather. The Python driver on the identical geometry incurs 24.26 s for NVT and 49.73 s for NPT, because its barostat recomputes the stress (Figure 2). Those times are 2.9× and 5.8× the engine, and the integrators are not identical (Nosé–Hoover here, Berendsen-style cell coupling in the driver), so this is a cost-per-step comparison, not a trajectory one. We also timed FairChem’s Python-wrapped LAMMPS path

(fairchem-lammps, fix external) [9] on a smaller cell: NaCl $N=6$ (2,000 atoms) on one A100, ten NPT steps. There, pair_style uma loops in 4.63 s (pair 99.98%), while the wrap needs 2.96 s to load, 7.36 s for a first ten steps that include LAMMPS construction, and 3.90 s for a second ten steps. On one H200, the same cell loops in 2.19 s with pair_style uma (pair 99.96%), while the wrap needs 2.60 s to load, 5.90 s for a first ten steps, and 1.60 s for a second ten steps. That wrap puts the force in LAMMPS Modify (Pair 0%), so it is not the same timer split. On this $W=1$ cell the wrap is not slower. On the 22,000-atom $W=12$ cell the wrap itself loops ten NPT steps in 8.07 s (0.107 ns/day) against 8.54 s for pair_style uma. pair_style uma does not reduce wall time relative to FairChem’s LAMMPS wrap. It instead keeps the force and energy path in double precision and forms the strain-gradient virial inside the pair call. pair_style uma fits this cell only at $W=12$. The wrap also produces forces at $W=8$ and $W=10$, where pair_style uma runs out of memory, so on the Intel node the wrap fits this cell at more layouts; we claim no capacity advantage over it. Table I lists the energy and per-atom force deviations of both LAMMPS paths against ASE, including the 8,000-atom $W=4$ cells on 4×A100-40GB and 4×H200 and the 22,000-atom $W=12$ cell on 12-tile Intel.

Capacity

With checkpointing off, retained activations set the cell size. On one 8×H200 node (141 GB/device) the ceiling is 47,000 NaCl atoms ($N=18$), with $N=19$ (55,000 atoms) exhausting device memory at artifact export (137.7 GiB committed when a 512 MiB allocation failed); the cell that fits finishes at 105.75 GiB allocated. That is the largest strain-gradient NPT cell we report, 5.8× the tight-memory A100 ceiling and 2.1× the 12-tile Intel ceiling. On one 4×A100-40GB node, 8,000 NaCl atoms finish a ten-step 300 K loop and run out of memory at 11,000. Export traces a fixed graph, so a longer or hotter NPT can exceed that cap at the same atom count; a production cell is therefore smaller: 8,000 NaCl atoms exceed the cap on the first step of an 800 K ramp, while 6,000 complete the 100 ps ramp. Neighbor density depends on chemistry, so atom count is not the right unit. CUDA managed memory thrashed on this path, and edge chunking does not raise NPT capacity, because retained activations bind it.

Throughput. At full-node ceilings the ten-step loop rates, which are not production throughputs, are 0.210 ns/day on one 4×A100-40GB node (8,000 atoms), 0.10 ns/day on one 12-tile Intel node (22,000 atoms), and 0.157 ns/day on 8×H200 (47,000 atoms) (Figure 3), with >99.9% of loop time in the pair call. The same H200 cell over 10,000 steps at 300 K is 0.149 ns/day. The production runs are the nine finished 100 ps ramps of Sec. VII, on cells smaller than these ten-step ceilings. A wall-day at the ten-step rates is tens of picoseconds, which is sufficient for lattice parameters, radial distribution functions (RDFs) and coordination, and not for diffusion or long-time kinetics.

Device count governs capacity, not throughput

Within a node, device count $W$ sets the cell that fits, not the loop time. The $[N,C]$ node tensor is replicated and the backward pass carries an in-graph all-reduce (Sec. VI-F), so adding tiles raises the checkpoint-off NPT ceiling rather than shortening the step; retained activations set the per-device ceiling, and we choose $W$ to fit the cell.

The resulting node-local NPT ceilings are these. On one 4×A100-40GB node the one-GPU cell is NaCl $N=6$ (2,000 atoms). On one 12-tile Intel node the node ceiling is $N=14$ (22,000 atoms), which fits only at $W=12$ (step-0 energy −74141.205 eV, pressure 11,032.4 bar); the same atoms at $W=8$ run out of memory. On H200 a single device holds $N=9$ (6,000 atoms) and the full 8×H200 node reaches 47,000 atoms. Relative to one device, the regime of the ASE calculator, graph parallelism therefore raises the ceiling 4.6× on A100 (2,000 to 8,000 atoms) and 8.0× on H200. Step-0 and step-10 thermo agree across every $W$ that fits, with both NCCL and oneCCL collectives, so device count changes capacity without changing the trajectory.

Because the step is replicated, adding tiles does not shorten the loop on small cells (one tile is fastest at $N=6$). It begins to help once the cell is large enough to amortize the collective: on the 6,000-atom H200 cell, four devices cut the loop to 2.76 s from 5.27 s at $W=1$, which is 1.91×, 48% of ideal. Eight devices on that cell take 4.12 s, slower than four. The cost of each extra tile is per-rank work that grows with the cell, not the barostat; Sec. VI-F gives the phase split.

Where the step goes

On the 12-tile Intel node the profiled checkpoint-off cell is $N=14$, $W=12$. Its ten-step loop is 8.47 s and the steady step is 820 ms (mean of the last nine force calls): graph build 127 ms, of which 118 ms is the vesin neighbor list; forward 270 ms, which holds the node all-gathers; and backward 422 ms, which holds the in-graph all-reduces. The drain-and-barrier split of Sec. VI-B measures the collective share on the 4,000-atom cell: 0.5–2.1% of the step on the Intel node and 2.0–4.5% on A100. On that cell the step is compute, not communication.

We did not repeat the split on the 22,000-atom cell or on the multi-node cells, so we quote no collective share for them. The post-graph force and virial sums are one $[N,3]$ FP64 tensor and one $[3,3]$ cell gradient, 24$N$ and 72 bytes (about 0.53 MB at 22,000 atoms). These sums take about 1 ms on the Intel node and under 1 ms on A100 at every layout; the force all-reduce takes 0.12 ms on one 4×A100-40GB cell (8,000 atoms) and 0.21 ms on one 8×H200 cell (47,000 atoms). They do not account for the two-node cost (Sec. VI-B).

A FairChem graph-parallel flop counter (PyTorch FlopCounterMode on the 8,000-atom A100 cell ($N=10$, $W=4$) counts 2.006 × $10^{12}$ FLOPs per rank for one forward-plus-backward predict. Dividing by the pair_style uma last step on that cell, 408 ms, gives 4.92 × $10^{12}$ FLOPs per GPU, 51% of the documented 9.7 TFLOP/s FP64 peak of an A100-40GB SXM [31]. That quotient is a labeled lower bound on the Python graph-parallel predictor, not a hardware counter in pair_style uma: the 1-GPU profiler returns 0 counted FLOPs on $N=6$ because UMA’s custom kernels are missing from the flop table. The same counter on the 47,000-atom H200 cell ($N=18$, $W=8$) counts 5.847 × $10^{12}$ FLOPs per rank (measured 4.678 × $10^{13}$). The pair_style uma last step on that cell is 548 ms, so the quotient is 1.07 × $10^{13}$ FLOP/s per GPU, 31% of the documented 34 TFLOP/s FP64 peak of an H200 SXM [32], labeled the same way. The two cells are not the same size, so the count is not transferred to the 12-tile Intel ceiling.

The demonstration workload (100 ps NPT heating ramps of NaCl, equiatomic Pt–Ru, and equiatomic (HfNbTaTiZr)N from one omat checkpoint) is the subject of Sec. VII. Those ramps use the checkpoint-off NPT path whose capacity and cost are measured above; they are not melting-point determinations. Section VII-B states what the jumps are and what the experimental marks are.

Heating ramps of three chemistries

One omat checkpoint drives NPT MD for three chemistries that have no shared classical potential (a salt, a binary alloy, and a five-cation nitride) with no refit. That single-checkpoint coverage is the systems result; the heating ramps are how we exercise the NPT path of Sec. VI. This section reports the protocol and nine jumps. Section VII-B is the accuracy discussion: the jumps are not melting points.

Protocol and observed jumps

The same omat checkpoint heats three crystals on one isotropic NPT path: NaCl, equiatomic Pt–Ru, and equiatomic rocksalt (Hf$_{0.2}$Nb$_{0.2}$Ta$_{0.2}$Ti$_{0.2}$Zr$_{0.2}$)N. The Zhou–Johnson–Wadley embedded-atom table covers 16 metals and no salt [17]. The nitride cations are taken from special-quasirandom occupations of that composition [23]. Each production run is 100 ps (100,000 steps of 1 fs) at 1 bar, with thermostat and barostat damping of 0.1 ps and 1.0 ps. Checkpointing stays off, because the barostat uses the strain–gradient virial. The ramps are 6 K/ps for NaCl, 9 K/ps for Pt–Ru, and 20 K/ps for the nitride. The signature is the jump in potential energy per atom and in density. Figure 4 is density versus temperature for the nine ramps.

tbl:II lists the nine jumps: the density drop and the temperature at which it occurs, which stays within a window of at most 120 K over the drop.

The cells in Figure 4 are the production ramps. Neighbor density depends on chemistry, so each ramp uses a cube that stays under the export cap as the neighbor list grows: Pt–Ru 15$^{3}$ (14,000 atoms) aborted at step 0 on 4×H200 because two ranks held about 197,000 edges each, so that node ran a 15×15×14 cell (13,000 atoms). Three ramps stopped once on a filesystem write and were rerun to completion. The A100 cubes are NaCl 9$^{3}$, Pt–Ru 10$^{3}$, and nitride 7$^{3}$; the Intel cubes are NaCl 14$^{3}$, nitride 10$^{3}$, and Pt–Ru 15$^{3}$. All nine 100 ps ramps finished. The nitride cells use special-quasirandom equiatomic occupations of Hf, Nb, Ta, Ti, and Zr on the cation sublattice, built from twelve tiles at the rocksalt (B1) lattice constant $a = 4.4845$ Å, with no new fit and no new tiling.

System	Node	Atoms	$\rho$ (g/cm$^3$)	Jump (K)
NaCl	4×A100	6,000	1.65→1.28	1070
NaCl	4×H200	18,000	1.64→1.31	1090
NaCl	12-tile Intel	22,000	1.63→1.31	1090
Pt–Ru	4×A100	4,000	15.2→14.2	2470
Pt–Ru	4×H200	13,000	15.0→14.2	2500
Pt–Ru	12-tile Intel	14,000	15.2→14.2	2480
(HfNbTaTiZr)N	4×A100	3,000	7.82→6.99	3870
(HfNbTaTiZr)N	4×H200	8,000	7.69→7.03	3870
(HfNbTaTiZr)N	12-tile Intel	8,000	7.81→7.03	3870

Accuracy of the jumps

The nine jumps are where a defect-free NPT ramp fails. They are not melting points of the crystals, of the potential, or of experiment. Four facts about the protocol set that number, and the experimental marks are context, not a score.

A periodic, defect-free crystal with no solid–liquid interface superheats [26]. Homogeneous nucleation has to create the liquid inside the bulk; the barrier is large on molecular-dynamics timescales, so the solid survives above the equilibrium melting temperature of the potential. The accepted way to measure that equilibrium temperature is a two-phase coexistence simulation [27]. We run neither coexistence nor a cooling branch. The ramps are fast (6, 9, and 20 K/ps), and the heating rate extra-superheats the solid. The cells have no free surface, grain boundary, dislocation, or vacancy to nucleate against, and they are small (3,000–22,000 atoms); homogeneous nucleation scales with volume, so a smaller cube typically fails later. Each reported jump is the density drop over about 1 ps while temperature is still moving, not a free-energy crossing.

We have no coexistence $T_m$ for uma-s-1p2/omat [1], so we cannot split ramp superheat from model-versus-experiment error.

The experimental marks in Figure 4 are therefore not an accuracy score for the engine. NaCl melts at 1074 K (801 °C) [24]. The three NaCl jumps (1070–1090 K) lie near that number by proximity only. Platinum melts at 2041 K (1768 °C) and ruthenium at 2607 K (2334 °C) [24]. At 50 at.% Ru the Massalski experimental diagram, reproduced as Fig. 6(b) of Cornish et al. [28], is a solidus–liquidus lens: about 2337 K (2064 °C) solidus and 2413 K (2140 °C) liquidus, read from the published axes. The 1800 to 2700 K ramp brackets the interval; all three Pt–Ru jumps (2470–2500 K) lie above the liquidus, which is the expected sign for a superheated homogeneous failure.

Equiatomic rocksalt (HfNbTaTiZr)N has no directly measured melting point. Transition-metal nitrides lose N$_2$ and decompose before they melt congruently at ambient pressure; even uranium nitride needs 0.25 MPa of N$_2$ to melt at 3120 ± 30 K [25]. The MD cell is closed, stoichiometric, and

held at 1 bar, with no N-vacancy population, so a defect-free ramp cannot reproduce a nitride melting experiment even in principle.

The literature that does bound this composition is consistent with the jump being a high superheated failure, not a $T_m$. Bulk (HfNbTaTiZr)N survives spark-plasma-sintering densification at 2473 K (2200 °C), a solid-state lower bound [30]. A heating/cooling hysteresis simulation of this same five-cation nitride places its melting temperature at 2900 ± 50 K [29], and the equiatomic rule-of-mixtures of the five binary endmembers (HfN, ZrN, TiN, TaN, NbN) gives ≈3190 K [24]. All three MD jumps (≈3870 K) lie above every one of these anchors: above the hysteresis-derived $T_m$ by about 1000 K and above the binary average by about 700 K. That ordering is the expected sign of homogeneous failure in a small, defect-free, fast ramp; it is not a melting-point determination.

Limitations

Multi-node speed-up of a single configuration is not delivered, and Sec. VI-B is why: per-rank work is $O(N)$ in the global atom count, so a second node adds memory but not speed. Multi-node data parallelism over independent configurations is available and untested at scale. Checkpointing stays off on the NPT path because we have not validated the cell gradient through our recompute functions, so the cell size is set by retained activation memory. Because tracing bakes partition offsets, $N$ and $W$ are fixed at export: no dynamic load balancing, no resizing mid-run. The graph is rebuilt every step with no reuse, correct but wasteful for small displacements. Across nodes the step grows with the cell: every rank rebuilds the full-cell neighbor graph, which is 15% of the step on one node and 41% on 32 (Sec. VI-B), and forward and backward grow with $N$ as well. Each row of that sweep is one run on rank 0, and the 2% boundary cost and the 1% NPT-versus-NVT difference are single measurements per layout, with no run-to-run variance. The collective split exists only at 4,000 atoms, so the compute-bound reading does not extend to the multi-node cells. The multi-device path requires precision double, so precision cannot currently be traded for capacity. The released model offers only single or double precision, so raising the NPT ceiling that way would mean validating an FP32 path against the FP64 one, which we have not done; it is future work. We also make no per-atom speed comparison with the spatially decomposed MLIPs of § 3. Code and scripts will be released at a later date. Wall-day trajectories are tens of picoseconds: enough for lattice constants, RDFs and coordination, not for diffusive annealing or catalytic turnover. The heating-ramp jumps are not melting points; the protocol and experimental marks are in Sec. VII-B.

measurement: crossing a node boundary at equal world size adds about 2% to the NPT step on a 4,000-atom cell, and doubling devices does not shorten it. Three results hold within a node. The strain-gradient virial adds 0.2% per step against NVT. The results do not depend on the device count or on whether NCCL or oneCCL provides the collectives, and the FP64 path agrees with ASE to ∼ $10^{-14}$ eV/Å, eight orders of magnitude below the residual caused by the wrap’s FP32 coordinates. In our implementation NPT costs capacity rather than time: ten-step NPT loops hold 8,000, 22,000, and 47,000 NaCl atoms on the three node types (the Intel-node figure is 20 times below the checkpoint-on NVT ceiling), and capacity reaches 779,000 atoms on 32 Intel nodes at 0.007 ns/day. The pair style runs on NVIDIA and Intel GPUs from one source tree. It trades speed per atom for one checkpoint; with it, NPT ramps of a salt, an alloy, and a nitride run with no refit. Two directions follow: carrying the cell gradient through checkpoint recompute to close the 20-fold gap, and recovering multi-node speed through a shared neighbor build and, further out, a depth-$L$ halo exchange.

Acknowledgments

This work was supported by Laboratory Directed Research and Development (LDRD) funding from Argonne National Laboratory, provided by the Director, Office of Science, of the U.S. Department of Energy under Contract No. DE-AC02-06CH11357. This research used resources of the Argonne Leadership Computing Facility, which is a U.S. Department of Energy Office of Science User Facility operated under contract DE-AC02-06CH11357. This research used the Delta advanced computing and data resource which is supported by the National Science Foundation (award OAC 2005572) and the State of Illinois. Delta is a joint effort of the University of Illinois Urbana-Champaign and its National Center for Supercomputing Applications.

We used Claude Sonnet 5.5 [45] to support article writing and editing. The authors reviewed the text and are responsible for its content.

Conclusion

Foundation MLIPs occupy a different point in the parallel design space from the custom-made potentials that have reached $10^8$–$10^{10}$ atoms. Their generality comes from depth and global couplings, which force graph parallelism. That makes one node the unit of speed, which we support by