fix(pd): fix adamw for paddle backend

[HydrogenSulfate]

This pull request updates the optimizer handling in the deepmd/pd/train/training.py file, adding support for selecting between Adam and AdamW optimizers based on configuration, and improves profiling clarity.

Optimizer selection improvements:

• Refactored optimizer initialization to dynamically select between paddle.optimizer.Adam and paddle.optimizer.AdamW depending on the value of self.opt_type. This allows for more flexibility in choosing the optimizer.

Profiling and logging enhancements:

• Updated the profiling context label from "Adam update" to "Optimizer update" to accurately reflect the optimizer in use, improving the clarity of profiling output.

Summary by CodeRabbit

• New Features
  • Added support for selecting between Adam and AdamW optimizers during training.

[Copilot]

Pull request overview

This PR fixes Paddle backend optimizer selection so that optimizer/type: AdamW actually uses Paddle’s AdamW optimizer (instead of always constructing Adam), and updates a profiler label to be optimizer-agnostic.

Changes:

• Select paddle.optimizer.Adam vs paddle.optimizer.AdamW based on self.opt_type.
• Rename the nvprof step label from “Adam update” to “Optimizer update”.

Comments suppressed due to low confidence (1)

deepmd/pd/train/training.py:669

• Newly enabling AdamW for the Paddle backend changes behavior, but there doesn’t appear to be any Paddle-side test that asserts the trainer actually constructs paddle.optimizer.AdamW when optimizer/type is set to AdamW. Adding a small unit/integration test (similar to existing source/tests/pd/test_training.py coverage) would prevent regressions back to Adam and confirm weight-decay behavior for this backend.

            opt_cls = (
                paddle.optimizer.Adam
                if self.opt_type == "Adam"
                else paddle.optimizer.AdamW
            )
            self.optimizer = opt_cls(
                learning_rate=self.scheduler,
                parameters=self.wrapper.parameters(),
                beta1=float(self.opt_param["adam_beta1"]),
                beta2=float(self.opt_param["adam_beta2"]),
                weight_decay=float(self.opt_param["weight_decay"]),
            )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change introduces dynamic selection of the optimizer class (Adam or AdamW) based on configuration in the training code, replacing a hard-coded Adam instantiation. A profiling label is also updated from "Adam update" to "Optimizer update" to reflect the generic optimizer usage.

Changes

Cohort / File(s)	Summary
Optimizer Selection and Profiling deepmd/pd/train/training.py	Dynamic optimizer instantiation based on opt_type parameter; profiling label updated to generic "Optimizer update" for broader compatibility across optimizer types.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Suggested reviewers

• wanghan-iapcm
• njzjz

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix(pd): fix adamw for paddle backend' accurately describes the main change: adding dynamic optimizer selection to support AdamW alongside Adam in the paddle backend.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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 and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 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/pd/train/training.py`:
- Around line 658-663: Runtime now constructs paddle.optimizer.AdamW in
training.py via opt_cls when self.opt_type == "AdamW", but the optimizer
docs/schema still mark "AdamW" as doc_only_pt_supported; update the optimizer
schema to reflect Paddle support by removing "AdamW" from the
doc_only_pt_supported set (or adding a separate marker for paddle support), and
ensure the optimizer name "AdamW" appears in the supported backends list used by
the doc generation code (look for the doc_only_pt_supported symbol and the
optimizer support mapping in argcheck.py and adjust it so "AdamW" is not
PyTorch-only).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 682953ff-1a57-4e22-8140-71e3d03f9abc

📥 Commits

Reviewing files that changed from the base of the PR and between a8d2a8d and 9e0c1bb.

📒 Files selected for processing (1)

• deepmd/pd/train/training.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.28%. Comparing base (a8d2a8d) to head (9e0c1bb).
⚠️ Report is 5 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5308      +/-   ##
==========================================
- Coverage   82.28%   82.28%   -0.01%     
==========================================
  Files         773      773              
  Lines       77330    77332       +2     
  Branches     3659     3659              
==========================================
  Hits        63631    63631              
- Misses      12528    12529       +1     
- Partials     1171     1172       +1     

☔ View full report in Codecov by Sentry.
📢 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.