feat(pt_expt): auto-select O(N) NeighborGraph builder by default#5903
Open
Shaurya2k06 wants to merge 1 commit into
Open
feat(pt_expt): auto-select O(N) NeighborGraph builder by default#5903Shaurya2k06 wants to merge 1 commit into
Shaurya2k06 wants to merge 1 commit into
Conversation
Replace the hard-coded None→dense default with a shared availability- probing ladder (CUDA: nv→vesin→dense; CPU: vesin→dense) used by the model default-flip, DeepEval, and compiled training's eager builder. Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>
Contributor
📝 WalkthroughWalkthrough
ChangesNeighborGraph auto selection
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant DeepEval
participant CM
participant resolve_auto_graph_builder
participant NeighborGraphBuilder
DeepEval->>resolve_auto_graph_builder: resolve "auto" for device
CM->>resolve_auto_graph_builder: resolve None or "auto" for device
resolve_auto_graph_builder-->>DeepEval: nv, vesin, or dense
resolve_auto_graph_builder-->>CM: nv, vesin, or dense
DeepEval->>NeighborGraphBuilder: build graph with resolved method
CM->>NeighborGraphBuilder: build graph with resolved method
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt_expt/model/test_graph_builder_dispatch.py (1)
146-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert dispatch, not only parity.
All backends are intentionally value-equivalent, so this passes if
None/"auto"incorrectly falls back todense. Mock or spy on the resolver/concrete builder and assert that the resolved backend is invoked.🤖 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/pt_expt/model/test_graph_builder_dispatch.py` around lines 146 - 158, Update test_none_and_auto_match_resolved_builder to spy on or mock the resolver/concrete graph builder, then assert that the backend resolved for None and "auto" is actually invoked. Retain the existing output-parity assertions, but ensure the test fails if either input silently falls back to dense.
🤖 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_expt/utils/vesin_graph_builder.py`:
- Around line 11-15: Update the documentation describing the shared
resolve_auto_graph_builder default ladder to state that vesin is selected for
CPU or CUDA fallback only when vesin.torch is importable; otherwise document
that dense is selected.
---
Nitpick comments:
In `@source/tests/pt_expt/model/test_graph_builder_dispatch.py`:
- Around line 146-158: Update test_none_and_auto_match_resolved_builder to spy
on or mock the resolver/concrete graph builder, then assert that the backend
resolved for None and "auto" is actually invoked. Retain the existing
output-parity assertions, but ensure the test fails if either input silently
falls back to dense.
🪄 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 Plus
Run ID: d02887c0-8595-4f4e-a31f-c9bc7c3749a5
📒 Files selected for processing (9)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/neighbor_graph_method.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/model/test_graph_builder_dispatch.pysource/tests/pt_expt/utils/test_neighbor_graph_method.py
Comment on lines
+11
to
+15
| for ``nf == 1`` inference and CPU use. On CPU (and on CUDA when ``nv`` is | ||
| unavailable) it is selected by the shared | ||
| :func:`~deepmd.pt_expt.utils.neighbor_graph_method.resolve_auto_graph_builder` | ||
| default ladder; otherwise prefer ``nv`` (:mod:`.nv_graph_builder`) for batched | ||
| multi-frame GPU work, which batches all frames in one kernel. Export and |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
State the vesin availability condition.
The ladder selects vesin on CPU or CUDA fallback only when vesin.torch is importable; otherwise it selects dense. Add that condition so the documented policy matches resolve_auto_graph_builder.
🤖 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/pt_expt/utils/vesin_graph_builder.py` around lines 11 - 15, Update the
documentation describing the shared resolve_auto_graph_builder default ladder to
state that vesin is selected for CPU or CUDA fallback only when vesin.torch is
importable; otherwise document that dense is selected.
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.
Closes #5902
Summary
resolve_auto_graph_builder(device)policy for the carry-all NeighborGraph builders: CUDA prefersnvthenvesinthendense; CPU prefersvesinthendense(asestays explicit-only).None/"auto") from hard-coded"dense"to that ladder, and use the same helper in DeepEval (default"auto") and compiled training's eager_forward_graph..pt2artifacts are unchanged.Why existing tests missed this
Builder dispatch already had vesin/nv vs dense energy/force parity, but nothing asserted that the model/DeepEval default would pick an O(N) builder when available, and compiled training still hardcoded dense while claiming to match the eager default-flip. The new resolver unit tests pin the availability ladder; the extended dispatch and DeepEval tests pin value-transparency of
None/"auto"against the resolved concrete builder.Validation
ruff check/ruff format --checkon touched filespytestfocused suite: 20 passed, 1 skipped (nvCUDA-only)source/tests/pt_expt/utils/test_neighbor_graph_method.pysource/tests/pt_expt/model/test_graph_builder_dispatch.py.pt2parityneighbor_graph_methodSummary by CodeRabbit
New Features
Bug Fixes
Tests