deepmodeling / deepmodeling/deepmd-kit
Auto-select the optimal NeighborGraph builder (nv/vesin/dense) by device and availability
- 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
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