feat(pt_expt): graph-native dpa1 — .pt2 export + compiled training + C++ inference single & multi-rank (NeighborGraph PR-B) - #5604
Merged
wanghan-iapcm merged 33 commits intoJul 2, 2026
Conversation
added 8 commits
June 29, 2026 18:42
… graph .pt2 export - Add make_model.forward_common_lower_graph_exportable: make_fx trace of forward_common_lower_graph with edge_vec as the autograd leaf; symbolic tracing enabled via two data-dependent shape fixes. - Add EnergyModel.forward_lower_graph_exportable: wraps the above with a second make_fx pass that translates internal keys to public names (atom_energy, energy, force, virial, atom_virial). - Fix edge_transform_output.py: replace int(n_node.sum()) with next(iter(fit_ret.values())).shape[0] (static shape attr, trace-safe); thread node_capacity through edge_energy_deriv -> edge_force_virial so that segment_sum sizes are static under make_fx symbolic mode. - Add edge_energy_deriv node_capacity param (None = eager fallback). - Test: source/tests/pt_expt/model/test_graph_export.py — TDD RED->GREEN, verifies traced module reproduces eager energy_redu (rtol=1e-10).
… do_atomic_virial branches), dedup key-translation, drop redundant detach
…r_graph Adds test_graph_static_capacity.py (4 tests): shape (2,64), real-edge count matches dynamic, real prefix identity + masked tail, overflow raises. No source change needed -- builder.py already threads layout.edge_capacity into pad_and_guard_edges.
…raph .pt2 metadata for the C++ hub (B2)
…pbc, 1e-10) Add a graph-form .pt2 DeepEval dispatch and parity test (NeighborGraph PR-B Phase B1.4). DeepEval: route lower_input_kind=="graph" archives through a new _eval_model_graph that builds a carry-all NeighborGraph (padded to the static edge_capacity baked in metadata), feeds the positional schema (atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin) to the AOTI forward, and reshapes the LOCAL public outputs by category (no communicate_extended_output, since the graph path is ghost-free). The dense (nlist) path is untouched. Fix dynamic-shape generalization of the graph .pt2 export: under symbolic make_fx/torch.export, int() on a SymInt SPECIALIZES the axis, baking the trace-time sample size (N=14, nf=2) into the autograd-derived force/energy/ virial outputs. Drop int() on node_capacity (edge_force_virial) and on n_node.shape[0] (fit_output_to_model_output_graph), and derive nf-1 in frame_id_from_n_node as a runtime 0-d tensor instead of xp.asarray(shape-1). The exported graph now generalizes across arbitrary N and nf. Test: dpa1(attn_layer=0) energy model exported with lower_kind="graph", evaluated through the pt_expt DeepPot and compared to the eager dense (sel-capped, neighbor_graph_method="legacy") reference at rtol=atol=1e-10 for pbc and nopbc. The fixture is a sparse 8-atom cluster and asserts non-binding sel (max neighbors < 30) so the carry-all/sel-capped neighbor sets coincide (anti-vacuity).
…); drop force_legacy_descriptor Retarget _CompiledModel to compile forward_common_lower_graph for graph-eligible descriptors (dpa1 attn_layer==0), gated by the same mixed_types()+uses_graph_lower() predicate the eager default-flip uses; se_e2_a/dpa2/dpa3 keep compiling the dense forward_lower. _trace_and_compile_graph builds a synthetic NeighborGraph with prime-distinct nf/N/E axes (no make_fx duck-shape merge) and edge_vec as the autograd leaf; _forward_graph builds the carry-all graph eagerly and unravels flat (N,*) node outputs to (nf,nloc,*). cpp.simdlen=0 for the graph compile avoids an inductor CPU scatter-vectorizer crash on the per-frame virial atomic_add. Also fixes an eager autograd bug in dpa1 call_graph: xp.asarray(type_embedding, device=dev) DETACHES under torch, so the type-embedding weights never trained in the graph path (grad None despite a real finite-diff dependency). make_fx traced through it, so compiled != eager and the optimizer diverged after step 0. Use type_embedding directly (mirrors the dense path); the tebd net now trains and eager==compiled to 1e-10 across the varying-natoms trajectory. Drops the force_legacy_descriptor workaround + uses_graph_lower monkeypatch.
… + legacy-gate assumption in the graph compile path
for more information, see https://pre-commit.ci
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a graph-native NeighborGraph lowering path across export, inference, training compilation, C++ runtime, and tests. It also updates dpmodel and neighbor-graph helpers to preserve gradients and symbolic shapes during graph tracing. ChangesGraph-form lower-forward export and runtime
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pt_expt/utils/serialization.py (1)
750-789: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReject unknown
lower_kindvalues instead of falling back to nlist.A typo such as
lower_kind="grpah"currently reaches the dense path and recordslower_input_kind="nlist", producing the wrong artifact without any signal.Proposed fix
def deserialize_to_file( @@ """ + if lower_kind not in {"nlist", "graph"}: + raise ValueError( + f"Unknown lower_kind {lower_kind!r}; expected 'nlist' or 'graph'." + ) if model_file.endswith(".pt2"):🤖 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/utils/serialization.py` around lines 750 - 789, The deserialize entrypoint in serialization.py currently accepts any lower_kind value and silently falls back to the nlist path, which can produce the wrong artifact for typos. Update deserialize_to_file and the lower-path selection logic in _deserialize_to_file_pt2 / _deserialize_to_file_pte to validate lower_kind explicitly against the supported values ("nlist" and "graph") and raise a clear error for anything else instead of defaulting to nlist, so the recorded lower_input_kind always matches the requested schema.source/tests/pt_expt/test_training.py (1)
1355-1374: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert that the DPA1 no-attn case really takes the graph lower.
These changes now document graph-lower routing, but the test still only proves eager/compiled parity. If
_CompiledModelsilently falls back to the dense lower and the sampled systems stay non-binding, this case still passes, so the dispatch change introduced by this PR is untested. Please add an explicit assertion for the DPA1 no-attn branch that the compiled model is graph-eligible / using the graph lower, while keeping the dense expectation for the other descriptors.Also applies to: 1391-1392, 1466-1479
🤖 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/test_training.py` around lines 1355 - 1374, The current test only checks eager/compiled parity in _check_varying_natoms, so the DPA1 no-attn routing to the graph lower is not explicitly verified. Add an assertion in the DPA1 branch that the compiled model is graph-eligible and actually uses forward_common_lower_graph, while preserving the dense forward_lower expectation for the non-eligible descriptors. Extend the same explicit dispatch check in the related test sections that cover the same descriptor routing.
🤖 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/train/training.py`:
- Around line 968-970: The graph-compiled forward path is dropping the
do_atomic_virial request, so atom_virial cannot be returned from compiled
models. Update forward and _forward_graph in training.py to thread
do_atomic_virial through the graph-eligible dispatch, and make
_trace_and_compile_graph cache/trace separately based on that flag so the
compiled lower path matches forward_common_lower_graph behavior.
- Around line 1283-1293: The reshaping logic in the result loop currently relies
on the `N != nframes` heuristic, which fails for `nloc == 1` and leaves
node-level outputs flat. Update the `training.py` postprocessing in the
`result.items()` loop to use explicit key categories from the graph/lower-level
public keys instead of shape inference, and reshape known node-level outputs
like `atom_energy` and `force` to `(nf, 1, *)` even when `N == nframes`.
In `@deepmd/pt_expt/utils/serialization.py`:
- Around line 447-448: The dynamic-shape spec in _build_graph_dynamic_shapes
still only marks the frame axis for aparam, so its atom axis remains fixed to
the sample nloc. Update the aparam entry in the shape mapping to include both
the frame axis and the atom axis as dynamic, matching the intended (nf, nloc,
nda) behavior. Keep the change localized to the graph shape spec alongside the
existing fparam and aparam handling so exported graphs no longer specialize to a
single atom count.
---
Outside diff comments:
In `@deepmd/pt_expt/utils/serialization.py`:
- Around line 750-789: The deserialize entrypoint in serialization.py currently
accepts any lower_kind value and silently falls back to the nlist path, which
can produce the wrong artifact for typos. Update deserialize_to_file and the
lower-path selection logic in _deserialize_to_file_pt2 /
_deserialize_to_file_pte to validate lower_kind explicitly against the supported
values ("nlist" and "graph") and raise a clear error for anything else instead
of defaulting to nlist, so the recorded lower_input_kind always matches the
requested schema.
In `@source/tests/pt_expt/test_training.py`:
- Around line 1355-1374: The current test only checks eager/compiled parity in
_check_varying_natoms, so the DPA1 no-attn routing to the graph lower is not
explicitly verified. Add an assertion in the DPA1 branch that the compiled model
is graph-eligible and actually uses forward_common_lower_graph, while preserving
the dense forward_lower expectation for the non-eligible descriptors. Extend the
same explicit dispatch check in the related test sections that cover the same
descriptor routing.
🪄 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
Run ID: 58e2933c-1a05-4d30-bdfa-0cec147b8f6a
📒 Files selected for processing (14)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/utils/neighbor_graph/derivatives.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/edge_transform_output.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/serialization.pysource/tests/common/dpmodel/test_graph_static_capacity.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/model/test_graph_export.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5604 +/- ##
==========================================
- Coverage 82.37% 81.85% -0.52%
==========================================
Files 902 960 +58
Lines 101527 106256 +4729
Branches 4056 4134 +78
==========================================
+ Hits 83630 86976 +3346
- Misses 16432 17787 +1355
- Partials 1465 1493 +28 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
added 7 commits
June 30, 2026 01:21
… static edge_capacity
The graph-form .pt2 export now marks the edge axis dynamic
(Dim("nedge", min=2)) instead of baking a static E_max=ceil(1.25*nloc*nnei)
capacity. The AOTI artifact accepts any system size with no capacity ceiling.
- _build_graph_dynamic_shapes: edge_index dim1 / edge_vec dim0 / edge_mask dim0
are now dynamic; mirrors the dense Dim("nnei", min=...) precedent.
- _trace_and_export graph branch: drop the metadata["edge_capacity"] write;
apply _strip_shape_assertions to neutralise the SIGFPE-prone deferred
shape guards on the dynamic E axis (same handling the spin dense path uses).
- deep_eval._eval_model_graph: build the carry-all graph at its tight edge
count (no edge_capacity padding).
- test_graph_deepeval: eval TWO different-size systems (8- and 20-atom,
56 vs 380 real edges) through the SAME exported artifact; both match eager
dense dpa1 at 1e-10 pbc+nopbc, both non-binding. The 20-atom system (380 > 263)
would have overflowed the B1 static artifact -> the RED.
- test_graph_pt2_metadata: graph metadata no longer carries edge_capacity.
…e graph (dynamic-E) caller + edge-axis safety
…) generator Section B of gen_dpa1.py produces deeppot_dpa1_graph.pt2 (lower_kind="graph", do_atomic_virial=True) and the accompanying deeppot_dpa1_graph.expected sidecar. Reference values come from an independent nlist .pt2 eval (NOT the graph .pt2) so the C++ gtest (B2.5) validates the graph AOTI path against a known-good reference. Sanity check: graph .pt2 vs nlist ref force diff 1.3e-18 (machine precision; sel=30 >> actual neighbors, so both paths see identical neighbor sets). Forces non-degenerate: max |F| ~5.3e-4 (PBC). Config: type_map=[O,H], sel=30, rcut=6.0, attn_layer=0, neuron=[2,4,8], axis_neuron=4, fitting neuron=[5,5,5], resnet_dt=True, seed=1 — mirrors DPA1_CONFIG in source/tests/pt_expt/utils/test_graph_pt2_metadata.py.
B2.2: read lower_input_kind="graph" -> lower_input_is_graph_. B2.3: run_model_graph with NeighborGraph AOTI input order (atype, n_node, edge_index, edge_vec, edge_mask, [fparam], [aparam], [charge_spin]); no coord / edge_scatter_index. B2.4: GraphTensorPack + buildGraphTensors (delegates to createEdgeTensors+compactEdgeTensors for the rcut filter, dynamic edge count and 2 masked dummy edges; drops edge_index_ext, adds n_node=[nloc], node types from atype_ext[0:nloc]) + compute_inner & standalone dispatch branches. Multi-rank graph fails fast (PR-B3).
… path Local std::vector<int64_t> mapping was declared in compute_inner and populated only inside if(ago==0). The graph branch called buildGraphTensors with this local vector on every step, causing an OOB heap read on ago>0 (mapping.size()==0). Fix: promote mapping to a member mapping_ (parallel to mapping_tensor) so it persists across steps. Edge-path (createEdgeTensors) and dense-path (mapping_tensor) are unaffected in behavior; only the vector source changes from local to member.
…xtraction
Add test_deeppot_dpa1_graph_ptexpt.cc — the first runtime exercise of the
NeighborGraph .pt2 C++ ingestion (B2.2-4). Four cases x {double,float}:
build-nlist parity vs the committed .expected, a second 12-atom system through
the same model (dynamic edge axis), the LAMMPS InputNlist+ago=0/ago=1 path
(compute_inner + cached mapping_), and a tiny no-edge system (nedge_min=2 guard).
The first run surfaced a B2.2-4 bug: the graph forward emits LOCAL flat-N PUBLIC
keys (atom_energy/energy/force/virial/atom_virial) but compute()'s output
extraction read the dense INTERNAL keys (energy_redu/energy_derv_r/...), so
output_map["energy_redu"].view() threw on an undefined tensor. The graph branch
had only ever been compiled, never run.
Fix: remap_graph_outputs_to_dense_keys() in commonPT.h, called after
extract_outputs in both compute overloads (gated on lower_input_is_graph_).
Rewrites the public keys into the dense internal-key layout; per-atom
force/atom_virial are local (nloc) and zero-padded up to nall so the existing
fold-back is a no-op on ghost rows (local rows already carry the folded ghost
contributions). Dense/edge paths untouched.
gen_dpa1.py now persists the dense nlist-ref .pt2 (deeppot_dpa1_graph_nlist_ref.pt2)
as a live graph-vs-dense oracle for the dynamic-edge cases.
…ring (nlist .pt2, 1e-5)
for more information, see https://pre-commit.ci
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
source/api_cc/src/DeepPotPTExpt.cc (1)
508-539: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRequire
lmp_list.mappingfor graph artifacts with ghosts.The graph branch uses
mapping_inbuildGraphTensorsto fold ghost neighbors onto local owners, but the current guard only requires an atom map forhas_message_passing_models. A DPA1 graph artifact can therefore run PBC LAMMPS withoutatom_modify map yes, fall back to identity mapping, and feed invalid/wrong folded graph indices.Proposed fix
if (lower_input_is_graph_ && multi_rank) { throw deepmd::deepmd_exception( "Multi-rank graph (NeighborGraph) .pt2 inference is not yet " "supported (PR-B3). Run single-rank, or use a dense/edge .pt2 for " "multi-rank LAMMPS."); } + if (lower_input_is_graph_ && nghost > 0 && !atom_map_present) { + throw deepmd::deepmd_exception( + "Graph (NeighborGraph) .pt2 inference with ghost atoms requires " + "`atom_modify map yes` so ghost neighbors can be folded onto their " + "local owners. Re-run with an atom map enabled, or use a dense .pt2."); + }🤖 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/api_cc/src/DeepPotPTExpt.cc` around lines 508 - 539, The graph-artifact path in DeepPotPTExpt::buildGraphTensors is using mapping_ to fold ghost neighbors, but the current validation only enforces atom-map presence for has_message_passing_ models. Add a guard for lower_input_is_graph_ with ghosts so graph artifacts also require lmp_list.mapping (or atom_modify map yes) before compute(), and fail fast when it is missing instead of falling back to identity mapping.deepmd/pt_expt/utils/serialization.py (1)
793-799: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReject unknown
lower_kindvalues before dispatch.A typo currently falls through to the dense nlist export while recording nlist metadata, producing a valid-looking artifact with the wrong ABI.
Proposed fix
""" + if lower_kind not in {"nlist", "graph"}: + raise ValueError( + f"Unsupported lower_kind={lower_kind!r}; expected 'nlist' or 'graph'." + ) if model_file.endswith(".pt2"): _deserialize_to_file_pt2(🤖 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/utils/serialization.py` around lines 793 - 799, The dispatch in serialization.py currently lets unknown lower_kind values fall through to the dense nlist export path, which can produce a misleading artifact with incorrect ABI metadata. Add validation before the model_file.endswith(".pt2") branch so only the expected lower_kind values are accepted, and raise an error for anything else. Keep the check close to the dispatch logic in the deserialization/export flow that calls _deserialize_to_file_pt2 and _deserialize_to_file_pte.
🤖 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.
Outside diff comments:
In `@deepmd/pt_expt/utils/serialization.py`:
- Around line 793-799: The dispatch in serialization.py currently lets unknown
lower_kind values fall through to the dense nlist export path, which can produce
a misleading artifact with incorrect ABI metadata. Add validation before the
model_file.endswith(".pt2") branch so only the expected lower_kind values are
accepted, and raise an error for anything else. Keep the check close to the
dispatch logic in the deserialization/export flow that calls
_deserialize_to_file_pt2 and _deserialize_to_file_pte.
In `@source/api_cc/src/DeepPotPTExpt.cc`:
- Around line 508-539: The graph-artifact path in
DeepPotPTExpt::buildGraphTensors is using mapping_ to fold ghost neighbors, but
the current validation only enforces atom-map presence for has_message_passing_
models. Add a guard for lower_input_is_graph_ with ghosts so graph artifacts
also require lmp_list.mapping (or atom_modify map yes) before compute(), and
fail fast when it is missing instead of falling back to identity mapping.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: c4b21156-2b33-4392-b6b0-4819a492ea77
📒 Files selected for processing (9)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/serialization.pysource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.ccsource/tests/infer/gen_dpa1.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/pt_expt/infer/deep_eval.py
- source/tests/pt_expt/infer/test_graph_deepeval.py
wanghan-iapcm
changed the title
added 2 commits
June 30, 2026 08:36
…st (B2.5 follow-ups)
…-comm; no with-comm)
OutisLi
approved these changes
Jul 1, 2026
wanghan-iapcm
enabled auto-merge
July 1, 2026 07:09
This was referenced Jul 2, 2026
pull Bot
pushed a commit
to ishandutta2007/deepmd-kit
that referenced
this pull request
Jul 4, 2026
…emiops (PR-C) (deepmodeling#5714) ## NeighborGraph PR-C — O(N) on-device graph builders (vesin / nv) Third PR of the NeighborGraph series (after deepmodeling#5581 foundation, deepmodeling#5583 PR-A dpa1 graph forward, deepmodeling#5604 PR-B `.pt2`/C++). Adds two **O(N) carry-all** NeighborGraph builders behind the World-2 `neighbor_graph_method` dispatcher, replacing PR-A's O(N²) `dense` search / per-frame ASE stopgap with on-device cell lists. ### What - **`build_neighbor_graph_vesin`** (`deepmd/pt_expt/utils/vesin_graph_builder.py`) — `vesin.torch` cell list. **Device-following**: runs the search on the input tensor's device (CUDA kernel on CUDA input, CPU cell list otherwise). - **`build_neighbor_graph_nv`** (`deepmd/pt_expt/utils/nv_graph_builder.py`) — nvalchemiops GPU cell list, **frame-batched** (`batch_idx`/`batch_ptr`, one kernel for all frames — no Python loop). CUDA-only. - Both structurally clone `build_neighbor_graph_ase`: search → per-frame local `(i, j, S)` → `neighbor_graph_from_ijs(...)`, which recomputes `edge_vec` **differentiably** from the original grad-carrying coords. - Wired into the pt_expt make_model graph dispatch (`neighbor_graph_method ∈ {"legacy","dense","ase","vesin","nv"}`) and DeepEval `.pt2` graph inference (new `neighbor_graph_method` kwarg, default `"dense"` → existing inference byte-identical). dpmodel/jax fail-fast on `vesin`/`nv` (torch/CUDA-only). **Perf-only:** all builders emit the SAME neighbor set as `dense` (carry-all, `sel`=normalization-only), proven by exact set-equality; energy/force/virial are unchanged (parity 1e-12 CPU / 1e-10 CUDA). ### Layering - `dpmodel` stays torch-free: vesin/nv builders live in `pt_expt`; the dpmodel dispatch only carries a fail-fast message. - vesin/nv are **optional deps, NOT in pyproject** — lazy-imported, guarded by `is_vesin_torch_available()` / `is_nv_available()`, `ImportError` with an install hint on absence. - nv decode is a faithful transcription of the tested `deepmd/pt/utils/nv_nlist.py:_matrix_to_extended_inputs` Step-1 extraction. ### Testing - **Local (CPU):** 13 passed (vesin builder 5, vesin/reject dispatch 2, DeepEval graph 6), nv self-skips. - **Remote GPU (Tesla T4), commit 78f6c24:** - `test_nv_graph_builder.py`: 4 passed (set-equality vs dense periodic+non-periodic, frame-batch, differentiable `edge_vec`). - `test_graph_builder_dispatch.py`: 3 passed on CUDA (**vesin** parity 1e-10, dpmodel reject vesin+nv, **nv** parity 1e-10). - `test_graph_deepeval.py`: 6 passed on CUDA (`.pt2` graph dense parity + vesin, AOTI compile + `torch.as_tensor` extraction). ### Known limitations - **vesin per-frame Python loop** — `vesin.torch.compute` is single-system (no batch dim), so multi-frame vesin loops over frames (each an O(N) search; `nf=1` inference has zero loop cost). **nv batches natively** — the loop-free path for batched training. - vesin/nv not in `pyproject`; no `"auto"` selector (explicit strings only). - nv `search_capacity = max(64, nloc)` initial heuristic + 1.25× grow loop (`.item()` host-sync per grow). - nv passes the *normalized* coord to `from_ijs` (differentiable lattice translation, identity gradient; `edge_vec` is lattice-invariant) — consistent with the pt nv-nlist path. - jax O(N) graph builders (matscipy/jax-md) = PR-F; attention/angles/MP = PR-D/E/G. Implements `plan_neighbor_graph_prC_implementation`; design spec: discussion wanghan-iapcm#4. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added experimental neighbor-graph builders supporting NVIDIA CUDA (NV) and Vesin cell-list construction (vesin), including periodic shifts and frame-aware edge decoding. * **Bug Fixes** * Excluded virtual atoms (`type < 0`) from graph edges consistently for ASE-based neighbor graphs. * Improved behavior for empty/zero-neighbor inputs and ensured differentiable edge vectors. * **Tests** * Added backend parity tests (dense vs vesin/nv), NV decode regression coverage, gradient/device checks, virtual-atom exclusions, and strengthened plugin entry-point import validation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This was referenced Jul 7, 2026
wanghan-iapcm
pushed a commit
to wanghan-iapcm/deepmd-kit
that referenced
this pull request
Jul 13, 2026
… DEVICE
_trace_and_compile_graph built its synthetic NeighborGraph trace inputs on
the module-level DEVICE constant instead of the device the model's own
parameters/buffers actually live on. In real training these always match
(the trainer moves the model to DEVICE before compiling), but any caller
that intentionally holds the model on a different device (e.g. a test
pinning the model to CPU while running on a CUDA host, mirroring the .pt2
export's model.to("cpu") pattern) hits a device split between the traced
graph tensors and the model params.
DPA2.call_graph's tebd gather (dpmodel/descriptor/dpa2.py:1159) surfaced
this first: xp.take(tebd_table, atype_local, axis=0) requires tebd_table
(a model param, left un-asarray'd to avoid detaching its gradient -- the
documented dpa1 lesson) and atype_local (moved to device(graph.edge_vec))
to share a device. Under test_compiled_training_graph_smoke the model is
pinned to CPU but the trace sample was still built on the global DEVICE
(cuda:0 on a GPU box), so tebd_table stayed CPU while atype_local was CUDA.
DPA1's call_graph has the identical gather pattern and would fail the same
way, but no existing dpa1 test calls _trace_and_compile_graph directly with
a device-pinned model, so the bug was latent there since PR-B (deepmodeling#5604).
Fix: add _model_trace_device(), reading the device off the model's own
parameters/buffers (falling back to DEVICE only if the model exposes
neither), and use it to build the trace sample instead of the global
DEVICE. No change to the gather itself -- gradient flow through
type_embedding is untouched.
wanghan-iapcm
pushed a commit
to wanghan-iapcm/deepmd-kit
that referenced
this pull request
Jul 13, 2026
… DEVICE
_trace_and_compile_graph built its synthetic NeighborGraph trace inputs on
the module-level DEVICE constant instead of the device the model's own
parameters/buffers actually live on. In real training these always match
(the trainer moves the model to DEVICE before compiling), but any caller
that intentionally holds the model on a different device (e.g. a test
pinning the model to CPU while running on a CUDA host, mirroring the .pt2
export's model.to("cpu") pattern) hits a device split between the traced
graph tensors and the model params.
DPA2.call_graph's tebd gather (dpmodel/descriptor/dpa2.py:1159) surfaced
this first: xp.take(tebd_table, atype_local, axis=0) requires tebd_table
(a model param, left un-asarray'd to avoid detaching its gradient -- the
documented dpa1 lesson) and atype_local (moved to device(graph.edge_vec))
to share a device. Under test_compiled_training_graph_smoke the model is
pinned to CPU but the trace sample was still built on the global DEVICE
(cuda:0 on a GPU box), so tebd_table stayed CPU while atype_local was CUDA.
DPA1's call_graph has the identical gather pattern and would fail the same
way, but no existing dpa1 test calls _trace_and_compile_graph directly with
a device-pinned model, so the bug was latent there since PR-B (deepmodeling#5604).
Fix: add _model_trace_device(), reading the device off the model's own
parameters/buffers (falling back to DEVICE only if the model exposes
neither), and use it to build the trace sample instead of the global
DEVICE. No change to the gather itself -- gradient flow through
type_embedding is untouched.
pull Bot
pushed a commit
to ishandutta2007/deepmd-kit
that referenced
this pull request
Jul 19, 2026
…model on the NeighborGraph lower (PR-G-dpa2) (deepmodeling#5779) DPA-2 becomes graph-native full-stack — the first **message-passing** model on the NeighborGraph lower (PR-G-dpa2 of the [NeighborGraph series](wanghan-iapcm#4); follows dpa1's route from deepmodeling#5581/deepmodeling#5583/deepmodeling#5604/deepmodeling#5714/deepmodeling#5715/deepmodeling#5717/deepmodeling#5733). The dense path coexists untouched; the graph route is opt-in for inference (`--lower-kind graph`) and the default for pt_expt training/eager on graph-eligible configs, exactly like dpa1. ## What's new **dpmodel (backend-agnostic math)** - `DescrptBlockRepformers.call_graph`: every repformer op ported to the flat edge list — conv/grrg/drrd/g1g1/`LocalAtten` via `segment_*` over `dst`; the dense `nnei×nnei` gated attention (`Atten2Map`/`Atten2MultiHeadApply`/`Atten2EquiVarApply`) via `center_edge_pairs(ordered, include_self)` + per-head `segment_softmax` (extended to trailing feature dims). - `DescrptDPA2.call_graph` + `uses_graph_lower` + dense-call adapter. The multi-level nlist (`build_multiple_neighbor_list`) becomes per-block edge masks; in the shape-static adapter layout the per-block **slice** replicates the dense `nlist[:, :, :ns]` truncation, making the adapter **bit-exact vs the dense path at ANY sel** (verified to 5e-16 at deliberately binding sel with attention enabled) — this is what keeps the cross-backend consistency suites green after the routing flip. Ineligible configs (`use_three_body`, compression, spin) keep the legacy dense route. - Owned-node energy mask: `fit_output_to_model_output_graph` consumes `n_local`, excluding halo rows from the differentiated energy (prevents cross-rank double counting; force stays full-N, halo partials reverse-commed by LAMMPS). **pt_expt / export** - Routing, autograd force/virial, compiled training, and freeze eligibility are inherited generically from the dpa1 machinery. - Per-layer MPI halo refresh on the graph path: `_exchange_ghosts_graph` — identity on ghost-free single-rank graphs (`src` IS the owner), `deepmd_export::border_op` overwrite of halo rows on extended-region multi-rank graphs. - Message-passing graph `.pt2` archives embed a **with-comm AOTInductor artifact** (`model/extra/forward_lower_with_comm.pt2`), traced with the 8-tensor comm ABI shared with the dense flow. **C++ / LAMMPS** - `DeepPotPTExpt::run_model_graph_with_comm` + dispatch replacing the PR-G fail-fast: MP message-passing graph models run multi-rank on the extended-region graph + with-comm artifact; non-MP (dpa1) multi-rank unchanged; old graph archives without the artifact get a clear re-freeze error. - Fixtures (`gen_dpa2.py` section B), `dpa2_graph_ptexpt` universal gtest row, and `test_lammps_dpa2_graph_pt2.py` (single-rank vs reference, per-atom virial, `mpirun -n 2` ≡ `-n 1`, graph-vs-nlist cross-artifact, bounded empty-rank). ## Validation - GPU (Tesla T4, CUDA): dpa2-graph LAMMPS **5/5** incl. the MP==SP gate; dpa1-graph LAMMPS 6/6; dense dpa2 MPI 1/1; every `dpa2_graph` C++ gtest variant passed; python GPU suites green. - CPU: full dpmodel+pt_expt sweep 2112 passed; dpa2 consistency suite (pt/dp/jax/array-api-strict) green via the bit-exact adapter. - Real bugs found & fixed during GPU validation: CUDA device placement of the `nlocal`/`nghost` comm scalars (the graph route consumes them in on-device owned-mask index math); `_trace_and_compile_graph` hardcoded the global device for trace samples (latent since deepmodeling#5604). ## Known limitations / deliberate divergences (documented in `doc/model/dpa2.md` + code Notes) - Carry-all graph attention is sel-independent by design: at binding sel — and for smooth attention generally — the graph route diverges from dense (dpa1 precedent, `KNOWN_GRAPH_DENSE_DIVERGENT`). - Per-atom virial attribution differs elementwise from the dense decomposition for message-passing models (full-to-src vs autograd spread); per-frame totals agree (cross-checked at 1e-8 in fixture generation). - A truly empty MPI rank (zero owned+ghost atoms) fails loudly rather than running; the thrown error cannot propagate through peers blocked in `border_op` collectives, so the job stalls until MPI timeout — phantom-node support is a follow-up. The LAMMPS test bounds this with a hard timeout and asserts it never silently succeeds. - Pre-existing dense bug surfaced (NOT introduced here): `se_atten` at `attn_layer=0` leaks a deterministic `-davg/dstd` padding residual when `set_davg_zero=False` and `exclude_types == []` (`PairExcludeMask` all-ones short-circuit is the only padding mask on that path). The graph path masks padding correctly, so graph and dense deliberately differ in that regime (pinned by test + docstring). A dense-side fix will be proposed separately. ## Follow-ups (separate issues/PRs) - Dense `se_atten` padding-residual fix (above). - Wire the `BUILD_PT_EXPT` C++ gtest suite into CI (pre-existing gap — the universal gtest rows, including the new one, never run in CI today). - `TestDeepPotPTExptWithCommLoadFailure.multi_rank_compute_throws` is vacuous (sets `nswap` but not `nprocs`; pre-existing since deepmodeling#5450). - Empty-rank phantom-node support / clean collective abort; nested with-comm artifact schema assertion in the export test. PR-G-dpa3 (repflows + angle channel + `charge_spin`, reusing this PR's MP machinery) comes next. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Enabled graph-native DPA2 descriptor/repformer execution for eligible energy models, including multi-rank “with-comm” ghost exchange inside graph artifacts. * Added `comm_dict` plumbing and persistent graph-lower enable/disable controls across save/restart. * Improved multi-rank owned/halo reduction using `n_local`, including flat node-axis `aparam` handling for graph ABI and exports. * **Bug Fixes** * Refined multi-rank reduction correctness by excluding ghost contributions from differentiated reductions. * Improved graph attention/softmax stability with phantom-aware behavior for smoother continuity. * **Documentation** * Added pt_expt graph-native route documentation, eligibility rules, and graph-freeze with-comm requirements. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: OutisLi <137472077+OutisLi@users.noreply.github.com>
This was referenced Jul 20, 2026
This was referenced Jul 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
NeighborGraph PR-B — graph
.pt2export, compiled training, and C++ inference (single & multi-rank)This PR spans the full PR-B: B1 (Python: graph
.pt2export + compiled training on the graph lower), B2 (C++ single-rank inference of the graph.pt2, dynamic edge axis), and B3 (C++/LAMMPS multi-rank). Built on the merged PR-A (#5583). Scope: dpa1,attn_layer=0, pt_expt.B1 — graph
.pt2export + compiled training (Python)forward_common_lower_graph_exportabletrace target;serialization.pygraph export branch (lower_kind="graph",lower_input_kindmetadata);_eval_model_graphDeepEval dispatch (parity vs eager dpa1 1e-10 pbc+nopbc).force_legacy_descriptordeleted. Root cause was a real dpa1call_graphautograd detach bug (xp.asarray(tebd, device=)drops the tebd-net gradient under torch); fixed.B2 — C++ graph ingestion (dynamic edge axis, single-rank)
.pt2uses a dynamic edge axis (Dim("nedge", min=2)) — one artifact evals any system size (proven across 56- and 380-edge systems at 1e-10), no C++ capacity ceiling.DeepPotPTExpt:lower_input_is_graph_+run_model_graph(NeighborGraph ABI:atype, n_node, edge_index, edge_vec, edge_mask, …) +buildGraphTensors(mirrors the refactor(pt/dpa4): use edge schema for dpa4 and ignore sel #5562 edge path; node types fromatype_ext);remap_graph_outputs_to_dense_keys(single-rank).ago>0, tiny system, atomic-overload). The review process caught two bugs that would otherwise have shipped: anago>0heap-OOB (by inspection) and a public-vs-internal output-key mismatch (at runtime).B3 — multi-rank C++ / LAMMPS (non-MP)
border_op/with-comm artifact (that is a message-passing concern, deferred to PR-G). Multi-rank reuses the same single-rank graph.pt2, fed an extended-region graph (buildGraphTensors(fold_to_local=false),N=nall, ghost node types fromatype_extincl. halo), with owned energy =sum(atom_energy[0:nloc])and the extended force folded to owners through the existing denseselect_mapreverse-comm. The fail-fast forgraph && multi_rank && has_message_passingis retained.test_lammps_dpa1_graph_pt2.py— single-rank vs reference,mpirun -n 2≡ single-rank (energy + per-atom force + virial, atol 1e-8), plus an empty-subdomain (nloc=0) corner. Single-rank gtests stay 10/10 (multi-rank is purely additive). Multi-rank matched single-rank on the first run.Tests / known limitations
nonzero, eager-only), PR-D attention, PR-E angles, PR-F jax graph force, PR-G dpa2/3 message-passing (forward halo + with-comm). CUDA multi-rank unvalidated locally. Carried code-cleanup follow-ups: a ~60-line DRY duplication intraining.py; the multi-rank atomic output branch has no direct gtest (covered indirectly by the mpirun per-atom-virial assertion, since a single-process gtest can't setnprocs>1).Summary by CodeRabbit
Summary
New Features
lower_kind="graph"export path, including CLI support and new graph-form inference handling.Bug Fixes
Tests
.pt2, plus metadata checks forlower_input_kind.