fix(cli): fit downstream bias without reapplying checkpoint presets

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -822,6 +822,7 @@ def change_out_bias(
         stat_file_path: DPPath | None = None,
         bias_adjust_mode: str = "change-by-statistic",
         observed_type: list[str] | None = None,
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change the output bias according to the input data and the pretrained model.
 
@@ -845,6 +846,12 @@ def change_out_bias(
             ('set-by-statistic') or their stored bias ('change-by-statistic'); no
             statistics are computed, read or written for such an output, whose
             output std keeps its stored value.
+        preset_out_bias : dict, optional
+            An explicit downstream preset. Assigned outputs use this bias and
+            skip fitting; other outputs fit residuals against the stored bias.
+            The model preset is used when omitted, preserving training and
+            fine-tuning behavior. Pass an empty dict to fit from the data
+            without reapplying the model's stored preset.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
         observed_type : list[str], optional
@@ -859,17 +866,22 @@ def change_out_bias(
         if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"):
             raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)
         change = bias_adjust_mode == "change-by-statistic"
+        preset = (
+            normalize_preset_out_bias(preset_out_bias, self.type_map)
+            if preset_out_bias is not None
+            else self.preset_out_bias
+        )
         distinguish_types = self.get_compute_stats_distinguish_types()
-        check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types)
+        check_preset_out_bias(preset, self.bias_keys, distinguish_types)
         sampled = sample_merged
         # === Step 1. Outputs fixed by the preset ===
         fitted_keys = self.bias_keys
-        if self.preset_out_bias:
+        if preset:
             if observed_type is None:
                 sampled = sample_merged() if callable(sample_merged) else sample_merged
                 observed_type = collect_observed_types(sampled, self.type_map)
             rows = preset_out_bias_rows(
-                self.preset_out_bias,
+                preset,
                 self.type_map,
                 observed_type,
                 to_numpy_array(self.out_bias),

deepmd/dpmodel/model/make_model.py

@@ -926,6 +926,7 @@ def change_out_bias(
             self,
             merged: Any,
             bias_adjust_mode: str = "change-by-statistic",
+            preset_out_bias: dict | None = None,
         ) -> None:
             """Change the output bias according to the input data and the pretrained model.
 
@@ -937,7 +938,11 @@ def change_out_bias(
                 The mode for changing output bias:
                 'change-by-statistic' or 'set-by-statistic'.
             """
-            self.atomic_model.change_out_bias(merged, bias_adjust_mode=bias_adjust_mode)
+            self.atomic_model.change_out_bias(
+                merged,
+                bias_adjust_mode=bias_adjust_mode,
+                preset_out_bias=preset_out_bias,
+            )
 
         def _input_type_cast(
             self,

deepmd/dpmodel/model/spin_model.py

@@ -31,6 +31,10 @@
 from deepmd.dpmodel.output_def import (
     ModelOutputDef,
 )
+from deepmd.utils.preset_out_bias import (
+    normalize_preset_out_bias,
+    remap_preset_out_bias,
+)
 from deepmd.utils.spin import (
     Spin,
 )
@@ -603,6 +607,7 @@ def change_out_bias(
         self,
         merged: Callable[[], list[dict]] | list[dict],
         bias_adjust_mode: str = "change-by-statistic",
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change the output bias of atomic model according to the input data and the pretrained model.
 
@@ -621,12 +626,17 @@ def change_out_bias(
                     and do least square on the errors to obtain the target shift as bias.
             'set-by-statistic' : directly use the statistic output bias in the target dataset.
         """
+        preset_out_bias = remap_preset_out_bias(
+            normalize_preset_out_bias(preset_out_bias, self.get_type_map()),
+            list(range(self.get_ntypes())) + [-1] * self.get_ntypes(),
+        )
         spin_sampled_func = self._get_spin_sampled_func(
             merged if callable(merged) else lambda: merged
         )
         self.backbone_model.change_out_bias(
             spin_sampled_func,
             bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias=preset_out_bias,
         )
 
     def change_type_map(

deepmd/dpmodel/train/trainer.py

@@ -74,12 +74,17 @@ def change_model_out_bias(
     *,
     bias_adjust_mode: str = "change-by-statistic",
     recompute_input_stats: bool = False,
+    preset_out_bias: dict | None = None,
 ) -> Any:
     """Change one model's output bias and log the before/after values."""
     old_bias = deepcopy(model.get_out_bias())
+    preset_kwargs = (
+        {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {}
+    )
     model.change_out_bias(
         sample_func,
         bias_adjust_mode=bias_adjust_mode,
+        **preset_kwargs,
     )
     new_bias = deepcopy(model.get_out_bias())
 

deepmd/pd/entrypoints/main.py

@@ -552,7 +552,10 @@ def change_bias(
             nbatches,
         )
         updated_model = training.model_change_out_bias(
-            model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode
+            model_to_change,
+            sampled_data,
+            _bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias={},
         )
 
     if not multi_task:

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -501,6 +501,7 @@ def change_out_bias(
         stat_file_path: DPPath | None = None,
         bias_adjust_mode: str = "change-by-statistic",
         observed_type: list[str] | None = None,
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change the output bias according to the input data and the pretrained model.
 
@@ -524,6 +525,12 @@ def change_out_bias(
             ('set-by-statistic') or their stored bias ('change-by-statistic'); no
             statistics are computed, read or written for such an output, whose
             output std keeps its stored value.
+        preset_out_bias : dict, optional
+            An explicit downstream preset. Assigned outputs use this bias and
+            skip fitting; other outputs fit residuals against the stored bias.
+            The model preset is used when omitted, preserving training and
+            fine-tuning behavior. Pass an empty dict to fit from the data
+            without reapplying the model's stored preset.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
         observed_type : list[str], optional
@@ -533,17 +540,22 @@ def change_out_bias(
         if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"):
             raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)
         change = bias_adjust_mode == "change-by-statistic"
+        preset = (
+            normalize_preset_out_bias(preset_out_bias, self.type_map)
+            if preset_out_bias is not None
+            else self.preset_out_bias
+        )
         distinguish_types = self.get_compute_stats_distinguish_types()
-        check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types)
+        check_preset_out_bias(preset, self.bias_keys, distinguish_types)
         sampled = sample_merged
         # === Step 1. Outputs fixed by the preset ===
         fitted_keys = self.bias_keys
-        if self.preset_out_bias:
+        if preset:
             if observed_type is None:
                 sampled = sample_merged() if callable(sample_merged) else sample_merged
                 observed_type = collect_observed_types(sampled, self.type_map)
             rows = preset_out_bias_rows(
-                self.preset_out_bias,
+                preset,
                 self.type_map,
                 observed_type,
                 to_numpy_array(self.out_bias),

deepmd/pd/model/model/make_model.py

@@ -231,6 +231,7 @@ def change_out_bias(
             self,
             merged: list[dict] | Callable[[], list[dict]],
             bias_adjust_mode: str = "change-by-statistic",
+            preset_out_bias: dict | None = None,
         ) -> None:
             """Change the output bias of atomic model according to the input data and the pretrained model.
 
@@ -252,6 +253,7 @@ def change_out_bias(
             self.atomic_model.change_out_bias(
                 merged,
                 bias_adjust_mode=bias_adjust_mode,
+                preset_out_bias=preset_out_bias,
             )
 
         def forward_common_lower(

deepmd/pd/train/training.py

@@ -1455,11 +1455,16 @@ def model_change_out_bias(
     _model: Any,
     _sample_func: Any,
     _bias_adjust_mode: str = "change-by-statistic",
+    preset_out_bias: dict | None = None,
 ) -> None:
     old_bias = deepcopy(_model.get_out_bias())
+    preset_kwargs = (
+        {"preset_out_bias": preset_out_bias} if preset_out_bias is not None else {}
+    )
     _model.change_out_bias(
         _sample_func,
         bias_adjust_mode=_bias_adjust_mode,
+        **preset_kwargs,
     )
     new_bias = deepcopy(_model.get_out_bias())
 

deepmd/pt/entrypoints/main.py

@@ -637,7 +637,10 @@ def change_bias(
             min_pair_dist=0.0 if pair_filter is None else float(pair_filter.default),
         )
         updated_model = training.model_change_out_bias(
-            model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode
+            model_to_change,
+            sampled_data,
+            _bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias={},
         )
 
     if not multi_task:

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -629,6 +629,7 @@ def change_out_bias(
         bias_adjust_mode: str = "change-by-statistic",
         model_forward: Callable[..., dict[str, torch.Tensor]] | None = None,
         observed_type: list[str] | None = None,
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change the output bias according to the input data and the pretrained model.
 
@@ -654,6 +655,12 @@ def change_out_bias(
             ('set-by-statistic') or their stored bias ('change-by-statistic'); no
             statistics are computed, read or written for such an output, whose
             output std keeps its stored value.
+        preset_out_bias : dict, optional
+            An explicit downstream preset. Assigned outputs use this bias and
+            skip fitting; other outputs fit residuals against the stored bias.
+            The model preset is used when omitted, preserving training and
+            fine-tuning behavior. Pass an empty dict to fit from the data
+            without reapplying the model's stored preset.
         model_forward : Callable[..., dict[str, torch.Tensor]], optional
             Predictor of the complete atomic outputs used by 'change-by-statistic'.
             Defaults to the forward of this atomic model; a model that adds
@@ -665,17 +672,22 @@ def change_out_bias(
         if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"):
             raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)
         change = bias_adjust_mode == "change-by-statistic"
+        preset = (
+            normalize_preset_out_bias(preset_out_bias, self.type_map)
+            if preset_out_bias is not None
+            else self.preset_out_bias
+        )
         distinguish_types = self.get_compute_stats_distinguish_types()
-        check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types)
+        check_preset_out_bias(preset, self.bias_keys, distinguish_types)
         sampled = sample_merged
         # === Step 1. Outputs fixed by the preset ===
         fitted_keys = self.bias_keys
-        if self.preset_out_bias:
+        if preset:
             if observed_type is None:
                 sampled = sample_merged() if callable(sample_merged) else sample_merged
                 observed_type = collect_observed_types(sampled, self.type_map)
             rows = preset_out_bias_rows(
-                self.preset_out_bias,
+                preset,
                 self.type_map,
                 observed_type,
                 to_numpy_array(self.out_bias),

deepmd/pt/model/model/make_model.py

@@ -335,6 +335,7 @@ def change_out_bias(
             self,
             merged: Any,
             bias_adjust_mode: str = "change-by-statistic",
+            preset_out_bias: dict | None = None,
         ) -> None:
             """Change the output bias of atomic model according to the input data and the pretrained model.
 
@@ -359,6 +360,7 @@ def change_out_bias(
             self.atomic_model.change_out_bias(
                 merged,
                 bias_adjust_mode=bias_adjust_mode,
+                preset_out_bias=preset_out_bias,
                 model_forward=self.predict_atomic_outputs_for_stat,
             )
 

deepmd/pt/model/model/sezm_spin_model.py

@@ -43,6 +43,10 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.preset_out_bias import (
+    normalize_preset_out_bias,
+    remap_preset_out_bias,
+)
 from deepmd.utils.spin import (
     Spin,
 )
@@ -430,14 +434,20 @@ def change_out_bias(
         self,
         merged: Callable[[], list[dict[str, Any]]] | list[dict[str, Any]],
         bias_adjust_mode: str = "change-by-statistic",
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change output bias using spin-expanded sampled data."""
+        preset_out_bias = remap_preset_out_bias(
+            normalize_preset_out_bias(preset_out_bias, self.get_type_map()),
+            list(range(self.get_ntypes())) + [-1] * self.get_ntypes(),
+        )
         spin_sampled_func = self._get_spin_sampled_func(
             merged if callable(merged) else lambda: merged
         )
         super().change_out_bias(
             spin_sampled_func,
             bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias=preset_out_bias,
         )
 
     def change_type_map(

deepmd/pt/model/model/spin_model.py

@@ -27,6 +27,10 @@
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.preset_out_bias import (
+    normalize_preset_out_bias,
+    remap_preset_out_bias,
+)
 from deepmd.utils.spin import (
     Spin,
 )
@@ -518,6 +522,7 @@ def change_out_bias(
         self,
         merged: Callable[[], list[dict]] | list[dict],
         bias_adjust_mode: str = "change-by-statistic",
+        preset_out_bias: dict | None = None,
     ) -> None:
         """Change the output bias of atomic model according to the input data and the pretrained model.
 
@@ -536,12 +541,17 @@ def change_out_bias(
                     and do least square on the errors to obtain the target shift as bias.
             'set-by-statistic' : directly use the statistic output bias in the target dataset.
         """
+        preset_out_bias = remap_preset_out_bias(
+            normalize_preset_out_bias(preset_out_bias, self.get_type_map()),
+            list(range(self.get_ntypes())) + [-1] * self.get_ntypes(),
+        )
         spin_sampled_func = self._get_spin_sampled_func(
             merged if callable(merged) else lambda: merged
         )
         self.backbone_model.change_out_bias(
             spin_sampled_func,
             bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias=preset_out_bias,
         )
 
     def change_type_map(

deepmd/pt/train/training.py

@@ -2603,6 +2603,7 @@ def model_change_out_bias(
     _model: Any,
     _sample_func: Callable[[], Any],
     _bias_adjust_mode: str = "change-by-statistic",
+    preset_out_bias: dict | None = None,
 ) -> Any:
     from deepmd.pt.model.model.dp_model import (
         DPModelCommon,
@@ -2613,4 +2614,5 @@ def model_change_out_bias(
         _sample_func,
         bias_adjust_mode=_bias_adjust_mode,
         recompute_input_stats=isinstance(_model, DPModelCommon),
+        preset_out_bias=preset_out_bias,
     )

deepmd/pt_expt/entrypoints/main.py

@@ -852,7 +852,10 @@ def change_bias(
             0.0 if pair_filter is None else float(pair_filter.default),
         )
         model_to_change = model_change_out_bias(
-            model_to_change, sampled_data, _bias_adjust_mode=bias_adjust_mode
+            model_to_change,
+            sampled_data,
+            _bias_adjust_mode=bias_adjust_mode,
+            preset_out_bias={},
         )
 
     if input_file.endswith(".pt"):