Merge branch 'master' into pr/dpa1

Author: OutisLi

deepmd/backend/pt_expt.py

@@ -13,6 +13,9 @@
 from deepmd.backend.backend import (
     Backend,
 )
+from deepmd.utils.pt_checkpoint import (
+    detect_pt_checkpoint_backend,
+)
 
 if TYPE_CHECKING:
     from argparse import (
@@ -51,11 +54,8 @@ def match_filename(cls, filename: str) -> int:
         Returns
         -------
         - 1 for the regular `.pte` / `.pt2` suffixes (default behaviour).
-        - 2 for `.pt` files whose state-dict uses pt_expt's dpmodel
-            parameter naming (`.w`/`.b`); this outranks the legacy pt
-            backend's default suffix score (1) so pt_expt-trained `.pt`
-            checkpoints route here, while genuine pt-trained `.pt` files
-            (which use `.matrix`/`.bias`) keep going to the pt backend.
+        - 2 for `.pt` files whose state dictionary uses the pt_expt parameter
+            dialect. This outranks the pt backend's default suffix score (1).
         - 0 otherwise.
         """
         score = super().match_filename(filename)
@@ -69,21 +69,14 @@ def match_filename(cls, filename: str) -> int:
 
             # weights_only=True avoids unpickling arbitrary code from an
             # untrusted .pt — sniffing only needs the dict keys.
-            sd = torch.load(filename, map_location="cpu", weights_only=True)
+            checkpoint = torch.load(filename, map_location="cpu", weights_only=True)
         except Exception:
             # Not a valid torch archive (corrupt file, wrong format, or a
             # weights_only=True restriction trip).  Surrender the claim so
             # the dispatcher falls back to the default suffix match — pt's
             # default score (1) will pick up the file under `dp --pt`.
             return 0
-        if isinstance(sd, dict) and "model" in sd:
-            sd = sd["model"]
-        keys = list(sd.keys()) if hasattr(sd, "keys") else []
-        has_pt_expt = any(k.endswith(".w") or k.endswith(".b") for k in keys)
-        has_pt = any(k.endswith(".matrix") or k.endswith(".bias") for k in keys)
-        if has_pt_expt and not has_pt:
-            return 2
-        return 0
+        return 2 if detect_pt_checkpoint_backend(checkpoint) == "pt-expt" else 0
 
     def is_available(self) -> bool:
         """Check if the backend is available.

deepmd/pt_expt/infer/deep_eval.py

@@ -5,6 +5,38 @@
     Any,
 )
 
+import torch
+
+
+def graph_edge_dtype(model: Any, lower_kind: str) -> str:
+    """Return the graph edge-vector dtype encoded by a deployment artifact.
+
+    Parameters
+    ----------
+    model : Any
+        Model exposing an atomic-model descriptor.
+    lower_kind : str
+        Concrete lower-forward schema.
+
+    Returns
+    -------
+    str
+        ``"float32"`` for eligible geometrically compressed DPA1 graph
+        lowers, otherwise ``"float64"``.
+    """
+    atomic_model = getattr(model, "atomic_model", None)
+    descriptor = getattr(atomic_model, "descriptor", None)
+    descriptor_block = getattr(descriptor, "se_atten", None)
+    statistics = getattr(descriptor_block, "mean", None)
+    if (
+        lower_kind in ("graph", "dpa1_canonical")
+        and bool(getattr(descriptor, "geo_compress", False))
+        and isinstance(statistics, torch.Tensor)
+        and statistics.dtype == torch.float32
+    ):
+        return "float32"
+    return "float64"
+
 
 def model_uses_graph_lower(model: Any) -> bool:
     """Return whether a model's default lower uses ``NeighborGraph``.

deepmd/pt_expt/model/graph_lower.py

@@ -26,6 +26,9 @@
 from deepmd.dpmodel.utils.serialization import (
     traverse_model_dict,
 )
+from deepmd.pt_expt.model.graph_lower import (
+    graph_edge_dtype,
+)
 
 # ---------------------------------------------------------------------------
 # AOTInductor ``.pt2`` archive layout.
@@ -1015,28 +1018,6 @@ def _build_dynamic_shapes(
     return (*base, None, None, None, None, None, None, None, None)
 
 
-def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str:
-    """Return the graph edge-vector dtype encoded by the deployment artifact.
-
-    Geometrically compressed DPA1 with float32 descriptor statistics evaluates
-    both descriptor directions in float32 and therefore accepts float32
-    geometry directly. Other graph descriptors retain the model-agnostic
-    float64 geometry ABI.
-    """
-    atomic_model = getattr(model, "atomic_model", None)
-    descriptor = getattr(atomic_model, "descriptor", None)
-    descriptor_block = getattr(descriptor, "se_atten", None)
-    statistics = getattr(descriptor_block, "mean", None)
-    if (
-        lower_kind in ("graph", "dpa1_canonical")
-        and bool(getattr(descriptor, "geo_compress", False))
-        and isinstance(statistics, torch.Tensor)
-        and statistics.dtype == torch.float32
-    ):
-        return "float32"
-    return "float64"
-
-
 def _supports_graph_export(model: torch.nn.Module) -> bool:
     """Whether the model has an exportable graph-lower implementation.
 
