fix(pt_expt): preserve graph routing for raw checkpoints by OutisLi · Pull Request #5885 · deepmodeling/deepmd-kit

OutisLi · 2026-07-21T04:21:01Z

Summary

select the eager lower ABI for raw PT-experimental checkpoints with the restored model's model_uses_graph_lower capability
route graph-eligible energy models through the existing graph DeepEval contract while preserving dense and spin checkpoint behavior
automatically choose the scalable graph builder (nv on CUDA, then Vesin, then dense fallback) so dp test does not stall in the single-core NumPy O(N²) builder
add a deterministic nonzero-statistics DPA1 regression covering both plain and compiled checkpoint layouts

Test plan

pytest source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py source/tests/pt_expt/infer/test_graph_deepeval.py -v
CUDA auto route resolves to nv; CPU auto route resolves to Vesin; both pass the raw checkpoint parity regression
pytest source/tests/pt_expt/model/test_dpa1_graph_lower.py source/tests/pt_expt/model/test_graph_builder_dispatch.py -v
repository pre-commit hooks, including Ruff, formatting, Velin, and pylint
representative OMat24 benchmark (8 frames × 128 atoms): dense median 0.1769 s, auto/NV median 0.0092 s
cross-backend validation with the supplied OMat24 weights: graph checkpoint outputs match public forward within float32 numerical tolerance

Summary by CodeRabbit

New Features
- Added graph-routed (“graph-lower”) eager .pt checkpoint loading and evaluation support.
- neighbor_graph_method now defaults to auto and applies graph-lower routing automatically.
- Graph-routed models expose lower-input kind and graph edge dtype metadata.
Bug Fixes
- Enforced correct neighbor-graph behavior when inappropriate neighbor-list settings are provided (with clearer warnings).
- Ensured graph-routed inference outputs match the model’s public forward(..., do_atomic_virial=True).
Tests
- Expanded coverage for plain and compiled graph-routed DPA1 .pt checkpoints.

Select the eager lower ABI from the restored model capability instead of forcing every raw checkpoint through the padded dense neighbor-list path. Graph-eligible energy models now reuse the graph DeepEval contract, keeping energy, force, virial, and atomic outputs aligned with public forward semantics when descriptor statistics are nonzero. Retain the existing dense and spin paths, expose graph builder selection for raw graph checkpoints, and cover both plain and compiled checkpoint layouts with a deterministic DPA1 regression.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

coderabbitai · 2026-07-21T04:29:05Z

📝 Walkthrough

Walkthrough

PT checkpoint inference now detects graph-lower models, routes them through graph-native execution, records graph metadata, and validates output parity for plain and compiled DPA1 checkpoints.

Changes

Graph-lower PT inference

Layer / File(s)	Summary
Routing contracts and backend selection `deepmd/pt_expt/infer/deep_eval.py`	Graph-only neighbor graph settings are validated, while graph-routed artifacts reject explicit neighbor lists and ignore non-default neighbor-list backends with a warning.
Graph-lower PT runner `deepmd/pt_expt/infer/deep_eval.py`	PT loading detects graph-lower models, records `lower_input_kind` and `graph_edge_dtype`, and runs `forward_common_lower_graph` with translated public outputs.
Graph DPA1 checkpoint regression `source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py`	Plain and compiled graph DPA1 checkpoints are tested against the source model’s public forward outputs.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant DeepPot
  participant DeepEval
  participant EnergyModel
  DeepPot->>DeepEval: load PT checkpoint
  DeepEval->>EnergyModel: detect graph-lower capability
  DeepEval->>EnergyModel: run graph-native forward
  EnergyModel-->>DeepEval: return graph outputs
  DeepEval-->>DeepPot: return translated public outputs

Possibly related issues

[pt_expt] NeighborGraph compiled training is energy-only; generalize the trace to non-energy fittings #5805: Related through graph-lower routing and output translation in the PT experimental code, although it targets a different execution path.

Possibly related PRs

