deepmodeling / deepmodeling/deepmd-kit

Auto-select the optimal NeighborGraph builder (nv/vesin/dense) by device and availability

Open
#5,902 2 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

The pt_expt NeighborGraph route has four builders — `dense` (carry-all O(N²) all-pairs, always available, backend-agnostic), `ase`, `vesin` (O(N) cell list, torch bindings, works on CPU and CUDA, requires the `vesin-torch` package), and `nv` (nvalchemiops, CUDA-only, batches natively over frames) — but the model-level default (`_resolve_graph_method`, `deepmd/pt_expt/model/make_model.py`) hard-resolves `None -> "dense"` for training/eager, so the O(N) builders are opt-in only. `DeepEval` already has an availability-probing `"auto"` mode for inference; the same policy should drive the model-level default.

Proposed resolution ladder for `None`/`"auto"`:

- CUDA device: `nv` if importable, else `vesin` if importable, else `dense`.
- CPU device: `vesin` if importable, else `dense`.

Implementation sketch:
- One shared `resolve_auto_graph_builder(device) -> str` helper used by BOTH the model-level default-flip and `DeepEval`'s auto mode (single owner of the policy).
- Builder choice must be value-transparent: add a parity test building the same system through every available builder and asserting identical energies/forces (the builders differ only in edge enumeration order; the segment reductions are order-independent up to fp addition order — pick the tolerance per the established device-conditional convention).
- The builder runs EAGERLY outside traced/compiled regions (export uses synthetic graph inputs; the training compile traces the lower, not the builder), so the flip does not affect `.pt2` artifacts — assert that in the test notes.
- Perf note: the graph-attention benchmark memo found the dense builder was not the bottleneck at benchmark sizes; the win is for large systems, so benchmark before/after on a large system to document the crossover.

Contributor guide

Open the contributing guide

Research direction

Start in deepmd/pt_expt/model/make_model.py and compare its _resolve_graph_method policy with DeepEval's existing availability-probing auto mode. Trace the dense, vesin, and nv builder entry points and identify the existing device-conditional test conventions. Done means one shared policy serves both paths, available builders produce matching energies and forces, compiled/exported behavior is covered in test notes, and large-system before/after benchmarks document the crossover.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance, testing
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.