Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions deepmd/backend/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,14 @@ class Feature(Flag):
"""The supported suffixes of the saved model.

The first element is considered as the default suffix."""
preserves_lower_input_kind: ClassVar[bool] = False
"""Whether the IO hook preserves lower-ABI metadata without materializing it.

Schema-neutral model containers retain ``lower_input_kind`` as provenance
even though their deserializer does not accept a concrete ``lower_kind``.
Executable backends instead materialize a lower ABI and must expose that
choice through their deserializer signature.
"""

@abstractmethod
def is_available(self) -> bool:
Expand Down
6 changes: 4 additions & 2 deletions deepmd/backend/dpmodel.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ class DPModelBackend(Backend):
"""The features of the backend."""
suffixes: ClassVar[list[str]] = [".dp", ".yaml", ".yml"]
"""The suffixes of the backend."""
preserves_lower_input_kind: ClassVar[bool] = True
"""DPModel files retain lower provenance without binding an execution ABI."""

def is_available(self) -> bool:
"""Check if the backend is available.
Expand Down Expand Up @@ -106,10 +108,10 @@ def serialize_hook(self) -> Callable[[str], dict]:
The serialize hook of the backend.
"""
from deepmd.dpmodel.utils.serialization import (
load_dp_model,
serialize_from_file,
)

return load_dp_model
return serialize_from_file

@property
def deserialize_hook(self) -> Callable[[str, dict], None]:
Expand Down
23 changes: 23 additions & 0 deletions deepmd/dpmodel/utils/serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,29 @@ def convert_numpy_ndarray(x: Any) -> Any:
return model_dict


def serialize_from_file(filename: str) -> dict:
"""Serialize a DPModel container for backend conversion.

DPModel files store model parameters rather than an executable lower ABI.
Concrete provenance written by an earlier conversion is retained; a native
file without provenance reports ``"auto"`` so the executable target selects
a compatible lower from the model capabilities.

Parameters
----------
filename : str
The DPModel filename.

Returns
-------
dict
The serialized model data with declared lower-input semantics.
"""
model_dict = load_dp_model(filename)
model_dict.setdefault("lower_input_kind", "auto")
return model_dict


def format_big_number(x: int) -> str:
"""Format a big number with suffixes.

Expand Down
23 changes: 22 additions & 1 deletion deepmd/entrypoints/convert_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,14 @@ def convert_backend(
If True, export .pt2/.pte models with per-atom virial correction.
This adds ~2.5x inference cost. Default False. Silently ignored
(with a warning) for backends that don't support the flag.

Notes
-----
Backend conversion preserves an explicit ``lower_input_kind`` reported by
the source serializer. Sources without this metadata retain the target's
automatic lower selection for backward compatibility. A target backend
that cannot represent an explicit non-dense lower is rejected rather than
silently changing the model function.
"""
inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
Expand All @@ -40,8 +48,21 @@ def convert_backend(

sig = inspect.signature(out_hook)
hook_kwargs: dict[str, Any] = {}
lower_input_kind = data.get("lower_input_kind")
if "lower_kind" in sig.parameters:
hook_kwargs["lower_kind"] = "auto"
hook_kwargs["lower_kind"] = (
lower_input_kind if lower_input_kind is not None else "auto"
)
elif (
lower_input_kind not in (None, "auto", "nlist")
and not out_backend.preserves_lower_input_kind
):
raise ValueError(
f"Cannot preserve lower_input_kind {lower_input_kind!r} when "
f"converting to output backend {out_backend.name!r}: its "
"deserializer does not accept a lower_kind. Retrain or freeze the "
"model with that backend instead of converting this artifact."
)
if "do_atomic_virial" in sig.parameters:
hook_kwargs["do_atomic_virial"] = atomic_virial
elif atomic_virial:
Expand Down
2 changes: 2 additions & 0 deletions deepmd/jax/utils/serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
"jax_version": jax.__version__,
"model": model_dict,
"model_def_script": model_def_script,
"lower_input_kind": "nlist",
"@variables": {},
}
if min_nbor_dist is not None:
Expand All @@ -436,6 +437,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
data = load_dp_model(model_file)
data.pop("constants")
data["@variables"].pop("stablehlo")
data["lower_input_kind"] = "nlist"
return data
elif model_file.endswith(".savedmodel"):
raise ValueError(
Expand Down
11 changes: 11 additions & 0 deletions deepmd/pt/model/model/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,17 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
"min_nbor_dist", torch.tensor(-1.0, dtype=torch.float64, device=env.DEVICE)
)

