diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index 43cb901449..4ed9b9cfbb 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -10,6 +10,8 @@ log = logging.getLogger(__name__) +_LOWER_INPUT_KINDS = frozenset({"nlist", "graph", "dpa1_canonical"}) + def convert_backend( *, # Enforce keyword-only arguments @@ -30,6 +32,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 the source model's lower-input semantics. + Formats without explicit lower metadata are dense neighbor-list models; + graph artifacts must expose ``lower_input_kind`` through their serializer. + A target backend that cannot represent the source 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 +50,24 @@ def convert_backend( sig = inspect.signature(out_hook) hook_kwargs: dict[str, Any] = {} + # Existing backend serializers describe dense models and predate explicit + # lower metadata. Graph-capable artifact serializers override this default + # from their archive metadata. + lower_input_kind = data.get("lower_input_kind", "nlist") + if lower_input_kind not in _LOWER_INPUT_KINDS: + raise ValueError( + f"Unsupported source lower_input_kind {lower_input_kind!r}; " + f"expected one of {sorted(_LOWER_INPUT_KINDS)}." + ) if "lower_kind" in sig.parameters: - hook_kwargs["lower_kind"] = "auto" + hook_kwargs["lower_kind"] = lower_input_kind + elif lower_input_kind != "nlist": + raise ValueError( + f"Cannot preserve graph lower semantics when converting to output " + f"backend {out_backend.name!r}: its deserializer does not accept " + "a lower_kind. Retrain or freeze the model with that backend " + "instead of converting the graph artifact." + ) if "do_atomic_virial" in sig.parameters: hook_kwargs["do_atomic_virial"] = atomic_virial elif atomic_virial: diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 5d7be558bb..aed90fdd5d 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1169,7 +1169,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) @@ -1179,10 +1180,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"] @@ -1200,6 +1211,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: @@ -1210,8 +1222,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 @@ -1312,10 +1331,12 @@ 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 + instead, preserving its 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``. """ diff --git a/doc/backend.md b/doc/backend.md index e09f230eea..4bcfddcafe 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -154,3 +154,10 @@ 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 source artifact's lower-input semantics. A +dense neighbor-list model therefore remains a dense model when converted to a +`.pt2` file; conversion does not reinterpret it as a graph-native model. Graph +artifacts retain their graph lower when the target format supports it, and a +conversion is rejected when the target backend cannot preserve that lower. +Train and freeze through a graph-native backend to obtain graph semantics. 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..a8d4ea451e 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -24,6 +24,7 @@ _needs_with_comm_artifact, _supports_graph_export, deserialize_to_file, + serialize_from_file, ) # dpa1 with attn_layer == 0 — the energy model exercised by the graph path. @@ -86,6 +87,44 @@ def _read_metadata(pt2_path: str) -> dict: return json.loads(raw) +@pytest.mark.parametrize("lower_input_kind", ["nlist", "graph", "dpa1_canonical"]) +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"]) +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.fixture(scope="module") def dpa1_dpmodel_data() -> dict: return _build_dpa1_data() diff --git a/source/tests/test_convert_backend.py b/source/tests/test_convert_backend.py index 063caec868..3564784081 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_preserves_default_dense_lower_kind( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -46,5 +46,65 @@ def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: convert_backend(INPUT="model.input", OUTPUT="model.output") - assert captured["lower_kind"] == "auto" + assert captured["lower_kind"] == "nlist" assert captured["do_atomic_virial"] is False + + +def test_convert_backend_preserves_graph_lower_kind( + 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" + + @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"] == "graph" + + +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" + + @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 graph lower"): + convert_backend(INPUT="model.input", OUTPUT="model.output")