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

Abierto
  • ¿Algo no está bien?
#6,087 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
55/100
Tipo de issue
Error
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
backend, machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
2.1k
Forks
663
Merge medio
3 d 11 h
PR fusionados (30 d)
25

Preparar el entorno

Abrir en Codespaces

Inicia el contenedor de desarrollo del proyecto en tu navegador, con tu propia cuenta de GitHub.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de deepmodeling/deepmd-kit

Todos los issues de deepmodeling/deepmd-kit

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.