feat(finetune): initialize native spin from spin-free models

deepmd/pt/train/training.py

@@ -96,6 +96,10 @@
 from deepmd.pt.utils.learning_rate import (
     BaseLR,
 )
+from deepmd.pt.utils.spin_finetune import (
+    compute_native_spin_input_stats,
+    get_native_spin_finetune_state,
+)
 from deepmd.pt.utils.stat import (
     make_stat_input,
     pair_margin_frame_mask,
@@ -770,6 +774,24 @@ def collect_single_finetune_params(
                                 _pretrained_model_params["descriptor"],
                                 _model_key_from,
                             )
+                        if _finetune_rule_single.convert_to_native_spin:
+                            log.info(
+                                "Initialize native-spin branch '%s' from spin-free branch '%s'.",
+                                _model_key,
+                                _model_key_from,
+                            )
+                            converted = get_native_spin_finetune_state(
+                                self.wrapper.model[_model_key],
+                                pretrained_model_wrapper.model[_model_key_from],
+                                random_fitting=_new_fitting,
+                            )
+                            _new_state_dict.update(
+                                {
+                                    f"model.{_model_key}.{key}": value
+                                    for key, value in converted.items()
+                                }
+                            )
+                            return
                         target_keys = [
                             i
                             for i in _random_state_dict.keys()
@@ -846,11 +868,16 @@ def update_finetune_bias(
                         model = (
                             self.model[_model_key] if self.multi_task else self.model
                         )
-                        model = model_change_out_bias(
-                            model,
+                        sample_func = (
                             self.get_sample_func[_model_key]
                             if self.multi_task
-                            else self.get_sample_func,
+                            else self.get_sample_func
+                        )
+                        if _finetune_rule.convert_to_native_spin:
+                            compute_native_spin_input_stats(model, sample_func)
+                        model = model_change_out_bias(
+                            model,
+                            sample_func,
                             _bias_adjust_mode="change-by-statistic"
                             if not _finetune_rule.get_random_fitting()
                             else "set-by-statistic",

deepmd/pt/utils/finetune.py

@@ -21,6 +21,33 @@
 __all__ = ["get_finetune_rule_single", "get_finetune_rules"]
 
 
+def _learned_part_path(model: torch.nn.Module) -> str | None:
+    """Return the state-dict path of the learned descriptor-fitting part of a model.
+
+    A plain model holds that part as its atomic model and a composition with
+    one learned sub-model, such as a zone-bridging one, as that sub-model, so
+    fine-tuning between the two transfers the part across the paths. A model
+    without a single learned part returns ``None``.
+    """
+    atomic_model = model.atomic_model
+    if getattr(atomic_model, "descriptor", None) is not None:
+        return "atomic_model."
+    learned = [
+        index
+        for index, child in enumerate(getattr(atomic_model, "models", ()))
+        if getattr(child, "descriptor", None) is not None
+    ]
+    return f"atomic_model.models.{learned[0]}." if len(learned) == 1 else None
+
+
+def _move_learned_key(key: str, path: str, source_path: str) -> str:
+    """Move a descriptor or fitting-net key between learned atomic-model paths."""
+    for part in ("descriptor.", "fitting_net."):
+        if f".{path}{part}" in key:
+            return key.replace(f".{path}{part}", f".{source_path}{part}", 1)
+    return key
+
+
 def get_finetune_rules(
     finetune_model: str,
     model_config: dict[str, Any],