DeePMD-kit Model Formats#

In the pt_expt (PyTorch experimental) backend, three file formats are in use, each mapping to a distinct inference path, neighbor-list strategy, and deployment stack. Understanding which format a model is in is critical — loading the wrong format through the wrong code path can produce silently incorrect results.

The .pt and .pt2 formats are the active focus of current pt_expt development; .pth is the legacy TorchScript path from the stable pt backend.

.pt — Training Checkpoint#

A .pt file is a standard PyTorch state-dict checkpoint produced by the pt_expt trainer (torch.save({"model": wrapper.state_dict()})). It stores trainable weights and model hyperparameters in _extra_state["model_params"].

When loaded by deepmd/pt_expt/infer/deep_eval.py (_load_pt), the checkpoint is reconstructed into a live model, but the eager runner installs model.forward_common_lower(...) as its call target — the dense lower path 1.

For most models this is correct. However, for graph-eligible DPA1 (descriptor type: dpa1 or se_atten_v2 with attn_layer: 0), the model's public forward() instead routes through call_common → graph-native path in deepmd/pt_expt/model/ener_model.py. The two paths are not numerically equivalent when descriptor mean statistics (davg) are nonzero 1. See the DPA1 path bug section below.

.pt2 — AOTInductor Graph Export#

A .pt2 file is an Ahead-of-Time Inductor (AOTInductor) compiled artifact produced by torch.export + aot_compile, packaged as a ZIP archive containing compiled kernels and metadata (lower_input_kind, has_message_passing, etc.) 2.

Key properties:

• Dynamic edge axis — one artifact evaluates any system size (proven across multiple system sizes at ≤1e-10 parity) 2.
• Multi-rank variant — message-passing models (DPA2, DPA3) embed a model/extra/forward_lower_with_comm.pt2 sub-artifact with an 8-tensor comm ABI for MPI ghost exchange 3.
• Graph-native path — calls forward_lower_graph_exportable / forward_lower_graph_exportable_with_comm on EnergyModel, not the dense lower.

Export Gate#

Export is gated by model_uses_graph_lower in deepmd/pt_expt/utils/serialization.py, which requires "energy" in the model's atomic output definition 4. Non-energy models (dos, dipole, polar) currently have no deployment consumers for .pt2 and the gate blocks export rather than produce an unusable file 4.

Loading#

• Python: torch._inductor.aoti_load_package()
• C++/LAMMPS: DeepPotPTExpt via torch::inductor::AOTIModelContainerRunnerCpu

Neighbor-List Strategy by Format#

Format choice determines which neighbor-list subsystem is used 5:

nlist-form (.pt / Python/ASE)#

Uses the NeighborList strategy, injected at forward_common/call_common. Configured via nlist_backend:

The O(N²) native baseline extends coords into ~27 periodic-image regions via extend_coord_with_ghosts in deepmd/dpmodel/utils/nlist.py 5.

graph-form (.pt2 / AOTI / C++)#

Uses the NeighborGraph subsystem, controlled by neighbor_graph_method 5:

There is no "auto" selector for the graph path — callers must specify explicitly 5.

C++ LAMMPS Bottleneck#

The C++ DeepPotPTExpt::compute() path builds edges via createEdgeTensors() — a CPU C++ loop over all neighbor pairs — then copies results to GPU with multiple .clone().to(device) calls. At 65,536 atoms this overhead accounts for 3–6× the model inference time itself, negating the AOTI compilation speedup 6. The Python path avoids this by building edges GPU-natively via VesinNeighborList or the nv backend. See issue #5595 and issue #5574 for ongoing fix work.

DPA1 Path Bug (.pt Checkpoint)#

Issue #5862 (open as of 2026-07-18): DeepPot(model.ckpt.pt) on a graph-eligible DPA1 model silently routes through the dense lower path instead of the graph-native path, producing catastrophically wrong outputs 1.

Trigger conditions (all required):

• pt_expt backend
• Descriptor type: dpa1 or se_atten_v2 with attn_layer: 0
• set_davg_zero: false (default) — nonzero davg
• Loaded via DeepPot(*.pt) or dp --pt-expt test -m model.ckpt.pt

Root cause: The dense DPA1 body (deepmd/dpmodel/descriptor/dpa1.py) retains a phantom padding-neighbor residual of -davg/dstd for all vacant neighbor slots. Models with large sel (e.g., sel=416 for OMat24) have most dense slots filled by padding, massively amplifying the residual. The graph path omits padding neighbors entirely 1.

Observed magnitude on an OMat24 checkpoint:

1

Related fix: PR #5785 (merged 2026-07-16) fixed DescrptDPA1.call routing via _call_graph_adapter, but issue #5862 specifically tracks the checkpoint-loading failure where _load_pt unconditionally installs the dense lower runner 7. The structural fix must make DeepPot(.pt) select the graph-native path for graph-eligible DPA1 checkpoints. The recommended regression guard is a graph-eligible DPA1 fixture with nonzero davg in source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py.

Deployment Compatibility Summary#

Key limitations to be aware of 4:

• .pt2 is energy-only — the model_uses_graph_lower gate in serialization blocks export for dos/dipole/polar models because no C++ consumer exists for those outputs.
• DeepTensor has no pt_expt support — dipole/polar models cannot use either dense or graph .pt2 schema in pt_expt.
• LAMMPS consumes energy models only, via DeepPotPTExpt.

For lifting these restrictions, see companion issues #5805 (training-side generalization) and #5806 (deployment consumers by output type).

Key Source References#