diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 2e30b8574a..7b40e68ee6 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/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( diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index f67efe8a8e..d24ef2e07a 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/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