fix(pt_expt): distinguish raw checkpoint backends by weights

Author: OutisLi

deepmd/backend/pt_expt.py

@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from collections.abc import (
     Callable,
+    Mapping,
 )
 from importlib.util import (
     find_spec,
 )
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
 )
 
@@ -27,6 +29,46 @@
     )
 
 
+def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None:
+    """Detect the backend 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"
+
+    # Bias-only state dictionaries retain the original dialect fallback.
+    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 == has_pt_bias:
+        return None
+    return "pt-expt" if has_pt_expt_bias else "pt"
+
+
 @Backend.register("pt-expt")
 @Backend.register("pytorch-exportable")
 class PyTorchExportableBackend(Backend):
@@ -51,11 +93,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 +108,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

@@ -13,6 +13,9 @@
 import numpy as np
 import torch
 
+from deepmd.backend.pt_expt import (
+    detect_pt_checkpoint_backend,
+)
 from deepmd.dpmodel.model.transform_output import (
     communicate_extended_output,
 )
@@ -132,18 +135,6 @@ def _reshape_charge_spin(
         ) from err
 
 
-def _is_pt_backend_dpa4_params(model_params: dict[str, Any]) -> bool:
-    """Return whether a training checkpoint should be loaded by the pt backend."""
-    model_type = str(model_params.get("type", "")).lower()
-    if model_type in {"sezm", "dpa4", "sezm_spin"}:
-        return True
-    descriptor = model_params.get("descriptor")
-    if isinstance(descriptor, dict):
-        descriptor_type = str(descriptor.get("type", "")).lower()
-        return descriptor_type in {"sezm", "dpa4"}
-    return False
-
-
 def _warn_legacy_edge_vec(metadata: dict) -> None:
     """Warn once per model load when an edge_vec-schema artifact is opened.
 
@@ -556,9 +547,17 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None:
 
         # Match the training resume path (training.py:712) — weights_only=True
         # avoids unpickling arbitrary code from untrusted checkpoints.
-        state_dict = torch.load(model_file, map_location=DEVICE, weights_only=True)
+        checkpoint = torch.load(model_file, map_location=DEVICE, weights_only=True)
+        checkpoint_backend = detect_pt_checkpoint_backend(checkpoint)
+        state_dict = checkpoint
         if isinstance(state_dict, dict) and "model" in state_dict:
             state_dict = state_dict["model"]
+        if checkpoint_backend == "pt":
+            raise ValueError(
+                f"Checkpoint '{model_file}' uses the regular `pt` parameter "
+                "dialect. Load it with `dp --pt`, or export it to `.pt2` / "
+                "`.pte` before loading it with `pt_expt`."
+            )
         extra = state_dict.get("_extra_state") if isinstance(state_dict, dict) else None
         if not (isinstance(extra, dict) and "model_params" in extra):
             raise ValueError(
@@ -598,13 +597,6 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None:
             state_dict = head_state
             model_params = head_params
 
-        if _is_pt_backend_dpa4_params(model_params):
-            raise ValueError(
-                "DPA4/SeZM `.pt` checkpoints belong to the regular `pt` backend. "
-                "Use the `pt` backend for eager checkpoint inference, or export "
-                "the checkpoint to `.pt2` / `.pte` before loading it with `pt_expt`."
-            )
-
         model = get_model(deepcopy(model_params)).to(DEVICE)
 
         # Strip the `_CompiledModel` wrapper that pt_expt training applies

source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py

@@ -4,8 +4,9 @@
 Covers two pieces:
 
 1. ``Backend.detect_backend_by_model`` sniffs ``.pt`` content
-   (``.w``/``.b`` -> pt_expt, ``.matrix``/``.bias`` -> pt) so that
-   ``dp test -m foo.pt`` routes to the right backend.
+   (``.w`` weights -> pt_expt, ``.matrix`` weights -> pt, with bias names
+   used only as a fallback) so that ``dp test -m foo.pt`` routes to the
+   right backend.
 2. ``pt_expt.DeepEval._load_pt`` reconstructs the model from
    ``_extra_state["model_params"]``, loads ``state_dict``, and runs
    inference in eager mode, producing outputs that match a direct
@@ -710,6 +711,80 @@ def _spin_eager_reference(model, COORD, ATYPE, SPIN, BOX):
     return {k: v.detach().cpu().numpy() for k, v in ref.items()}
 
 
+class TestPtExptLoadPtNativeSpinDPA4(unittest.TestCase):
+    """A native-spin DPA4 checkpoint reaches the graph-spin eager runner."""
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        from ..model.test_dpa4_native_spin import (
+            NATIVE_SPIN_CONFIG,
+            _build_native_spin_model_cpu,
+        )
+        from .test_deep_eval_spin import (
+            ATYPE,
+            BOX,
+            COORD,
+            SPIN,
+        )
+
+        cls.ATYPE = ATYPE
+        cls.BOX = BOX
+        cls.COORD = COORD
+        cls.SPIN = SPIN
+        cls.model = _build_native_spin_model_cpu().to(DEVICE).eval()
+        cls.ref = _spin_eager_reference(
+            cls.model, cls.COORD, cls.ATYPE, cls.SPIN, cls.BOX
+        )
+        cls.pt_path = tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name
+        _save_pt_checkpoint(
+            cls.model,
+            copy.deepcopy(NATIVE_SPIN_CONFIG),
+            cls.pt_path,
+        )
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        if os.path.exists(cls.pt_path):
+            os.unlink(cls.pt_path)
+
+    def test_auto_dispatch_and_eval_match_eager(self) -> None:
+        backend = Backend.detect_backend_by_model(self.pt_path)
+        self.assertIs(backend, Backend.get_backend("pt-expt"))
+
+        dp = DeepPot(
+            self.pt_path,
+            auto_batch_size=False,
+            neighbor_graph_method="dense",
+        )
+        self.assertTrue(dp.has_spin)
+        self.assertEqual(dp.deep_eval.metadata["lower_input_kind"], "graph")
+
+        energy, force, virial, force_mag, mask_mag = dp.eval(
+            self.COORD,
+            self.BOX,
+            self.ATYPE,
+            atomic=False,
+            spin=self.SPIN,
+        )
+        for name, actual in (
+            ("energy", energy),
+            ("force", force),
+            ("virial", virial),
+            ("force_mag", force_mag),
+        ):
+            np.testing.assert_allclose(
+                actual.reshape(-1),
+                self.ref[name].reshape(-1),
+                rtol=1e-10,
+                atol=1e-10,
+                err_msg=name,
+            )
+        np.testing.assert_array_equal(
+            mask_mag.reshape(-1),
+            self.ref["mask_mag"].reshape(-1),
+        )
+
+
 class _SpinFilesMixin:
     """Build .pt + .pte for the chosen ``spin_config`` once per class."""