fix(pt_expt): preserve graph routing for raw checkpoints - #5912
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.
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.
📝 WalkthroughWalkthrough
ChangesGraph-routed PT inference
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DeepEval
participant PTCheckpoint
participant NeighborBackend
participant PTModel
participant PublicEvaluation
DeepEval->>PTCheckpoint: detect checkpoint dialect
DeepEval->>PTModel: inspect graph-lower capability
DeepEval->>NeighborBackend: resolve and configure neighbor construction
DeepEval->>PTModel: call forward_common_lower_graph
PTModel-->>DeepEval: return graph-lower outputs
DeepEval->>PublicEvaluation: translate and expose evaluation results
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}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py (1)
342-423: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest duplicates the production resolution order.
setUpClassre-implements theauto→nv/vesin/denseladder, so a regression inDeepEval._resolve_neighbor_graph_methodthat changes the order would be mirrored by the test and stay green. Consider asserting againstDeepEval._resolve_neighbor_graph_method("auto")and covering the explicit-method path (e.g.neighbor_graph_method="dense") for the deterministic parity check.🤖 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/infer/test_deep_eval_pt_checkpoint.py` around lines 342 - 423, Update TestPtExptLoadPtGraphDPA1.setUpClass to derive expected_graph_method by calling DeepEval._resolve_neighbor_graph_method("auto") instead of duplicating the nv/vesin/dense selection logic. Add coverage for an explicit neighbor_graph_method="dense" configuration and assert its deterministic method selection and evaluation parity, reusing the existing checkpoint/test flow.
🤖 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/infer/deep_eval.py`:
- Around line 728-777: The use_graph_lower branch in _load_pt incorrectly
selects the dense _eager_runner_graph for native-spin checkpoints. Exclude
native-spin models from this branch or provide a native-spin graph runner
matching NativeSpinEnergyModel.forward_lower_graph_exportable’s 14-argument
signature, ensuring exported_module receives spin separately rather than
treating it as fparam.
---
Nitpick comments:
In `@source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py`:
- Around line 342-423: Update TestPtExptLoadPtGraphDPA1.setUpClass to derive
expected_graph_method by calling DeepEval._resolve_neighbor_graph_method("auto")
instead of duplicating the nv/vesin/dense selection logic. Add coverage for an
explicit neighbor_graph_method="dense" configuration and assert its
deterministic method selection and evaluation parity, reusing the existing
checkpoint/test flow.
🪄 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: fffcc8fd-f028-48f2-9e8b-e2448c4e79d2
📒 Files selected for processing (2)
deepmd/pt_expt/infer/deep_eval.pysource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5912 +/- ##
==========================================
- Coverage 79.06% 79.00% -0.07%
==========================================
Files 1067 1070 +3
Lines 123805 124123 +318
Branches 4510 4527 +17
==========================================
+ Hits 97891 98064 +173
- Misses 24304 24440 +136
- Partials 1610 1619 +9 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for the inline raw-checkpoint routing issue.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/backend/pt_expt.py (1)
32-71: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd focused tests for every detector outcome.
This helper now controls both backend auto-dispatch and
.ptloading. Cover wrapped checkpoints,.w,.matrix, bias-only, mixed, and non-string-key cases to prevent silent misrouting.🤖 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/backend/pt_expt.py` around lines 32 - 71, Add focused tests for detect_pt_checkpoint_backend covering every outcome: wrapped checkpoints under “model”, unambiguous .w and .matrix weights, bias-only .b and .bias states, mixed/ambiguous keys returning None, and mappings containing non-string keys. Assert the expected backend or None for each case, including both wrapped and unwrapped payload behavior.
🤖 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 `@source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py`:
- Around line 754-761: Update the test around DeepPot initialization to assert
that the resolved neighbor-graph method is “dense,” not only that metadata
reports graph input. Use the DeepPot-resolved method symbol exposed by the
existing API, preserving the explicit neighbor_graph_method="dense" setup and
current assertions.
---
Nitpick comments:
In `@deepmd/backend/pt_expt.py`:
- Around line 32-71: Add focused tests for detect_pt_checkpoint_backend covering
every outcome: wrapped checkpoints under “model”, unambiguous .w and .matrix
weights, bias-only .b and .bias states, mixed/ambiguous keys returning None, and
mappings containing non-string keys. Assert the expected backend or None for
each case, including both wrapped and unwrapped payload behavior.
🪄 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: 8280e232-c6d8-4b66-9991-cfefe05b349b
📒 Files selected for processing (3)
deepmd/backend/pt_expt.pydeepmd/pt_expt/infer/deep_eval.pysource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt_expt/infer/deep_eval.py
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The latest commit distinguishes raw checkpoint backends by parameter dialect, makes the pt_expt native-spin DPA4 runner reachable, and adds an end-to-end dispatch and inference regression test. The earlier blocker is resolved.
Focused validation: backend-dispatch tests and native-spin DPA4 checkpoint inference passed.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The routing fix is the right diagnosis and the weight-name dispatch is a genuine improvement. Raw .pt was hardcoded to lower_input_kind="nlist" since #5562, so a model trained on the graph lower was being evaluated through the dense path - and the old has_pt_expt and not has_pt rule mis-handed pt_expt DPA4 checkpoints to the pt backend, because they legitimately carry torch-native .bias alongside dpmodel .w. Making weight names decisive fixes that; I confirmed pt's SeZM/DPA4 routes through MLPLayer and so carries .matrix, and that freshly built state dicts for pt and pt_expt across dpa2 and DPA4-native-spin all classify correctly.
Two blocking comments, on the dispatch inconsistency and on the untested new guards. Eight further comments inline.
Worth stating plainly, because it frames several of the others: this PR changes the default evaluation path for every graph-eligible artifact, not only the raw .pt files it set out to fix. neighbor_graph_method flips from "dense" to a device-aware "auto", so the builder now depends on the host and on which optional packages happen to be installed. The builders are asserted to be edge-set equivalent, but the graph lower accumulates through segment_sum over edges, so ordering differences alone move fp64 results - and nothing in the PR asserts that dense, vesin and nv agree. The one test that touches this is tautological (see the comment on _resolve_neighbor_graph_method).
A couple of things I checked and want to record as clear, so they are not re-raised:
- Virtual-atom spin models do not misroute into the native-spin graph runner.
SpinModel.__init__callsdescriptor.disable_graph_lower()and_graph_lower_disabledis a persisted buffer restored byload_state_dictbefore the check, somodel_uses_graph_lowercorrectly returns False for them. - The earlier bot concern that the native-spin fix was dead code - because
_load_ptrejected DPA4/SeZM before that branch - was true at7c48a6cbut is resolved at head by the final commit, which does exactly what that comment asked. The branch is reachable and has a regression test. - CI did run on this PR: 11 passing, 2 skipping, and the single failure is a Read the Docs build, not a test. The "Checks: not run per request" note refers to local runs.
Smaller notes:
-
When no
.wand no.matrixkey exists, the bias-only fallback returns"pt"for anything carrying.bias, and_load_ptthen tells the user to load it withdp --pt. That is safe today only because every pt_expt layer is aNativeLayerusing.w/.b. A pt_expt model that became fully torch-native would be refused by its own backend, pointed at the wrong one. ReturningNonefor "unknown" would fail soft instead. -
.ptfiles are fullytorch.loaded twice during dispatch - once inmatch_filenameand again in_load_pt. Pre-existing, but this PR adds a seconddetect_pt_checkpoint_backendpass on top, and on multi-GB checkpoints backend selection alone pays a full deserialization. -
Codecov reports 81.43% patch coverage with 13 newly-added lines uncovered, 7 of them in
deepmd/backend/pt_expt.py- consistent with the coverage gaps in the two blocking comments.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
One readability note on the eager-runner dispatch.
|
Follow-up to the review summary, implemented in
The PR description now also records the graph/nlist compatibility contract and directs ASE graph users to |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt_expt/infer/deep_eval.py (1)
246-259: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a single
_is_graph_routedhelper.The
lower_input_kind in ("graph", "dpa1_canonical")classification is now encoded here and again in_setup_neighbor_backend(Line 318), and_eval_model/_eval_model_spinuse slightly different variants (_eval_model_spinchecks only"graph"). A single accessor would keep these in sync.🤖 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/infer/deep_eval.py` around lines 246 - 259, Add a shared _is_graph_routed helper and use it for all lower_input_kind classification. Replace the inline check in the current validation and the duplicate logic in _setup_neighbor_backend, _eval_model, and _eval_model_spin, preserving graph and dpa1_canonical as graph-routed kinds consistently.
🤖 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.
Nitpick comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 246-259: Add a shared _is_graph_routed helper and use it for all
lower_input_kind classification. Replace the inline check in the current
validation and the duplicate logic in _setup_neighbor_backend, _eval_model, and
_eval_model_spin, preserving graph and dpa1_canonical as graph-routed kinds
consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 9d6be660-788d-4192-a41d-a93af6bf7bbc
📒 Files selected for processing (7)
deepmd/backend/pt_expt.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/graph_lower.pydeepmd/pt_expt/utils/serialization.pydeepmd/utils/pt_checkpoint.pysource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Verified all eleven at 5f37d97b4. The nine you agreed with are genuinely fixed, and I accept both of the ones you pushed back on.
The routing foundation is now right. Moving the detector into deepmd/utils/pt_checkpoint.py so dispatch and loading share one verdict was the important part, and _load_pt rejects the explicit pt verdict and the indeterminate one separately, before _extra_state parsing, so a forced dp --pt-expt foo.pt no longer reaches load_state_dict with an opaque failure. The bias asymmetry is better than what I asked for: treating a lone .b as pt_expt-specific while leaving a lone .bias unclassified is conservative in the safe direction, since pt_expt models legitimately contain torch-native modules.
The rest check out as described: Vesin is out of the auto cascade and auto now resolves to NV only on a CUDA host with the package present; an explicit nlist_backend on a graph-routed model raises rather than being ignored; dense is no longer silently accepted on an nlist artifact now that auto is the default; the unreachable is_native_spin_graph condition is gone; and the runner dispatch nests on lower ABI then spin, with all four quadrants covered by test classes. The mocked CPU / CUDA+NV / CUDA-without-NV resolution test replaces the tautological host-dependent assertion, which was the actual problem there.
On the two you did not take:
Keeping neighbor_list= as a hard failure is the right call. My concern was that a working call becomes a construction error, but the old call could produce badly wrong energies for graph-native DPA1 with nonzero descriptor statistics, and failing loudly beats returning plausible garbage. The error now names neighbor_graph_method='ase' and the exported nlist artifact as the two real alternatives, and the break is written down in the PR body, which is what I was actually asking for.
I also accept not reusing _supports_graph_export for eager routing. Exportability and reproducing the source model's public forward are different questions, and applying the export gate here would push raw checkpoints back onto the nlist lower and reintroduce #5862. What remains is that a compressed DPA1 can be frozen to an nlist artifact while eager evaluation of the same model takes the graph lower; I read that as a limitation of the export path rather than something this PR should carry.
On coverage: TestPtExptLoadPtGraphDPA1::test_eval_matches_public_forward is the test that pins the original defect, since it compares a raw .pt evaluation against the source model's public forward for a graph-native descriptor -- the comparison nothing made before, which is why the unconditional lower_input_kind = "nlist" went unnoticed.
Summary
Compatibility
Graph-eligible raw
.ptcheckpoints intentionally keep the source model's graph-forward semantics. Passingneighbor_list=or a non-defaultnlist_backendto such a checkpoint now fails fast because switching to the nlist lower can change predictions (the root cause of #5862). Useneighbor_graph_method="ase"for ASE-based carry-all graph construction. Explicit nlist controls remain supported for nlist-routed checkpoints and deployment artifacts.Checks
Summary by CodeRabbit
neighbor_graph_method="auto"(now the default) with improved neighbor-graph metadata reporting.