feat(dpmodel): NeighborGraph support for dpa1 tebd_input_mode="strip"#5747
feat(dpmodel): NeighborGraph support for dpa1 tebd_input_mode="strip"#5747wanghan-iapcm wants to merge 7 commits into
Conversation
…gate keeps exclude_types/compression dense)
|
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 (1)
📝 WalkthroughWalkthroughThis PR adds ChangesStrip mode graph routing
Estimated code review effort: 3 (Moderate) | ~25 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py (1)
231-256: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd an explicit NaN guard before
assert_allclose.Same concern as the equivalent strip-parity test in
test_dpa1_call_graph_descriptor.py:np.testing.assert_allclosetreatsNaN == NaNas equal by default (equal_nan=True), so sharedNaNs betweenroutedanddensewould pass unnoticed here.🧪 Proposed guard
routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) assert len(routed) == len(dense) for r, d in zip(routed, dense, strict=True): if r is None: assert d is None continue + assert not np.any(np.isnan(r)) np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12)🤖 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_dpa1_graph_attention_parity.py` around lines 231 - 256, The parity check in test_se_atten_v2_graph_equals_dense needs an explicit NaN guard before np.testing.assert_allclose, because matching NaNs could otherwise pass unnoticed. Update the loop over routed and dense to assert that neither value contains NaNs before comparing them, using the existing test_dpa1_graph_attention_parity test structure and the DescrptSeAttenV2 call/_call_dense outputs as the comparison points.source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py (1)
255-286: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd an explicit NaN guard before
assert_allclose.
np.testing.assert_allclosedefaults toequal_nan=True, so ifroutedanddenseboth produceNaNat the same position (e.g. a shared division-by-zero in the strip/attention path), this loop would pass silently instead of catching the regression. The sibling block-level test (test_dpa1_call_graph_block.py) guards against exactly this withassert not np.any(np.isnan(graph_g))before itsassert_allclosecall — worth mirroring here for consistency.🧪 Proposed guard
routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) assert len(routed) == len(dense) for r, d in zip(routed, dense, strict=True): if r is None: assert d is None continue + assert not np.any(np.isnan(r)) np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12)🤖 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_dpa1_call_graph_descriptor.py` around lines 255 - 286, The `test_call_strip_graph_equals_dense` comparison in `test_dpa1_call_graph_descriptor.py` should explicitly reject NaNs before calling `np.testing.assert_allclose`, since matching NaNs can otherwise hide regressions in the routed graph path. Add a NaN check on the per-item outputs right before the `assert_allclose` inside the `routed`/`dense` loop, mirroring the existing guard used in `test_dpa1_call_graph_block.py`, and keep the rest of the equality comparison unchanged.
🤖 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/common/dpmodel/test_dpa1_call_graph_descriptor.py`:
- Around line 255-286: The `test_call_strip_graph_equals_dense` comparison in
`test_dpa1_call_graph_descriptor.py` should explicitly reject NaNs before
calling `np.testing.assert_allclose`, since matching NaNs can otherwise hide
regressions in the routed graph path. Add a NaN check on the per-item outputs
right before the `assert_allclose` inside the `routed`/`dense` loop, mirroring
the existing guard used in `test_dpa1_call_graph_block.py`, and keep the rest of
the equality comparison unchanged.
In `@source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py`:
- Around line 231-256: The parity check in test_se_atten_v2_graph_equals_dense
needs an explicit NaN guard before np.testing.assert_allclose, because matching
NaNs could otherwise pass unnoticed. Update the loop over routed and dense to
assert that neither value contains NaNs before comparing them, using the
existing test_dpa1_graph_attention_parity test structure and the
DescrptSeAttenV2 call/_call_dense outputs as the comparison points.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 09bc9b5d-ea2d-4462-ac04-f52ab6788c4d
📒 Files selected for processing (5)
deepmd/dpmodel/descriptor/dpa1.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_call_graph_descriptor.pysource/tests/common/dpmodel/test_dpa1_graph_attention_parity.pysource/tests/pt_expt/descriptor/test_dpa1.py
| if self.tebd_input_mode not in ["concat", "strip"]: | ||
| raise NotImplementedError( | ||
| "graph path supports tebd_input_mode='concat' only (NeighborGraph PR-A)" | ||
| f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}" |
There was a problem hiding this comment.
if tebd_input_mode not in ["concat", "strip"] it will not pass argcheck?
There was a problem hiding this comment.
Good catch — confirmed no, it will pass argcheck. tebd_input_mode is declared as a plain str Argument (no choices=) in argcheck.py, so any string value is accepted at input-validation time; the docstrings there just prose-describe 'concat'/'strip' as the supported modes without enforcing it. So this NotImplementedError guard is genuinely reachable, not dead code — it's needed because an invalid tebd_input_mode can reach call_graph at runtime. For reference, the dense call() path (dpa1.py around line 1575, if ... in ["concat"]: ... elif ... in ["strip"]: ... with no else) has no equivalent guard and would fail with an unclear UnboundLocalError on gg for the same bad input — that's a pre-existing gap outside this PR's scope, so I've left it as-is here rather than expanding scope. Tightening argcheck.py with an explicit choices=["concat", "strip"] would be the real fix for this whole class of issue; happy to file that separately if useful.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5747 +/- ##
==========================================
- Coverage 80.87% 79.07% -1.80%
==========================================
Files 1003 1014 +11
Lines 112492 115382 +2890
Branches 4236 4274 +38
==========================================
+ Hits 90981 91242 +261
- Misses 19986 22599 +2613
- Partials 1525 1541 +16 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…close np.testing.assert_allclose treats NaN == NaN as equal by default, so a shared NaN between the graph and dense outputs would silently pass the parity check. Applies to the two strip/attention parity tests CodeRabbit flagged: test_call_strip_graph_equals_dense (test_dpa1_call_graph_descriptor.py) and test_se_atten_v2_graph_equals_dense (test_dpa1_graph_attention_parity.py).
…spatch test CI failure: test_default_fallback[se_atten_v2] failed (energy off by 1.2e-07, rtol=1e-10). Root cause is this PR's own change: uses_graph_lower() now admits tebd_input_mode='strip', which makes se_atten_v2 (always strip, always smooth_type_embedding=True) graph-eligible. call_common's dispatch (make_model.py) forces the dense route whenever an explicit neighbor_list is passed, but pt_expt's default-flip (decision deepmodeling#17) routes neighbor_list=None to the carry-all graph for graph-eligible mixed_types descriptors -- so 'None' and 'explicit DefaultNeighborList()' stop being the same computation for this model. model_dpa1 already works around the same issue by pinning smooth_type_embedding=False (with a comment explaining why); se_atten_v2 can't do that since DescrptSeAttenV2 hardcodes smooth_type_embedding=True, so the divergence is unavoidable and documented (NeighborGraph PR-D: dense keeps sel-padding phantom softmax terms the graph route doesn't). Exclude it from this specific parametrize with a docstring explaining why; it is still covered by test_pt_expt_equivalence/test_dpmodel_equivalence, which always pass an explicit neighbor_list on both sides and so never hit this route split.
Follow-up to 66089f1. That commit fixed the CI failure by excluding se_atten_v2 from test_default_fallback's parametrize, but a plain exclusion silently drops coverage rather than documenting the expected behavior. Replace it with a positive assertion, mirroring the precedent already set by test_block_compact_graph_smooth_clean_divergence (test_dpa1_graph_attention_parity.py): for models where tebd_input_mode in {concat, strip} + attn_layer > 0 + smooth_type_embedding=True, dense and the carry-all graph default (neighbor_list=None) intentionally diverge (NeighborGraph PR-D: dense keeps sel-padding phantom terms in the attention softmax denominator, the graph route does not). Verified empirically this is not strip-specific: a plain concat se_atten with smooth_type_embedding=True shows the same ~1e-7 order-of-magnitude divergence at attn_layer>0, and exactly zero divergence at attn_layer=0 (no softmax involved there at all) -- confirming the gap lives entirely in the shared attention code, not in anything this PR added. Add model_dpa1_smooth (concat's counterpart to model_se_atten_v2) so both tebd_input_mode values are covered by the same mechanism, and a KNOWN_GRAPH_DENSE_DIVERGENT set used only by test_default_fallback to widen its tolerance for exactly these two models (atol=3e-5, rtol=1e-3 -- measured empirically: energy ~1e-7, force ~1e-6, virial up to ~1.3e-5, since virial is a derivative and amplifies the softmax-denominator perturbation more than the value itself). All other models keep the original tight tolerance.
2 participants
Summary
Adds NeighborGraph (graph-native lower) support for the dpa1 descriptor with
tebd_input_mode="strip", closing the last descriptor-level gap that forced strip-mode models (andse_atten_v2, which is strip-by-construction) onto the legacy dense path.The dense strip branch factorizes the per-neighbor feature as
gg = gg_s*gg_t + gg_s— a radial-only geometric embedding times a type-pair strip embedding (optionally switch-smoothed). Because this has no neighbor-axis coupling, it maps to the graph path edge-for-edge. The change is:DescrptBlockSeAtten): a new per-edge helper_graph_edge_gg_strip(op-for-op mirror of the dense strip branch, including thecenter*ntypes_pad + neinei-fastest two-side table layout andint64gather indices), selected by aconcat/stripbranch incall_graph.DescrptDPA1.uses_graph_lower): admitsstrip, while keeping compressed descriptors andexclude_typeson the dense path (they have no graph kernel here).se_atten_v2inherits this and becomes graph-eligible for free.make_fxexport tests are parametrized overtebd_input_modeto prove the strip kernel is fx-traceable.No new op, no attention/
segment_sumchange, no C++/serialization change.Scope
This PR is deliberately independent of #5733 (graph
exclude_types). It does not changeexclude_typeseligibility — theuses_graph_lowerand not exclude_typesgate and thecall_graphexclude_typesraise are both kept. When both land, whichever merges second resolves a small (2–3 line) mechanical conflict atuses_graph_lower/ thecall_graphguard.Test plan
rtol=atol=1e-12overtype_one_side × smooth(attn=0) andtype_one_side(attn=2, non-smooth).callvs_call_denseparity overtype_one_side × smooth × attn_layer(incl. attn=2 + smooth=True, bit-exact via thestatic_nneiadapter), plus a negative-contract gate test (compressed → dense, strip+exclude_types→ dense).se_atten_v2eligibility + graph-vs-dense parity (replaces the obsolete "strip stays dense" test).make_fxexport; cross-backend consistency strip cases now route pt_expt through the graph adapter.Validated on CPU and on GPU (Tesla T4, cuda:0): pt_expt dpa1 50 passed, consistency strip 22 passed +
se_atten_v2110 passed, dpmodel strip suites 46 passed. No tolerances relaxed, no tests skipped.Known limitations
self.compress.exclude_typesstays dense (out of scope — owned by refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18) #5733).segment_sum→index_addis atomic/non-deterministic on CUDA (1–2 fp64 ULP), inherent to atomic scatter; GPU parity validated within tolerance.RuntimeWarningon the shared attention path (max over fully-masked segments), also present on the concat path.Summary by CodeRabbit
Summary
New Features
Bug Fixes
Tests