feat(pt/dpa4): add nvalchemi-toolkit warpper for DPA4#5568
Conversation
|
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 (3)
💤 Files with no reviewable changes (1)
✅ Files skipped from review due to trivial changes (2)
📝 WalkthroughWalkthroughAdds ChangesDPA4Wrapper nvalchemi adapter
Sequence Diagram(s)sequenceDiagram
participant User as User Script
participant DPA4Wrapper
participant NeighborListHook
participant SeZMModel
User->>DPA4Wrapper: from_checkpoint(path, device, compute_stress=True)
DPA4Wrapper-->>User: wrapper instance
User->>NeighborListHook: register BEFORE_COMPUTE hook
User->>DPA4Wrapper: forward(AtomicData/Batch)
DPA4Wrapper->>DPA4Wrapper: _atype(atomic_numbers) → type indices
DPA4Wrapper->>DPA4Wrapper: _edge_schema(neighbors, shifts, cell) → edge_index, edge_vec
DPA4Wrapper->>DPA4Wrapper: adapt_input() → coord, atype, edges, charge_spin
DPA4Wrapper->>SeZMModel: forward_lower(coord, atype, edges, ...)
SeZMModel-->>DPA4Wrapper: atom_energy, extended_force, extended_virial
DPA4Wrapper->>DPA4Wrapper: adapt_output() → energy, forces, Cauchy stress
DPA4Wrapper-->>User: ModelOutputs (energy, forces, stress)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 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}
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 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/pt/nvalchemi/wrapper.py`:
- Line 485: Before accessing model_params["model_dict"][head] at line 485,
validate that the head parameter exists as a key in the model_dict dictionary.
If the head is not found, raise a ValueError (instead of letting a raw KeyError
propagate) that includes a descriptive message with the list of available heads
to provide better debugging information and a cleaner API contract.
- Around line 274-283: The conditional check on line 274 that evaluates `if
neighbor_list_shifts is not None and cell is not None` allows incomplete
periodic metadata to pass through silently, causing incorrect physics when one
field is provided without the other. Add validation logic before this
conditional to check that both neighbor_list_shifts and cell are either both
present or both absent, and raise an informative error if they are inconsistent
(one is None while the other is not). This ensures that periodic boundary
condition calculations fail fast with a clear error message rather than silently
using incorrect zero shifts.
- Around line 366-367: The stress computation at the line where virial is
divided by volume does not validate whether the volume is valid before
performing the division, which can result in inf or nan values for singular or
degenerate cells. Add a validation check before computing the stress output that
ensures the volume tensor contains only valid (non-zero, non-NaN, non-inf)
values, and raise a clear and informative error if any invalid volumes are
detected. This check should occur after computing the volume using
torch.det(cell.to(virial.dtype)).abs().view(n_graph, 1, 1) but before the
division in the stress calculation.
In `@examples/water/dpa4/nvalchemi/relax.py`:
- Around line 95-96: The argument parser for --log-every lacks validation and
the optimization loop does not properly handle remaining steps. Add validation
to the argument parser around the --log-every definition to ensure the value is
greater than 0. Then examine the stepping logic around lines 141-151 where the
loop increments by log_every chunks, and modify it to ensure all remaining steps
up to --max-steps are executed even when --max-steps is not perfectly divisible
by log_every, preventing under-run of the optimization budget.
In `@examples/water/dpa4/nvalchemi/run_nve.py`:
- Around line 114-115: Add validation for the --log-every argument to ensure it
is greater than zero, preventing runtime crashes from zero or negative values.
Additionally, fix the stepping logic in the simulation loop to handle cases
where the total number of steps is not evenly divisible by log_every. Instead of
only advancing in fixed log_every chunks, explicitly run any remainder steps at
the end of the loop to ensure all requested steps are executed (e.g., if 205
steps are requested with log_every 20, ensure 5 additional steps are run after
the main loop completes). Update the argument parser to include validation for
--log-every and modify the stepping logic around the main simulation iteration
to calculate and execute remaining steps.
In `@examples/water/dpa4/nvalchemi/run_nvt.py`:
- Around line 104-105: Add validation for the log_every argument to ensure it is
a positive integer (greater than zero) in the argument parser definition.
Additionally, fix the NVT integration loop around line 168-174 to properly
handle cases where the total number of steps is not evenly divisible by
log_every. Ensure the remainder steps are executed after the main chunked loop
completes so that the simulation runs for the full requested number of steps
rather than truncating early.
🪄 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: 8c562042-bbdf-4ddd-806e-688cf8ca0c38
📒 Files selected for processing (11)
deepmd/pt/nvalchemi/__init__.pydeepmd/pt/nvalchemi/wrapper.pydoc/inference/index.rstdoc/inference/nvalchemi.mdexamples/water/dpa4/nvalchemi/README.mdexamples/water/dpa4/nvalchemi/relax.pyexamples/water/dpa4/nvalchemi/run_nve.pyexamples/water/dpa4/nvalchemi/run_nvt.pyexamples/water/dpa4/nvalchemi/single_point.pypyproject.tomlsource/tests/pt/model/test_sezm_nvalchemi.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5568 +/- ##
==========================================
- Coverage 82.18% 82.01% -0.18%
==========================================
Files 898 901 +3
Lines 103576 103799 +223
Branches 4434 4435 +1
==========================================
+ Hits 85129 85130 +1
- Misses 17056 17278 +222
Partials 1391 1391 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
OutisLi
force-pushed
the
pr/nvtoolkit-dpa4
branch
from
11ffa62 to
7ea521f
Compare
PR Summary by QodoAdd nvalchemi-toolkit adapter for DPA-4/SeZM PyTorch models Description
Diagram
High-Level Assessment
Files changed (11)
|
Code Review by Qodo
1. Sparse edges passed as nlist
|
![qodo-free-for-open-source-projects[bot]](https://avatars.githubusercontent.com/in/1068000?s=60&v=4){kind=small}
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Thanks for adding this integration. I think this needs changes before merge: I verified the current head (f62b88775cfb00ca616bea857463c13002c48320) and several main paths in the wrapper are still incompatible with the SeZM lower/AOTI interfaces, so the nvalchemi adapter would fail at runtime rather than just needing polish.
Blocking issues:
-
DPA4Wrapper.forward()passes sparse-edge tensors intoSeZMModel.forward_lower()with the wrong interface.
Indeepmd/pt/nvalchemi/dpa4wrapper.py, the wrapper buildsedge_indexas a COO tensor with shape(2, E)and passesedge_index,edge_vec,edge_scatter_index, andedge_maskpositionally toforward_lower. But the currentSeZMModel.forward_lower()signature expects DeePMD lower-interface inputs:(extended_coord, extended_atype, nlist, mapping, fparam, aparam, ...), wherenlistis a padded 3D neighbor list(nf, nloc, nsel). The model immediately routes throughformat_nlist()/build_edge_list_from_nlist(), so a(2, E)COO tensor will not work. Please either convert the nvalchemi COO list into the standard paddednlist+mapping, or add a supported sparse-edge SeZM entrypoint and call that explicitly. -
The
.pt2backend has the same callable-signature mismatch.
_run_pt2()calls the AOTI runner with 9 sparse-edge-style arguments, but the existing export path/test for DPA-4.pt2artifacts uses the 7-argumentforward_common_lowerinterface:(ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin). A standarddp --pt freezepackage will therefore not match this wrapper call. This should be fixed consistently with item 1, and ideally covered with a.pt2wrapper test. -
compute_embeddings()calls an unsupported API.
self.model.forward_common_lower(..., embedding_only=True)is not a valid call in this repository:forward_common_lower()has noembedding_onlykeyword, and it does not return adescriptorkey. Please implement embeddings via a supported descriptor path (for example, usingatomic_model.descriptor.forward_with_edges(...)with the prepared edge tensors), or add a dedicated SeZM method for descriptor-only output. -
The new optional extra is unconditional but
nvalchemi-toolkitis not available on all DeePMD-supported environments.
pyproject.tomladdsnvalchemi = ["nvalchemi-toolkit"], whilenvalchemi-toolkit==0.1.0currently declares Python>=3.11,<3.14; DeePMD-kit still supports Python 3.10. The existing PyTorch extra already guardsnvalchemi-toolkit-opswith environment markers for Python/Linux. Please add appropriate markers for this extra and document the supported Python/platform constraints indoc/third-party/nvalchemi.mdand the example README. -
The documented example install path misses an example dependency.
The example scripts import ASE (from ase.data import atomic_numbers as ASE_Z), butpip install deepmd-kit[nvalchemi]only pullsnvalchemi-toolkit. A user following the docs/README will hitModuleNotFoundError: No module named 'ase'before reaching the wrapper. Either remove the ASE dependency from these scripts by using an already-required element mapping, or include/document ASE (or an appropriate nvalchemi extra) in the prerequisites.
Smaller follow-up: the stress path divides by abs(det(cell)) without validating finite/non-zero volume. It would be safer to fail fast with a clear ValueError for singular/invalid cells, especially before using the output in barostats.
The new tests are a good direction, but because nvalchemi-toolkit is optional and absent from the current CI environment, the wrapper tests appear to be skipped and these runtime interface issues are not exercised. It would be worth adding at least one CI-covered path (or a lightweight mocked/interface test) that catches the SeZM lower-interface call shape.
Reviewed by OpenClaw 2026.6.8 (844f405)
model: custom-chat-jinzhezeng-group/gpt-5.5
|
Should be merged after PR #5571 is merged, this PR relies on that |
| ret = self.model.forward_common_lower( | ||
| model_inputs["coord"], | ||
| model_inputs["atype"], | ||
| model_inputs["edge_index"], | ||
| model_inputs["edge_vec"], | ||
| model_inputs["edge_scatter_index"], | ||
| model_inputs["edge_mask"], | ||
| charge_spin=model_inputs["charge_spin"], | ||
| embedding_only=True, | ||
| ) | ||
| n_node = data.num_nodes | ||
| node_embeddings = ret["descriptor"].reshape(n_node, self._descriptor_dim) |
There was a problem hiding this comment.
This call will crash at runtime. forward_common_lower has no embedding_only parameter (and no **kwargs), and its return dict does not contain a "descriptor" key. The current signature in deepmd/pt/model/model/sezm_model.py is (coord, atype, edge_index, edge_vec, edge_scatter_index, edge_mask, fparam, aparam, comm_dict, extended_coord_corr, charge_spin, input_prec, use_compile) and it returns {energy, energy_redu, energy_derv_r, energy_derv_c, energy_derv_c_redu, mask}. The token embedding_only does not exist anywhere in deepmd/. So compute_embeddings (and embedding_shapes/_descriptor_dim, which feed the same path) raise TypeError: ... unexpected keyword argument 'embedding_only' on first use, then KeyError: 'descriptor' if that were patched. The embeddings feature needs to be wired to whatever the model actually exposes for descriptor output.
| delattr(cls, "_default_device_ctx") | ||
|
|
||
|
|
||
| @unittest.skipUnless(NVALCHEMI_AVAILABLE, "nvalchemi-toolkit is not installed") |
There was a problem hiding this comment.
The entire suite is gated on NVALCHEMI_AVAILABLE, but nvalchemi-toolkit is added only as a user-facing optional extra in pyproject.toml -- it is not in any test/CI dependency group (and is not installable from PyPI). So in upstream CI every test here silently skips and this ~665-line module gets zero executed coverage. Concrete evidence this hurts: the compute_embeddings runtime crash (forward_common_lower(embedding_only=True) / ret["descriptor"]) is exercised by test_compute_embeddings_shapes yet went uncaught, i.e. the suite was never actually run with nvalchemi installed. Please add nvalchemi-toolkit to the CI test dependencies (or otherwise ensure these tests run), so the parity tests provide real coverage.
| n_graph, 9, dtype=atom_virial.dtype, device=atom_virial.device | ||
| ).index_add_(0, batch_idx, atom_virial) | ||
| volume = torch.det(cell.to(virial.dtype)).abs().view(n_graph, 1, 1) | ||
| output["stress"] = (virial.view(n_graph, 3, 3) / volume).to(out_dtype) |
There was a problem hiding this comment.
Warning -- please double-check the stress sign convention. This emits stress = +virial / |det(cell)| (unsymmetrized), whereas deepmd's own ASE calculator uses the opposite sign and symmetrizes: stress = -0.5 * (v + v.T) / volume (deepmd/calculator.py:160-161, "tensile stress is positive"), where v is the same virial quantity. If nvalchemi follows the standard ASE/tensile convention, this stress is sign-inverted (and missing symmetrization), which would drive NPT/NPH barostats the wrong way. The parity test only asserts stress*volume == native_virial, which is circular and cannot catch a wrong sign. Please confirm nvalchemi's expected stress convention and add a non-circular check (e.g. a finite-difference stress, or a known-pressure reference).
4 participants
Summary by CodeRabbit
Release Notes
New Features
.pt2packagesDocumentation