@@ -1173,7 +1154,7 @@ def _probe_has_message_passing(obj: object) -> bool | None:
     #   "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask)
     # The C++ loader branches on this to build the matching inputs.
     meta["lower_input_kind"] = lower_kind
-    meta["graph_edge_dtype"] = _graph_edge_dtype(model, lower_kind)
+    meta["graph_edge_dtype"] = graph_edge_dtype(model, lower_kind)
 
     # Model-level pair-type exclusion (``pair_exclude_types``): a list of
     # ``[ti, tj]`` type pairs whose interaction is dropped.  Exclusion is a

deepmd/pt_expt/utils/serialization.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Utilities shared by PyTorch checkpoint backends."""
+
+from collections.abc import (
+    Mapping,
+)
+from typing import (
+    Any,
+)
+
+
+def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None:
+    """Detect the parameter dialect of a raw PyTorch checkpoint.
+
+    Parameters
+    ----------
+    checkpoint : Any
+        A checkpoint payload or its unwrapped model state dictionary.
+
+    Returns
+    -------
+    str or None
+        ``"pt-expt"`` or ``"pt"`` when the parameter names identify one
+        backend unambiguously, otherwise ``None``.
+    """
+    state_dict = checkpoint
+    if isinstance(state_dict, Mapping) and "model" in state_dict:
+        state_dict = state_dict["model"]
+    if not isinstance(state_dict, Mapping):
+        return None
+
+    keys = tuple(key for key in state_dict if isinstance(key, str))
+
+    # Weight names are decisive. pt_expt DPA4 also contains ordinary
+    # torch-native ``.bias`` parameters, so bias names cannot override a
+    # clear ``.w`` versus ``.matrix`` distinction.
+    has_pt_expt_weight = any(key.endswith(".w") for key in keys)
+    has_pt_weight = any(key.endswith(".matrix") for key in keys)
+    if has_pt_expt_weight or has_pt_weight:
+        if has_pt_expt_weight == has_pt_weight:
+            return None
+        return "pt-expt" if has_pt_expt_weight else "pt"
+
+    # A lone ``.b`` is specific to pt_expt's NativeLayer. A lone ``.bias``
+    # is not specific to pt because pt_expt models can contain torch-native
+    # modules with that suffix, so it remains deliberately unclassified.
+    has_pt_expt_bias = any(key.endswith(".b") for key in keys)
+    has_pt_bias = any(key.endswith(".bias") for key in keys)
+    if has_pt_expt_bias and not has_pt_bias:
+        return "pt-expt"
+    return None

deepmd/utils/pt_checkpoint.py

@@ -27,8 +27,11 @@
 )
 
 sys.path.append(os.path.dirname(__file__))
+import github_linkcode
 import sphinx_contrib_exhale_multiproject  # noqa: F401
 
+linkcode_resolve = github_linkcode.linkcode_resolve
+
 # -- Project information -----------------------------------------------------
 
 project = "DeePMD-kit"
@@ -59,7 +62,7 @@
     "myst_nb",
     "sphinx.ext.autosummary",
     "sphinx.ext.mathjax",
-    "sphinx.ext.viewcode",
+    "sphinx.ext.linkcode",
     "sphinx.ext.imgconverter",
     "sphinx.ext.intersphinx",
     "sphinx.ext.napoleon",
@@ -277,5 +280,13 @@ def _cap_cli_secnumbers(app: Sphinx, doctree: nodes.document, docname: str) -> N
 
 
 def setup(app: Sphinx) -> dict[str, bool]:
+    # AutoAPI records exact source locations without importing backend modules.
+    # Reuse that metadata for commit-pinned GitHub links after AutoAPI's default
+    # priority-500 ``builder-inited`` callback has populated the environment.
+    app.connect(
+        "builder-inited",
+        github_linkcode.collect_autoapi_source_locations,
+        priority=600,
+    )
     app.connect("doctree-resolved", _cap_cli_secnumbers)
     return {"parallel_read_safe": True, "parallel_write_safe": True}