fix(train): dispatch registered learning rate schedules - #5776
Conversation
Use the shared BaseLR registry in JAX, TF2, and pt_expt so cosine and WSD schedules are honored consistently with the input schema. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughA shared ChangesLearning-rate schedule construction
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant make_learning_rate_schedule
participant BaseLR
Trainer->>make_learning_rate_schedule: pass learning-rate parameters and num_steps
make_learning_rate_schedule->>BaseLR: construct registered schedule
BaseLR-->>Trainer: return selected schedule
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5776 +/- ##
==========================================
- Coverage 79.69% 78.31% -1.39%
==========================================
Files 1020 1050 +30
Lines 116359 120550 +4191
Branches 4303 4353 +50
==========================================
+ Hits 92736 94404 +1668
- Misses 22076 24575 +2499
- Partials 1547 1571 +24 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix is correct across all three backends: routing through make_learning_rate_schedule -> BaseLR(**params) dispatches by type (BaseLR.new resolves the registered subclass and init absorbs the extra type key via **kwargs), so cosine/wsd now work where JAX previously raised and TF2/pt_expt silently forced exp. Copying the config into a fresh dict before adding num_steps also fixes the in-place config mutation. The new regression test genuinely locks the fix: it fails on the buggy behavior (an exp-always factory returns the wrong class for cosine/wsd; a mutating factory leaves num_steps in the input).
Non-blocking: the regression test covers the shared factory but not the trainer call sites that were actually broken. A trainer-level test (construct a JAX/TF2/pt_expt trainer with type: cosine/wsd and assert self.lr_schedule is the expected class) would guard against a future re-hardcoding of a backend to exp. The trainers are now one-line delegations to the tested factory, so this is a coverage note rather than a correctness concern.
Approving.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Requesting changes for one pt_expt integration blocker: valid cosine/WSD warmup schedules crash while constructing LambdaLR, before training starts.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional compatibility issue in the shared factory: its documented/schema default for an omitted learning-rate type is not implemented.
Use each schedule's nonzero start_lr as the PyTorch LambdaLR base, preserve the documented exponential default when type is omitted, and add trainer-level cosine/WSD warmup regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/tests/pt_expt/test_training.py`:
- Around line 314-315: Add the pytest timeout marker to
test_zero_start_warmup_schedulers_construct, using a 60-second limit while
preserving the test’s existing behavior.
- Around line 314-315: Add the pytest timeout decorator to
test_zero_start_warmup_schedulers_construct, using a 60-second limit while
preserving the test’s existing behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2f99ac77-9b95-4fb4-be64-6c0e932d052a
📒 Files selected for processing (4)
deepmd/dpmodel/utils/learning_rate.pydeepmd/pt_expt/train/training.pysource/tests/pt_expt/test_training.pysource/tests/universal/dpmodel/utils/test_learning_rate.py
🚧 Files skipped from review as they are similar to previous changes (3)
- deepmd/dpmodel/utils/learning_rate.py
- source/tests/universal/dpmodel/utils/test_learning_rate.py
- deepmd/pt_expt/train/training.py
Keep the trainer-level warmup regression within the repository's validation timeout. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Also addressed the approved review's non-blocking trainer-coverage note: the pt_expt trainer now has an integration regression for cosine and WSD warmup schedules, while the shared factory remains covered for all registered variants. Coding agent: Codex |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Re-reviewed the latest head. The zero-start warmup LambdaLR issue and omitted-type compatibility issue are both fixed. The shared schedule semantics now match the PT backend for dispatch, warmup, step progression, and restart. Targeted learning-rate, pt_expt warmup/restart, and cross-backend checks pass. Approving.
3 participants
Summary
BaseLRregistrypt_expttraining soexp,cosine, andwsdfollow the common schemaAddresses the learning-rate checklist items in #5755, #5756, and #5757. The remaining checklist items in those issues are intentionally left open.
Tests
venv/bin/pytest source/tests/universal/dpmodel/utils/test_learning_rate.py -qvenv/bin/pytest source/tests/pt_expt/test_training.py::TestTraining::test_training_loop -qDP_TEST_TF2_ONLY=1 venv/bin/pytest source/tests/tf2/test_training.py -qvenv/bin/ruff check .venv/bin/ruff format .JAX integration tests were not run locally because the existing virtual environment does not include JAX/Optax; the shared dispatch regression test covers the schedule factory used by the JAX trainer.
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Bug Fixes
Tests