deepmodeling/deepmd-kit#5714: Adds graph-builder routing that overlaps with this PR’s neighbor_graph_method handling.
deepmodeling/deepmd-kit#5715: Defines graph-native DPA1 eligibility used by this PR’s PT loading path.
deepmodeling/deepmd-kit#5747: Changes DPA1 graph-lower eligibility directly used by model_uses_graph_lower.

Suggested labels: bug

Suggested reviewers: njzjz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes route graph-eligible PT checkpoints through the graph-native path and add regression coverage matching `#5862`.
Out of Scope Changes check	✅ Passed	The refactors and doc updates are directly tied to the checkpoint routing fix and its regression tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: preserving graph routing for raw PT-experimental checkpoints.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Resolve the graph neighbor backend once when DeepEval loads the model. Prefer the batched nvalchemiops builder on CUDA, fall back to Vesin when available, and retain the dense all-pairs implementation as the dependency-free final fallback. This removes the single-core NumPy O(N²) graph construction bottleneck from the default dp test path while preserving every explicit builder selection. Cover the resolved backend in the raw DPA1 checkpoint regression.

codecov · 2026-07-21T07:23:34Z

Codecov Report

❌ Patch coverage is 87.50000% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.57%. Comparing base (83e7245) to head (7f73a80).

Files with missing lines	Patch %	Lines
deepmd/pt_expt/infer/deep_eval.py	87.50%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5885      +/-   ##
==========================================
- Coverage   78.87%   78.57%   -0.30%     
==========================================
  Files        1054     1054              
  Lines      121770   121806      +36     
  Branches     4413     4412       -1     
==========================================
- Hits        96046    95715     -331     
- Misses      24159    24510     +351     
- Partials     1565     1581      +16

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

wanghan-iapcm

The routing fix is correct by construction (it keys on the same model_uses_graph_lower flag that drives the model's own forward(), so .pt eval goes graph iff forward() does), and the DPA1 nonzero-stats regression is a genuine test for #5862. One test-coverage note inline.

wanghan-iapcm · 2026-07-22T13:51:32Z

                f"Unknown nlist_backend '{nlist_backend}'; "
                "expected 'auto', 'vesin', or 'native'."
            )
+        if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"):


This new graph-routed branch adds two reachable, user-facing guardrails that no test exercises: (1) the raise ValueError when an explicit neighbor_list is passed for a graph artifact, and (2) the warnings.warn when nlist_backend != "auto" for a graph artifact. Both are cheap to cover — construct DeepEval/DeepPot on the new graph DPA1 checkpoint with neighbor_list=<something> (assertRaises) and with nlist_backend="vesin" (assertWarns). Worth adding so every reachable branch here is tested.

Copilot AI review requested due to automatic review settings July 21, 2026 04:21

Copilot AI reviewed Jul 21, 2026

dosubot Bot added the enhancement label Jul 21, 2026

github-actions Bot added the Python label Jul 21, 2026

OutisLi requested review from njzjz and wanghan-iapcm July 21, 2026 04:39

wanghan-iapcm reviewed Jul 22, 2026

View reviewed changes

coderabbitai Bot mentioned this pull request Jul 24, 2026

feat(pt_expt): auto-select O(N) NeighborGraph builder by default #5903

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(pt_expt): preserve graph routing for raw checkpoints#5885

fix(pt_expt): preserve graph routing for raw checkpoints#5885
OutisLi wants to merge 2 commits into
deepmodeling:masterfrom
OutisLi:pr/exptest

OutisLi commented Jul 21, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

coderabbitai Bot commented Jul 21, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

codecov Bot commented Jul 21, 2026

Uh oh!

wanghan-iapcm left a comment

Uh oh!

wanghan-iapcm Jul 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

OutisLi commented Jul 21, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Summary by CodeRabbit

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

codecov Bot commented Jul 21, 2026

Codecov Report

Uh oh!

wanghan-iapcm left a comment

Choose a reason for hiding this comment

Uh oh!

wanghan-iapcm Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

OutisLi commented Jul 21, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 21, 2026 •

edited

Loading