fix(train): share post-training bias adjustment (#5745)

## Summary
- add dpmodel shared helpers for post-training output bias adjustment
- route TF2/JAX/PT/PT-expt change-bias paths through the shared helper
- enable change_bias_after_training in JAX and pt_expt training and
preserve PT wrapper APIs

## Tests
- ruff format . && ruff check .
- pytest source/tests/common/test_dpmodel_train.py -q
- pytest
source/tests/pt_expt/test_change_bias.py::TestChangeBiasFittingStats -q
- pytest
source/tests/pt/test_training.py::TestModelChangeOutBiasFittingStat::test_fitting_stat_consistency
-q

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an optional post-training output-bias adjustment step for
supported training workflows.
  * Extended bias updates across both single-task and multi-task models.
* Exposed shared bias-update helpers as part of the public
backend-independent training API.
* **Bug Fixes**
* Ensured updated bias state is propagated correctly in distributed JAX
runs and checkpoints are saved after the post-training adjustment.
* **Tests**
* Added coverage for bias-update helper behavior and JAX broadcast
during post-training.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: njzjz

deepmd/dpmodel/train/__init__.py

@@ -20,6 +20,8 @@
     TrainingTask,
     TrainingTaskCollection,
     TrainStepResult,
+    change_model_out_bias,
+    change_model_out_bias_by_task,
 )
 
 __all__ = [
@@ -34,6 +36,8 @@
     "TrainingTask",
     "TrainingTaskCollection",
     "TrainingTaskConfig",
+    "change_model_out_bias",
+    "change_model_out_bias_by_task",
     "iter_training_task_configs",
     "make_task_maps",
     "print_data_summaries",

deepmd/dpmodel/train/trainer.py

@@ -22,8 +22,12 @@
     Callable,
     Iterator,
     Mapping,
+    MutableMapping,
     Sequence,
 )
