From 8b5e3f85fecee18283263fbb0151b1da280ae2b8 Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 13:59:45 +0000 Subject: [PATCH 1/7] fix(pt): preserve legacy property checkpoint semantics --- deepmd/pt/entrypoints/freeze_pt2.py | 5 +- deepmd/pt/entrypoints/main.py | 11 +++- deepmd/pt/infer/deep_eval.py | 7 +- deepmd/pt/infer/inference.py | 7 +- deepmd/pt/nvalchemi/dpa4wrapper.py | 7 +- deepmd/pt/train/training.py | 7 +- deepmd/pt/utils/finetune.py | 5 +- deepmd/pt/utils/serialization.py | 7 +- deepmd/pt_expt/entrypoints/main.py | 11 +++- deepmd/pt_expt/infer/deep_eval.py | 3 +- deepmd/pt_expt/train/training.py | 7 +- deepmd/pt_expt/utils/finetune.py | 5 +- deepmd/utils/pt_checkpoint.py | 78 +++++++++++++++++++++++ source/tests/common/test_pt_checkpoint.py | 61 ++++++++++++++++++ 14 files changed, 206 insertions(+), 15 deletions(-) create mode 100644 source/tests/common/test_pt_checkpoint.py diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..68caf09a90 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/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: diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index da6b890d13..357c064852 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -94,6 +94,9 @@ from deepmd.utils.model_preset import ( expand_model_preset, ) +from deepmd.utils.pt_checkpoint import ( + normalize_pt_model_params, +) from deepmd.utils.stat_file import ( StatFileSpec, ) @@ -375,7 +378,9 @@ 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( @@ -541,7 +546,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() diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..c1693e81e0 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/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: diff --git a/deepmd/pt/infer/inference.py b/deepmd/pt/infer/inference.py index b026cd54c5..51518215ad 100644 --- a/deepmd/pt/infer/inference.py +++ b/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!" diff --git a/deepmd/pt/nvalchemi/dpa4wrapper.py b/deepmd/pt/nvalchemi/dpa4wrapper.py index 720d1b5472..96308f8cba 100644 --- a/deepmd/pt/nvalchemi/dpa4wrapper.py +++ b/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: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7a7210962e..2aa702e43e 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -139,6 +139,9 @@ from deepmd.utils.out_stat import ( ReduStatScanner, ) +from deepmd.utils.pt_checkpoint import ( + normalize_pt_model_params, +) if torch.__version__.startswith("2"): import torch._dynamo @@ -834,7 +837,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 diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index bf1c328408..165842a3f7 100644 --- a/deepmd/pt/utils/finetune.py +++ b/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, diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index 6cdd65712f..022599fc77 100644 --- a/deepmd/pt/utils/serialization.py +++ b/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) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7a381b4938..76899d66d6 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -39,6 +39,9 @@ get_data, process_systems, ) +from deepmd.utils.pt_checkpoint import ( + normalize_pt_model_params, +) from deepmd.utils.stat_file import ( StatFileSpec, ) @@ -350,7 +353,9 @@ 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"] + ) return config @@ -578,7 +583,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 +750,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 ( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..4c4382c120 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/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). diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..a94ca97c6c 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -151,6 +151,9 @@ from deepmd.utils.finetune import ( warn_configuration_mismatch_during_finetune, ) +from deepmd.utils.pt_checkpoint import ( + normalize_pt_model_params, +) from deepmd.utils.stat_file import ( StatFileSpec, open_stat_file, @@ -2150,7 +2153,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: diff --git a/deepmd/pt_expt/utils/finetune.py b/deepmd/pt_expt/utils/finetune.py index c44de996dc..2788e6e29a 100644 --- a/deepmd/pt_expt/utils/finetune.py +++ b/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( diff --git a/deepmd/utils/pt_checkpoint.py b/deepmd/utils/pt_checkpoint.py index de71c50a09..a04a186751 100644 --- a/deepmd/utils/pt_checkpoint.py +++ b/deepmd/utils/pt_checkpoint.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Utilities shared by PyTorch checkpoint backends.""" +import logging from collections.abc import ( Mapping, + MutableMapping, ) from typing import ( Any, ) +log = logging.getLogger(__name__) + def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None: """Detect the parameter dialect of a raw PyTorch checkpoint. @@ -49,3 +53,77 @@ def detect_pt_checkpoint_backend(checkpoint: Any) -> str | None: if has_pt_expt_bias and not has_pt_bias: return "pt-expt" return None + + +def normalize_pt_model_params(model_params: dict[str, Any]) -> dict[str, Any]: + """Apply compatibility fixes to model parameters loaded from a ``.pt`` file. + + Property fitting networks used to omit ``distinguish_types`` from their + serialized model parameters. The historical checkpoint behavior for + those networks is ``False`` even though the constructor default remains + ``True`` for newly-created models. This helper fills in the omitted value + while loading a checkpoint and leaves all other model parameters untouched. + + The model parameter layout can contain nested ``model_dict`` branches, + ``shared_dict`` entries, or child models in ``models``. Walking every + mapping/list makes the compatibility fix apply to each of these layouts, + including a shared fitting-network definition, without changing the + behavior of regular model construction. + + Parameters + ---------- + model_params : dict[str, Any] + The ``_extra_state["model_params"]`` dictionary from a PyTorch + checkpoint. The dictionary is normalized in place. + + Returns + ------- + dict[str, Any] + ``model_params`` itself, after applying compatibility fixes. + """ + seen: set[int] = set() + changed = False + + def visit(value: Any, shared_dict: Mapping[str, Any] | None = None) -> None: + nonlocal changed + if isinstance(value, MutableMapping): + # A shared config can be reachable through more than one path. + # Avoid revisiting it (and emitting duplicate warnings) when the + # object is referenced multiple times in the loaded parameters. + value_id = id(value) + if value_id in seen: + return + seen.add(value_id) + + local_shared_dict = value.get("shared_dict") + if not isinstance(local_shared_dict, Mapping): + local_shared_dict = shared_dict + + fitting_net = value.get("fitting_net") + if isinstance(fitting_net, str) and local_shared_dict is not None: + # Shared references may carry a parameter-sharing level (for + # example ``"property_head:1"``); the level does not change + # the fitting-net configuration to normalize. + fitting_net = local_shared_dict.get(fitting_net.split(":", 1)[0]) + if ( + isinstance(fitting_net, MutableMapping) + and str(fitting_net.get("type", "")).lower() == "property" + and "distinguish_types" not in fitting_net + ): + fitting_net["distinguish_types"] = False + changed = True + + for child in value.values(): + visit(child, local_shared_dict) + elif isinstance(value, (list, tuple)): + for child in value: + visit(child, shared_dict) + + visit(model_params) + if changed: + log.warning( + "The `distinguish_types` option is missing from a property fitting " + "network in a PyTorch checkpoint; defaulting to False for backward " + "compatibility." + ) + return model_params diff --git a/source/tests/common/test_pt_checkpoint.py b/source/tests/common/test_pt_checkpoint.py new file mode 100644 index 0000000000..212d1dd762 --- /dev/null +++ b/source/tests/common/test_pt_checkpoint.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for PyTorch checkpoint compatibility helpers.""" + +import logging + +from deepmd.utils.pt_checkpoint import ( + normalize_pt_model_params, +) + + +def test_normalize_property_fitting_net_defaults_distinguish_types(caplog) -> None: + model_params = { + "type_map": ["H"], + "fitting_net": {"type": "property", "property_name": "foo"}, + } + + with caplog.at_level(logging.WARNING, logger="deepmd.utils.pt_checkpoint"): + result = normalize_pt_model_params(model_params) + + assert result is model_params + assert model_params["fitting_net"]["distinguish_types"] is False + assert "missing from a property fitting network" in caplog.text + + +def test_normalize_pt_model_params_handles_multitask_and_shared_models(caplog) -> None: + shared_fitting = {"type": "property"} + model_params = { + "shared_dict": {"property_fitting": shared_fitting}, + "model_dict": { + "head_a": {"fitting_net": "property_fitting:1"}, + "head_b": {"fitting_net": {"type": "property"}}, + }, + "models": [{"fitting_net": {"type": "property"}}], + } + + with caplog.at_level(logging.WARNING, logger="deepmd.utils.pt_checkpoint"): + normalize_pt_model_params(model_params) + + assert shared_fitting["distinguish_types"] is False + assert ( + model_params["model_dict"]["head_b"]["fitting_net"]["distinguish_types"] + is False + ) + assert model_params["models"][0]["fitting_net"]["distinguish_types"] is False + assert caplog.text.count("missing from a property fitting network") == 1 + + +def test_normalize_pt_model_params_preserves_explicit_and_non_property_values( + caplog, +) -> None: + model_params = { + "fitting_net": {"type": "property", "distinguish_types": True}, + "other_fitting": {"type": "ener"}, + } + + with caplog.at_level(logging.WARNING, logger="deepmd.utils.pt_checkpoint"): + normalize_pt_model_params(model_params) + + assert model_params["fitting_net"]["distinguish_types"] is True + assert "distinguish_types" not in model_params["other_fitting"] + assert not caplog.records From b8652c5b712ac59ad8187d452b632ab977c6bb92 Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 14:22:43 +0000 Subject: [PATCH 2/7] test(pt): cover legacy property checkpoint loading --- source/tests/pt/test_checkpoint_compat.py | 51 +++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 source/tests/pt/test_checkpoint_compat.py diff --git a/source/tests/pt/test_checkpoint_compat.py b/source/tests/pt/test_checkpoint_compat.py new file mode 100644 index 0000000000..e68efd8689 --- /dev/null +++ b/source/tests/pt/test_checkpoint_compat.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for PyTorch checkpoint compatibility.""" + +from copy import ( + deepcopy, +) + +import torch + +from deepmd.pt.infer.inference import Tester as PTTester +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.train.wrapper import ( + ModelWrapper, +) + + +def test_legacy_property_checkpoint_defaults_to_type_agnostic_stats(tmp_path) -> None: + """Legacy property checkpoints retain their pre-#5322 statistics behavior.""" + model_params = { + "type_map": ["A", "B"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "rcut": 3.0, + "rcut_smth": 0.5, + "neuron": [4], + "axis_neuron": 2, + "seed": 1, + }, + "fitting_net": { + "type": "property", + "property_name": "property", + "task_dim": 1, + "intensive": True, + "neuron": [4], + "seed": 1, + "distinguish_types": True, + }, + } + model = get_model(deepcopy(model_params)).to("cpu") + wrapper = ModelWrapper(model, model_params=deepcopy(model_params)) + state_dict = wrapper.state_dict() + state_dict["_extra_state"]["model_params"]["fitting_net"].pop("distinguish_types") + checkpoint = tmp_path / "legacy.pt" + torch.save({"model": state_dict}, checkpoint) + + loaded = PTTester(str(checkpoint)) + + assert loaded.model.atomic_model.fitting_net.get_distinguish_types() is False From 8581488a455b1836f4fa50e6d086ee66c568663a Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 15:01:34 +0000 Subject: [PATCH 3/7] fix(pt): normalize legacy restart model parameters --- deepmd/pt/entrypoints/main.py | 18 +++++ deepmd/pt/train/training.py | 12 +++ deepmd/pt_expt/entrypoints/main.py | 17 ++++ deepmd/pt_expt/train/training.py | 12 +++ deepmd/utils/pt_checkpoint.py | 99 +++++++++++++++++++++++ source/tests/common/test_pt_checkpoint.py | 28 +++++++ 6 files changed, 186 insertions(+) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 357c064852..8ba2a78bb6 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -95,6 +95,7 @@ expand_model_preset, ) from deepmd.utils.pt_checkpoint import ( + apply_pt_checkpoint_defaults, normalize_pt_model_params, ) from deepmd.utils.stat_file import ( @@ -388,6 +389,23 @@ def train( ).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: + 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) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 2aa702e43e..68fe31bc55 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -140,6 +140,7 @@ ReduStatScanner, ) from deepmd.utils.pt_checkpoint import ( + apply_pt_checkpoint_defaults, normalize_pt_model_params, ) @@ -210,6 +211,17 @@ 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) training_params = config["training"] optimizer_params = config.get("optimizer", {}) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 76899d66d6..8afff64bfa 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -40,6 +40,7 @@ process_systems, ) from deepmd.utils.pt_checkpoint import ( + apply_pt_checkpoint_defaults, normalize_pt_model_params, ) from deepmd.utils.stat_file import ( @@ -357,6 +358,22 @@ def preprocess_config( 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 def update_neighbor_stat( diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index a94ca97c6c..cce8d44055 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -152,6 +152,7 @@ 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 ( @@ -1798,6 +1799,17 @@ 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) training_params = config["training"] optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating", {}) or {} diff --git a/deepmd/utils/pt_checkpoint.py b/deepmd/utils/pt_checkpoint.py index a04a186751..516de96965 100644 --- a/deepmd/utils/pt_checkpoint.py +++ b/deepmd/utils/pt_checkpoint.py @@ -127,3 +127,102 @@ def visit(value: Any, shared_dict: Mapping[str, Any] | None = None) -> None: "compatibility." ) return model_params + + +def apply_pt_checkpoint_defaults( + model_params: dict[str, Any], + checkpoint_model_params: Mapping[str, Any], +) -> dict[str, Any]: + """Fill omitted compatibility fields in a model config from a checkpoint. + + This is used before constructing a model for ``init_model`` and + ``restart_model``. A checkpoint that predates the property fitting-net + metadata field represents ``distinguish_types=False``. Only fields that + are omitted from both the current config and the checkpoint are filled, so + an explicit value in the current training input remains authoritative. + + Parameters + ---------- + model_params : dict[str, Any] + The model parameters that will be used to construct the new model. + checkpoint_model_params : Mapping[str, Any] + The raw ``_extra_state["model_params"]`` from the checkpoint. + + Returns + ------- + dict[str, Any] + ``model_params`` itself, after applying compatible defaults. + """ + changed = False + seen: set[tuple[int, int]] = set() + + def resolve_fitting_net( + fitting_net: Any, + shared_dict: Mapping[str, Any] | None, + ) -> Any: + if isinstance(fitting_net, str) and shared_dict is not None: + return shared_dict.get(fitting_net.split(":", 1)[0]) + return fitting_net + + def visit( + value: Any, + checkpoint_value: Any, + shared_dict: Mapping[str, Any] | None = None, + checkpoint_shared_dict: Mapping[str, Any] | None = None, + ) -> None: + nonlocal changed + if isinstance(value, MutableMapping) and isinstance(checkpoint_value, Mapping): + pair = (id(value), id(checkpoint_value)) + if pair in seen: + return + seen.add(pair) + + local_shared_dict = value.get("shared_dict") + if not isinstance(local_shared_dict, Mapping): + local_shared_dict = shared_dict + local_checkpoint_shared_dict = checkpoint_value.get("shared_dict") + if not isinstance(local_checkpoint_shared_dict, Mapping): + local_checkpoint_shared_dict = checkpoint_shared_dict + + fitting_net = resolve_fitting_net( + value.get("fitting_net"), local_shared_dict + ) + checkpoint_fitting_net = resolve_fitting_net( + checkpoint_value.get("fitting_net"), local_checkpoint_shared_dict + ) + if ( + isinstance(fitting_net, MutableMapping) + and isinstance(checkpoint_fitting_net, Mapping) + and str(fitting_net.get("type", "")).lower() == "property" + and str(checkpoint_fitting_net.get("type", "")).lower() == "property" + and "distinguish_types" not in fitting_net + and "distinguish_types" not in checkpoint_fitting_net + ): + fitting_net["distinguish_types"] = False + changed = True + + for key, child in value.items(): + if key in checkpoint_value: + visit( + child, + checkpoint_value[key], + local_shared_dict, + local_checkpoint_shared_dict, + ) + elif isinstance(value, list) and isinstance(checkpoint_value, list): + for child, checkpoint_child in zip(value, checkpoint_value, strict=False): + visit( + child, + checkpoint_child, + shared_dict, + checkpoint_shared_dict, + ) + + visit(model_params, checkpoint_model_params) + if changed: + log.warning( + "The `distinguish_types` option is missing from a property fitting " + "network in the model config and checkpoint; defaulting to False " + "for backward compatibility." + ) + return model_params diff --git a/source/tests/common/test_pt_checkpoint.py b/source/tests/common/test_pt_checkpoint.py index 212d1dd762..15d4bd3d1c 100644 --- a/source/tests/common/test_pt_checkpoint.py +++ b/source/tests/common/test_pt_checkpoint.py @@ -4,6 +4,7 @@ import logging from deepmd.utils.pt_checkpoint import ( + apply_pt_checkpoint_defaults, normalize_pt_model_params, ) @@ -59,3 +60,30 @@ def test_normalize_pt_model_params_preserves_explicit_and_non_property_values( assert model_params["fitting_net"]["distinguish_types"] is True assert "distinguish_types" not in model_params["other_fitting"] assert not caplog.records + + +def test_apply_pt_checkpoint_defaults_preserves_explicit_config_values(caplog) -> None: + checkpoint_params = { + "fitting_net": {"type": "property", "property_name": "foo"}, + } + model_params = { + "fitting_net": {"type": "property", "property_name": "foo"}, + } + + with caplog.at_level(logging.WARNING, logger="deepmd.utils.pt_checkpoint"): + apply_pt_checkpoint_defaults(model_params, checkpoint_params) + + assert model_params["fitting_net"]["distinguish_types"] is False + assert "model config and checkpoint" in caplog.text + + explicit_params = { + "fitting_net": { + "type": "property", + "property_name": "foo", + "distinguish_types": True, + }, + } + with caplog.at_level(logging.WARNING, logger="deepmd.utils.pt_checkpoint"): + apply_pt_checkpoint_defaults(explicit_params, checkpoint_params) + + assert explicit_params["fitting_net"]["distinguish_types"] is True From 5335e2bb5fb73ee567f2ad76e3ac525fdfe215a7 Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 15:08:38 +0000 Subject: [PATCH 4/7] perf(pt): reuse init checkpoint metadata --- deepmd/pt/entrypoints/main.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 8ba2a78bb6..4c73ae50ed 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -394,12 +394,15 @@ def train( # constructor default, so restart and init-model construction preserve the # checkpoint behavior. if init_model is not None or restart is not None: - 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"] + 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" ) From 8befcc3e18f9cd08a221a1ed23aafc38bba5ff23 Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 15:12:22 +0000 Subject: [PATCH 5/7] fix(pt-expt): preserve metadata when resuming --- deepmd/pt_expt/train/training.py | 7 ++++ source/tests/pt/test_checkpoint_compat.py | 45 +++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index cce8d44055..f5602a186b 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2296,6 +2296,13 @@ 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. + state_dict["_extra_state"] = self._unwrapped.state_dict()[ + "_extra_state" + ] self._unwrapped.load_state_dict(state_dict) if shared_links is not None: diff --git a/source/tests/pt/test_checkpoint_compat.py b/source/tests/pt/test_checkpoint_compat.py index e68efd8689..f72a26f050 100644 --- a/source/tests/pt/test_checkpoint_compat.py +++ b/source/tests/pt/test_checkpoint_compat.py @@ -49,3 +49,48 @@ def test_legacy_property_checkpoint_defaults_to_type_agnostic_stats(tmp_path) -> loaded = PTTester(str(checkpoint)) assert loaded.model.atomic_model.fitting_net.get_distinguish_types() is False + + +def test_legacy_checkpoint_load_keeps_current_model_metadata() -> None: + """Loading legacy weights does not replace explicit current metadata.""" + model_params = { + "type_map": ["A", "B"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "rcut": 3.0, + "rcut_smth": 0.5, + "neuron": [4], + "axis_neuron": 2, + "seed": 1, + }, + "fitting_net": { + "type": "property", + "property_name": "property", + "task_dim": 1, + "intensive": True, + "neuron": [4], + "seed": 1, + "distinguish_types": True, + }, + } + source = ModelWrapper( + get_model(deepcopy(model_params)).to("cpu"), + model_params=deepcopy(model_params), + ) + legacy_state = source.state_dict() + legacy_state["_extra_state"]["model_params"]["fitting_net"].pop("distinguish_types") + + target = ModelWrapper( + get_model(deepcopy(model_params)).to("cpu"), + model_params=deepcopy(model_params), + ) + legacy_state["_extra_state"] = target.state_dict()["_extra_state"] + target.load_state_dict(legacy_state) + + assert ( + target.state_dict()["_extra_state"]["model_params"]["fitting_net"][ + "distinguish_types" + ] + is True + ) From 50d1c0532e113818b42327ec25078acbc16968d9 Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Fri, 16 Oct 2026 15:25:51 +0000 Subject: [PATCH 6/7] perf(pt): release preliminary resume checkpoint --- deepmd/pt/train/training.py | 1 + deepmd/pt_expt/train/training.py | 1 + 2 files changed, 2 insertions(+) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 68fe31bc55..f3b2c8260f 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -222,6 +222,7 @@ def __init__( ) 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", {}) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f5602a186b..35c183adde 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1810,6 +1810,7 @@ def __init__( ) 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 {} From 3ade174d3a4767ece37b7b048aebe1fdf89f1bfa Mon Sep 17 00:00:00 2001 From: Ming-Yu Guo Date: Sat, 17 Oct 2026 06:49:23 +0000 Subject: [PATCH 7/7] fix(pt-expt): preserve validation state on restart --- deepmd/pt_expt/train/training.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 35c183adde..fdc3b131aa 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2301,9 +2301,14 @@ def update_finetune_bias( # construction. Legacy checkpoints omit property # ``distinguish_types``; restoring their extra state here # would make a subsequent save lose an explicit config value. - state_dict["_extra_state"] = self._unwrapped.state_dict()[ - "_extra_state" - ] + 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: