fix(pt): preserve legacy property checkpoint semantics

deepmd/pt/entrypoints/freeze_pt2.py

@@ -85,6 +85,9 @@
 from deepmd.utils.model_branch_dict import (
     get_model_dict,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 log = logging.getLogger(__name__)
 
@@ -155,7 +158,7 @@ def _extract_state_and_params(
     params = extra.get("model_params")
     if not isinstance(params, dict):
         raise ValueError("Unsupported checkpoint: missing '_extra_state.model_params'.")
-    return inner, params
+    return inner, normalize_pt_model_params(params)
 
 
 def is_sezm_checkpoint(ckpt_path: str) -> bool:

deepmd/pt/entrypoints/main.py

@@ -94,6 +94,10 @@
 from deepmd.utils.model_preset import (
     expand_model_preset,
 )
+from deepmd.utils.pt_checkpoint import (
+    apply_pt_checkpoint_defaults,
+    normalize_pt_model_params,
+)
 from deepmd.utils.stat_file import (
     StatFileSpec,
 )
@@ -375,14 +379,36 @@ def train(
             )
             if "model" in init_state_dict:
                 init_state_dict = init_state_dict["model"]
-            config["model"] = init_state_dict["_extra_state"]["model_params"]
+            config["model"] = normalize_pt_model_params(
+                init_state_dict["_extra_state"]["model_params"]
+            )
         else:
             config["model"] = json.loads(
                 torch.jit.load(
                     init_frz_model, map_location=DEVICE
                 ).get_model_def_script()
             )
 
+    # A legacy checkpoint may omit property ``distinguish_types``.  Apply its
+    # compatibility default before argcheck fills omitted input fields with the
+    # constructor default, so restart and init-model construction preserve the
+    # checkpoint behavior.
+    if init_model is not None or restart is not None:
+        if init_model is not None and use_pretrain_script:
+            resume_state_dict = init_state_dict
+        else:
+            resume_path = init_model or restart
+            resume_state_dict = torch.load(
+                resume_path, map_location=DEVICE, weights_only=True
+            )
+            if "model" in resume_state_dict:
+                resume_state_dict = resume_state_dict["model"]
+        resume_model_params = resume_state_dict.get("_extra_state", {}).get(
+            "model_params"
+        )
+        if isinstance(resume_model_params, dict):
+            apply_pt_checkpoint_defaults(config["model"], resume_model_params)
+
     # argcheck
     config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
     config = normalize(config, multi_task=multi_task)
@@ -541,7 +567,9 @@ def change_bias(
             input_file, map_location=env.DEVICE, weights_only=True
         )
         model_state_dict = old_state_dict.get("model", old_state_dict)
-        model_params = model_state_dict["_extra_state"]["model_params"]
+        model_params = normalize_pt_model_params(
+            model_state_dict["_extra_state"]["model_params"]
+        )
     elif input_file.endswith(".pth"):
         old_model = torch.jit.load(input_file, map_location=env.DEVICE)
         model_params_string = old_model.get_model_def_script()

deepmd/pt/infer/deep_eval.py

@@ -93,6 +93,9 @@
 from deepmd.utils.model_branch_dict import (
     get_model_dict,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 if TYPE_CHECKING:
     import ase.neighborlist
@@ -170,7 +173,9 @@ def __init__(
             )
             if "model" in state_dict:
                 state_dict = state_dict["model"]
-            self.input_param = state_dict["_extra_state"]["model_params"]
+            self.input_param = normalize_pt_model_params(
+                state_dict["_extra_state"]["model_params"]
+            )
             self.model_def_script = self.input_param
             self.multi_task = "model_dict" in self.input_param
             if self.multi_task:

deepmd/pt/infer/inference.py

@@ -19,6 +19,9 @@
     DEVICE,
     JIT,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 if torch.__version__.startswith("2"):
     import torch._dynamo
@@ -40,7 +43,9 @@ def __init__(
         state_dict = torch.load(model_ckpt, map_location=DEVICE, weights_only=True)
         if "model" in state_dict:
             state_dict = state_dict["model"]
-        model_params = state_dict["_extra_state"]["model_params"]
+        model_params = normalize_pt_model_params(
+            state_dict["_extra_state"]["model_params"]
+        )
         self.multi_task = "model_dict" in model_params
         if self.multi_task:
             assert head is not None, "Head must be specified in multitask mode!"

deepmd/pt/nvalchemi/dpa4wrapper.py

@@ -62,6 +62,9 @@
     ELEMENT_TO_Z,
     SeZMModel,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 if TYPE_CHECKING:
     from pathlib import (
@@ -584,7 +587,9 @@ def _from_pt(
         state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
         if "model" in state_dict:
             state_dict = state_dict["model"]
-        model_params = state_dict["_extra_state"]["model_params"]
+        model_params = normalize_pt_model_params(
+            state_dict["_extra_state"]["model_params"]
+        )
 
         if "model_dict" in model_params:
             if head is None:

deepmd/pt/train/training.py

@@ -139,6 +139,10 @@
 from deepmd.utils.out_stat import (
     ReduStatScanner,
 )
+from deepmd.utils.pt_checkpoint import (
+    apply_pt_checkpoint_defaults,
+    normalize_pt_model_params,
+)
 
 if torch.__version__.startswith("2"):
     import torch._dynamo
@@ -207,6 +211,18 @@ def __init__(
         has_initial_state = resuming or init_frz_model is not None
         self.restart_training = restart_model is not None
         model_params = config["model"]
+        if resuming:
+            resume_state = torch.load(
+                resume_model, map_location="cpu", weights_only=True
+            )
+            if "model" in resume_state:
+                resume_state = resume_state["model"]
+            resume_model_params = resume_state.get("_extra_state", {}).get(
+                "model_params"
+            )
+            if isinstance(resume_model_params, dict):
+                apply_pt_checkpoint_defaults(model_params, resume_model_params)
+            del resume_state
         training_params = config["training"]
         optimizer_params = config.get("optimizer", {})
 
@@ -834,7 +850,9 @@ def epoch_length(model_key: str) -> int:
                     new_state_dict = {}
                     target_state_dict = self.wrapper.state_dict()
                     # pretrained_model
-                    pretrained_model_params = state_dict["_extra_state"]["model_params"]
+                    pretrained_model_params = normalize_pt_model_params(
+                        state_dict["_extra_state"]["model_params"]
+                    )
                     with scoped_env_defaults(eval_env_defaults):
                         pretrained_model = get_model_for_wrapper(
                             pretrained_model_params

deepmd/pt/utils/finetune.py

@@ -17,6 +17,9 @@
 from deepmd.utils.finetune import (
     get_finetune_rules_from_model_params,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 __all__ = ["get_finetune_rule_single", "get_finetune_rules"]
 
@@ -32,7 +35,7 @@ def get_finetune_rules(
     if "model" in state_dict:
         state_dict = state_dict["model"]
     return get_finetune_rules_from_model_params(
-        state_dict["_extra_state"]["model_params"],
+        normalize_pt_model_params(state_dict["_extra_state"]["model_params"]),
         model_config,
         model_branch=model_branch,
         change_model_params=change_model_params,

deepmd/pt/utils/serialization.py

@@ -15,6 +15,9 @@
 from deepmd.pt.utils import (
     env,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 
 def serialize_from_file(model_file: str) -> dict:
@@ -39,7 +42,9 @@ def serialize_from_file(model_file: str) -> dict:
         state_dict = torch.load(model_file, map_location="cpu", weights_only=True)
         if "model" in state_dict:
             state_dict = state_dict["model"]
-        model_def_script = state_dict["_extra_state"]["model_params"]
+        model_def_script = normalize_pt_model_params(
+            state_dict["_extra_state"]["model_params"]
+        )
         model = get_model(model_def_script)
         modelwrapper = ModelWrapper(model)
         modelwrapper.load_state_dict(state_dict)

deepmd/pt_expt/entrypoints/main.py

@@ -39,6 +39,10 @@
     get_data,
     process_systems,
 )
+from deepmd.utils.pt_checkpoint import (
+    apply_pt_checkpoint_defaults,
+    normalize_pt_model_params,
+)
 from deepmd.utils.stat_file import (
     StatFileSpec,
 )
@@ -350,7 +354,25 @@ def preprocess_config(
             )
             if "model" in init_state_dict:
                 init_state_dict = init_state_dict["model"]
-            config["model"] = init_state_dict["_extra_state"]["model_params"]
+            config["model"] = normalize_pt_model_params(
+                init_state_dict["_extra_state"]["model_params"]
+            )
+
+        # Preserve legacy property-checkpoint semantics before argcheck adds
+        # the constructor default for an omitted ``distinguish_types`` field.
+        if options.init_model is not None or options.restart is not None:
+            resume_path = options.init_model or options.restart
+            if not resume_path.endswith((".pte", ".pt2")):
+                resume_state_dict = torch.load(
+                    resume_path, map_location=DEVICE, weights_only=True
+                )
+                if "model" in resume_state_dict:
+                    resume_state_dict = resume_state_dict["model"]
+                resume_model_params = resume_state_dict.get("_extra_state", {}).get(
+                    "model_params"
+                )
+                if isinstance(resume_model_params, dict):
+                    apply_pt_checkpoint_defaults(config["model"], resume_model_params)
 
         return config
 
@@ -578,7 +600,7 @@ def freeze(
             f"Unsupported checkpoint format at '{model}': missing "
             "'_extra_state.model_params' in model state dict."
         )
-    model_params = extra_state["model_params"]
+    model_params = normalize_pt_model_params(extra_state["model_params"])
 
     multi_task = "model_dict" in model_params
     if multi_task:
@@ -745,7 +767,7 @@ def change_bias(
                 f"Unsupported checkpoint format at '{input_file}': missing "
                 "'_extra_state.model_params' in model state dict."
             )
-        model_params = extra_state["model_params"]
+        model_params = normalize_pt_model_params(extra_state["model_params"])
     elif input_file.endswith((".pte", ".pt2")):
         pte_data = serialize_from_file(input_file)
         from deepmd.pt_expt.model.model import (

deepmd/pt_expt/infer/deep_eval.py

@@ -75,6 +75,7 @@
 )
 from deepmd.utils.pt_checkpoint import (
     detect_pt_checkpoint_backend,
+    normalize_pt_model_params,
 )
 
 if TYPE_CHECKING:
@@ -593,7 +594,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None:
                 "load it with `dp --pt` instead. If this is an exported model, "
                 "use a `.pte` or `.pt2` artifact."
             )
-        model_params = deepcopy(extra["model_params"])
+        model_params = normalize_pt_model_params(deepcopy(extra["model_params"]))
 
         if "model_dict" in model_params:
             # Multi-task: pick the requested head (defaults to "Default" if present).

deepmd/pt_expt/train/training.py

@@ -151,6 +151,10 @@
 from deepmd.utils.finetune import (
     warn_configuration_mismatch_during_finetune,
 )
+from deepmd.utils.pt_checkpoint import (
+    apply_pt_checkpoint_defaults,
+    normalize_pt_model_params,
+)
 from deepmd.utils.stat_file import (
     StatFileSpec,
     open_stat_file,
@@ -1795,6 +1799,18 @@ def __init__(
         self.restart_training = restart_model is not None
 
         model_params = config["model"]
+        if resuming and not resume_model.endswith((".pte", ".pt2")):
+            resume_state = torch.load(
+                resume_model, map_location="cpu", weights_only=True
+            )
+            if "model" in resume_state:
+                resume_state = resume_state["model"]
+            resume_model_params = resume_state.get("_extra_state", {}).get(
+                "model_params"
+            )
+            if isinstance(resume_model_params, dict):
+                apply_pt_checkpoint_defaults(model_params, resume_model_params)
+            del resume_state
         training_params = config["training"]
         optimizer_params = config.get("optimizer", {})
         validating_params = config.get("validating", {}) or {}
@@ -2150,7 +2166,9 @@ def initialize_statistics(
                     pretrained_model_params = data["model_def_script"]
                     pretrained_model = BaseModel.deserialize(data["model"]).to(DEVICE)
                 else:
-                    pretrained_model_params = state_dict["_extra_state"]["model_params"]
+                    pretrained_model_params = normalize_pt_model_params(
+                        state_dict["_extra_state"]["model_params"]
+                    )
 
                 # Build pretrained model (single-task or multi-task)
                 if "model_dict" not in pretrained_model_params:
@@ -2279,6 +2297,18 @@ def update_finetune_bias(
                 )
             else:
                 # --- Normal resume (init_model / restart) --------------------
+                # Keep the effective model parameters from the current
+                # construction.  Legacy checkpoints omit property
+                # ``distinguish_types``; restoring their extra state here
+                # would make a subsequent save lose an explicit config value.
+                current_extra_state = self._unwrapped.state_dict()["_extra_state"]
+                if self.restart_training:
+                    state_dict["_extra_state"] = {
+                        **state_dict.get("_extra_state", {}),
+                        "model_params": current_extra_state["model_params"],
+                    }
+                else:
+                    state_dict["_extra_state"] = current_extra_state
                 self._unwrapped.load_state_dict(state_dict)
 
             if shared_links is not None:

deepmd/pt_expt/utils/finetune.py

@@ -20,6 +20,9 @@
 from deepmd.utils.finetune import (
     get_finetune_rules_from_model_params,
 )
+from deepmd.utils.pt_checkpoint import (
+    normalize_pt_model_params,
+)
 
 __all__ = ["get_finetune_rule_single", "get_finetune_rules"]
 
@@ -41,7 +44,7 @@ def _load_model_params(finetune_model: str) -> dict[str, Any]:
         state_dict = torch.load(finetune_model, map_location=DEVICE, weights_only=True)
         if "model" in state_dict:
             state_dict = state_dict["model"]
-        return state_dict["_extra_state"]["model_params"]
+        return normalize_pt_model_params(state_dict["_extra_state"]["model_params"])
 
 
 def get_finetune_rules(