[BUG] pt_expt DeepPot(.pt) uses dense lower for graph-native DPA1 and returns incorrect outputs

[OutisLi]

Bug summary

dp --pt-expt test -m model.ckpt.pt can return catastrophically incorrect energy, force, and virial predictions for a graph-eligible PT-experimental DPA1 checkpoint, although the checkpoint's direct public model.forward(...) predictions and in-training validation are correct.

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(...), forcing the legacy dense neighbor-list path. A graph-eligible DPA1 model's public forward(...) instead reaches the graph-native call_common path by default. These paths are not numerically equivalent when descriptor statistics are nonzero: the dense DPA1 implementation retains a padding-neighbor -davg/dstd residual, as already documented in deepmd/dpmodel/descriptor/dpa1.py. Large sel values amplify the discrepancy because most dense slots are padding.

This violates the stated .pt inference contract and the intent of source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py: DeepPot(.pt) should reproduce the source model's public forward outputs. That test currently uses only se_e2_a, so it does not exercise graph-eligible DPA1 routing.

On a trained OMat24 checkpoint (se_atten_v2, attn_layer=0, sel=416), the same five validation frames gave:

Direct public model.forward (graph): Energy MAE 0.03955 eV, Force MAE 0.02755 eV/Å
DeepPot/checkpoint dp test (dense):   Energy MAE 2600.81 eV, Force MAE 12.97899 eV/Å

Type maps matched. Single-frame and batched inference agreed with each other, and native and Vesin neighbor-list builders also agreed, isolating the discrepancy to graph public forward versus the dense lower selected by raw-checkpoint DeepEval.

DeePMD-kit Version

• Clean Python source checkout: 1be0082cceff7934a88874af6daa0fb8fa60b9f3
• Installed package metadata: v3.2.0b1.dev110+g8399520c5.d20260704

The minimal reproducer below sets DP_CUDA_INFER=0 and exercises only eager Python/PyTorch paths; it does not invoke compiled graph custom operators.

Backend and its version

• Backend: PyTorch Experimental (pt_expt)
• PyTorch: 2.11.0+cu128
• CUDA: 12.8
• Python: 3.13.13

How did you download the software?

Built from source.

Input Files, Running Commands, Error Log, etc.

The production command was:

dp --pt-expt test \
  -m models/model.ckpt.pt \
  -s /path/to/val.lmdb

The following standalone reproducer requires no external model or dataset. It constructs a small graph-eligible DPA1 model, sets deterministic nonzero descriptor statistics, saves a normal PT-experimental training checkpoint, and compares its public forward with DeepPot(.pt):

import copy
import os
import tempfile

import numpy as np
import torch
from deepmd.infer import DeepPot
from deepmd.pt_expt.model import get_model
from deepmd.pt_expt.train.wrapper import ModelWrapper
from deepmd.pt_expt.utils.env import DEVICE

params = {
    "type_map": ["H", "O"],
    "descriptor": {
        "type": "dpa1",
        "sel": 20,
        "rcut_smth": 0.5,
        "rcut": 4.0,
        "neuron": [3, 6],
        "axis_neuron": 2,
        "attn": 4,
        "attn_layer": 0,
        "smooth_type_embedding": True,
        "set_davg_zero": False,
        "type_one_side": True,
        "precision": "float64",
        "seed": 1,
    },
    "fitting_net": {
        "type": "ener",
        "neuron": [8, 8],
        "precision": "float64",
        "seed": 1,
    },
}

model = get_model(copy.deepcopy(params)).to(torch.float64).to(DEVICE).eval()
with torch.no_grad():
    model.atomic_model.descriptor.se_atten.mean.fill_(0.01)
    model.atomic_model.descriptor.se_atten.stddev.fill_(0.1)

coords = np.array(
    [[[1.0, 1.0, 1.0], [2.0, 1.0, 1.0], [1.0, 2.0, 1.0]]]
)
cells = np.eye(3).reshape(1, 9) * 10.0
atom_types = np.array([0, 1, 0], dtype=np.int32)

coord_t = torch.tensor(
    coords, dtype=torch.float64, device=DEVICE
).requires_grad_(True)
atype_t = torch.tensor(
    atom_types.reshape(1, -1), dtype=torch.int64, device=DEVICE
)
cell_t = torch.tensor(cells, dtype=torch.float64, device=DEVICE)
reference = model.forward(coord_t, atype_t, cell_t)

fd, checkpoint = tempfile.mkstemp(suffix=".pt")
os.close(fd)
try:
    wrapper = ModelWrapper(model, model_params=params)
    torch.save({"model": wrapper.state_dict()}, checkpoint)
    energy, force, _ = DeepPot(checkpoint, auto_batch_size=False).eval(
        coords, cells, atom_types
    )
finally:
    os.unlink(checkpoint)

reference_energy = reference["energy"].detach().cpu().numpy()
reference_force = reference["force"].detach().cpu().numpy()
print("public graph energy:", reference_energy.reshape(-1))
print("DeepPot .pt energy:", energy.reshape(-1))
print("energy max abs delta:", np.max(np.abs(energy - reference_energy)))
print("force max abs delta:", np.max(np.abs(force - reference_force)))

Run with:

DP_CUDA_INFER=0 python reproduce.py

Observed output:

public graph energy: [2.42767467]
DeepPot .pt energy: [2.33223403]
energy max abs delta: 0.09544063670727354
force max abs delta: 0.18991794401885587

Expected behavior: all outputs from DeepPot(checkpoint) match the source model's public forward(...) within numerical tolerance.

Steps to Reproduce

1. Use a current clean deepmd-kit checkout with the PT-experimental backend.
2. Save the script above as reproduce.py.
3. Run DP_CUDA_INFER=0 python reproduce.py.
4. Observe that raw-checkpoint DeepPot does not match the source model's public graph forward.

Further Information, Files, and Links

The relevant control flow is:

• deepmd/pt_expt/infer/deep_eval.py::_load_pt defines _eager_runner using model.forward_common_lower(...), which is the dense lower.
• deepmd/pt_expt/model/ener_model.py::EnergyModel.forward calls call_common(...); PT-experimental graph-eligible DPA1 resolves this to the graph-native route.
• deepmd/dpmodel/descriptor/dpa1.py documents that the dense body leaks a phantom padding-neighbor -davg/dstd residual when davg != 0, whereas the graph path omits it.

A structural fix should make raw .pt DeepEval preserve the source model's public-forward semantics, selecting the graph-native route for graph-eligible DPA1 checkpoints instead of unconditionally adapting every .pt checkpoint to the dense lower ABI. A regression test can extend test_deep_eval_pt_checkpoint.py with a graph-eligible DPA1 fixture whose descriptor mean is nonzero and assert energy/force/virial parity with model.forward(...).

[OutisLi]

Resolved by #5912, which preserves graph routing for pt_expt raw checkpoints. #5975 addresses the complementary case by preserving dense lower semantics during backend conversion.