feat(pt_expt): full model and refact the module output names of dpmodel backend#5243
feat(pt_expt): full model and refact the module output names of dpmodel backend#5243wanghan-iapcm wants to merge 77 commits intodeepmodeling:masterfrom
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, torch.Tensor]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| return | ||
| return super().__setattr__(name, value) | ||
|
|
||
| def call(self, x: torch.Tensor) -> torch.Tensor: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| ) | ||
| # Compare the common keys | ||
| common_keys = set(dp_ret.keys()) & set(pt_ret.keys()) | ||
| self.assertTrue(len(common_keys) > 0) |
Check notice
Code scanning / CodeQL
Imprecise assert Note test
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the full model support for the experimental PyTorch backend (pt_expt), including hooks for descriptor and fitting layer evaluation, and autograd-based derivative calculations. The changes also improve backend-agnostic array operations in the core dpmodel and ensure compatibility with torch.fx tracing. I have identified a few issues regarding potential runtime errors in the new evaluation hooks and unintended side effects in the output definition translation logic.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds PT-Expt PyTorch-exportable model layers and autodiff-based output transformations, introduces descriptor and bias accessors, refactors input/output casting across backends, enables middle-output capture in fitting, updates device-aware tensor creation, replaces in-place network ops, and adds extensive pt_expt tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client Code
participant EM as EnergyModel (PT-Expt)
participant Desc as Descriptor
participant Fit as Fitting Net
participant TO as TransformOutput
participant Autograd as PyTorch Autograd
Client->>EM: forward(coord, atype, box, ...)
EM->>Desc: compute descriptors
Desc-->>EM: descriptor tensor
EM->>Fit: evaluate fitting network
Fit-->>EM: fitting outputs (per-atom/reducible)
EM->>TO: fit_output_to_model_output(fit_ret, coord_ext, ...)
TO->>Autograd: enable grad on extended coords
Autograd->>TO: compute ∇energy -> forces, virial, atom_virial
TO-->>EM: model outputs (energy, atom_energy, force, virial, ...)
EM-->>Client: return outputs
sequenceDiagram
participant GP as GeneralFitting
participant PT as Per-type Net(s)
participant Acc as MiddleOutput Accumulator
participant Res as Results Dict
GP->>GP: eval_return_middle_output?
alt mixed_types False
GP->>Acc: init per-type accumulation
loop per-type
GP->>PT: evaluate net(xx)
PT-->>GP: output (+ middle_output)
GP->>Acc: accumulate middle_output
end
GP->>Res: store Acc as "middle_output"
else mixed_types True
GP->>PT: call_until_last(xx)
PT-->>GP: middle_output
GP->>Res: store middle_output
end
GP-->>Res: set var_name -> output tensor
GP-->>Client: return Res
(Note: colored rectangles not used; flows kept minimal.) Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 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}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@source/tests/consistent/model/test_ener.py`:
- Around line 1184-1190: The variable nloc is assigned but never used—remove the
unused assignment or replace it with a meaningful use; specifically in the test
block where nframes, coords_2f, atype_2f, box_2f, natoms_data, and energy_data
are set, delete the declaration "nloc = 6" (or if intended, use nloc to drive
array shapes/validation) so there are no unused local variables in the test_ener
setup.
- Around line 124-126: The test contains a duplicate assignment to pd_class
(assigned to EnergyModelPD twice); remove the redundant second assignment so
pd_class is only set once and leave the other class assignments (pt_expt_class =
EnergyModelPTExpt and jax_class = EnergyModelJAX) unchanged.
- Around line 979-988: In test_change_out_bias the local variable nloc is
assigned but never used; remove the unused assignment (nloc = 6) from the
test_change_out_bias function to satisfy static analysis (or if the intended
intent was to use it, replace references to hardcoded 6 with nloc where
appropriate) so that only necessary variables (e.g. coords_2f, atype_2f, box_2f,
natoms_data, energy_data) remain.
🧹 Nitpick comments (6)
deepmd/dpmodel/fitting/general_fitting.py (1)
599-639: Minor: double forward pass wheneval_return_middle_outputis enabled.When the hook is active, the network is evaluated twice per type (or once extra for mixed types) — once for the full output (line 609/635) and again via
call_until_last(line 621/639). This duplicates computation of all layers except the last.Since the dpmodel backend is primarily for reference/testing rather than production, this is acceptable. However, if performance becomes a concern, you could refactor the network to return both the final and penultimate outputs in a single pass.
deepmd/pt_expt/model/transform_output.py (1)
109-110:zip()withoutstrict=True(ruff B905).Both
split_vv1andsplit_svv1are produced from the samesizesplit, so they're guaranteed to have equal length. Addingstrict=Truemakes this invariant explicit and guards against future refactors that might change one without the other.Suggested fix
- for vvi, svvi in zip(split_vv1, split_svv1): + for vvi, svvi in zip(split_vv1, split_svv1, strict=True):source/tests/pt_expt/model/test_ener_model.py (1)
207-250: Consider extending DP-consistency test to also verify force values.The consistency test validates
energyandatom_energybut doesn't checkforce. Since the dpmodel sets derivative outputs toNone, this is understandable, but you could compute a numerical finite-difference force from the dpmodel to cross-validate the autograd-based force from pt_expt. This would strengthen confidence in the derivative pathway.(Note: the autodiff test file mentioned in the summary may already cover this — feel free to disregard if so.)
deepmd/pt_expt/model/ener_model.py (1)
100-147:forward_lowerbakesfparam,aparam,do_atomic_virialinto the traced graph.The closure captures these values as constants during
make_fxtracing, so the returned module is specialized for the specificfparam/aparam/do_atomic_virialvalues passed at trace time. This is appropriate for export workflows but means the caller must re-trace for different parameter configurations. The docstring correctly documents this ("Sample inputs with representative shapes"), but it might be worth explicitly noting thatfparam/aparam/do_atomic_virialare baked in as well.source/tests/pt_expt/model/test_autodiff.py (2)
49-76: Docstring sayscellshape is[nf, 3, 3], but the body reshapes to[nframes, 9].The function accepts
cellwith shape[nf, 3, 3]per the docstring, which is consistent with the callers (e.g.,cell.unsqueeze(0)producing[1, 3, 3]). However, line 72 reshapes it to[nframes, 9]. This works, but the docstring could note that either[nf, 3, 3]or[nf, 9]is accepted, sincereshape(nframes, 9)handles both shapes silently. Very minor — just a documentation clarity nit.
154-181: DuplicatedsetUpbetweenTestEnergyModelSeAForceandTestEnergyModelSeAVirial.Both test classes have identical
setUpmethods. Consider extracting a shared helper or a common base mixin to reduce duplication.That said, this pattern (identical setUp in mixin-based test classes) is common in this codebase and the duplication is small, so this is a minor nit.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1fa1eb27b2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5243 +/- ##
==========================================
- Coverage 82.12% 82.06% -0.06%
==========================================
Files 736 740 +4
Lines 74237 74667 +430
Branches 3615 3616 +1
==========================================
+ Hits 60966 61277 +311
- Misses 12107 12225 +118
- Partials 1164 1165 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
How do we handle model output keys in the new backend? |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@deepmd/dpmodel/atomic_model/dp_atomic_model.py`:
- Around line 217-221: The runtime check using assert in the block guarded by
self.enable_eval_fitting_last_layer_hook is unsafe because asserts are removed
with -O; replace it with an explicit runtime exception (e.g., raise RuntimeError
or ValueError) that checks "middle_output" in ret and raises a clear error
message, then pop the key and append to self.eval_fitting_last_layer_list as
before (refer to the symbols self.enable_eval_fitting_last_layer_hook,
ret.pop("middle_output"), and self.eval_fitting_last_layer_list to locate the
code).
🧹 Nitpick comments (1)
deepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
140-148: Consider clearing the cache after retrieval to prevent stale data and unbounded growth.Both
eval_descriptor()andeval_fitting_last_layer()leave the cache intact after returning the concatenated result. If a caller invokes them, then runs more forward passes, a subsequent call silently includes data from both the old and new passes. If clearing isn't desired, at minimum document the accumulation semantics so callers know to callset_eval_*_hook(True)again to reset.Also applies to: 156-164
Shouldn't it always use the output keys in the dpmodel backend? |
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| def call( | ||
| self, | ||
| coord: Array, | ||
| atype: Array, | ||
| box: Array | None = None, | ||
| fparam: Array | None = None, | ||
| aparam: Array | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, Array]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
wanghan-iapcm
changed the title
2 participants
Summary by CodeRabbit
New Features
Improvements
Refactor
Tests