[Q&A] PyTorch backend compression of converted TF se_e2_a models causes unstable LAMMPS MD #5438

Hi @Dm0216, thanks for the very detailed report. I tried to reproduce on a fresh fp64 se_e2_a model (TF type_one_side=False, 2 types, 100-step training on the example water dataset) — dp convert-backend → dp --pt compress works cleanly there, with compressed vs uncompressed forces matching to ~3e-14 across configurations whose minimum pair distance is 0.9 Å. So the pt compression code is not categorically broken; the failure looks specific to either your 5-type weight set or to a runtime configuration LAMMPS exercises that pure-Python eval does not.

Suggested quick workaround first: please verify that your model is type_one_side=True, and if it isn't, re-train (or fine-tune) with type_one_side=True. Convert-backend and compression both work much more reliably on the type_one_side=True path — there is one embedding network per neighbor type instead of one per (center, neighbor) pair, so the tabulation builds fewer, larger-coverage tables that are less sensitive to per-pair boundary effects. Many TF legacy models default to type_one_side=False; flipping it for the new train usually has negligible accuracy impact for energy/force prediction. This may be the cleanest fix in your situation.

If you'd like to understand the root cause regardless, the most likely structural cause for the unstable MD is runtime ss exceeding the table's trained upper bound:

• During compression, the table upper bound is computed from the model's stored min_nbor_dist (set at training time):
  s_upper = ((1/min_nbor_dist) * sw - davg_s) / dstd_s
• At runtime, ss = (1/r * sw - davg_s) / dstd_s. If LAMMPS hands the model an atom pair with r < min_nbor_dist (initial config, bad starting structure, or a thermal spike at the very first step), then ss > s_upper and the lookup falls into the extrapolation region.
• The polynomial spline in the extrapolation region is fit to embedding-network output at xx values the network never saw during training, so both the spline values and their derivatives are essentially unconstrained — forces of arbitrary magnitude are possible. The 75 eV/Å your LAMMPS log shows is consistent with this failure mode but doesn't by itself prove it; the Python diagnostic below distinguishes.

The 7.6 GB file from -s 0.001 -e 10 failing to load is an unrelated LAMMPS-side libtorch deserialization issue (file too large for the legacy load path); the bound issue is what's causing the unstable MD with the default-stride compressed model.

Quick diagnostic to confirm and localize — please run before LAMMPS, in pure Python, against the exact starting structure LAMMPS reads. Important: if your LAMMPS run uses periodic boundary conditions (the default for boundary p p p), pass the cell to DeepPot.eval and compute the minimum distance with PBC accounted for; otherwise omit the box and use plain distances. Could you also confirm whether the run is PBC or non-periodic?

import numpy as np
from ase import Atoms
from deepmd.infer import DeepPot
# coord (natoms, 3), atype (natoms,), box (3, 3) from your LAMMPS data file
# Set pbc=True/False matching your LAMMPS `boundary` setting.
USE_PBC = True
dp_un = DeepPot("model5.1.pth")
dp_co = DeepPot("model5.1_compressed.pth")
box_arg = box.reshape(1, 9) if USE_PBC else None
e_un, f_un, _ = dp_un.eval(coord.reshape(1, -1), box_arg, atype)
e_co, f_co, _ = dp_co.eval(coord.reshape(1, -1), box_arg, atype)
print("uncomp |F| max:", float(np.max(np.abs(f_un))))
print("comp |F| max:", float(np.max(np.abs(f_co))))
print("max force diff:", float(np.max(np.abs(f_un - f_co))))
# PBC-aware minimum-image pair distance (ase handles the cell properly)
atoms = Atoms(
 positions=coord.reshape(-1, 3),
 cell=box.reshape(3, 3) if USE_PBC else None,
 pbc=USE_PBC,
)
d = atoms.get_all_distances(mic=USE_PBC)
np.fill_diagonal(d, np.inf) # mask self-distances
print("min pair distance (mic):", float(d.min()))
# And the min_nbor_dist baked into the .pth itself
import torch
m = torch.jit.load("model5.1.pth", map_location="cpu")
print("model min_nbor_dist: ", float(m.min_nbor_dist))

Three outcomes tell us where to look:

1. comp |F| max matches the 75 eV/Å number and is much larger than uncomp |F| max → the bug is reproducible from Python eval on this single frame. If min pair distance (mic) is below model min_nbor_dist, the over-extrapolation hypothesis above is the cause. Workarounds: pre-relax with the uncompressed model, or freeze the model with a smaller min_nbor_dist (use dp neighbor-stat on a wider training set or set it manually before compress).

2. Diffs are O(1e-3) or larger but forces are reasonable → genuine compression accuracy regression for your specific weights. In that case it would be very useful if you could attach model5.1.pth and the read_data data file so we can reproduce locally; the compressed model is fully self-contained and reproduces deterministically.

3. Diffs are O(1e-10) or smaller → the bug is LAMMPS-side (pair_deepmd.cpp interpretation of compressed PT models), not in the compression itself. We'd then investigate the C++ dispatch path.

For reference, the internal indexing/bound calculations I traced in pt-side compression appear correct (table net-name ↔ embedding_net_nodes mapping, runtime (center, neighbor) ↔ embedding_idx mapping, and _get_env_mat_range extracting the s-component bound after the min/max collapse). The cosmetic over-extension of upper to the sx/sy/sz axis bounds (visible in the compression log as a 13.13 upper instead of the s-component's ~7.7) does not actually corrupt runtime values, because the table polynomials at queried indices are still computed from in-distribution vv/dd/d2.