fix(train): dispatch registered learning rate schedules

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)

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)

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/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."""