DPA1 Inference — Correctness Bugs in pt_expt Backend#

Two distinct bugs affect DPA1 inference in the PT-experimental (pt_expt) backend. Both involve graph-eligible DPA1 models (i.e., attn_layer: 0, se_atten_v2/dpa1 descriptor type) but manifest in different scenarios.

Bug 1 — Wrong Outputs from .pt Checkpoint Inference (Issue #5862)#

Root cause: DeepPot(model.ckpt.pt) routes inference through the dense lower (forward_common_lower) instead of the graph-native path used by the model's public forward() 1.

The raw .pt loader in deepmd/pt_expt/infer/deep_eval.py::_load_pt reconstructs the model but installs an eager runner that calls model.forward_common_lower(...). For graph-eligible DPA1, model.forward(...) instead reaches call_common → graph-native route. These paths are not numerically equivalent when descriptor statistics are nonzero 1.

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

Observed magnitude on a real OMat24 checkpoint:

1

Trigger conditions:

• PT-experimental backend
• type: dpa1 or se_atten_v2 descriptor with attn_layer: 0
• set_davg_zero: false (default) — nonzero descriptor mean statistics
• Loaded via DeepPot(*.pt) or dp --pt-expt test -m model.ckpt.pt

Key files:

• deepmd/pt_expt/infer/deep_eval.py — _load_pt / _eager_runner: installs dense lower path
• deepmd/pt_expt/model/ener_model.py — EnergyModel.forward → call_common: graph-native route
• deepmd/dpmodel/descriptor/dpa1.py — documents the dense padding-neighbor residual

Status: PR #5785 (merged 2026-07-16) fixed a related routing issue (DescrptDPA1.call no longer routes through _call_graph_adapter). Issue #5862 specifically targets the checkpoint-loading failure; the structural fix must make DeepPot(.pt) select the graph-native path for graph-eligible DPA1 checkpoints. The recommended regression guard is extending source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py with a graph-eligible DPA1 fixture whose davg is nonzero 1.

Bug 2 — f_use_norm Training Failure on NPY Data (Issue #5813)#

Root cause: Shape mismatch between force predictions and force labels when f_use_norm: true is used with an NPY training system and a graph-eligible DPA1 model 2.

The graph lower returns force as shape (nf, nloc, 3). Legacy NPY (DeepmdDataSystem) supplies force labels as shape (nf, nloc * 3) (flat). Most force-loss paths in deepmd/dpmodel/loss/ener.py::EnergyLoss.call first flatten both tensors into diff_f, making the shape difference invisible. The f_use_norm branches instead subtract the original (unflattened) tensors before reshaping, triggering a shape broadcast error 2:

RuntimeError: The size of tensor a (576) must match the size of tensor b (3)
    at non-singleton dimension 2
  File "deepmd/dpmodel/loss/ener.py", line 445
    diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))

Trigger conditions (all must be true):

• PT-experimental backend
• type: dpa1 with attn_layer: 0 (graph-eligible)
• NPY training data — does not reproduce with LMDB (LMDB normalizes force shape to (nloc, 3) on read)
• loss.f_use_norm: true combined with loss_func: "mae" or use_huber: true

Key files:

• deepmd/dpmodel/loss/ener.py — EnergyLoss.call, f_use_norm branches subtract original tensors before reshape
• deepmd/pt_expt/train/training.py — _CompiledModel._forward_graph exposes force as (nf, nloc, 3)
• deepmd/dpmodel/utils/batch.py — normalize_batch preserves the legacy flat NPY label shape

Fix direction: Derive every f_use_norm force difference from the already-flattened diff_f (reshaped to (_nf, _nloc, 3) where needed) rather than re-computing force_hat - force on the original layouts 2.

Relationship Between the Two Bugs#

Both bugs involve graph-eligible DPA1 (attn_layer: 0) in pt_expt and are rooted in the dense ↔ graph-native path transition:

Key References#

• Issue #5862 — .pt checkpoint wrong-output report with standalone reproducer 1
• Issue #5813 — f_use_norm + NPY crash report 2
• PR #5785 — Related merged fix: DPA1 dense-call routing via _call_graph_adapter
• PR #5294 — Introduced f_use_norm to EnergyLoss 3
• Force Loss Computation (KB) — Explains f_use_norm, relative_f, and EnergyStdLoss shape flows