[BUG] pt_expt .pt2 of a DPZBL model (nlist lower) rejects inputs without ghost atoms: Expected Ne(s54, s96)

Open
  • Is something wrong?
#6,087 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
55/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
backend, machine-learning

Research direction

Start by running the provided repro.py from the repository root with CUDA_VISIBLE_DEVICES= python repro.py; it freezes a DPZBL model and fails when nall == nloc. Then inspect deepmd.pt_expt.utils.serialization._trace_and_export and _strip_shape_assertions, which the issue identifies as relevant entry points. Done means the exported model accepts inputs without ghost atoms while preserving correct inference; add a regression check for the supplied reproduction.

Written by the indexing model from the issue text.

Description

Bug summary

A DPZBL model (DPA1 with a pair table) frozen to .pt2 through the pt_expt backend, with the nlist lower, fails at inference whenever the extended system has no ghost atom within the cutoff, so that nall == nloc:

RuntimeError: Expected Ne(s54, s96) to be True but received 6

That is the normal situation for a non-periodic input (a molecule or cluster), and for a periodic cell in which no image lies within the cutoff.

The exported program carries runtime assertions that require nall != nloc. A plain DPA1 model exported the same way carries none of them, and it evaluates the same input correctly. So on current master, the guard comes with the pair-table composition. #5502 reported the same class of guard for DPA4/SeZM.

DeePMD-kit Version

master (668b60cc2)

Backend and its version

PyTorch exportable backend (pt_expt); PyTorch 2.13.0+cu130, Python 3.13; reproduced on CPU

How did you download the software?

Built from source

Input Files, Running Commands, Error Log, etc.

Save the script below as repro.py and run it from the repository root, so that the relative use_srtab path resolves. Use CUDA_VISIBLE_DEVICES= python repro.py:

import numpy as np

from deepmd.infer import DeepPot
from deepmd.pt_expt.model import get_model
from deepmd.pt_expt.utils.serialization import deserialize_to_file

config = {
    "type_map": ["O", "H", "B"],
    "use_srtab": "source/tests/pt/model/water/data/zbl_tab_potential/H2O_tab_potential.txt",
    "smin_alpha": 0.1,
    "sw_rmin": 0.2,
    "sw_rmax": 4.0,
    "descriptor": {
        "type": "dpa1",
        "rcut_smth": 0.5,
        "rcut": 4.0,
        "sel": 20,
        "neuron": [3, 6],
        "axis_neuron": 2,
        "attn": 4,
        "attn_layer": 2,
        "attn_dotr": True,
        "attn_mask": False,
        "activation_function": "tanh",
        "set_davg_zero": True,
        "type_one_side": True,
        "seed": 1,
    },
    "fitting_net": {"type": "ener", "neuron": [8], "seed": 1},
}
model = get_model(config).to("cpu").eval()
deserialize_to_file("dpzbl.pt2", {"model": model.serialize()}, lower_kind="nlist")

coord = np.array(
    [
        [2.5181844868, 1.7816201484, 3.2653708448],
        [3.4681844868, 1.7816201484, 3.2653708448],
        [2.1069247227, 2.3984421871, 2.5468661677],
        [3.9483803343, 3.7817931063, 3.0107818782],
        [2.3386765694, 2.7552686174, 4.0262878436],
        [3.2453983914, 3.7819603024, 1.6852189190],
    ]
)
atype = np.array([0, 1, 2, 0, 1, 2])
DeepPot("dpzbl.pt2").eval(coord.reshape(1, -1), None, atype)

The last call fails with:

RuntimeError: Expected Ne(s54, s96) to be True but received 6
Steps to Reproduce
  1. Run the script above. It freezes a small, randomly initialized DPZBL model and evaluates six atoms without a cell. DeepEval passes them with nall == nloc == 6.
  2. Each of these controls evaluates without error:
    • the same archive with cell = np.eye(3) * 5.0, where every atom has periodic images within the 4 Å cutoff;
    • the same model without the pair table (drop use_srtab, smin_alpha, sw_rmin and sw_rmax from the config), on the same six atoms, with and without a cell.
Further Information, Files, and Links

We listed the aten._assert_scalar nodes of deepmd.pt_expt.utils.serialization._trace_and_export(..., lower_kind="nlist")[0].graph_module. The DPZBL export includes four assertions that the DPA1 export lacks, where s54 is extended_coord.shape[1] (nall), s96 is nlist.shape[1] (nloc) and s77 is an internal size of the program:

s96 <= s77
Ne(s77, s96)
s96 <= s54
Ne(s54, s96)

Our reading, which we have not traced to a source line: these are relations that torch.export records from the trace's sample inputs, which have ghost atoms (nall > nloc), as described in #5502. The DPZBL archive evaluates without error when ghost atoms lie within the cutoff, which is consistent with that reading. _strip_shape_assertions already neutralizes the same kind of guard for spin models; its docstring states that those models are correct when nall == nloc.

Dominant language
Python
Stars
2.1k
Forks
663
Avg merge
3d 11h
Merged PRs (30d)
25

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepmodeling/deepmd-kit

All issues in deepmodeling/deepmd-kit

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.