From a09e557c109619610892181d6b4cfbe9cbd600d4 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 12 Jul 2026 20:48:32 +0800 Subject: [PATCH 1/3] fix(train): dispatch registered learning rate schedules Use the shared BaseLR registry in JAX, TF2, and pt_expt so cosine and WSD schedules are honored consistently with the input schema. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/learning_rate.py | 27 +++++++++++++++++++ deepmd/jax/train/trainer.py | 11 ++++---- deepmd/pt_expt/train/training.py | 8 +++--- deepmd/tf2/train/trainer.py | 8 +++--- .../dpmodel/utils/test_learning_rate.py | 25 +++++++++++++++++ 5 files changed, 65 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index d26432528c..05c6c9121e 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -692,3 +692,30 @@ def _decay_value(self, step: int | Array) -> Array: # Clip to min_lr for steps beyond decay_num_steps step_lr = xp.where(step >= self.decay_num_steps, min_lr, step_lr) return step_lr + + +def make_learning_rate_schedule( + lr_params: dict[str, Any], + num_steps: int, +) -> BaseLR: + """Build a registered learning-rate schedule for a training run. + + The input schema selects schedules through ``learning_rate.type``. Keep + backend trainers on the shared :class:`BaseLR` registry so every backend + accepts the same registered variants without mutating its input config. + + Parameters + ---------- + lr_params : dict[str, Any] + Learning-rate configuration, including the optional ``type`` key. + num_steps : int + Total number of training steps used to parameterize the schedule. + + Returns + ------- + BaseLR + The schedule selected by ``lr_params["type"]`` (``exp`` by default). + """ + params = dict(lr_params) + params["num_steps"] = num_steps + return BaseLR(**params) diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c19267250e..41c0628e69 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -54,7 +54,8 @@ resolve_best_checkpoint_dir, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + BaseLR, + make_learning_rate_schedule, ) from deepmd.dpmodel.utils.multi_task import ( apply_shared_links, @@ -277,11 +278,9 @@ def _deserialize_models(model_data: dict[str, Any]) -> dict[str, BaseModel]: } return {DEFAULT_TASK_KEY: BaseModel.deserialize(model_data["model"])} - def _get_lr_and_coef(self, lr_param: dict[str, Any]) -> LearningRateExp: - lr_type = lr_param.get("type", "exp") - if lr_type == "exp": - return LearningRateExp(**lr_param, num_steps=self.num_steps) - raise RuntimeError("unknown learning_rate type " + lr_type) + def _get_lr_and_coef(self, lr_param: dict[str, Any]) -> BaseLR: + """Construct the schema-selected shared learning-rate schedule.""" + return make_learning_rate_schedule(lr_param, self.num_steps) def _build_losses( self, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f8dc2d9a1f..e2fce14bb2 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -40,7 +40,7 @@ split_batch, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + make_learning_rate_schedule, ) from deepmd.pt.train.utils import ( resolve_best_checkpoint_dir, @@ -1476,9 +1476,9 @@ def _make_sample( self.model_prob = None # Learning rate ------------------------------------------------------- - lr_params = config["learning_rate"].copy() - lr_params["num_steps"] = self.num_steps - self.lr_schedule = LearningRateExp(**lr_params) + self.lr_schedule = make_learning_rate_schedule( + config["learning_rate"], self.num_steps + ) # Gradient clipping self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 8cfa12cbda..1c1205e1ca 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -47,7 +47,7 @@ split_batch, ) from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, + make_learning_rate_schedule, ) from deepmd.dpmodel.utils.training_utils import ( resolve_model_prob, @@ -385,9 +385,9 @@ def sample( resume=init_model is not None or restart_model is not None ) - lr_params = dict(config["learning_rate"]) - lr_params["num_steps"] = self.num_steps - self.lr_schedule = LearningRateExp(**lr_params) + self.lr_schedule = make_learning_rate_schedule( + config["learning_rate"], self.num_steps + ) self.optimizer = self._build_optimizer(config.get("optimizer", {})) self.model_container = _TaskModelContainer(self.models) self.step = tf.Variable(0, dtype=tf.int64, trainable=False, name="step") diff --git a/source/tests/universal/dpmodel/utils/test_learning_rate.py b/source/tests/universal/dpmodel/utils/test_learning_rate.py index 8406f7f352..19551d2165 100644 --- a/source/tests/universal/dpmodel/utils/test_learning_rate.py +++ b/source/tests/universal/dpmodel/utils/test_learning_rate.py @@ -7,9 +7,11 @@ to_numpy_array, ) from deepmd.dpmodel.utils.learning_rate import ( + BaseLR, LearningRateCosine, LearningRateExp, LearningRateWSD, + make_learning_rate_schedule, ) @@ -369,6 +371,29 @@ def test_array_input_wsd_cosine(self) -> None: np.testing.assert_allclose(lrs[3], 1e-5, rtol=1e-10) +class TestMakeLearningRateSchedule(unittest.TestCase): + """Test the shared factory used by backend trainers.""" + + def test_dispatches_all_schema_variants_without_mutating_config(self) -> None: + """Select exp, cosine, and WSD schedules through ``type``.""" + schedule_types: list[tuple[str, type[BaseLR]]] = [ + ("exp", LearningRateExp), + ("cosine", LearningRateCosine), + ("wsd", LearningRateWSD), + ] + for schedule_type, expected_class in schedule_types: + with self.subTest(schedule_type=schedule_type): + params = { + "type": schedule_type, + "start_lr": 1e-3, + "stop_lr": 1e-5, + } + schedule = make_learning_rate_schedule(params, num_steps=100) + + self.assertIsInstance(schedule, expected_class) + self.assertNotIn("num_steps", params) + + class TestLearningRateBeyondStopSteps(unittest.TestCase): """Test learning rate behavior beyond num_steps.""" From 28dfd4e11d853d4d9302dc2bb12a447e9e6dfc22 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 17:56:36 +0800 Subject: [PATCH 2/3] fix(train): initialize warmup schedulers safely Use each schedule's nonzero start_lr as the PyTorch LambdaLR base, preserve the documented exponential default when type is omitted, and add trainer-level cosine/WSD warmup regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/learning_rate.py | 3 ++ deepmd/pt_expt/train/training.py | 11 ++++++- source/tests/pt_expt/test_training.py | 32 +++++++++++++++++++ .../dpmodel/utils/test_learning_rate.py | 10 ++++++ 4 files changed, 55 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index 05c6c9121e..a720da81b7 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -717,5 +717,8 @@ def make_learning_rate_schedule( The schedule selected by ``lr_params["type"]`` (``exp`` by default). """ params = dict(lr_params) + # ``type`` is optional in the public schema. Supply the documented + # exponential default before dispatching through the strict registry. + params.setdefault("type", "exp") params["num_steps"] = num_steps return BaseLR(**params) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index e2fce14bb2..3ff2348870 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1535,7 +1535,11 @@ def _make_sample( # Optimiser ----------------------------------------------------------- opt_type = training_params.get("opt_type", "Adam") - initial_lr = float(self.lr_schedule.value(self.start_step)) + # LambdaLR multiplies each param group's initial learning rate by the + # lambda value. Warmup schedules legitimately return zero at step 0, + # so use the nonzero schedule base as the denominator and let the + # lambda initialize the optimizer to the requested warmup value. + initial_lr = float(self.lr_schedule.start_lr) if opt_type == "Adam": self.optimizer = torch.optim.Adam(self.wrapper.parameters(), lr=initial_lr) @@ -1549,6 +1553,9 @@ def _make_sample( else: raise ValueError(f"Unsupported optimizer type: {opt_type}") + for param_group in self.optimizer.param_groups: + param_group["initial_lr"] = initial_lr + self.scheduler = torch.optim.lr_scheduler.LambdaLR( self.optimizer, lambda step: self.lr_schedule.value(step) / initial_lr, @@ -1733,6 +1740,8 @@ def _make_sample( if optimizer_state_dict is not None: self.optimizer.load_state_dict(optimizer_state_dict) + for param_group in self.optimizer.param_groups: + param_group["initial_lr"] = initial_lr # rebuild scheduler from the resumed step. # last_epoch handles the step offset; the lambda must NOT # add self.start_step again (that would double-count). diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index c8a21a8a59..5c4f28deba 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -311,6 +311,38 @@ def test_training_loop(self) -> None: config = normalize(config) self._run_training(config) + def test_zero_start_warmup_schedulers_construct(self) -> None: + """Cosine and WSD warmup must initialize LambdaLR without division by zero.""" + for schedule_type in ("cosine", "wsd"): + with self.subTest(schedule_type=schedule_type): + config = _make_config(self.data_dir, numb_steps=4) + config["learning_rate"] = { + "type": schedule_type, + "start_lr": 1e-3, + "stop_lr": 1e-5, + "warmup_steps": 1, + } + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + tmpdir = tempfile.mkdtemp(prefix=f"pt_expt_{schedule_type}_warmup_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(config) + + self.assertEqual(trainer.lr_schedule.value(0), 0.0) + self.assertEqual(trainer.scheduler.get_last_lr(), [0.0]) + self.assertTrue( + all( + group["initial_lr"] == trainer.lr_schedule.start_lr + for group in trainer.optimizer.param_groups + ) + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") def test_full_validation_loop(self, mocked_eval) -> None: """Run pt_expt full validation and verify best-checkpoint outputs.""" diff --git a/source/tests/universal/dpmodel/utils/test_learning_rate.py b/source/tests/universal/dpmodel/utils/test_learning_rate.py index 19551d2165..2a0d616577 100644 --- a/source/tests/universal/dpmodel/utils/test_learning_rate.py +++ b/source/tests/universal/dpmodel/utils/test_learning_rate.py @@ -393,6 +393,16 @@ def test_dispatches_all_schema_variants_without_mutating_config(self) -> None: self.assertIsInstance(schedule, expected_class) self.assertNotIn("num_steps", params) + def test_omitted_type_uses_exponential_default(self) -> None: + """Honor the public schema default without mutating the input.""" + params = {"start_lr": 1e-3, "stop_lr": 1e-5} + + schedule = make_learning_rate_schedule(params, num_steps=100) + + self.assertIsInstance(schedule, LearningRateExp) + self.assertNotIn("type", params) + self.assertNotIn("num_steps", params) + class TestLearningRateBeyondStopSteps(unittest.TestCase): """Test learning rate behavior beyond num_steps.""" From 5d5804d6fcf36a40b131f79f0b17e8148485e7f7 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 18:14:03 +0800 Subject: [PATCH 3/3] test(pt): bound warmup scheduler regression Keep the trainer-level warmup regression within the repository's validation timeout. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/pt_expt/test_training.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 5c4f28deba..c41a73640b 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -18,6 +18,7 @@ patch, ) +import pytest import torch from deepmd.loggers.training import ( @@ -311,6 +312,7 @@ def test_training_loop(self) -> None: config = normalize(config) self._run_training(config) + @pytest.mark.timeout(60) def test_zero_start_warmup_schedulers_construct(self) -> None: """Cosine and WSD warmup must initialize LambdaLR without division by zero.""" for schedule_type in ("cosine", "wsd"):