Merge remote-tracking branch 'origin/master' into fix/pr5978-maintainer-review-20260902

Author: atulcthakur

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:

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]:

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.
 

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:

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(

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,

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."""

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:

deepmd/pt_expt/utils/serialization.py

@@ -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.
@@ -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)
@@ -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"]
@@ -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:
@@ -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
@@ -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,
     )
@@ -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,
         )
@@ -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(
@@ -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"]
     )