fix

Author: OutisLi

deepmd/dpmodel/utils/learning_rate.py

@@ -511,18 +511,10 @@ def __init__(
         # === Derive stable and decay phase lengths ===
         self.decay_phase_ratio = decay_phase_ratio
         self.decay_type = decay_type
-        self.decay_phase_steps = int(self.decay_phase_ratio * self.num_steps)
-        if self.decay_phase_steps <= 0:
-            raise ValueError(
-                "decay_phase_ratio results in zero decay steps. "
-                "Increase num_steps or decay_phase_ratio."
-            )
-        if self.decay_phase_steps > self.decay_num_steps:
-            raise ValueError(
-                "decay phase steps must not exceed the post-warmup steps. "
-                f"Got decay_phase_steps={self.decay_phase_steps}, "
-                f"post_warmup_steps={self.decay_num_steps}."
-            )
+        # Clamp decay_phase_steps to valid range [1, decay_num_steps]
+        self.decay_phase_steps = max(
+            1, min(int(self.decay_phase_ratio * self.num_steps), self.decay_num_steps)
+        )
         self.stable_steps = self.decay_num_steps - self.decay_phase_steps
 
     def _decay_value(self, step: int | Array) -> Array:
@@ -556,7 +548,6 @@ def _decay_value(self, step: int | Array) -> Array:
         stop_lr = xp.asarray(self.stop_lr, dtype=step_dtype)
         stable_steps = xp.asarray(self.stable_steps, dtype=step_dtype)
         decay_phase_steps = xp.asarray(self.decay_phase_steps, dtype=step_dtype)
-        decay_num_steps = xp.asarray(self.decay_num_steps, dtype=step_dtype)
 
         # === Step 2. Keep a constant learning rate in the stable phase ===
         decay_progress = (typed_step - stable_steps) / decay_phase_steps

source/tests/universal/dpmodel/utils/test_learning_rate.py

@@ -164,15 +164,17 @@ def test_invalid_decay_phase_ratio(self) -> None:
             )
 
     def test_decay_phase_exceeds_post_warmup_steps(self) -> None:
-        """Test WSD rejects decay phases longer than post-warmup steps."""
-        with self.assertRaises(ValueError):
-            LearningRateWSD(
-                start_lr=1e-3,
-                stop_lr=1e-5,
-                num_steps=10,
-                warmup_steps=9,
-                decay_phase_ratio=0.2,
-            )
+        """Test WSD clamps decay_phase_steps to post-warmup steps when ratio is too large."""
+        lr = LearningRateWSD(
+            start_lr=1e-3,
+            stop_lr=1e-5,
+            num_steps=10,
+            warmup_steps=9,
+            decay_phase_ratio=0.2,
+        )
+        # decay_num_steps = 1, so decay_phase_steps should be clamped to 1
+        self.assertEqual(lr.decay_phase_steps, 1)
+        self.assertEqual(lr.stable_steps, 0)
 
 
 class TestLearningRateWarmup(unittest.TestCase):