diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 3a8965284e..943f9ffb05 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/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", diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 83d3900e3f..74f61f486a 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/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. diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c77bc944b5..99661531e6 100644 --- a/deepmd/jax/train/trainer.py +++ b/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, *, diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index df883c8ee6..26e114e4c8 100644 --- a/deepmd/pt/train/training.py +++ b/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, @@ -2609,24 +2612,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 diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index d2868a4082..497e15713d 100644 --- a/deepmd/pt_expt/train/training.py +++ b/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( diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index ada601b077..d7572abe7c 100644 --- a/deepmd/tf2/train/trainer.py +++ b/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, diff --git a/source/tests/common/test_dpmodel_train.py b/source/tests/common/test_dpmodel_train.py new file mode 100644 index 0000000000..22149f8e99 --- /dev/null +++ b/source/tests/common/test_dpmodel_train.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.train import ( + change_model_out_bias, + change_model_out_bias_by_task, +) + + +class FakeFittingNet: + def __init__(self) -> None: + self.input_stats_sample_func = None + + def compute_input_stats(self, sample_func): + self.input_stats_sample_func = sample_func + + +class FakeModel: + def __init__(self, bias: float = 0.0) -> None: + self.out_bias = np.full((1, 2, 1), bias) + self.change_calls = [] + self.fitting_net = FakeFittingNet() + + def get_out_bias(self): + return self.out_bias + + def change_out_bias(self, sample_func, bias_adjust_mode): + self.change_calls.append((sample_func, bias_adjust_mode)) + self.out_bias = self.out_bias + 1.0 + + def get_type_map(self): + return ["O", "H"] + + def get_fitting_net(self): + return self.fitting_net + + +class TestChangeModelOutBias(unittest.TestCase): + def test_change_model_out_bias_recomputes_input_stats_when_requested(self): + model = FakeModel() + + def sample_func(): + return [{"atype": np.array([[0, 1]])}] + + returned = change_model_out_bias( + model, + sample_func, + bias_adjust_mode="set-by-statistic", + recompute_input_stats=True, + ) + + self.assertIs(returned, model) + self.assertEqual(model.change_calls, [(sample_func, "set-by-statistic")]) + self.assertIs(model.fitting_net.input_stats_sample_func, sample_func) + np.testing.assert_allclose(model.get_out_bias(), np.ones((1, 2, 1))) + + def test_change_model_out_bias_by_task_updates_all_models(self): + models = { + "task_a": FakeModel(0.0), + "task_b": FakeModel(2.0), + } + sample_funcs = { + "task_a": lambda: [{"atype": np.array([[0]])}], + "task_b": lambda: [{"atype": np.array([[1]])}], + } + + returned = change_model_out_bias_by_task( + models, + sample_funcs, + ["task_a", "task_b"], + ) + + self.assertIs(returned, models) + np.testing.assert_allclose(models["task_a"].get_out_bias(), np.ones((1, 2, 1))) + np.testing.assert_allclose( + models["task_b"].get_out_bias(), + np.full((1, 2, 1), 3.0), + ) + self.assertEqual( + models["task_a"].change_calls, + [(sample_funcs["task_a"], "change-by-statistic")], + ) + self.assertEqual( + models["task_b"].change_calls, + [(sample_funcs["task_b"], "change-by-statistic")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 25d3ccdc49..c9145ad0bd 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -25,6 +25,8 @@ import optax from deepmd.dpmodel.train import ( + DEFAULT_TASK_KEY, + RankContext, TrainEntrypointOptions, ) from deepmd.jax.entrypoints.freeze import ( @@ -39,6 +41,7 @@ ) from deepmd.jax.env import ( jnp, + nnx, ) from deepmd.jax.train.trainer import ( DPTrainer, @@ -334,6 +337,73 @@ def save_checkpoint(path: Path, lr: float = 0.0, step: int = 0) -> None: assert save_calls == [(Path("best.jax"), 0.25, 99)] +class _BiasModel(nnx.Module): + def __init__(self, value: float) -> None: + self.bias = nnx.Param(jnp.asarray([value])) + + +def _bias_sync_trainer(rank: int) -> DPTrainer: + trainer = DPTrainer.__new__(DPTrainer) + trainer.rank_context = RankContext(rank=rank, world_size=2) + trainer.models = {DEFAULT_TASK_KEY: _BiasModel(0.0)} + trainer._sample_funcs = {DEFAULT_TASK_KEY: object()} + trainer.model_keys = [DEFAULT_TASK_KEY] + return trainer + + +def test_jax_change_bias_after_training_broadcasts_chief_state() -> None: + """Rank 0 recomputes post-training bias and broadcasts the resulting state.""" + trainer = _bias_sync_trainer(rank=0) + + def change_bias(models, *args, **kwargs) -> None: + del args, kwargs + nnx.update(models[DEFAULT_TASK_KEY], {"bias": jnp.asarray([3.0])}) + + with ( + patch( + "deepmd.jax.train.trainer.change_model_out_bias_by_task", + side_effect=change_bias, + ) as change_model_out_bias_by_task, + patch( + "jax.experimental.multihost_utils.broadcast_one_to_all", + side_effect=lambda state, **kwargs: state, + ) as broadcast_one_to_all, + ): + trainer._change_bias_after_training() + + change_model_out_bias_by_task.assert_called_once() + broadcast_one_to_all.assert_called_once() + assert broadcast_one_to_all.call_args.kwargs["is_source"] is True + np.testing.assert_allclose( + np.asarray(trainer.models[DEFAULT_TASK_KEY].bias.value), + [3.0], + ) + + +def test_jax_change_bias_after_training_uses_broadcast_on_peer_rank() -> None: + """Peer ranks receive rank-0 post-training bias instead of recomputing it.""" + trainer = _bias_sync_trainer(rank=1) + + with ( + patch( + "deepmd.jax.train.trainer.change_model_out_bias_by_task", + ) as change_model_out_bias_by_task, + patch( + "jax.experimental.multihost_utils.broadcast_one_to_all", + return_value={"bias": jnp.asarray([5.0])}, + ) as broadcast_one_to_all, + ): + trainer._change_bias_after_training() + + change_model_out_bias_by_task.assert_not_called() + broadcast_one_to_all.assert_called_once() + assert broadcast_one_to_all.call_args.kwargs["is_source"] is False + np.testing.assert_allclose( + np.asarray(trainer.models[DEFAULT_TASK_KEY].bias.value), + [5.0], + ) + + class TestJAXTraining(unittest.TestCase): """Regression tests for complete JAX training runs."""