Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new warmup-stable-decay (WSD) learning-rate scheduler to the backend-agnostic BaseLR registry, wires it into argument checking and backend exports, and expands cross-backend test coverage to validate behavior and consistency.
Changes:
- Implement
LearningRateWSD(type="wsd") with warmup support plus configurable decay modes (inverse_linear,cosine,linear). - Extend CLI/config arg validation to accept and validate WSD-specific parameters.
- Add/extend tests across universal, TF, PT, PD, and consistency suites for WSD behavior.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
deepmd/dpmodel/utils/learning_rate.py |
Adds the LearningRateWSD scheduler implementation and registers it under BaseLR. |
deepmd/utils/argcheck.py |
Adds WSD-specific validation and exposes WSD arguments via the lr args plugin registry. |
deepmd/pt/utils/learning_rate.py |
Re-exports LearningRateWSD for the PyTorch backend utilities. |
deepmd/pd/utils/learning_rate.py |
Re-exports LearningRateWSD for the Paddle backend utilities. |
source/tests/universal/dpmodel/utils/test_learning_rate.py |
Adds extensive unit tests for WSD (modes, warmup, array input, boundary conditions). |
source/tests/tf/test_lr.py |
Adds TF wrapper build/value tests for WSD (default + cosine). |
source/tests/pt/test_lr.py |
Adds PT-side WSD curve tests for all decay modes. |
source/tests/pd/test_lr.py |
Adds PD-side WSD curve tests for all decay modes. |
source/tests/consistent/test_learning_rate.py |
Extends consistency parameterization to include WSD and adjusts the sampling step to hit the decay phase. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
📝 WalkthroughWalkthroughAdds a new plugin-registered learning-rate schedule Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
source/tests/tf/test_lr.py (1)
106-140: These assertions don't exercise the TF graph path.
lr_schedule.value()just delegates tobase_lr.value(), so both new tests still pass ifLearningRateSchedule.build()is wrong fortype="wsd". Please evaluate the tensor returned bybuild()at a mid-decay step and compare that result instead.🧪 Suggested coverage improvement
- global_step = tf.constant(0, dtype=tf.int64) - lr_schedule.build(global_step, num_steps=10000) + g = tf.Graph() + with g.as_default(): + global_step = tf.placeholder(shape=[], dtype=tf.int64) + lr_tensor = lr_schedule.build(global_step, num_steps=10000) self.assertIsInstance(lr_schedule.base_lr, LearningRateWSD) - np.testing.assert_allclose( - lr_schedule.value(9500), lr_schedule.base_lr.value(9500), rtol=1e-10 - ) + with tf.Session(graph=g) as sess: + tensor_value = sess.run(lr_tensor, feed_dict={global_step: 9500}) + np.testing.assert_allclose( + tensor_value, lr_schedule.base_lr.value(9500), rtol=1e-10 + )Apply the same pattern to the cosine WSD variant as well.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/tf/test_lr.py` around lines 106 - 140, The tests are only exercising the Python path because lr_schedule.value(9500) passes a plain int; update both tests to pass a TF tensor so the built TF graph is evaluated: after calling LearningRateSchedule.build(global_step, num_steps=10000), call lr_schedule.value(tf.constant(9500, dtype=tf.int64)) and compare it to lr_schedule.base_lr.value(tf.constant(9500, dtype=tf.int64)) (and do the same change in test_wsd_cosine_build_and_value) to ensure LearningRateSchedule.build, base_lr and LearningRateWSD TF-paths are exercised.source/tests/universal/dpmodel/utils/test_learning_rate.py (1)
142-165: Consider splitting mixed validation assertions into separate tests.
test_invalid_decay_phase_ratioalso checks invaliddecay_type; splitting improves failure localization.♻️ Suggested test split
- def test_invalid_decay_phase_ratio(self) -> None: - """Test invalid WSD decay_phase_ratio values.""" + def test_invalid_decay_phase_ratio(self) -> None: + """Test invalid WSD decay_phase_ratio values.""" with self.assertRaises(ValueError): LearningRateWSD( start_lr=1e-3, stop_lr=1e-5, num_steps=10000, decay_phase_ratio=0.0, ) with self.assertRaises(ValueError): LearningRateWSD( start_lr=1e-3, stop_lr=1e-5, num_steps=10000, decay_phase_ratio=1.1, ) + + def test_invalid_decay_type(self) -> None: + """Test invalid WSD decay_type values.""" with self.assertRaises(ValueError): LearningRateWSD( start_lr=1e-3, stop_lr=1e-5, num_steps=10000, decay_type="bad_mode", )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/universal/dpmodel/utils/test_learning_rate.py` around lines 142 - 165, test_invalid_decay_phase_ratio currently asserts multiple unrelated invalid cases (invalid decay_phase_ratio and invalid decay_type) in one test; split it into separate tests so each assertion checks a single validation rule. Create one test (e.g., test_invalid_decay_phase_ratio_values) that calls LearningRateWSD with decay_phase_ratio=0.0 and decay_phase_ratio=1.1 and asserts ValueError, and another test (e.g., test_invalid_decay_type) that calls LearningRateWSD with decay_type="bad_mode" and asserts ValueError; keep function/class names LearningRateWSD, decay_phase_ratio, and decay_type to locate the code. Ensure each new test has a clear name and only one responsibility for better failure localization.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/utils/learning_rate.py`:
- Around line 514-525: The current int(self.decay_phase_ratio * self.num_steps)
can produce 0 or exceed post-warmup steps; change the computation of
decay_phase_steps to derive from and clamp against decay_num_steps instead of
raising: compute e.g. desired = max(0, int(round(self.decay_phase_ratio *
self.num_steps))) (or 0 if you prefer), then set self.decay_phase_steps =
min(desired, self.decay_num_steps) and if self.decay_phase_steps == 0 and
self.decay_num_steps > 0 set it to 1 (or otherwise handle the zero-decay case),
replacing the two ValueError checks; update code around decay_phase_steps,
decay_phase_ratio, num_steps and decay_num_steps to use this clamped value so
short runs or heavy-warmup configs don’t raise at runtime.
---
Nitpick comments:
In `@source/tests/tf/test_lr.py`:
- Around line 106-140: The tests are only exercising the Python path because
lr_schedule.value(9500) passes a plain int; update both tests to pass a TF
tensor so the built TF graph is evaluated: after calling
LearningRateSchedule.build(global_step, num_steps=10000), call
lr_schedule.value(tf.constant(9500, dtype=tf.int64)) and compare it to
lr_schedule.base_lr.value(tf.constant(9500, dtype=tf.int64)) (and do the same
change in test_wsd_cosine_build_and_value) to ensure LearningRateSchedule.build,
base_lr and LearningRateWSD TF-paths are exercised.
In `@source/tests/universal/dpmodel/utils/test_learning_rate.py`:
- Around line 142-165: test_invalid_decay_phase_ratio currently asserts multiple
unrelated invalid cases (invalid decay_phase_ratio and invalid decay_type) in
one test; split it into separate tests so each assertion checks a single
validation rule. Create one test (e.g., test_invalid_decay_phase_ratio_values)
that calls LearningRateWSD with decay_phase_ratio=0.0 and decay_phase_ratio=1.1
and asserts ValueError, and another test (e.g., test_invalid_decay_type) that
calls LearningRateWSD with decay_type="bad_mode" and asserts ValueError; keep
function/class names LearningRateWSD, decay_phase_ratio, and decay_type to
locate the code. Ensure each new test has a clear name and only one
responsibility for better failure localization.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: a283f5a5-f445-4ee4-858e-4f089ee4287d
📒 Files selected for processing (9)
deepmd/dpmodel/utils/learning_rate.pydeepmd/pd/utils/learning_rate.pydeepmd/pt/utils/learning_rate.pydeepmd/utils/argcheck.pysource/tests/consistent/test_learning_rate.pysource/tests/pd/test_lr.pysource/tests/pt/test_lr.pysource/tests/tf/test_lr.pysource/tests/universal/dpmodel/utils/test_learning_rate.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5326 +/- ##
==========================================
- Coverage 82.42% 82.42% -0.01%
==========================================
Files 784 784
Lines 79125 79189 +64
Branches 3676 3676
==========================================
+ Hits 65220 65268 +48
- Misses 12732 12749 +17
+ Partials 1173 1172 -1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
deepmd/dpmodel/utils/learning_rate.py (1)
514-518:⚠️ Potential issue | 🟡 MinorHandle edge case when
decay_num_stepsis zero.The clamping logic at line 515-516 ensures
decay_phase_steps >= 1, but whendecay_num_steps = 0(i.e.,num_steps = warmup_steps = 0), this results instable_steps = -1. While the_decay_valuelogic handles this gracefully by returningstop_lrimmediately, having negativestable_stepsis semantically incorrect.Consider adding early validation or adjusting the clamping:
🛡️ Suggested fix
# === Derive stable and decay phase lengths === self.decay_phase_ratio = decay_phase_ratio self.decay_type = decay_type + # When decay_num_steps is 0, there's no room for decay + if self.decay_num_steps == 0: + self.decay_phase_steps = 0 + self.stable_steps = 0 + return # 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🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/learning_rate.py` around lines 514 - 518, The current clamping forces decay_phase_steps >= 1 which yields stable_steps negative when decay_num_steps == 0; update the logic so decay_phase_steps is clamped to the range [0, decay_num_steps] (allowing zero) or explicitly handle the decay_num_steps == 0 case by setting decay_phase_steps = 0 and stable_steps = 0 and/or validating decay_num_steps >= 0; adjust the block that sets self.decay_phase_steps and self.stable_steps (and any callers like _decay_value) to rely on non-negative stable_steps and ensure no negative values are produced.
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/learning_rate.py (1)
559-562: Redundantxp.asarraywrapping.Since
tauis already an array withstep_dtype, the expressionxp.pi * tauwill produce an array of the same dtype. Thexp.asarray(..., dtype=step_dtype)call is redundant.♻️ Suggested simplification
elif self.decay_type == "cosine": decay_lr = stop_lr + (start_lr - stop_lr) * 0.5 * ( - one + xp.cos(xp.asarray(xp.pi * tau, dtype=step_dtype)) + one + xp.cos(xp.pi * tau) )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/learning_rate.py` around lines 559 - 562, In the cosine decay branch (where self.decay_type == "cosine") the expression xp.asarray(xp.pi * tau, dtype=step_dtype) is redundant because xp.pi * tau already yields an array with the correct dtype; replace that subexpression with xp.pi * tau so decay_lr is computed as stop_lr + (start_lr - stop_lr) * 0.5 * (one + xp.cos(xp.pi * tau)), keeping references to decay_lr, start_lr, stop_lr, tau, xp.cos and step_dtype unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@deepmd/dpmodel/utils/learning_rate.py`:
- Around line 514-518: The current clamping forces decay_phase_steps >= 1 which
yields stable_steps negative when decay_num_steps == 0; update the logic so
decay_phase_steps is clamped to the range [0, decay_num_steps] (allowing zero)
or explicitly handle the decay_num_steps == 0 case by setting decay_phase_steps
= 0 and stable_steps = 0 and/or validating decay_num_steps >= 0; adjust the
block that sets self.decay_phase_steps and self.stable_steps (and any callers
like _decay_value) to rely on non-negative stable_steps and ensure no negative
values are produced.
---
Nitpick comments:
In `@deepmd/dpmodel/utils/learning_rate.py`:
- Around line 559-562: In the cosine decay branch (where self.decay_type ==
"cosine") the expression xp.asarray(xp.pi * tau, dtype=step_dtype) is redundant
because xp.pi * tau already yields an array with the correct dtype; replace that
subexpression with xp.pi * tau so decay_lr is computed as stop_lr + (start_lr -
stop_lr) * 0.5 * (one + xp.cos(xp.pi * tau)), keeping references to decay_lr,
start_lr, stop_lr, tau, xp.cos and step_dtype unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 1363b61c-f8fb-4804-b834-420b56c86723
📒 Files selected for processing (2)
deepmd/dpmodel/utils/learning_rate.pysource/tests/universal/dpmodel/utils/test_learning_rate.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/universal/dpmodel/utils/test_learning_rate.py
wanghan-iapcm
approved these changes
Mar 20, 2026
iProzd
reviewed
Mar 23, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/learning_rate.py (1)
476-483: TheRaisessection is stale.The constructor now clamps oversized/empty derived decay phases instead of raising for them, so this docstring no longer matches the implementation.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/learning_rate.py` around lines 476 - 483, Update the stale "Raises" docstring in the constructor of the learning rate class (referenced by decay_phase_ratio and the derived decay phase logic in deepmd/dpmodel/utils/learning_rate.py) to reflect current behavior: remove the claim that an empty or oversized derived decay phase raises ValueError and instead document that such cases are clamped to valid bounds; keep the existing notes about non-positive learning rates and invalid decay_type raising ValueError. Mention the clamping behavior and the valid range for decay_phase_ratio (0, 1] so the docstring matches the implementation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/utils/learning_rate.py`:
- Around line 514-518: The WSD scheduler currently forces decay_phase_steps to
at least 1 even when decay_num_steps == 0, producing negative stable_steps and
making value(0) return stop_lr; modify the LearningRateWSD initialization to
special-case zero-decay runs (or explicitly reject num_steps == 0): if
self.decay_num_steps == 0 then set self.decay_phase_steps = 0 and
self.stable_steps = 0 (and ensure value(index) treats all indices as
pre-decay/start_lr), or raise a ValueError when self.num_steps == 0; update
references to decay_phase_steps, decay_num_steps, stable_steps and the
value(...) method so zero-decay behavior matches
LearningRateExp/LearningRateCosine.
---
Nitpick comments:
In `@deepmd/dpmodel/utils/learning_rate.py`:
- Around line 476-483: Update the stale "Raises" docstring in the constructor of
the learning rate class (referenced by decay_phase_ratio and the derived decay
phase logic in deepmd/dpmodel/utils/learning_rate.py) to reflect current
behavior: remove the claim that an empty or oversized derived decay phase raises
ValueError and instead document that such cases are clamped to valid bounds;
keep the existing notes about non-positive learning rates and invalid decay_type
raising ValueError. Mention the clamping behavior and the valid range for
decay_phase_ratio (0, 1] so the docstring matches the implementation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 3c30e1f9-f232-4bbc-8eac-79c3c14cfbe3
📒 Files selected for processing (10)
deepmd/dpmodel/utils/learning_rate.pydeepmd/pd/utils/learning_rate.pydeepmd/pt/utils/learning_rate.pydeepmd/utils/argcheck.pydoc/train/learning-rate.mdsource/tests/consistent/test_learning_rate.pysource/tests/pd/test_lr.pysource/tests/pt/test_lr.pysource/tests/tf/test_lr.pysource/tests/universal/dpmodel/utils/test_learning_rate.py
✅ Files skipped from review due to trivial changes (3)
- source/tests/tf/test_lr.py
- doc/train/learning-rate.md
- source/tests/consistent/test_learning_rate.py
🚧 Files skipped from review as they are similar to previous changes (3)
- deepmd/pd/utils/learning_rate.py
- deepmd/pt/utils/learning_rate.py
- source/tests/pd/test_lr.py
Comment on lines
+514
to
+518
| # 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 |
Contributor
There was a problem hiding this comment.
Handle num_steps == 0 before forcing a one-step decay.
Lines 515-517 turn decay_num_steps == 0 into decay_phase_steps == 1, which makes stable_steps negative and causes Lines 565-566 to return stop_lr immediately for value(0). LearningRateExp and LearningRateCosine already special-case zero-decay runs, so WSD should do the same or reject num_steps == 0 explicitly.
💡 Possible fix
# === Derive stable and decay phase lengths ===
self.decay_phase_ratio = decay_phase_ratio
self.decay_type = decay_type
+ if self.decay_num_steps == 0:
+ self.decay_phase_steps = 0
+ self.stable_steps = 0
+ return
# 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
@@
step_dtype = (
step.dtype
if xp.isdtype(step.dtype, "real floating")
else get_xp_precision(xp, "global")
)
+ if self.decay_num_steps == 0:
+ return xp.full_like(step, self._start_lr, dtype=step_dtype)Also applies to: 565-566
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/dpmodel/utils/learning_rate.py` around lines 514 - 518, The WSD
scheduler currently forces decay_phase_steps to at least 1 even when
decay_num_steps == 0, producing negative stable_steps and making value(0) return
stop_lr; modify the LearningRateWSD initialization to special-case zero-decay
runs (or explicitly reject num_steps == 0): if self.decay_num_steps == 0 then
set self.decay_phase_steps = 0 and self.stable_steps = 0 (and ensure
value(index) treats all indices as pre-decay/start_lr), or raise a ValueError
when self.num_steps == 0; update references to decay_phase_steps,
decay_num_steps, stable_steps and the value(...) method so zero-decay behavior
matches LearningRateExp/LearningRateCosine.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The doc will be added through PR #5276
Summary by CodeRabbit
New Features
Tests
Documentation