fix(pd): add default chg spin parameter#5469
Open
HydrogenSulfate wants to merge 7 commits into
Open
Conversation
… charge_spin Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
…t_tebd descriptors Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
…e_spin keyword Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughUniform charge-spin interface methods are added to multiple descriptors; DPA3 gains default value support, validation, serialization, and uses ChangesCharge-spin descriptor interface and DPA3 implementation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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
🤖 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 `@deepmd/pd/model/atomic_model/dp_atomic_model.py`:
- Line 335: When constructing the call that sets charge_spin (the line currently
using "charge_spin=fparam if self.add_chg_spin_ebd else None"), ensure that when
self.add_chg_spin_ebd is True but fparam is None you fall back to the configured
default_chg_spin instead of passing None; change the expression to use fparam if
fparam is not None else self.default_chg_spin (only when self.add_chg_spin_ebd
is True), so charge_spin becomes (fparam if fparam is not None else
self.default_chg_spin) when self.add_chg_spin_ebd else None and avoids
triggering the DPA3 assertion path.
🪄 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: 6949ec6c-cc26-4f0f-a800-8b4381111ee6
📒 Files selected for processing (6)
deepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/model/descriptor/se_a.pydeepmd/pd/model/descriptor/se_t_tebd.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5469 +/- ##
==========================================
+ Coverage 81.41% 81.89% +0.47%
==========================================
Files 871 892 +21
Lines 96952 101609 +4657
Branches 4240 4242 +2
==========================================
+ Hits 78938 83217 +4279
- Misses 16711 17087 +376
- Partials 1303 1305 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/consistent/descriptor/common.py (1)
272-282: ⚡ Quick winRemove dead
fparam_pdassignment in PD descriptor helper
fparam_pdis computed fromfparambut never used/added tokwargson the PD path; it’s safe to remove for clarity (Ruff ignoresF841in this repo, so it won’t fail CI).♻️ Proposed cleanup
- fparam_pd = ( - paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None - ) charge_spin_pd = ( paddle.to_tensor(charge_spin).to(PD_DEVICE) if charge_spin is not None else None ) kwargs = {"nlist": nlist, "mapping": mapping} if charge_spin_pd is not None: kwargs["charge_spin"] = charge_spin_pd🤖 Prompt for 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. In `@source/tests/consistent/descriptor/common.py` around lines 272 - 282, The temporary paddle tensor fparam_pd is created but never used; remove the unused assignment "fparam_pd = (paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None)" from the PD descriptor helper so we don't allocate an unused tensor, and ensure only charge_spin_pd is computed/added to kwargs (leave PD_DEVICE, paddle usage and the kwargs mapping/nlist logic intact).
🤖 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.
Nitpick comments:
In `@source/tests/consistent/descriptor/common.py`:
- Around line 272-282: The temporary paddle tensor fparam_pd is created but
never used; remove the unused assignment "fparam_pd =
(paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None)" from
the PD descriptor helper so we don't allocate an unused tensor, and ensure only
charge_spin_pd is computed/added to kwargs (leave PD_DEVICE, paddle usage and
the kwargs mapping/nlist logic intact).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 7a52faf4-1efb-473c-9cc2-1c3928995a53
📒 Files selected for processing (3)
source/tests/consistent/descriptor/common.pysource/tests/consistent/descriptor/test_dpa3.pysource/tests/pd/model/test_dpa3.py
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Collaborator
Author
|
@iProzd please help review this PR. Thanks |
iProzd
requested changes
Jun 5, 2026
Member
There was a problem hiding this comment.
Thanks for you commit, but it seems that the charge_spin in model level still depends on the fparam. In #5431 , charge_spin is decoupled as an end to end independent param, including training data (charge_spin.npy) and forward. Please refer to this PR and change the corresponding implementation. Thanks!
| mapping=mapping, | ||
| comm_dict=comm_dict, | ||
| fparam=fparam if self.add_chg_spin_ebd else None, | ||
| charge_spin=fparam if self.add_chg_spin_ebd else None, |
Member
There was a problem hiding this comment.
charge_spin and fparam should be two seperated params, instead of getting charge_spin from fparam.
…r PT backend Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
iProzd
requested changes
Jun 11, 2026
iProzd
left a comment
iProzd
left a comment
Member
There was a problem hiding this comment.
It seems that the training parts and the end-to-end training test of charge_spin are still missing. (From reading charge_spin.npy to model forward. Now it only changes the model forward parts.)
…hods Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pd/model/model/make_model.py (1)
139-140:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winNormalize
charge_spindtype/device via_input_type_castbefore forwardingLine 195 and Line 321 forward raw
charge_spin, while other frame-level inputs are normalized through_input_type_cast(Line 168 and Line 309). This can introduce mixed-precision dtype mismatches in the downstream embedding path.💡 Proposed fix
- cc, bb, fp, ap, input_prec = self._input_type_cast( - coord, box=box, fparam=fparam, aparam=aparam + cc, bb, fp, ap, cs, input_prec = self._input_type_cast( + coord, box=box, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) @@ - charge_spin=charge_spin, + charge_spin=cs, ) @@ - cc_ext, _, fp, ap, input_prec = self._input_type_cast( - extended_coord, fparam=fparam, aparam=aparam + cc_ext, _, fp, ap, cs, input_prec = self._input_type_cast( + extended_coord, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) @@ - charge_spin=charge_spin, + charge_spin=cs, )def _input_type_cast( self, coord: paddle.Tensor, box: paddle.Tensor | None = None, fparam: paddle.Tensor | None = None, aparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, paddle.Tensor | None, paddle.Tensor | None, paddle.Tensor | None, + paddle.Tensor | None, str, ]: @@ - _lst: list[paddle.Tensor | None] = [ + _lst: list[paddle.Tensor | None] = [ vv.astype(coord.dtype) if vv is not None else None - for vv in [box, fparam, aparam] + for vv in [box, fparam, aparam, charge_spin] ] - box, fparam, aparam = _lst + box, fparam, aparam, charge_spin = _lst @@ - return coord, box, fparam, aparam, input_prec + return coord, box, fparam, aparam, charge_spin, input_prec else: pp = self.global_pd_float_precision return ( coord.to(pp), box.to(pp) if box is not None else None, fparam.to(pp) if fparam is not None else None, aparam.to(pp) if aparam is not None else None, + charge_spin.to(pp) if charge_spin is not None else None, input_prec, )Also applies to: 195-196, 268-269, 321-322
🤖 Prompt for 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. In `@deepmd/pd/model/model/make_model.py` around lines 139 - 140, The charge_spin parameter is being forwarded raw at multiple locations without dtype/device normalization, while other frame-level inputs are normalized through _input_type_cast. In deepmd/pd/model/model/make_model.py at lines 195-196, 268-269, and 321-322, apply the same _input_type_cast normalization to charge_spin before forwarding it, consistent with how other frame-level inputs are processed at lines 168 and 309. This ensures uniform dtype/device handling across all inputs to prevent mixed-precision mismatches in the embedding path.deepmd/pd/model/atomic_model/base_atomic_model.py (1)
248-249:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
charge_spinis not propagated in the bias-stat forward wrapperLine 248/Line 304 add
charge_spinto atomic inference, but_get_forward_wrapper_func.model_forward(Line 562 onward) still callsforward_common_atomic(...)withoutcharge_spin. Inchange-by-statisticbias updates, that can drop per-frame charge/spin and produce incorrect bias estimates (or break when defaults are unavailable).💡 Proposed fix
def model_forward( coord: paddle.Tensor, atype: paddle.Tensor, box: paddle.Tensor | None, fparam: paddle.Tensor | None = None, aparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> dict[str, paddle.Tensor]: @@ atomic_ret = self.forward_common_atomic( extended_coord, extended_atype, nlist, mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, )Also applies to: 304-305, 562-592
🤖 Prompt for 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. In `@deepmd/pd/model/atomic_model/base_atomic_model.py` around lines 248 - 249, The `charge_spin` parameter is added to the method signatures around lines 248-249 and 304-305, but it is not being propagated through the forward wrapper function. In the `_get_forward_wrapper_func` function (around lines 562-592), the `model_forward` method needs to be updated to accept the `charge_spin` parameter and pass it to the `forward_common_atomic` method call. Ensure that `charge_spin` is properly threaded through from the wrapper's input to the `forward_common_atomic` invocation so that charge and spin information is preserved during bias-statistic updates and not lost in the inference pipeline.
🤖 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.
Outside diff comments:
In `@deepmd/pd/model/atomic_model/base_atomic_model.py`:
- Around line 248-249: The `charge_spin` parameter is added to the method
signatures around lines 248-249 and 304-305, but it is not being propagated
through the forward wrapper function. In the `_get_forward_wrapper_func`
function (around lines 562-592), the `model_forward` method needs to be updated
to accept the `charge_spin` parameter and pass it to the `forward_common_atomic`
method call. Ensure that `charge_spin` is properly threaded through from the
wrapper's input to the `forward_common_atomic` invocation so that charge and
spin information is preserved during bias-statistic updates and not lost in the
inference pipeline.
In `@deepmd/pd/model/model/make_model.py`:
- Around line 139-140: The charge_spin parameter is being forwarded raw at
multiple locations without dtype/device normalization, while other frame-level
inputs are normalized through _input_type_cast. In
deepmd/pd/model/model/make_model.py at lines 195-196, 268-269, and 321-322,
apply the same _input_type_cast normalization to charge_spin before forwarding
it, consistent with how other frame-level inputs are processed at lines 168 and
309. This ensures uniform dtype/device handling across all inputs to prevent
mixed-precision mismatches in the embedding path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 09db3ca1-14d5-48c3-ba61-61530cc3d526
📒 Files selected for processing (5)
deepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/model/model/make_model.pydeepmd/pd/train/training.pydeepmd/pd/train/wrapper.pysource/tests/consistent/model/test_ener.py
💤 Files with no reviewable changes (1)
- source/tests/consistent/model/test_ener.py
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.
This pull request introduces support for handling
charge_spininformation in atomic and descriptor models. The main changes add new methods to query and provide defaultcharge_spinvalues, update theforwardmethods to accept acharge_spinargument, and refactor code to use this argument consistently instead of the previousfparamfield. Additionally, the DPA3 descriptor now supports defaultcharge_spinvalues and serializes them.Support for charge_spin in descriptor and atomic models:
dpa1.py,dpa2.py,dpa3.py,se_a.py,se_t_tebd.py) now implementget_dim_chg_spin,has_default_chg_spin, andget_default_chg_spinmethods to standardize how charge and spin information is queried. [1] [2] [3] [4] [5]forwardmethods in all descriptor classes, as well as in the atomic model (dp_atomic_model.py), now accept acharge_spinargument instead of (or in addition to)fparam, and all internal logic is updated to usecharge_spin. [1] [2] [3] [4] [5] [6]DPA3 descriptor enhancements:
default_chg_spinparameter, validates its shape, and serializes it as part of its configuration. [1] [2] [3]charge_spinargument.These changes standardize how charge and spin information is handled across descriptors and atomic models, making the codebase more extensible and robust for future features involving charge and spin.
Summary by CodeRabbit
New Features
charge_spin, shape nf×2) support across atomic, energy, and model forward pipelines, including model wrapper interfaces.charge_spinbased on model support and configuration.Tests
charge_spininput and improved backend coverage checks.