fix(pt-expt): add key remapping for cross-backend checkpoint compatibility

deepmd/pt/infer/deep_eval.py

@@ -2,6 +2,7 @@
 import io
 import json
 import logging
+import re
 from collections.abc import (
     Callable,
 )
@@ -84,6 +85,44 @@
 log = logging.getLogger(__name__)
 
 
+def _remap_state_dict_keys_for_pt(state_dict: dict[str, Any]) -> dict[str, Any]:
+    """Remap state dict keys from pt_expt naming to PT naming for compatibility.
+
+    The pt_expt backend uses different naming conventions than the PT backend:
+    - pt_expt uses "_min_nbor_dist" → pt uses "min_nbor_dist"
+    - pt_expt uses ".w" (weights) → pt uses ".matrix"
+    - pt_expt uses ".b" (bias) → pt uses ".bias"
+
+    This function remaps pt_expt keys to PT format when loading pt_expt checkpoints
+    into the PT backend.
+
+    Parameters
+    ----------
+    state_dict : dict
+        The state dict to remap.
+
+    Returns
+    -------
+    dict
+        The remapped state dict.
+    """
+    remapped = {}
+    for key, value in state_dict.items():
+        new_key = key
+        # Remap _min_nbor_dist → min_nbor_dist
+        new_key = new_key.replace("._min_nbor_dist", ".min_nbor_dist")
+        # Remap layer weights: .w → .matrix (must be at end of key or before a dot)
+        # Match ".w" that ends the key or is followed by a dot (for nested keys)
+        new_key = re.sub(r"\.w$", ".matrix", new_key)
+        new_key = re.sub(r"\.w\.", ".matrix.", new_key)
+        # Remap layer bias: .b → .bias (must be at end of key or before a dot)
+        new_key = re.sub(r"\.b$", ".bias", new_key)
+        new_key = re.sub(r"\.b\.", ".bias.", new_key)
+        remapped[new_key] = value
+
+    return remapped
+
+
 class DeepEval(DeepEvalBackend):
     """PyTorch backend implementation of DeepEval.
 
@@ -170,6 +209,8 @@ def __init__(
             if not self.input_param.get("hessian_mode") and not no_jit:
                 model = torch.jit.script(model)
             self.dp = ModelWrapper(model)
+            # Remap state dict keys for compatibility with pt_expt checkpoints
+            state_dict = _remap_state_dict_keys_for_pt(state_dict)
             missing, unexpected = self.dp.load_state_dict(state_dict, strict=False)
             if missing:
                 log.warning(

deepmd/pt_expt/train/wrapper.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
+import re
 from typing import (
     Any,
 )
@@ -214,3 +215,67 @@ def get_extra_state(self) -> dict:
             "model_params": self.model_params,
             "train_infos": self.train_infos,
         }
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        strict: bool = True,
+        assign: bool = False,
+    ) -> torch.nn.modules.module._IncompatibleKeys:
+        """Load state dict with key remapping for PT backend compatibility.
+
+        This method handles loading checkpoints from the PT backend, which uses
+        different naming conventions:
+        - PT uses "min_nbor_dist" → pt_expt uses "_min_nbor_dist"
+        - PT uses ".matrix" (weights) → pt_expt uses ".w"
+        - PT uses ".bias" (bias) → pt_expt uses ".b"
+
+        Parameters
+        ----------
+        state_dict : dict
+            The state dict to load.
+        strict : bool
+            Whether to strictly enforce that the keys in state_dict match.
+        assign : bool
+            Whether to assign tensors in-place (PyTorch 2.1+ feature).
+
+        Returns
+        -------
+        _IncompatibleKeys
+            Named tuple with missing_keys and unexpected_keys.
+        """
+        # Remap keys from PT backend naming to pt_expt naming
+        remapped_state_dict = {}
+        for key, value in state_dict.items():
+            new_key = key
+            # Remap min_nbor_dist → _min_nbor_dist
+            new_key = new_key.replace(".min_nbor_dist", "._min_nbor_dist")
+            # Remap layer weights: .matrix → .w (must be at end of key or before a dot)
+            new_key = re.sub(r"\.matrix$", ".w", new_key)
+            new_key = re.sub(r"\.matrix\.", ".w.", new_key)
+            # Remap layer bias: .bias → .b (must be at end of key or before a dot)
+            # Note: only match ".bias" when it's a parameter, not when it's part of
+            # a module name. We detect this by checking if it ends the key or
+            # is followed by another dot (indicating it's a parameter name).
+            new_key = re.sub(r"\.bias$", ".b", new_key)
+            new_key = re.sub(r"\.bias\.", ".b.", new_key)
+            remapped_state_dict[new_key] = value
+
+        # Call parent's load_state_dict with remapped keys
+        result = super().load_state_dict(
+            remapped_state_dict, strict=strict, assign=assign
+        )
+
+        # Log warnings for missing/unexpected keys (matching PT backend behavior)
+        if result.missing_keys:
+            log.warning(
+                "Checkpoint loaded with missing keys (likely from an older version): %s",
+                result.missing_keys,
+            )
+        if result.unexpected_keys:
+            log.warning(
+                "Checkpoint loaded with unexpected keys: %s",
+                result.unexpected_keys,
+            )
+
+        return result