def export_lower_input_kind(self) -> str:
"""Return the lower-input ABI that preserves this model's semantics.

Returns
-------
str
``"nlist"`` for the standard PyTorch model contract. Models with
a graph-native deployment ABI override this method.
"""
return "nlist"

def compute_or_load_stat(
self,
sampled_func: Any,
Expand Down
11 changes: 11 additions & 0 deletions deepmd/pt/model/model/spin_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -454,6 +454,17 @@ def has_spin(self) -> bool:
"""Returns whether it has spin input and output."""
return True

def export_lower_input_kind(self) -> str:
"""Return the dense ABI used by the virtual-atom spin model.

Returns
-------
str
``"nlist"``, because virtual atoms are expanded inside the
bounded neighbor-list contract.
"""
return "nlist"

@torch.jit.export
def has_message_passing(self) -> bool:
"""Returns whether the model has message passing."""
Expand Down
1 change: 1 addition & 0 deletions deepmd/pt/utils/serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ def serialize_from_file(model_file: str) -> dict:
"pt_version": str(torch.__version__),
"model": model_dict,
"model_def_script": model_def_script,
"lower_input_kind": model.export_lower_input_kind(),
"@variables": {},
}
if model.get_min_nbor_dist() is not None:
Expand Down
160 changes: 119 additions & 41 deletions deepmd/pt_expt/utils/serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,21 @@
# ---------------------------------------------------------------------------
PT2_EXTRA_PREFIX = "model/extra/"

# Backend conversion supplies the source artifact's lower ABI. Concrete target
# schemas pass through unchanged. PT SeZM's ``edge_vec`` identifies an edge-list
# source contract rather than a pt_expt schema; the target model capabilities
# determine whether that contract is materialized as NeighborGraph or dense
# nlist input.
_LOWER_INPUT_KINDS = frozenset(
{
"nlist",
"graph",
"dpa1_canonical",
"dpa4c_canonical",
"edge_vec",
}
)


def _strip_shape_assertions(graph_module: torch.nn.Module) -> None:
"""Neutralise deferred shape-guard assertion nodes in an exported graph.
Expand Down Expand Up @@ -1284,7 +1299,8 @@ def serialize_from_file(model_file: str) -> dict:
dict
The serialized model data. If the archive contains
``model_def_script.json`` (training config), it is included
under the ``"model_def_script"`` key.
under the ``"model_def_script"`` key. ``lower_input_kind`` records
the concrete lower ABI from the artifact metadata.
"""
if model_file.endswith(".pt2"):
return _serialize_from_file_pt2(model_file)
Expand All @@ -1294,10 +1310,20 @@ def serialize_from_file(model_file: str) -> dict:

