fix(train): dispatch registered learning rate schedules

deepmd/dpmodel/utils/learning_rate.py

@@ -692,3 +692,33 @@ 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)
+    # ``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)

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,

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)
@@ -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).

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")

source/tests/pt_expt/test_training.py

@@ -18,6 +18,7 @@
     patch,
 )
 
+import pytest
 import torch
 
 from deepmd.loggers.training import (
@@ -311,6 +312,39 @@ 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"):
+            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."""

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,39 @@ 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)
+
+    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."""