diff --git a/deepmd/backend/backend.py b/deepmd/backend/backend.py index ecd132ad9f..1d5bb36c76 100644 --- a/deepmd/backend/backend.py +++ b/deepmd/backend/backend.py @@ -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: diff --git a/deepmd/backend/dpmodel.py b/deepmd/backend/dpmodel.py index 0e6e0964f3..23a2a701e4 100644 --- a/deepmd/backend/dpmodel.py +++ b/deepmd/backend/dpmodel.py @@ -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. @@ -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]: diff --git a/deepmd/dpmodel/utils/serialization.py b/deepmd/dpmodel/utils/serialization.py index ce65ff7784..b0c70d7a33 100644 --- a/deepmd/dpmodel/utils/serialization.py +++ b/deepmd/dpmodel/utils/serialization.py @@ -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. diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index 43cb901449..3f1f624fab 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -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)() @@ -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: diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 62a6851160..33801e4149 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -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: @@ -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( diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index 5f89ff50ad..02ecedbce7 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -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, diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index c0adf618c8..029f6c1dd0 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -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.""" diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index db23eef4dc..6cdd65712f 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -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: diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 2ad89a22c4..0cbf164e00 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -52,6 +52,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. @@ -1242,7 +1257,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) @@ -1252,10 +1268,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"] @@ -1273,6 +1299,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: @@ -1283,8 +1310,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 @@ -1318,19 +1352,22 @@ def _cuda_infer_at_least_2() -> Iterator[None]: os.environ["DP_CUDA_INFER"] = saved -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, ) @@ -1339,7 +1376,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.kernels.cuda.dpa1.canonical import ( canonical_model_eligible as dpa1_canonical_eligible, ) @@ -1351,8 +1390,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( @@ -1393,31 +1486,17 @@ 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 always preserves the fused - inference operators (``DP_CUDA_INFER >= 2``) and the per-atom virial. + ``"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 always preserves + the fused inference operators (``DP_CUDA_INFER >= 2``) and the + per-atom virial. 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) # A graph lower deploys the fused inference pipeline. The trace runs at # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom # operators, while the per-atom virial is mandatory for the LAMMPS Kokkos diff --git a/deepmd/tf/utils/serialization.py b/deepmd/tf/utils/serialization.py index 1d2f1b597f..81d46c1e2d 100644 --- a/deepmd/tf/utils/serialization.py +++ b/deepmd/tf/utils/serialization.py @@ -50,6 +50,7 @@ def serialize_from_file(model_file: str) -> dict: "tf_version": tf.__version__, "model": model_dict, "model_def_script": jdata["model"], + "lower_input_kind": "nlist", } # neighbor stat information try: diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 63bbdf6583..b859d16b10 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -611,6 +611,7 @@ def serialize_from_file(model_file: str) -> dict: "backend": "TensorFlow2", "model": model_payload, "model_def_script": model_def_script, + "lower_input_kind": "nlist", "shared_links": state.get("shared_links"), "@variables": { "current_step": int(state.get("current_step", 0)), diff --git a/doc/backend.md b/doc/backend.md index e09f230eea..c5155cc991 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -154,3 +154,16 @@ then selects its `edge_vec`, dense `nlist`, NeighborGraph, or compact ## Convert model files between backends If a model is supported by two backends, one can use [`dp convert-backend`](./cli.rst) to convert the model file between these two backends. + +Backend conversion preserves the concrete `lower_input_kind` reported by the +source serializer. Dense TensorFlow, TensorFlow 2, JAX, and standard PyTorch +models therefore remain dense `nlist` models when converted to `.pt2`; model +families with a graph-native deployment ABI report their corresponding kind. +Compiled `.pt2` and `.pte` artifacts retain the kind recorded in their metadata. +Native `.dp`, `.yaml`, and `.yml` files are schema-neutral parameter containers: they +retain concrete source provenance across conversion but report `auto` when no +lower kind was stored. An executable target then selects a compatible schema +from the model capabilities. A conversion is rejected only when an executable +target cannot represent an explicit source kind. Legacy model files without +lower metadata retain target-specific automatic selection for backward +compatibility. diff --git a/source/tests/common/dpmodel/test_serialization.py b/source/tests/common/dpmodel/test_serialization.py index 049fcef9bd..1b26ea01b6 100644 --- a/source/tests/common/dpmodel/test_serialization.py +++ b/source/tests/common/dpmodel/test_serialization.py @@ -13,6 +13,7 @@ from deepmd.dpmodel.utils.serialization import ( load_dp_model, save_dp_model, + serialize_from_file, ) @@ -70,3 +71,27 @@ def test_save_dp_model_accepts_hdf5_dataset_without_mutation(tmp_path: Path) -> np.testing.assert_equal( loaded_model["model"]["layers"][0]["@variables"]["weights"], values ) + + +@pytest.mark.parametrize( + ("stored_lower_kind", "expected_lower_kind"), + [ + pytest.param(None, "auto", id="unbound"), + pytest.param("graph", "graph", id="preserved"), + ], +) +def test_serialize_from_file_declares_lower_semantics( + tmp_path: Path, + stored_lower_kind: str | None, + expected_lower_kind: str, +) -> None: + """The conversion hook distinguishes unbound and concrete lower ABIs.""" + data: dict = {"model": {}} + if stored_lower_kind is not None: + data["lower_input_kind"] = stored_lower_kind + model_file = tmp_path / "model.dp" + save_dp_model(str(model_file), data) + + serialized = serialize_from_file(str(model_file)) + + assert serialized["lower_input_kind"] == expected_lower_kind diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 2a34e2bbe5..913aca4d51 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -275,6 +275,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -319,6 +320,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -375,6 +377,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: @@ -422,6 +425,7 @@ def setUp(self) -> None: "model": model.serialize(), "backend": "test", "model_def_script": model_def_script, + "lower_input_kind": "nlist", } def tearDown(self) -> None: diff --git a/source/tests/jax/test_hlo.py b/source/tests/jax/test_hlo.py index 34488f4629..2e923205bd 100644 --- a/source/tests/jax/test_hlo.py +++ b/source/tests/jax/test_hlo.py @@ -1,9 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Regression tests for metadata exposed by serialized JAX HLO models.""" +from types import ( + SimpleNamespace, +) + +import pytest +from typing_extensions import ( + Self, +) + from deepmd.jax.model.hlo import ( HLO, ) +from deepmd.jax.utils import ( + serialization, +) def test_hlo_get_nnei_uses_stored_selection() -> None: @@ -17,3 +29,47 @@ def test_hlo_get_nnei_uses_stored_selection() -> None: model.sel = [6, 12, 1] assert model.get_nnei() == sum(model.sel) + + +def test_hlo_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A JAX HLO artifact exposes its dense source execution semantics.""" + stored_data = { + "model": {}, + "constants": {}, + "@variables": {"stablehlo": b"module"}, + } + monkeypatch.setattr(serialization, "load_dp_model", lambda _path: stored_data) + + data = serialization.serialize_from_file("model.hlo") + + assert data["lower_input_kind"] == "nlist" + + +def test_checkpoint_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A JAX training checkpoint exposes its dense source semantics.""" + + class Checkpointer: + def __init__(self, _handler: object) -> None: + pass + + def __enter__(self) -> Self: + return self + + def __exit__(self, *_args: object) -> None: + pass + + def restore(self, *_args: object, **_kwargs: object) -> SimpleNamespace: + return SimpleNamespace( + state={}, + model_def_script={"model_dict": {}}, + ) + + monkeypatch.setattr(serialization.ocp, "Checkpointer", Checkpointer) + + data = serialization.serialize_from_file("model.jax") + + assert data["lower_input_kind"] == "nlist" diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index 49864375e2..80199934d2 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -111,6 +111,10 @@ def test_output_shape( torch.testing.assert_close(result["force"].shape, [nframes, nloc, 3]) torch.testing.assert_close(result["force_mag"].shape, [nframes, nloc, 3]) + def test_export_lower_input_kind(self) -> None: + """Virtual-atom spin models retain the dense export ABI.""" + self.assertEqual(self.model.export_lower_input_kind(), "nlist") + def test_input_output_process(self) -> None: nframes, nloc = self.coord.shape[:2] self.real_ntypes = self.model.spin.get_ntypes_real() diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index 7ee61a1526..4d3964b77b 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -14,16 +14,22 @@ import tempfile import zipfile +import numpy as np import pytest import torch +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, ) from deepmd.pt_expt.utils.serialization import ( _needs_with_comm_artifact, + _resolve_target_lower_kind, _supports_graph_export, deserialize_to_file, + serialize_from_file, ) # dpa1 with attn_layer == 0 — the energy model exercised by the graph path. @@ -86,11 +92,255 @@ def _read_metadata(pt2_path: str) -> dict: return json.loads(raw) +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_pt2_serialization_preserves_lower_input_kind( + tmp_path, lower_input_kind: str +) -> None: + """The interchange dictionary exposes the artifact's lower semantics.""" + model_file = tmp_path / "model.pt2" + with zipfile.ZipFile(model_file, "w") as zf: + zf.writestr("model/extra/model.json", json.dumps({"model": {}})) + zf.writestr( + "model/extra/metadata.json", + json.dumps({"lower_input_kind": lower_input_kind}), + ) + + data = serialize_from_file(str(model_file)) + + assert data["lower_input_kind"] == lower_input_kind + + +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_pte_serialization_preserves_lower_input_kind( + tmp_path, monkeypatch: pytest.MonkeyPatch, lower_input_kind: str +) -> None: + """PTE extra metadata has the same interchange contract as PT2.""" + + def load_exported_program(_model_file: str, *, extra_files: dict[str, str]) -> None: + extra_files["model.json"] = json.dumps({"model": {}}) + extra_files["model_def_script.json"] = "" + extra_files["metadata.json"] = json.dumps( + {"lower_input_kind": lower_input_kind} + ) + + monkeypatch.setattr(torch.export, "load", load_exported_program) + + data = serialize_from_file(str(tmp_path / "model.pte")) + + assert data["lower_input_kind"] == lower_input_kind + + +@pytest.mark.parametrize( + ("embedded_lower_input_kind", "expected"), + [("graph", "graph"), (None, "nlist")], +) +def test_pt2_serialization_legacy_lower_input_kind_fallback( + tmp_path, + embedded_lower_input_kind: str | None, + expected: str, +) -> None: + """Legacy PT2 archives use embedded metadata, then dense fallback.""" + model_file = tmp_path / "model.pt2" + model_data: dict[str, object] = {"model": {}} + if embedded_lower_input_kind is not None: + model_data["lower_input_kind"] = embedded_lower_input_kind + with zipfile.ZipFile(model_file, "w") as zf: + zf.writestr("model/extra/model.json", json.dumps(model_data)) + + data = serialize_from_file(str(model_file)) + + assert data["lower_input_kind"] == expected + + +@pytest.mark.parametrize( + ("embedded_lower_input_kind", "expected"), + [("graph", "graph"), (None, "nlist")], +) +def test_pte_serialization_legacy_lower_input_kind_fallback( + tmp_path, + monkeypatch: pytest.MonkeyPatch, + embedded_lower_input_kind: str | None, + expected: str, +) -> None: + """Legacy PTE archives use embedded metadata, then dense fallback.""" + model_data: dict[str, object] = {"model": {}} + if embedded_lower_input_kind is not None: + model_data["lower_input_kind"] = embedded_lower_input_kind + + def load_exported_program(_model_file: str, *, extra_files: dict[str, str]) -> None: + extra_files["model.json"] = json.dumps(model_data) + extra_files["model_def_script.json"] = "" + extra_files["metadata.json"] = "" + + monkeypatch.setattr(torch.export, "load", load_exported_program) + + data = serialize_from_file(str(tmp_path / "model.pte")) + + assert data["lower_input_kind"] == expected + + @pytest.fixture(scope="module") def dpa1_dpmodel_data() -> dict: return _build_dpa1_data() +def test_convert_regular_pt_dpa1_preserves_dense_semantics(tmp_path) -> None: + """A nonzero-davg PT artifact remains numerically dense after conversion.""" + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt.utils.serialization import deserialize_to_file as deserialize_to_pt + from deepmd.pt.utils.serialization import serialize_from_file as serialize_from_pt + + data = _build_dpa1_data() + descriptor_variables = data["model"]["descriptor"]["@variables"] + descriptor_variables["davg"] = np.full_like(descriptor_variables["davg"], 0.01) + source_model = tmp_path / "model.pth" + converted_model = tmp_path / "model.pt2" + deserialize_to_pt(str(source_model), copy.deepcopy(data)) + + source_data = serialize_from_pt(str(source_model)) + assert source_data["lower_input_kind"] == "nlist" + + convert_backend(INPUT=str(source_model), OUTPUT=str(converted_model)) + assert _read_metadata(str(converted_model))["lower_input_kind"] == "nlist" + + coord = np.array( + [ + [0.0, 0.0, 0.0], + [1.1, 0.2, 0.1], + [0.3, 1.4, 0.2], + [1.2, 1.1, 0.8], + ], + dtype=np.float64, + )[None, ...] + atype = np.array([[0, 1, 0, 1]], dtype=np.int32) + source_result = DeepPot(str(source_model), auto_batch_size=False).eval( + coord, None, atype + ) + converted_result = DeepPot(str(converted_model), auto_batch_size=False).eval( + coord, None, atype + ) + np.testing.assert_allclose( + converted_result[0], source_result[0], rtol=1e-10, atol=1e-10 + ) + np.testing.assert_allclose( + converted_result[1], source_result[1], rtol=1e-10, atol=1e-10 + ) + + +def test_convert_pt_dpa4_through_dp_maps_edge_vec_to_graph(tmp_path) -> None: + """A schema-neutral DPModel container preserves PT SeZM's edge-list ABI.""" + from deepmd.dpmodel.utils.serialization import ( + load_dp_model, + ) + from deepmd.pt.model.model import get_model as get_pt_model + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + from deepmd.pt.utils.serialization import serialize_from_file as serialize_from_pt + + from ..model.test_dpa4_export import ( + _DPA4_CONFIG, + ) + + config = copy.deepcopy(_DPA4_CONFIG) + source_model = tmp_path / "model.pt" + intermediate_model = tmp_path / "model.dp" + converted_model = tmp_path / "model.pte" + model = get_pt_model(config) + wrapper = ModelWrapper(model, model_params=config) + torch.save({"model": wrapper.state_dict()}, source_model) + + source_data = serialize_from_pt(str(source_model)) + assert source_data["lower_input_kind"] == "edge_vec" + + convert_backend(INPUT=str(source_model), OUTPUT=str(intermediate_model)) + assert load_dp_model(str(intermediate_model))["lower_input_kind"] == "edge_vec" + + convert_backend(INPUT=str(intermediate_model), OUTPUT=str(converted_model)) + converted_data = serialize_from_file(str(converted_model)) + assert converted_data["lower_input_kind"] == "graph" + + +def test_edge_vec_uses_dense_lower_for_non_energy_target(tmp_path) -> None: + """Target model capabilities, not the source spelling, select graph export.""" + from deepmd.dpmodel.utils.serialization import ( + save_dp_model, + ) + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + + from ..model.test_dpa4_export import ( + _DPA4_CONFIG, + ) + + config = copy.deepcopy(_DPA4_CONFIG) + config.pop("type") + config["fitting_net"] = { + "type": "property", + "task_dim": 2, + "neuron": [16], + "precision": "float64", + "seed": 1, + } + model = get_model(config) + source_model = tmp_path / "property.dp" + converted_model = tmp_path / "property.pte" + save_dp_model( + str(source_model), + {"model": model.serialize(), "lower_input_kind": "edge_vec"}, + ) + + convert_backend(INPUT=str(source_model), OUTPUT=str(converted_model)) + + assert serialize_from_file(str(converted_model))["lower_input_kind"] == "nlist" + + +def test_native_spin_auto_pte_reports_target_constraint(tmp_path) -> None: + """An unbound native-spin container reports why automatic PTE export fails.""" + from deepmd.dpmodel.utils.serialization import ( + save_dp_model, + ) + + from ..model.test_dpa4_native_spin import ( + _build_native_spin_model_cpu, + ) + + source_model = tmp_path / "native_spin.dp" + data = {"model": _build_native_spin_model_cpu().serialize()} + assert _resolve_target_lower_kind("native_spin.pt2", data, "auto") == "graph" + save_dp_model( + str(source_model), + data, + ) + + with pytest.raises( + ValueError, + match=r"automatic lower selection for native-spin models requires a \.pt2 output", + ): + convert_backend( + INPUT=str(source_model), OUTPUT=str(tmp_path / "native_spin.pte") + ) + + +def test_deserialize_rejects_unknown_lower_kind(dpa1_dpmodel_data, tmp_path) -> None: + """The target serializer owns validation of its supported lower ABIs.""" + with pytest.raises(ValueError, match="Unsupported lower_kind 'unknown'"): + deserialize_to_file( + str(tmp_path / "model.pt2"), + copy.deepcopy(dpa1_dpmodel_data), + lower_kind="unknown", + ) + + def test_graph_pt2_has_lower_input_kind_graph(dpa1_dpmodel_data) -> None: """``lower_kind="graph"`` -> metadata ``lower_input_kind == "graph"``.""" with tempfile.TemporaryDirectory() as d: diff --git a/source/tests/test_convert_backend.py b/source/tests/test_convert_backend.py index 063caec868..17624006a3 100644 --- a/source/tests/test_convert_backend.py +++ b/source/tests/test_convert_backend.py @@ -9,7 +9,7 @@ ) -def test_convert_backend_automatically_selects_lower_kind( +def test_convert_backend_uses_auto_for_unannotated_source( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -23,6 +23,7 @@ def serialize_hook(path: str) -> dict[str, str]: class OutputBackend: name = "output" + preserves_lower_input_kind = False @staticmethod def deserialize_hook( @@ -48,3 +49,135 @@ def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: assert captured["lower_kind"] == "auto" assert captured["do_atomic_virial"] is False + + +@pytest.mark.parametrize( + "lower_input_kind", + ["nlist", "graph", "dpa1_canonical", "dpa4c_canonical", "edge_vec"], +) +def test_convert_backend_preserves_explicit_lower_kind( + monkeypatch: pytest.MonkeyPatch, + lower_input_kind: str, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path, "lower_input_kind": lower_input_kind} + + class OutputBackend: + name = "output" + preserves_lower_input_kind = False + + @staticmethod + def deserialize_hook( + path: str, + data: dict[str, str], + *, + lower_kind: str = "nlist", + ) -> None: + captured.update(path=path, data=data, lower_kind=lower_kind) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["lower_kind"] == lower_input_kind + + +@pytest.mark.parametrize("lower_input_kind", [None, "auto", "nlist"]) +def test_convert_backend_allows_dense_compatible_source_for_plain_output( + monkeypatch: pytest.MonkeyPatch, + lower_input_kind: str | None, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + data = {"path": path} + if lower_input_kind is not None: + data["lower_input_kind"] = lower_input_kind + return data + + class OutputBackend: + name = "output" + preserves_lower_input_kind = False + + @staticmethod + def deserialize_hook(path: str, data: dict[str, str]) -> None: + captured.update(path=path, data=data) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["path"] == "model.output" + + +def test_convert_backend_rejects_graph_for_dense_only_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path, "lower_input_kind": "graph"} + + class OutputBackend: + name = "output" + preserves_lower_input_kind = False + + @staticmethod + def deserialize_hook(path: str, data: dict[str, str]) -> None: + raise AssertionError("dense-only output hook must not be called") + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + with pytest.raises(ValueError, match="Cannot preserve lower_input_kind 'graph'"): + convert_backend(INPUT="model.input", OUTPUT="model.output") + + +def test_convert_backend_preserves_graph_for_schema_neutral_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path, "lower_input_kind": "graph"} + + class OutputBackend: + name = "output" + preserves_lower_input_kind = True + + @staticmethod + def deserialize_hook(path: str, data: dict[str, str]) -> None: + captured.update(path=path, data=data) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["data"]["lower_input_kind"] == "graph" diff --git a/source/tests/tf2/test_serialization.py b/source/tests/tf2/test_serialization.py new file mode 100644 index 0000000000..a8ac2af15b --- /dev/null +++ b/source/tests/tf2/test_serialization.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""TensorFlow 2 backend serialization contracts.""" + +import os + +import pytest + +if os.environ.get("DP_TEST_TF2_ONLY") != "1": + pytest.skip( + "TF2 tests require DP_TEST_TF2_ONLY=1", + allow_module_level=True, + ) + +from deepmd.tf2.utils import ( + serialization, +) + + +def test_checkpoint_serialization_declares_dense_lower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A TF2 training checkpoint exposes its dense source semantics.""" + state = { + "backend": "TensorFlow2", + "model_def_script": {}, + "current_step": 0, + } + monkeypatch.setattr( + serialization, + "_load_checkpoint_state", + lambda _path: ("checkpoint", state), + ) + monkeypatch.setattr( + serialization, + "_restore_models_from_checkpoint", + lambda _checkpoint, _script, _state: {}, + ) + monkeypatch.setattr( + serialization, + "_serialize_models", + lambda _models, _script: {}, + ) + + data = serialization.serialize_from_file("model.tf2") + + assert data["lower_input_kind"] == "nlist"