chore(pd): sync get_lr from pt to pd (#5144)

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

## Summary by CodeRabbit

* **Refactor**
* Updated learning rate scheduling implementation with enhanced type
annotations and more flexible parameter handling for improved code
clarity.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz

deepmd/pd/train/training.py

@@ -30,6 +30,9 @@
 from deepmd.common import (
     symlink_prefix_files,
 )
+from deepmd.dpmodel.utils.learning_rate import (
+    BaseLR,
+)
 from deepmd.loggers.training import (
     format_training_message,
     format_training_message_per_task,
@@ -62,9 +65,6 @@
     SAMPLER_RECORD,
     enable_prim,
 )
-from deepmd.pd.utils.learning_rate import (
-    LearningRateExp,
-)
 from deepmd.pd.utils.stat import (
     make_stat_input,
 )
@@ -238,13 +238,10 @@ def get_sample():
                     _stat_file_path.root.close()
             return get_sample
 
-        def get_lr(lr_params):
-            assert lr_params.get("type", "exp") == "exp", (
-                "Only learning rate `exp` is supported!"
-            )
+        def get_lr(lr_params: dict[str, Any]) -> BaseLR:
             lr_params["stop_steps"] = self.num_steps - self.warmup_steps
-            lr_exp = LearningRateExp(**lr_params)
-            return lr_exp
+            lr_schedule = BaseLR(**lr_params)
+            return lr_schedule
 
         # Optimizer
         if self.multi_task and training_params.get("optim_dict", None) is not None: