fix(train): dispatch registered learning rate schedules

[njzjz-bot]

Summary

• add a shared learning-rate schedule factory backed by the BaseLR registry
• use it in JAX, TF2, and pt_expt training so exp, cosine, and wsd follow the common schema
• preserve the caller's learning-rate configuration while injecting the training step count
• add regression coverage for all registered schema variants

Addresses 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 -q
• venv/bin/pytest source/tests/pt_expt/test_training.py::TestTraining::test_training_loop -q
• DP_TEST_TF2_ONLY=1 venv/bin/pytest source/tests/tf2/test_training.py -q
• venv/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

  • Added a shared learning-rate schedule builder supporting exponential, cosine, and WSD schedules.
  • Training workflows now select learning-rate schedules from configuration settings consistently across supported backends.

• Bug Fixes

  • Prevented learning-rate configuration dictionaries from being modified during schedule creation.

• Tests

  • Added coverage for schedule selection and configuration immutability.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 603458a4-4854-4f7a-954b-517b70ec4ac1

📥 Commits

Reviewing files that changed from the base of the PR and between 0c5a914 and a09e557.

📒 Files selected for processing (5)

• deepmd/dpmodel/utils/learning_rate.py
• deepmd/jax/train/trainer.py
• deepmd/pt_expt/train/training.py
• deepmd/tf2/train/trainer.py
• source/tests/universal/dpmodel/utils/test_learning_rate.py

---

📝 Walkthrough

Walkthrough

A shared make_learning_rate_schedule factory now constructs registered schedules without mutating input parameters. JAX, PyTorch, and TensorFlow trainers use it instead of direct exponential schedule construction, with tests covering exponential, cosine, and WSD schedules.

Changes

Learning-rate schedule construction

Layer / File(s)	Summary
Schedule factory and validation deepmd/dpmodel/utils/learning_rate.py, source/tests/universal/dpmodel/utils/test_learning_rate.py	Adds a factory that copies parameters, injects num_steps, constructs the registered schedule, and tests exp, cosine, wsd, and input immutability.
Trainer schedule integration deepmd/jax/train/trainer.py, deepmd/pt_expt/train/training.py, deepmd/tf2/train/trainer.py	Updates all three trainers to build schedules through the shared factory and removes direct exponential-schedule dispatch.รร

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

Loading

Possibly related PRs

• deepmodeling/deepmd-kit#5142: Introduced the registered BaseLR schedule machinery used by this factory.

Suggested labels: new feature, Python

Suggested reviewers: wanghan-iapcm, njzjz, OutisLi

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: training now dispatches registered learning-rate schedules through a shared factory.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 79.57%. Comparing base (0c5a914) to head (a09e557).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/tf2/train/trainer.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5776      +/-   ##
==========================================
- Coverage   79.69%   79.57%   -0.12%     
==========================================
  Files        1020     1020              
  Lines      116359   116356       -3     
  Branches     4303     4306       +3     
==========================================
- Hits        92736    92594     -142     
- Misses      22076    22219     +143     
+ Partials     1547     1543       -4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.