From a51ca75308956918f6e7272e86e280ccb89686f5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 26 Apr 2026 04:41:40 +0000 Subject: [PATCH 1/3] fix(pt-expt): add key remapping for cross-backend checkpoint loading Fix checkpoint key mismatches between pt and pt_expt backends during dptest. The backends use different naming conventions for some buffers: - pt uses "min_nbor_dist" - pt_expt uses "_min_nbor_dist" Add key remapping in both directions: - pt_expt/train/wrapper.py: remap PT keys to pt_expt format - pt/infer/deep_eval.py: remap pt_expt keys to PT format Also add logging of missing/unexpected keys to pt_expt wrapper for better debugging. Agent-Logs-Url: https://github.com/deepmodeling/deepmd-kit/sessions/ee9417b6-69d2-497c-81d7-a9afff7c7d6c Co-authored-by: anyangml <137014849+anyangml@users.noreply.github.com> --- deepmd/pt/infer/deep_eval.py | 39 ++++++++++++++++++++ deepmd/pt_expt/train/wrapper.py | 63 +++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 2e30b8574a..f14d80667c 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -84,6 +84,43 @@ 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 slightly different naming conventions for some buffers: + - pt_expt uses "_min_nbor_dist" (with underscore prefix) + - pt uses "min_nbor_dist" (without underscore prefix) + + 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. + """ + # Mapping from pt_expt naming to PT naming + key_mappings = { + "._min_nbor_dist": ".min_nbor_dist", + } + + remapped = {} + for key, value in state_dict.items(): + new_key = key + for old_pattern, new_pattern in key_mappings.items(): + if old_pattern in key: + new_key = key.replace(old_pattern, new_pattern) + break + remapped[new_key] = value + + return remapped + + class DeepEval(DeepEvalBackend): """PyTorch backend implementation of DeepEval. @@ -170,6 +207,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..044f1f3e75 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -2,6 +2,7 @@ import logging from typing import ( Any, + ClassVar, ) import torch @@ -214,3 +215,65 @@ def get_extra_state(self) -> dict: "model_params": self.model_params, "train_infos": self.train_infos, } + + # Key mappings from PT backend to pt_expt backend for checkpoint compatibility + # PT backend uses different naming conventions for some buffers/parameters + _PT_TO_PT_EXPT_KEY_MAP: ClassVar[dict[str, str]] = { + # Buffer name difference: pt uses "min_nbor_dist", pt_expt uses "_min_nbor_dist" + ".min_nbor_dist": "._min_nbor_dist", + } + + 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 + slightly different naming conventions for some buffers (e.g., "min_nbor_dist" + vs "_min_nbor_dist"). + + 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 + for old_suffix, new_suffix in self._PT_TO_PT_EXPT_KEY_MAP.items(): + if old_suffix in key: + new_key = key.replace(old_suffix, new_suffix) + break + 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 From ea795c4d69a4a50bde26cdc3a664cfb990dc6e67 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 26 Apr 2026 05:18:19 +0000 Subject: [PATCH 2/3] fix(pt,pt_expt): expand state_dict key remapping for cross-backend compatibility Add complete key remapping between PT and pt_expt backends: - PT uses .matrix/.bias, pt_expt uses .w/.b for layer weights/biases - PT uses .min_nbor_dist, pt_expt uses ._min_nbor_dist This allows loading pt_expt checkpoints (.pt files) in PT backend and vice versa. Agent-Logs-Url: https://github.com/deepmodeling/deepmd-kit/sessions/287e8580-f3ed-4b4e-b31f-1c25935791fd Co-authored-by: anyangml <137014849+anyangml@users.noreply.github.com> --- deepmd/pt/infer/deep_eval.py | 25 ++++++++++++++----------- deepmd/pt_expt/train/wrapper.py | 23 +++++++++++++++++------ 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f14d80667c..e92c9e544c 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -87,9 +87,10 @@ 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 slightly different naming conventions for some buffers: - - pt_expt uses "_min_nbor_dist" (with underscore prefix) - - pt uses "min_nbor_dist" (without underscore prefix) + 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. @@ -104,18 +105,20 @@ def _remap_state_dict_keys_for_pt(state_dict: dict[str, Any]) -> dict[str, Any]: dict The remapped state dict. """ - # Mapping from pt_expt naming to PT naming - key_mappings = { - "._min_nbor_dist": ".min_nbor_dist", - } + import re remapped = {} for key, value in state_dict.items(): new_key = key - for old_pattern, new_pattern in key_mappings.items(): - if old_pattern in key: - new_key = key.replace(old_pattern, new_pattern) - break + # 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 diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index 044f1f3e75..6669c8b2e9 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -232,8 +232,10 @@ def load_state_dict( """Load state dict with key remapping for PT backend compatibility. This method handles loading checkpoints from the PT backend, which uses - slightly different naming conventions for some buffers (e.g., "min_nbor_dist" - vs "_min_nbor_dist"). + 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 ---------- @@ -249,14 +251,23 @@ def load_state_dict( _IncompatibleKeys Named tuple with missing_keys and unexpected_keys. """ + import re + # Remap keys from PT backend naming to pt_expt naming remapped_state_dict = {} for key, value in state_dict.items(): new_key = key - for old_suffix, new_suffix in self._PT_TO_PT_EXPT_KEY_MAP.items(): - if old_suffix in key: - new_key = key.replace(old_suffix, new_suffix) - break + # 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 From 22e3535863fe2a1b20c3a1a07a4409232f28b0a8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 26 Apr 2026 05:19:37 +0000 Subject: [PATCH 3/3] style: move re import to module level and remove unused ClassVar Agent-Logs-Url: https://github.com/deepmodeling/deepmd-kit/sessions/287e8580-f3ed-4b4e-b31f-1c25935791fd Co-authored-by: anyangml <137014849+anyangml@users.noreply.github.com> --- deepmd/pt/infer/deep_eval.py | 3 +-- deepmd/pt_expt/train/wrapper.py | 11 +---------- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index e92c9e544c..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, ) @@ -105,8 +106,6 @@ def _remap_state_dict_keys_for_pt(state_dict: dict[str, Any]) -> dict[str, Any]: dict The remapped state dict. """ - import re - remapped = {} for key, value in state_dict.items(): new_key = key diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index 6669c8b2e9..d24ef2e07a 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +import re from typing import ( Any, - ClassVar, ) import torch @@ -216,13 +216,6 @@ def get_extra_state(self) -> dict: "train_infos": self.train_infos, } - # Key mappings from PT backend to pt_expt backend for checkpoint compatibility - # PT backend uses different naming conventions for some buffers/parameters - _PT_TO_PT_EXPT_KEY_MAP: ClassVar[dict[str, str]] = { - # Buffer name difference: pt uses "min_nbor_dist", pt_expt uses "_min_nbor_dist" - ".min_nbor_dist": "._min_nbor_dist", - } - def load_state_dict( self, state_dict: dict[str, Any], @@ -251,8 +244,6 @@ def load_state_dict( _IncompatibleKeys Named tuple with missing_keys and unexpected_keys. """ - import re - # Remap keys from PT backend naming to pt_expt naming remapped_state_dict = {} for key, value in state_dict.items():