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

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -51,6 +51,7 @@
     check_preset_out_bias,
     normalize_preset_out_bias,
     preset_out_bias_rows,
+    recover_nonfinite_preset_bias,
     remap_preset_out_bias,
 )
 
@@ -822,6 +823,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 +847,14 @@ 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.
+            Before residual fitting, non-finite stored cells covered by the
+            checkpoint preset are recovered; finite adjusted cells are retained.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
         observed_type : list[str], optional
@@ -859,17 +869,32 @@ 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
+        # A per-call override must not propagate a corrupt checkpoint baseline.
+        if change and preset_out_bias is not None and self.preset_out_bias:
+            recovered = recover_nonfinite_preset_bias(
+                self.preset_out_bias,
+                self.type_map,
+                to_numpy_array(self.out_bias),
+                self.bias_keys,
+                [self.atomic_output_def()[kk].size for kk in self.bias_keys],
+            )
+            self._store_out_stat(recovered)
         # === 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.
 
@@ -936,8 +937,17 @@ def change_out_bias(
             bias_adjust_mode : str
                 The mode for changing output bias:
                 'change-by-statistic' or 'set-by-statistic'.
+            preset_out_bias : dict, optional
+                Per-call preset override. Omit to use the configured preset;
+                pass an empty dict to fit without fixing preset outputs.
+                Non-finite stored cells covered by the checkpoint preset are
+                recovered before residual fitting; finite values are retained.
             """
-            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.
 
@@ -620,13 +625,24 @@ def change_out_bias(
             'change-by-statistic' : perform predictions on labels of target dataset,
                     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 : dict, optional
+            Preset on the public real-atom type map. It is normalized and
+            extended with unassigned virtual types before forwarding. Omit
+            to retain the model preset; an empty dict fits without fixing
+            preset outputs. Non-finite stored cells covered by a checkpoint
+            preset are recovered before residual fitting.
         """
+        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

@@ -48,6 +48,7 @@
     check_preset_out_bias,
     normalize_preset_out_bias,
     preset_out_bias_rows,
+    recover_nonfinite_preset_bias,
     remap_preset_out_bias,
 )
 
@@ -501,6 +502,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 +526,14 @@ 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.
+            Before residual fitting, non-finite stored cells covered by the
+            checkpoint preset are recovered; finite adjusted cells are retained.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
         observed_type : list[str], optional
@@ -533,17 +543,34 @@ 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
+        # A per-call override must not propagate a corrupt checkpoint baseline.
+        if change and preset_out_bias is not None and self.preset_out_bias:
+            recovered = recover_nonfinite_preset_bias(
+                self.preset_out_bias,
+                self.type_map,
+                to_numpy_array(self.out_bias),
+                self.bias_keys,
+                [self.atomic_output_def()[kk].size for kk in self.bias_keys],
+            )
+            self._store_out_stat(
+                {kk: to_paddle_tensor(vv) for kk, vv in recovered.items()}
+            )
         # === 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.
 
@@ -248,10 +249,16 @@ def change_out_bias(
                 'change-by-statistic' : perform predictions on labels of target dataset,
                         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 : dict, optional
+                Per-call preset override. Omit to use the configured preset;
+                pass an empty dict to fit without fixing preset outputs.
+                Non-finite stored cells covered by the checkpoint preset are
+                recovered before residual fitting; finite values are retained.
             """
             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

@@ -52,6 +52,7 @@
     check_preset_out_bias,
     normalize_preset_out_bias,
     preset_out_bias_rows,
+    recover_nonfinite_preset_bias,
     remap_preset_out_bias,
 )
 
@@ -629,6 +630,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 +656,14 @@ 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.
+            Before residual fitting, non-finite stored cells covered by the
+            checkpoint preset are recovered; finite adjusted cells are retained.
         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 +675,34 @@ 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
+        # A per-call override must not propagate a corrupt checkpoint baseline.
+        if change and preset_out_bias is not None and self.preset_out_bias:
+            recovered = recover_nonfinite_preset_bias(
+                self.preset_out_bias,
+                self.type_map,
+                to_numpy_array(self.out_bias),
+                self.bias_keys,
+                [self.atomic_output_def()[kk].size for kk in self.bias_keys],
+            )
+            self._store_out_stat(
+                {kk: to_torch_tensor(vv) for kk, vv in recovered.items()}
+            )
         # === 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.
 
@@ -355,10 +356,16 @@ def change_out_bias(
                 'change-by-statistic' : perform predictions on labels of target dataset,
                         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 : dict, optional
+                Per-call preset override. Omit to use the configured preset;
+                pass an empty dict to fit without fixing preset outputs.
+                Non-finite stored cells covered by the checkpoint preset are
+                recovered before residual fitting; finite values are retained.
             """
             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,
             )