feat(dpmodel): NeighborGraph 3-body angle machinery (PR-E)#5717
Merged
OutisLi merged 14 commits intoJul 9, 2026
Conversation
Contributor
|
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:
📝 WalkthroughWalkthroughThis PR adds angle-graph support to ChangesAngle graph construction and aggregation
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![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
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_angle_builder.py (1)
326-364: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a padded-
angle_indexcase to the aggregation tests.Both aggregation tests here only use angle data where every slot is real (no padding). Given the padding-mask precondition gap noted in
angles.py(angle_to_edge_sum/angle_to_node_sum), a test with a paddedangle_indexand non-zero paddingdatavalues would catch regressions if that precondition is ever violated.🤖 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/common/dpmodel/test_angle_builder.py` around lines 326 - 364, The aggregation tests in test_angle_builder only cover fully real angle slots, so add a new padded-angle case for angle_to_edge_sum and angle_to_node_sum using a padded angle_index with non-zero data in the padded positions. Reuse the existing test_angle_aggregation and test_angle_aggregation_torch_namespace patterns, but assert that padding is ignored in both the numpy and torch-namespace paths so regressions around the padding-mask precondition are caught.
🤖 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/dpmodel/utils/neighbor_graph/segment.py`:
- Around line 59-89: The masked softmax in segment_softmax can overflow because
shifted is computed from raw data instead of data_for_max, allowing masked
entries to produce inf and then nan when multiplied by the zero mask. Update
segment_softmax to base the subtraction on data_for_max so masked elements
remain -inf before exp, keeping segment_sum and the denom_e guard from being
poisoned; use the existing segment_max, segment_sum, and mask handling paths to
locate the fix.
---
Nitpick comments:
In `@source/tests/common/dpmodel/test_angle_builder.py`:
- Around line 326-364: The aggregation tests in test_angle_builder only cover
fully real angle slots, so add a new padded-angle case for angle_to_edge_sum and
angle_to_node_sum using a padded angle_index with non-zero data in the padded
positions. Reuse the existing test_angle_aggregation and
test_angle_aggregation_torch_namespace patterns, but assert that padding is
ignored in both the numpy and torch-namespace paths so regressions around the
padding-mask precondition are caught.
🪄 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: b038075f-6809-4579-8ce1-5cf5d1d41c90
📒 Files selected for processing (17)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/angles.pydeepmd/dpmodel/utils/neighbor_graph/env.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/dpmodel/utils/neighbor_graph/pairs.pydeepmd/dpmodel/utils/neighbor_graph/segment.pysource/tests/common/dpmodel/test_angle_builder.pysource/tests/common/dpmodel/test_center_edge_pairs.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_graph_attention_parity.pysource/tests/common/dpmodel/test_graph_angle_cos_parity.pysource/tests/common/dpmodel/test_segment_softmax.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_linear_model.pysource/tests/pt_expt/utils/test_neighbor_list.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5717 +/- ##
==========================================
- Coverage 79.61% 79.50% -0.11%
==========================================
Files 1014 1015 +1
Lines 115359 115483 +124
Branches 4272 4274 +2
==========================================
- Hits 91839 91813 -26
- Misses 21978 22127 +149
- Partials 1542 1543 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
force-pushed
the
feat-graph-angles-prE
branch
from
July 7, 2026 08:13
ad73de0 to
dd58c02
Compare
- Remove unused locals flagged by CodeQL (ai_def/ai_full in the ordered-superset test, ev in the angle-cos parity test); strengthen the ordered+include_self test to assert the default pairs are an actual subset of the full pair set (using the previously-dead ai_def/ai_full). - Add masked-padding aggregation tests (CodeRabbit nitpick): pin the precondition that angle_to_edge_sum/angle_to_node_sum ignore guard angles only once data is zeroed via angle_mask, in numpy and torch.
iProzd
reviewed
Jul 7, 2026
iProzd
reviewed
Jul 7, 2026
…nd mask precondition - Document that build_angle_index's strict '<' a_rcut gate intentionally mirrors dpa3's dense angle gate (repflows.py:598 a_dist_mask), unlike the edge channel's '<= rcut' (builder.py:284). Answers iProzd's review question: yes, boundary divergence is intentional dpa3 parity. - Document the angle_to_edge_sum/angle_to_node_sum 'caller must zero data at padded slots first' contract directly in their docstrings (previously only pinned by a test, per iProzd's second review comment).
OutisLi
requested changes
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the careful machinery and the thorough oracle/parity tests. The forward parity is solid; my one change request is about autograd safety of the norm used in graph_angle_cos (details inline). A couple of minor, non-blocking notes will follow separately.
OutisLi
reviewed
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
One non-blocking performance follow-up on the unordered angle enumeration (details inline). The change itself lands in the shared center_edge_pairs primitive.
OutisLi
reviewed
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
One more non-blocking readability note on the pair_mask fold (inline).
Address OutisLi's blocking review on PR deepmodeling#5717. graph_angle_cos normalized with xp.linalg.vector_norm, but the dpa3 dense channel it mirrors uses safe_for_vector_norm for the same step (repflows.py:642-643), and the a_rcut gate mirrors repflows.py:598 which also uses safe_for_vector_norm. The two are value-identical for non-zero vectors (fp64 dense/se_t parity still passes at 1e-12), but differ in gradient at ||v||==0: plain vector_norm back-props NaN under jax, whereas safe_for_vector_norm yields a 0 gradient. edge_vec is the sole autograd leaf, so this was a latent NaN-gradient landmine on the geometry path (not yet triggered: padding angles point at a real, non-zero edge, and no gradients flow until the dpa3 graph descriptor lands in PR-G). Switching both call sites removes the landmine and makes the 'mirrors dpa3 exactly' docstring claims literally true. Add test_graph_angle_cos_zero_edge_grad_is_finite: a jax-autograd test with a zero-length edge_vec asserting the geometry-path gradient stays finite, plus a discrimination check that plain vector_norm DOES produce a NaN gradient under jax on the same construction (torch's vector_norm is already safe at zero, so jax is the discriminating backend). Guarded by pytest.importorskip('jax').
iProzd
approved these changes
Jul 8, 2026
OutisLi
approved these changes
Jul 9, 2026
HydrogenSulfate
pushed a commit
to HydrogenSulfate/deepmd-kit
that referenced
this pull request
Jul 10, 2026
deepmodeling#5728) ## Summary `source/tests/pt_expt/test_plugin.py` (from deepmodeling#5559) pops `deepmd.pt_expt` from `sys.modules` without restoring it, leaving the package's cached submodules bound to a dead parent. Any later import of a cached submodule (e.g. `deepmd.pt_expt.infer.deep_eval`) re-creates a BARE parent package whose `utils`/`infer` attributes are never rebound, and `mock.patch("deepmd.pt_expt.utils...")` in `test_deep_eval_serialize_api.py` then fails with `AttributeError: module 'deepmd.pt_expt' has no attribute 'utils'` under py3.10's mock target resolution (py3.13 tolerates it). The failure is shard-order dependent: it needs (1) something to import `deepmd.pt_expt.infer.deep_eval` before `test_plugin` runs, and (2) the serialize-API test to run after — so it appears/disappears as PRs add test files and reshuffle the duration-based shards. It currently fails `Test Python (4, 3.10)` on deepmodeling#5714 and `Test Python (8, 3.10)` on deepmodeling#5717 while master stays green by ordering luck. ## Fix Snapshot the whole `deepmd.pt_expt.*` module tree before the re-import and restore it (including the `deepmd` parent-package attribute binding) in the `finally`. Verified: fixed test passes together with the serialize-API tests; an inline emulation of the worst-case ordering (child cached → pop/reimport/pop → `mock.patch` target resolution) resolves cleanly. Same one-file fix is cherry-picked on deepmodeling#5714 (as 6422007) and deepmodeling#5717; whichever lands first, the others resolve trivially. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Improved import handling in the plugin loading test to better isolate module state between test runs. * Reduced the chance of leftover cached imports affecting later tests in the same session. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
pull Bot
pushed a commit
to ishandutta2007/deepmd-kit
that referenced
this pull request
Jul 19, 2026
…model on the NeighborGraph lower (PR-G-dpa2) (deepmodeling#5779) DPA-2 becomes graph-native full-stack — the first **message-passing** model on the NeighborGraph lower (PR-G-dpa2 of the [NeighborGraph series](wanghan-iapcm#4); follows dpa1's route from deepmodeling#5581/deepmodeling#5583/deepmodeling#5604/deepmodeling#5714/deepmodeling#5715/deepmodeling#5717/deepmodeling#5733). The dense path coexists untouched; the graph route is opt-in for inference (`--lower-kind graph`) and the default for pt_expt training/eager on graph-eligible configs, exactly like dpa1. ## What's new **dpmodel (backend-agnostic math)** - `DescrptBlockRepformers.call_graph`: every repformer op ported to the flat edge list — conv/grrg/drrd/g1g1/`LocalAtten` via `segment_*` over `dst`; the dense `nnei×nnei` gated attention (`Atten2Map`/`Atten2MultiHeadApply`/`Atten2EquiVarApply`) via `center_edge_pairs(ordered, include_self)` + per-head `segment_softmax` (extended to trailing feature dims). - `DescrptDPA2.call_graph` + `uses_graph_lower` + dense-call adapter. The multi-level nlist (`build_multiple_neighbor_list`) becomes per-block edge masks; in the shape-static adapter layout the per-block **slice** replicates the dense `nlist[:, :, :ns]` truncation, making the adapter **bit-exact vs the dense path at ANY sel** (verified to 5e-16 at deliberately binding sel with attention enabled) — this is what keeps the cross-backend consistency suites green after the routing flip. Ineligible configs (`use_three_body`, compression, spin) keep the legacy dense route. - Owned-node energy mask: `fit_output_to_model_output_graph` consumes `n_local`, excluding halo rows from the differentiated energy (prevents cross-rank double counting; force stays full-N, halo partials reverse-commed by LAMMPS). **pt_expt / export** - Routing, autograd force/virial, compiled training, and freeze eligibility are inherited generically from the dpa1 machinery. - Per-layer MPI halo refresh on the graph path: `_exchange_ghosts_graph` — identity on ghost-free single-rank graphs (`src` IS the owner), `deepmd_export::border_op` overwrite of halo rows on extended-region multi-rank graphs. - Message-passing graph `.pt2` archives embed a **with-comm AOTInductor artifact** (`model/extra/forward_lower_with_comm.pt2`), traced with the 8-tensor comm ABI shared with the dense flow. **C++ / LAMMPS** - `DeepPotPTExpt::run_model_graph_with_comm` + dispatch replacing the PR-G fail-fast: MP message-passing graph models run multi-rank on the extended-region graph + with-comm artifact; non-MP (dpa1) multi-rank unchanged; old graph archives without the artifact get a clear re-freeze error. - Fixtures (`gen_dpa2.py` section B), `dpa2_graph_ptexpt` universal gtest row, and `test_lammps_dpa2_graph_pt2.py` (single-rank vs reference, per-atom virial, `mpirun -n 2` ≡ `-n 1`, graph-vs-nlist cross-artifact, bounded empty-rank). ## Validation - GPU (Tesla T4, CUDA): dpa2-graph LAMMPS **5/5** incl. the MP==SP gate; dpa1-graph LAMMPS 6/6; dense dpa2 MPI 1/1; every `dpa2_graph` C++ gtest variant passed; python GPU suites green. - CPU: full dpmodel+pt_expt sweep 2112 passed; dpa2 consistency suite (pt/dp/jax/array-api-strict) green via the bit-exact adapter. - Real bugs found & fixed during GPU validation: CUDA device placement of the `nlocal`/`nghost` comm scalars (the graph route consumes them in on-device owned-mask index math); `_trace_and_compile_graph` hardcoded the global device for trace samples (latent since deepmodeling#5604). ## Known limitations / deliberate divergences (documented in `doc/model/dpa2.md` + code Notes) - Carry-all graph attention is sel-independent by design: at binding sel — and for smooth attention generally — the graph route diverges from dense (dpa1 precedent, `KNOWN_GRAPH_DENSE_DIVERGENT`). - Per-atom virial attribution differs elementwise from the dense decomposition for message-passing models (full-to-src vs autograd spread); per-frame totals agree (cross-checked at 1e-8 in fixture generation). - A truly empty MPI rank (zero owned+ghost atoms) fails loudly rather than running; the thrown error cannot propagate through peers blocked in `border_op` collectives, so the job stalls until MPI timeout — phantom-node support is a follow-up. The LAMMPS test bounds this with a hard timeout and asserts it never silently succeeds. - Pre-existing dense bug surfaced (NOT introduced here): `se_atten` at `attn_layer=0` leaks a deterministic `-davg/dstd` padding residual when `set_davg_zero=False` and `exclude_types == []` (`PairExcludeMask` all-ones short-circuit is the only padding mask on that path). The graph path masks padding correctly, so graph and dense deliberately differ in that regime (pinned by test + docstring). A dense-side fix will be proposed separately. ## Follow-ups (separate issues/PRs) - Dense `se_atten` padding-residual fix (above). - Wire the `BUILD_PT_EXPT` C++ gtest suite into CI (pre-existing gap — the universal gtest rows, including the new one, never run in CI today). - `TestDeepPotPTExptWithCommLoadFailure.multi_rank_compute_throws` is vacuous (sets `nswap` but not `nprocs`; pre-existing since deepmodeling#5450). - Empty-rank phantom-node support / clean collective abort; nested with-comm artifact schema assertion in the export test. PR-G-dpa3 (repflows + angle channel + `charge_spin`, reusing this PR's MP machinery) comes next. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Enabled graph-native DPA2 descriptor/repformer execution for eligible energy models, including multi-rank “with-comm” ghost exchange inside graph artifacts. * Added `comm_dict` plumbing and persistent graph-lower enable/disable controls across save/restart. * Improved multi-rank owned/halo reduction using `n_local`, including flat node-axis `aparam` handling for graph ABI and exports. * **Bug Fixes** * Refined multi-rank reduction correctness by excluding ghost contributions from differentiated reductions. * Improved graph attention/softmax stability with phantom-aware behavior for smoother continuity. * **Documentation** * Added pt_expt graph-native route documentation, eligibility rules, and graph-freeze with-comm requirements. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: OutisLi <137472077+OutisLi@users.noreply.github.com>
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
4 participants
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.
Summary
NeighborGraph PR-E: the optional 3-body angle extension of the edge-graph neighbor-list contract (design discussion wanghan-iapcm#4). An angle is a pair of edges sharing a center (
dst(edge_a) == dst(edge_b)), stored asangle_index (2, A)into[0, E)+angle_mask (A,)—edge_vecstays the ONLY geometry leaf, so force/virial assembly is untouched (proven by an invariance test).What's added (all in
deepmd/dpmodel/utils/neighbor_graph/)pad_and_guard_angles(graph.py) — angle-axis padder mirroringpad_and_guard_edges(dynamic guard append / staticangle_capacitywith overflow ValueError).angles.py(new):build_angle_index(edge_index, edge_vec, edge_mask, n_total, a_rcut, *, ordered=False, include_self=False, layout=None)— unordered, no-self pairs of edges sharing a center where BOTH edges are withina_rcut; built on PR-D'scenter_edge_pairs;pair_maskfolded intoangle_mask(never discarded).attach_angles(graph, a_rcut, ...)— post-hoc: edge graph in, graph with angle fields out (dataclasses.replace); default builders keep anglesNone.angle_to_edge_sum/angle_to_node_sum— segment-sum aggregation to the query edge / shared center.graph_angle_cos(angle_index, edge_vec, eps=1e-6)— per-angle cos θ mirroring dpa3 repflowscosine_ijeps placement exactly (+epsin norm denominators,*(1-eps)on the product).angle_padding_fraction(graph)— mask-derived padding-waste report for static capacities.Semantics / decisions
a_sel x a_selsquare including thej==kdiagonal; the graph set keeps one entry per unordered{j,k}pair and moves the degenerate diagonal to the (a_rcut-filtered) edge channel. Dense parity is therefore asserted against the OFF-DIAGONALcosine_ij[j,k](j != k) at rtol/atol 1e-12 (same-math fp64), at non-bindinga_sel. The ordered+self full square stays available via flags.a_sel= normalization-only (carry-all withina_rcut), consistent with the edge-seldecision.sw == 1regime (rtol 1e-12).Tests
source/tests/common/dpmodel/test_angle_builder.py(21) +test_graph_angle_cos_parity.py(6): brute-force triplet oracle (all flag combinations, multi-center, static layout,node_capacitybranch), dpa3 dense-parity + no-self-angle assertion, se_t coordinate oracle, force/virial bit-exact invariance with/without angles, padding-fraction (incl.total==0), torch-namespace smoke tests for every new function. Full neighbor-graph suite: 54 passed.Known limitations
mixed_types=False), used as oracle only.nonzeroincenter_edge_pairs) even when a staticlayoutis passed —angle_capacityfixes the output shape only; shape-static enumeration for export is deferred to PR-G.angle_maskbefore summing (padding angles point at edge 0).A ~ sum(deg^2)capacity overhead mitigated bya_rcut < rcutand reported byangle_padding_fraction.Summary by CodeRabbit
New Features
Tests