def _serialize_from_file_pte(model_file: str) -> dict:
"""Serialize a .pte model file to a dictionary."""
extra_files = {"model.json": "", "model_def_script.json": ""}
extra_files = {
"model.json": "",
"model_def_script.json": "",
"metadata.json": "",
}
torch.export.load(model_file, extra_files=extra_files)
model_dict = json.loads(extra_files["model.json"])
model_dict = _json_to_numpy(model_dict)
metadata = (
json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {}
)
model_dict["lower_input_kind"] = metadata.get(
"lower_input_kind", model_dict.get("lower_input_kind", "nlist")
)
if extra_files["model_def_script.json"]:
model_dict["model_def_script"] = json.loads(
extra_files["model_def_script.json"]
Expand All @@ -1315,6 +1341,7 @@ def _serialize_from_file_pt2(model_file: str) -> dict:

model_json_entry = PT2_EXTRA_PREFIX + "model.json"
model_def_script_entry = PT2_EXTRA_PREFIX + "model_def_script.json"
metadata_entry = PT2_EXTRA_PREFIX + "metadata.json"
with zipfile.ZipFile(model_file, "r") as zf:
names = zf.namelist()
if model_json_entry not in names:
Expand All @@ -1325,8 +1352,15 @@ def _serialize_from_file_pt2(model_file: str) -> dict:
model_def_script_json = ""
if model_def_script_entry in names:
model_def_script_json = zf.read(model_def_script_entry).decode("utf-8")
metadata_json = ""
if metadata_entry in names:
metadata_json = zf.read(metadata_entry).decode("utf-8")
model_dict = json.loads(model_json)
model_dict = _json_to_numpy(model_dict)
metadata = json.loads(metadata_json) if metadata_json else {}
model_dict["lower_input_kind"] = metadata.get(
"lower_input_kind", model_dict.get("lower_input_kind", "nlist")
)
if model_def_script_json:
model_dict["model_def_script"] = json.loads(model_def_script_json)
return model_dict
Expand Down Expand Up @@ -1506,19 +1540,22 @@ def _dpa4_kernel_levels_for_target(
os.environ[name] = value


def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
"""Resolve ``lower_kind="auto"`` to a concrete lower-forward schema.
def _select_graph_lower_kind(data: dict, *, allow_canonical: bool) -> str | None:
"""Select the graph schema supported by the target model.

``"auto"`` selects the graph lower for a graph-lower model whose graph
implementation is exportable to ``.pt2`` and the dense nlist lower for
everything else. Eligible compressed DPA1 and DPA4C energy models select
their compact canonical graph schemas. Any explicit lower kind is returned
unchanged.
Parameters
----------
data : dict
Serialized model data.
allow_canonical : bool
Whether an eligible compact canonical schema may replace NeighborGraph.

Returns
-------
str or None
The supported graph schema, or ``None`` when the model uses the dense
lower.
"""
if lower_kind != "auto":
return lower_kind
if not model_file.endswith(".pt2") or data["model"].get("type") == "spin_ener":
return "nlist"
from deepmd.pt_expt.model.graph_lower import (
model_uses_graph_lower,
)
Expand All @@ -1527,7 +1564,9 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
)

model = BaseModel.deserialize(data["model"])
if model_uses_graph_lower(model) and _supports_graph_export(model):
if not (model_uses_graph_lower(model) and _supports_graph_export(model)):
return None
if allow_canonical:
from deepmd.pt_expt.kernels.cuda.dpa1.canonical import (
canonical_model_eligible as dpa1_canonical_eligible,
)
Expand All @@ -1539,8 +1578,62 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
return "dpa4c_canonical"
if dpa1_canonical_eligible(model):
return "dpa1_canonical"
return "graph"
return "nlist"
return "graph"


def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
"""Resolve ``lower_kind="auto"`` to a concrete lower-forward schema.

``"auto"`` selects the graph lower for a graph-lower model whose graph
implementation is exportable to ``.pt2`` and the dense nlist lower for
everything else. Eligible compressed DPA1 and DPA4C energy models select
their compact canonical graph schemas. Any explicit lower kind is returned
unchanged.
"""
if lower_kind != "auto":
return lower_kind
if not model_file.endswith(".pt2") or data["model"].get("type") == "spin_ener":
return "nlist"
return _select_graph_lower_kind(data, allow_canonical=True) or "nlist"


def _resolve_target_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
"""Resolve a source lower ABI to a concrete pt_expt export schema."""
source_lower_kind = _resolve_lower_kind(model_file, data, lower_kind)
if source_lower_kind not in _LOWER_INPUT_KINDS:
raise ValueError(
f"Unsupported lower_kind {source_lower_kind!r}; expected one of "
f"{sorted(_LOWER_INPUT_KINDS)}."
)
target_lower_kind = source_lower_kind
if source_lower_kind == "edge_vec":
target_lower_kind = (
_select_graph_lower_kind(data, allow_canonical=False) or "nlist"
)

if data["model"].get("type") == "native_spin" and target_lower_kind not in (
"graph",
"dpa4c_canonical",
):
if lower_kind == "auto":
if not model_file.endswith(".pt2"):
raise ValueError(
"automatic lower selection for native-spin models requires "
"a .pt2 output because native-spin models do not implement "
"the dense nlist lower"
)
raise ValueError(
"automatic lower selection found no exportable graph lower for "
"this native-spin model, which does not implement the dense "
"nlist lower"
)
raise ValueError(
"native-spin models implement only the NeighborGraph and compact "
f"canonical lowers (got lower_kind={target_lower_kind!r}); use "
"lower_kind='graph', or lower_kind='dpa4c_canonical' for an "
"eligible compressed DPA4C model, with a .pt2 output."
)
return target_lower_kind


def deserialize_to_file(
Expand Down Expand Up @@ -1581,34 +1674,19 @@ def deserialize_to_file(
(``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask`` and
the destination/source CSR views) with a DYNAMIC edge axis ``E``
(``Dim("nedge", min=2)``), so the artifact accepts any system size.
``"auto"`` (used by ``convert-backend``) resolves to ``"graph"`` for an
exportable graph-lower ``.pt2`` and ``"nlist"`` otherwise (see
:func:`_resolve_lower_kind`). A graph lower preserves the selected
inference operators and always includes the per-atom virial. DPA1 and
DPA4C graph pipelines use ``DP_CUDA_INFER >= 2``; DPA4 ``.pt2`` follows
its PT freeze defaults unless the environment explicitly selects other
levels.
``"auto"`` resolves to ``"graph"`` for an exportable graph-lower
``.pt2`` and ``"nlist"`` otherwise (see :func:`_resolve_lower_kind`).
Backend conversion passes the source artifact's concrete lower kind;
compatible source ABIs are mapped to the target's native schema while
preserving their execution semantics. A graph lower preserves the
selected inference operators and always includes the per-atom virial.
DPA1 and DPA4C graph pipelines use ``DP_CUDA_INFER >= 2``; DPA4
``.pt2`` follows its PT freeze defaults unless the environment
explicitly selects other levels.
The selected schema is recorded as ``lower_input_kind`` in
``metadata.json``.
"""
lower_kind = _resolve_lower_kind(model_file, data, lower_kind)
if data["model"].get("type") == "native_spin" and lower_kind not in (
"graph",
"dpa4c_canonical",
):
# Native-spin models implement the NeighborGraph lower and, for an
# eligible compressed DPA4C, the compact canonical one; the dense/nlist
# trace branch does not exist for them. The public freeze layer
# resolves this before calling here (see
# deepmd.pt_expt.entrypoints.main.freeze); this guard pins the
# contract for direct programmatic callers with a clear error instead
# of an opaque trace-time failure.
raise ValueError(
"native-spin models implement only the NeighborGraph and compact "
f"canonical lowers (got lower_kind={lower_kind!r}); use "
"lower_kind='graph', or lower_kind='dpa4c_canonical' for an "
"eligible compressed DPA4C model, with a .pt2 output."
)
lower_kind = _resolve_target_lower_kind(model_file, data, lower_kind)
uses_dpa4_defaults = model_file.endswith(".pt2") and _uses_dpa4_kernel_defaults(
data["model"]
)
Expand Down
Loading
Loading