+from copy import (
+    deepcopy,
+)
 from dataclasses import (
     dataclass,
     field,
@@ -38,6 +42,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
 from deepmd.loggers.training import (
     format_training_message,
     format_training_message_per_task,
@@ -52,6 +59,53 @@
 DisplayResults = LossResults | TaskResults
 
 
+def change_model_out_bias(
+    model: Any,
+    sample_func: Callable[[], Any],
+    *,
+    bias_adjust_mode: str = "change-by-statistic",
+    recompute_input_stats: bool = False,
+) -> Any:
+    """Change one model's output bias and log the before/after values."""
+    old_bias = deepcopy(model.get_out_bias())
+    model.change_out_bias(
+        sample_func,
+        bias_adjust_mode=bias_adjust_mode,
+    )
+    new_bias = deepcopy(model.get_out_bias())
+
+    if recompute_input_stats and bias_adjust_mode == "set-by-statistic":
+        model.get_fitting_net().compute_input_stats(sample_func)
+
+    model_type_map = model.get_type_map()
+    log.info(
+        f"Change output bias of {model_type_map!s} "
+        f"from {to_numpy_array(old_bias).reshape(-1)[: len(model_type_map)]!s} "
+        f"to {to_numpy_array(new_bias).reshape(-1)[: len(model_type_map)]!s}."
+    )
+    return model
+
+
+def change_model_out_bias_by_task(
+    models: MutableMapping[str, Any],
+    sample_funcs: Mapping[str, Callable[[], Any]],
+    model_keys: Sequence[str],
+    *,
+    bias_adjust_mode: str = "change-by-statistic",
+    recompute_input_stats: bool = False,
+) -> MutableMapping[str, Any]:
+    """Change output bias for all requested training-task models."""
+    log.info("Changing output bias after training.")
+    for model_key in model_keys:
+        models[model_key] = change_model_out_bias(
+            models[model_key],
+            sample_funcs[model_key],
+            bias_adjust_mode=bias_adjust_mode,
+            recompute_input_stats=recompute_input_stats,
+        )
+    return models
+
+
 @dataclass(frozen=True)
 class RankContext:
     """Rank metadata used by a trainer.

deepmd/jax/train/trainer.py

@@ -7,6 +7,7 @@
 import os
 import platform
 import shutil
+import time
 from collections.abc import (
     Mapping,
 )
@@ -41,6 +42,7 @@
     TrainingTask,
     TrainingTaskCollection,
     TrainStepResult,
+    change_model_out_bias_by_task,
 )
 from deepmd.dpmodel.train.validation import (
     resolve_best_checkpoint_dir,
@@ -197,8 +199,8 @@ def __init__(
         self.tensorboard_log_dir = tr_data.get("tensorboard_log_dir", "log")
         self.tensorboard_freq = tr_data.get("tensorboard_freq", 1)
         self.mixed_prec = tr_data.get("mixed_precision", None)
-        self.change_bias_after_training = tr_data.get(
-            "change_bias_after_training", False
+        self.change_bias_after_training = bool(
+            tr_data.get("change_bias_after_training", False)
         )
         self.numb_fparam = (
             {key: model.get_dim_fparam() for key, model in self.models.items()}
@@ -731,6 +733,39 @@ def save_checkpoint(self, step: int) -> None:
         """Persist a JAX checkpoint for a one-based step."""
         self._save_checkpoint(step)
 
+    def run(self, tasks: TrainingTaskCollection) -> None:
+        """Run JAX training through the backend-independent trainer loop."""
+        log.info("Start to train %d steps.", self.num_steps)
+        wall_start = time.time()
+        super().run(tasks)
+        if self.change_bias_after_training and self.num_steps > self.start_step:
+            self._change_bias_after_training()
+            if self.rank_context.is_chief:
+                self.save_checkpoint(self.num_steps)
+        log.info("Training finished. Total wall time: %.2fs", time.time() - wall_start)
+
+    def _change_bias_after_training(self) -> None:
+        if self.rank_context.is_chief:
+            change_model_out_bias_by_task(
+                self.models,
+                self._sample_funcs,
+                self.model_keys,
+                bias_adjust_mode="change-by-statistic",
+            )
+        if self.rank_context.world_size <= 1:
+            return
+        from jax.experimental import (
+            multihost_utils,
+        )
+
+        for model_key in self.model_keys:
+            _, state = nnx.split(self.models[model_key])
+            state = multihost_utils.broadcast_one_to_all(
+                state.to_pure_dict(),
+                is_source=self.rank_context.is_chief,
+            )
+            nnx.update(self.models[model_key], state)
+
     def run_full_validation(
         self,
         *,

deepmd/pt/train/training.py

@@ -28,6 +28,9 @@
 from deepmd.common import (
     symlink_prefix_files,
 )
+from deepmd.dpmodel.train import (
+    change_model_out_bias,
+)
 from deepmd.dpmodel.utils import (
     compute_total_numb_batch,
     resolve_model_prob,
@@ -2618,24 +2621,13 @@ def model_change_out_bias(
     _sample_func: Callable[[], Any],
     _bias_adjust_mode: str = "change-by-statistic",
 ) -> Any:
-    old_bias = deepcopy(_model.get_out_bias())
-    _model.change_out_bias(
-        _sample_func,
-        bias_adjust_mode=_bias_adjust_mode,
-    )
-    new_bias = deepcopy(_model.get_out_bias())
-
     from deepmd.pt.model.model.dp_model import (
         DPModelCommon,
     )
 
-    if isinstance(_model, DPModelCommon) and _bias_adjust_mode == "set-by-statistic":
-        _model.get_fitting_net().compute_input_stats(_sample_func)
-
-    model_type_map = _model.get_type_map()
-    log.info(
-        f"Change output bias of {model_type_map!s} "
-        f"from {to_numpy_array(old_bias).reshape(-1)[: len(model_type_map)]!s} "
-        f"to {to_numpy_array(new_bias).reshape(-1)[: len(model_type_map)]!s}."
+    return change_model_out_bias(
+        _model,
+        _sample_func,
+        bias_adjust_mode=_bias_adjust_mode,
+        recompute_input_stats=isinstance(_model, DPModelCommon),
     )
-    return _model

deepmd/pt_expt/train/training.py

@@ -24,9 +24,6 @@
 import torch
 import torch.distributed as dist
 
-from deepmd.dpmodel.common import (
-    to_numpy_array,
-)
 from deepmd.dpmodel.train import (
     DEFAULT_TASK_KEY,
     AbstractTrainer,
@@ -35,6 +32,8 @@
     TrainingTask,
     TrainingTaskCollection,
     TrainStepResult,
+    change_model_out_bias,
+    change_model_out_bias_by_task,
 )
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
@@ -1350,6 +1349,9 @@ def __init__(
         self.max_ckpt_keep = int(training_params.get("max_ckpt_keep", 5))
         self.display_in_training = training_params.get("disp_training", True)
         self.timing_in_training = training_params.get("time_training", True)
+        self.change_bias_after_training = bool(
+            training_params.get("change_bias_after_training", False)
+        )
 
         # Model ---------------------------------------------------------------
         self.models: dict[str, torch.nn.Module] = {}
@@ -2139,8 +2141,25 @@ def run(self) -> None:
         log.info("Start to train %d steps.", self.num_steps)
         wall_start = time.time()
         super().run(self.training_tasks)
+        if self.change_bias_after_training and self.num_steps > self.start_step:
+            self._change_bias_after_training()
+            if self.rank_context.is_chief:
+                self.save_checkpoint(self.num_steps)
         log.info("Training finished. Total wall time: %.2fs", time.time() - wall_start)
 
+    def _change_bias_after_training(self) -> None:
+        if self.rank == 0:
+            change_model_out_bias_by_task(
+                self.models,
+                self._sample_funcs,
+                self.model_keys,
+                bias_adjust_mode="change-by-statistic",
+            )
+        if self.is_distributed:
+            for model_key in self.model_keys:
+                self._broadcast_model_stat(self.models[model_key])
+        self.model = self.models if self.multi_task else self.models[DEFAULT_TASK_KEY]
+
     def run_full_validation(
         self,
         *,
@@ -2326,27 +2345,16 @@ def model_change_out_bias(
     -------
     The model with updated bias.
     """
-    old_bias = deepcopy(_model.get_out_bias())
-    _model.change_out_bias(
-        _sample_func,
-        bias_adjust_mode=_bias_adjust_mode,
-    )
-    new_bias = deepcopy(_model.get_out_bias())
-
     from deepmd.dpmodel.model.dp_model import (
         DPModelCommon,
     )
 
-    if isinstance(_model, DPModelCommon) and _bias_adjust_mode == "set-by-statistic":
-        _model.get_fitting_net().compute_input_stats(_sample_func)
-
-    model_type_map = _model.get_type_map()
-    log.info(
-        f"Change output bias of {model_type_map!s} "
-        f"from {to_numpy_array(old_bias).reshape(-1)[: len(model_type_map)]!s} "
-        f"to {to_numpy_array(new_bias).reshape(-1)[: len(model_type_map)]!s}."
+    return change_model_out_bias(
+        _model,
+        _sample_func,
+        bias_adjust_mode=_bias_adjust_mode,
+        recompute_input_stats=isinstance(_model, DPModelCommon),
     )
-    return _model
 
 
 def _get_case_embd_config(

deepmd/tf2/train/trainer.py

@@ -40,6 +40,7 @@
     TrainingTask,
     TrainingTaskCollection,
     TrainStepResult,
+    change_model_out_bias_by_task,
 )
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
@@ -1261,12 +1262,12 @@ def _write_tensorboard_step(
             self.summary_writer.flush()
 
     def _change_bias_after_training(self) -> None:
-        log.info("Changing output bias after training.")
-        for model_key in self.model_keys:
-            self.models[model_key].change_out_bias(
-                self._sample_funcs[model_key],
-                bias_adjust_mode="change-by-statistic",
-            )
+        change_model_out_bias_by_task(
+            self.models,
+            self._sample_funcs,
+            self.model_keys,
+            bias_adjust_mode="change-by-statistic",
+        )
 
     def get_data(
         self,