fix(pt_expt): preserve graph routing for raw checkpoints - #5885
Conversation
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.
📝 WalkthroughWalkthroughPT 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. ChangesGraph-lower PT inference
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
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 |
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 Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
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.
| f"Unknown nlist_backend '{nlist_backend}'; " | ||
| "expected 'auto', 'vesin', or 'native'." | ||
| ) | ||
| if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"): |
There was a problem hiding this comment.
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.
njzjz
left a comment
njzjz
left a comment
There was a problem hiding this comment.
Nice piece of work — routing raw .pt checkpoints through the graph lower and auto-selecting the O(N) builder is a real usability win, and the parity test against model.forward(..., do_atomic_virial=True) at 1e-10 is the right assertion to anchor it on. A few things I would want resolved before this lands.
1. neighbor_list= now raises for graph-eligible .pt checkpoints (backward-incompatible)
if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"):
if self.neighbor_list is not None:
raise ValueError("neighbor_list only applies to nlist-routed artifacts; ...")Before this PR _load_pt always set lower_input_kind = "nlist", so neighbor_list= was accepted for every .pt. After it, any checkpoint where model_uses_graph_lower(model) is true takes the new branch and the same call raises. That reaches user code that never opted into anything: DP in deepmd/calculator.py forwards neighbor_list straight through (calculator.py:100), so an existing ASE script pairing a DPA1 .pt with a custom NewPrimitiveNeighborList goes from working to ValueError on upgrade.
Since the PR's own premise is that all builders emit the same neighbor set, the ASE list is not wrong here — it is just redundant. I would either honor it (route to neighbor_graph_method="ase", which already exists) or downgrade to a warning, and in either case call the change out in the PR description. Raising is the one option that breaks callers silently at import-of-model time.
2. The fail-fast guard stops catching an explicit "dense"
if neighbor_graph_method not in ("auto", "dense") and ... != "graph":
raise ValueError(...)The guard exists because "the knob would silently do nothing there". With the default moved to "auto", "dense" is no longer the default but is still exempt, so DeepEval(nlist_model, neighbor_graph_method="dense") now silently does nothing — exactly the case the check was written to prevent. Comparing against a sentinel rather than a value set fixes this cleanly: default the parameter to None, treat None as "auto" after routing is known, and raise for any non-None value on a non-graph artifact.
3. The two knobs disagree on policy in mirrored situations
neighbor_graph_method on an nlist artifact → ValueError. nlist_backend on a graph artifact → UserWarning. Same mistake, same consequence (the knob does nothing), opposite outcome. Worth picking one. If the warning stays, stacklevel=2 points at DeepEval.__init__ rather than the user's DeepPot(...) call — the warning as emitted names deep_eval.py, so it is not actionable. It needs to account for __init__ and the DeepEval wrapper frame.
4. auto resolution is untested for the branches CI does not hit
cls.expected_graph_method = "nv" if (cuda and is_nv_available()) else "vesin" if ... else "dense"The test computes the expected method with the same predicate the implementation uses, then asserts they match — so it pins "the resolver agrees with itself", not "the resolver picks a builder that is correct". The parity assertion against model.forward does check correctness, but only for whichever builder happens to be available on that runner; on a CPU-only CI box nv is never exercised at all, and the "performance-only" claim in the docstring is never verified.
Since that claim is what makes the default change safe, I would add a loop over the builders that are available (skipping the rest) asserting they produce identical outputs, rather than only checking the one auto chose. That also gives ase — which is still reachable explicitly and currently has no coverage on this path — a home.
5. Smaller points
_resolve_neighbor_graph_methodimportsdeepmd.pt.utils.nv_nlistfrom insidept_expt. Presumably deliberate reuse, but apt_expt→ptdependency is worth a one-line comment saying so, otherwise it reads like a typo forpt_expt.utils.- The device is sampled once at construction (
DEVICE.type == "cuda"), which is fine, but the resolved value is stored in the same attribute that held the user's request. Keeping the request and the resolution in separate attributes would make the_neighbor_graph_methodassertion in the test unambiguous about which one it is checking. - The parity fixture is 3 atoms in a 10 Å cell with
rcut=4.0, so every atom sees every other and no periodic image matters. An atom with an empty neighbor row, and a cell small enough to need images, would exercise the parts of the graph builders most likely to differ between backends. docs/readthedocs.org:deepmdis red on this PR. That is the shared 40-minute RTD timeout affecting many open PRs, not something in this diff — #5888 is the fix.
Summary
model_uses_graph_lowercapabilitynvon CUDA, then Vesin, then dense fallback) sodp testdoes not stall in the single-core NumPy O(N²) builderCloses #5862
Test plan
pytest source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py source/tests/pt_expt/infer/test_graph_deepeval.py -vnv; CPU auto route resolves to Vesin; both pass the raw checkpoint parity regressionpytest source/tests/pt_expt/model/test_dpa1_graph_lower.py source/tests/pt_expt/model/test_graph_builder_dispatch.py -vSummary by CodeRabbit
New Features
.ptcheckpoint loading and evaluation support.neighbor_graph_methodnow defaults toautoand applies graph-lower routing automatically.Bug Fixes
forward(..., do_atomic_virial=True).Tests
.ptcheckpoints.