feat(pt_expt): dpa1(attn_layer=0) graph-native NeighborGraph forward - #5583
Merged
wanghan-iapcm merged 69 commits intoJun 29, 2026
Merged
Conversation
added 26 commits
June 25, 2026 17:26
…_graph The dense path masks excluded type pairs; the graph path does not yet, so raise NotImplementedError instead of silently diverging.
…pter (attn_layer=0)
…k kernel; sw to dense adapter
…ph_method (Option B)
…gy_deriv) + parity
…d/pt_expt/model/)
serialize roundtrip + dpmodel->pt_expt interop on the attn_layer=0 graph path are already covered by test_dpa1.py::test_consistency (lines 86-113), which routes through the graph forward via the Task-3 dense-call adapter.
…all back to dense Task 3's adapter routed ALL attn_layer==0 through the graph, but the graph only supports tebd_input_mode='concat', no exclude_types, and needs mapping for ghosts. strip-mode / exclude / mapping-None-with-ghosts attn_layer=0 models raised/IndexError'd. uses_graph_lower() now encodes full eligibility and ineligible configs fall back to the legacy dense body unchanged. Fixes test_compressed_forward (attn_layer=0 strip).
…ping-ghosts) dense fallback
…pt graph mask key; legacy opt-out in Option-B test - _resolve_graph_method/_call_common_graph use getattr(atomic_model,'descriptor',None) so Linear/ZBL models (no descriptor) fall back to dense instead of AttributeError - pt_expt _call_common_graph override adds the all-ones mask key for dense parity - test_dpa1_graph_model_energy dense refs use neighbor_graph_method='legacy' to opt out of the now-default carry-all graph (decision deepmodeling#17 default-flip)
…nse default dpmodel/jax compute force/virial analytically inside call_common (energy_derv_r); the energy-only graph lower drops it -> KeyError when force is requested. Only pt_expt has the autograd graph force/virial path, so only pt_expt defaults eligible models to the graph. dpmodel base _resolve_graph_method no longer auto-routes; pt_expt overrides it to re-enable AUTO.
…nse call jit/export-traceable (decision deepmodeling#16)
…x int-sum, Array typing) - swap dangling memory/spec_unified_edge_nlist.md refs -> public design discussion (#4) so the references resolve - edge_force_virial: short-circuit n_out=int(node_capacity) when supplied so the static jax/export path never calls int() on a traced sum(n_node) - derivatives.py: move Array import under TYPE_CHECKING (+ from __future__ import annotations) for subpackage uniformity
for more information, see https://pre-commit.ci
…tion Add the missing Parameters/Returns sections (and fill incomplete ones) on the NeighborGraph / graph-lower functions so they match the package numpydoc style: - dpa1: _call_graph_adapter, _call_dense (Parameters+Returns) - general_fitting.call_graph: add missing g2, h2 params - neighbor_graph: pad_and_guard_edges, node_validity_mask (Parameters+Returns); from_dense_quartet, build_neighbor_graph_ase (Returns); edge_force_virial (add g_e/edge_vec/edge_index/edge_mask params) - dpmodel/pt_expt make_model: _resolve_graph_method, _call_common_graph (Parameters+Returns); call_common_lower_graph (replace "Parameters mirror ..." cross-ref with an explicit Parameters section) - pt_expt edge_transform_output: edge_energy_deriv (Parameters+Returns); fit_output_to_model_output_graph (Returns) Docstring-only; no behavior change.
iProzd
requested changes
Jun 28, 2026
- call_common: an explicit `neighbor_list` (a dense-nlist strategy) is no longer silently ignored by the graph default. Raise on `neighbor_list` + explicit `neighbor_graph_method`; otherwise honor the nlist by taking the dense route. - frame_id_from_n_node: accept an optional static `n_total` (jax/export trace-friendly, avoids `int(sum(n_node))`); clamp padding nodes to the last frame so a padded node axis stays in range for segment_sum. - thread `charge_spin` (accept-for-ABI-stability, like comm_dict/n_local) through the graph interface: forward_atomic_graph, forward_common_atomic_graph, call_common_lower_graph, forward_common_lower_graph. - docs: list neighbor_graph_method options one per line incl. "legacy", clarify "dense"/"ase" are carry-all GRAPH builders (not the dense nlist lower); contrast from_dense_quartet (legacy-quartet adapter, keeps sel truncation) vs the carry-all builders. Tests: neighbor_list conflict-raise + dense-route fallback; frame_id static n_total (exact + padded).
dpa1 does not consume charge_spin (get_dim_chg_spin()==0; the dense atomic model passes None to the descriptor since add_chg_spin_ebd is False). charge_spin is accepted on the graph lower only for ABI stability with charge/spin-conditioned descriptors (dpa3/dpa4, PR-G). Pin that the dpa1 graph lower output is INVARIANT to charge_spin: - dpmodel call_common_lower_graph: energy/atom_energy/mask unchanged. - pt_expt forward_common_lower_graph: energy/force/virial/atom_virial unchanged. With the existing graph==dense parity at non-binding sel this gives the full claim graph(charge_spin) == graph(None) == dense. Guards against a future regression where charge_spin leaks into the dpa1 graph path.
for more information, see https://pre-commit.ci
OutisLi
approved these changes
Jun 28, 2026
CodeQL flagged the unused local `N = nf * nloc`; fold it into the comment.
wanghan-iapcm
enabled auto-merge
June 28, 2026 13:15
iProzd
approved these changes
Jun 28, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jun 28, 2026
wccc-phys
pushed a commit
to wccc-phys/deepmd-kit
that referenced
this pull request
Jul 2, 2026
…C++ inference single & multi-rank (NeighborGraph PR-B) (deepmodeling#5604) ## NeighborGraph PR-B — graph `.pt2` export, compiled training, and C++ inference (single & multi-rank) This PR spans the full PR-B: **B1** (Python: graph `.pt2` export + 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 (deepmodeling#5583). Scope: dpa1, `attn_layer=0`, pt_expt. ### B1 — graph `.pt2` export + compiled training (Python) - `forward_common_lower_graph_exportable` trace target; `serialization.py` graph export branch (`lower_kind="graph"`, `lower_input_kind` metadata); `_eval_model_graph` DeepEval dispatch (parity vs eager dpa1 **1e-10 pbc+nopbc**). - **Compiled training retargeted to the graph lower so eager == compiled** (the MUST-FIX) → `force_legacy_descriptor` deleted. Root cause was a real dpa1 `call_graph` autograd **detach** bug (`xp.asarray(tebd, device=)` drops the tebd-net gradient under torch); fixed. ### B2 — C++ graph ingestion (dynamic edge axis, single-rank) - Graph `.pt2` uses 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. - C++ `DeepPotPTExpt`: `lower_input_is_graph_` + `run_model_graph` (NeighborGraph ABI: `atype, n_node, edge_index, edge_vec, edge_mask, …`) + `buildGraphTensors` (mirrors the deepmodeling#5562 edge path; node types from `atype_ext`); `remap_graph_outputs_to_dense_keys` (single-rank). - gtest: 5 cases × {double,float} = 10/10 (build-nlist parity, dynamic-E 2nd size, `ago>0`, tiny system, atomic-overload). The review process caught two bugs that would otherwise have shipped: an `ago>0` heap-OOB (by inspection) and a public-vs-internal output-key mismatch (at runtime). ### B3 — multi-rank C++ / LAMMPS (non-MP) - **dpa1 is non-message-passing ⇒ multi-rank needs NO `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 from `atype_ext` incl. halo), with owned energy = `sum(atom_energy[0:nloc])` and the extended force folded to owners through the **existing dense `select_map` reverse-comm**. The fail-fast for `graph && multi_rank && has_message_passing` is retained. - **Validated locally on multi-CPU** (no GPU needed for correctness): `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 - Per-task + whole-phase reviews all Ready-to-merge. - **pt_expt-only; dpa1 (non-MP) only.** Follow-ons: **PR-C** vesin/nv O(N) builders (carry-all builders still use `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 in `training.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 set `nprocs>1`). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary * **New Features** * Added support for graph-schema (NeighborGraph) model archives with a selectable `lower_kind="graph"` export path, including CLI support and new graph-form inference handling. * Added static edge-capacity support during graph construction. * **Bug Fixes** * Improved gradient continuity for type embeddings in graph mode. * Enhanced trace/export stability by preventing out-of-range graph indices/frame IDs and making scatter/frame sizing more consistent. * **Tests** * Added/extended parity, export metadata, training, and LAMMPS single-/multi-rank validation for graph-form `.pt2`, plus metadata checks for `lower_input_kind`. <!-- 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 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>
HydrogenSulfate
pushed a commit
to HydrogenSulfate/deepmd-kit
that referenced
this pull request
Jul 6, 2026
…eepmodeling#5715) Implements NeighborGraph PR-D: the graph path now supports `attn_layer > 0` for dpa1/se_atten, removing the attn_layer=0-only restriction shipped in deepmodeling#5583. ## What - **Segment toolkit**: `segment_max` + numerically-stable, mask-aware `segment_softmax` (`deepmd/dpmodel/utils/neighbor_graph/segment.py`), built on the existing `xp_maximum_at`. - **`center_edge_pairs`** (`neighbor_graph/pairs.py`): pairs of edges sharing a center — the edge-pair axis shared with the upcoming angle machinery (PR-E). Segment-based enumeration (a global `(E,E)` boolean is deliberately avoided: `O(N²·nnei²)` memory). Two forms: compact eager (dynamic `P`, carry-all graphs) and **shape-static** (`P = n_center·nnei²`, pure arange/reshape arithmetic, no `nonzero`) for the center-major static layout — this keeps the traced/compiled/export path traceable. - **`DescrptBlockSeAtten._graph_attention`**: op-for-op ragged mirror of `GatedAttentionLayer`/`NeighborGatedAttention` — per-center `q@kᵀ` becomes per-pair `q_m·k_n`, softmax over keys becomes `segment_softmax` grouped by the query edge; head_dim QKV slicing, q/k/v normalize, temperature/scaling, smooth shift trick, post-softmax `sw` and `dotr` weighting, residual + LayerNorm per layer. - `edge_env_mat(return_sw=True)` exposes the per-edge switch (zeroed on padding) for the smooth branch. - `uses_graph_lower` widened: attention configs (concat tebd, no exclude_types) are now graph-eligible — pt_expt eager/compiled/exported paths route them through the graph lower by default. ## Numerical semantics (reviewed decision) - **Shape-static adapter path** (the dense `call` adapter, `from_dense_quartet(compact=False)` + `static_nnei`): **bit-exact vs the dense body, rtol 1e-12**, full flag matrix (attn_layer 1/2 × dotr × smooth × normalize × temperature, binding AND non-binding sel). - **Carry-all graphs**: exact for non-smooth attention. For `smooth_type_embedding=True`, the dense branch keeps sel-padding slots in the attention softmax **denominator** (weight `exp(-attnw_shift)`), which makes the dense output *depend on sel itself* (measured up to ~1e-4 with an identical physical neighbor set). The carry-all form **drops those phantom terms by design** — the sel-independent math. Pinned by a clean-divergence test; route-equivalence fixtures pin `smooth_type_embedding=False`. - se_atten_v2 (`tebd_input_mode="strip"`) remains graph-ineligible (strip mode is a later PR) — pinned by test. ## Testing - 38 new dpmodel tests (segment toolkit, pairs incl. random-vs-oracle + static-vs-compact equality, attention parity matrix, binding-sel divergence sanity). - pt_expt: `test_make_fx_graph_attn` (graph forward + autograd at attn_layer=2 traces under make_fx, both smooth branches — required since compiled training uses the graph lower); model-level graph-vs-legacy force/virial/atom-virial parity parametrized over attn_layer {0,2}. - Local CPU: common/dpmodel 583, consistent dpa1+se_atten_v2 209, pt_expt descriptor/model/utils 701 (2 failures: dpa4 export inductor error **pre-existing on upstream/master**, and a route-parity fixture fixed in-branch). - **GPU-validated (Tesla T4, cuda:0)**: dpmodel suites 38, pt_expt graph-lower/make_fx/consistency 44 (CUDA 1e-10), route-parity 6, attention AOTI export pipeline + dpa1 cross-backend consistency 105 — all passed. ## Known limitations - Strip-mode (se_atten_v2) attention stays on the dense path. - Carry-all smooth attention diverges from dense by design (see above); old behavior reachable via `neighbor_graph_method="legacy"` / explicit World-1 builders. - `num_heads == 1` assumed (dpa1 never exposes num_heads); fail-fast otherwise. - Compact `center_edge_pairs` is eager-only (`nonzero`); traced paths use the shape-static form. - 3-body angles (PR-E), jax graph force (PR-F), dpa2/3 MP (PR-G) unchanged. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Expanded graph-native attention support for additional DPA1/se_atten configurations, enabling transformer-style graph execution suitable for tracing/export. * Added center-based neighbor edge-pair enumeration with shape-static control to improve graph layout consistency. * Improved graph tracing/export with optional dynamic-shape hinting. * **Bug Fixes** * Stabilized graph attention softmax under masking/padding and ensured correct behavior for empty/no-edge cases. * **Tests** * Added/updated parity, eligibility, FX traceability, export/graph-lower, and single-atom (no edges) coverage across attention settings. <!-- 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>
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>
pull Bot
pushed a commit
to ishandutta2007/deepmd-kit
that referenced
this pull request
Jul 26, 2026
… multi-rank, native spin, charge-spin/bridging (deepmodeling#5884) ## Summary Brings the DPA4/SeZM descriptor fully onto the NeighborGraph ("graph") lower — the sel-free, edge-native route that dpa1 (deepmodeling#5583) and dpa2 (deepmodeling#5779) already use — and adds the conditioning inputs DPA4 needs on that route. Stacked on the merged dpa2 graph work (deepmodeling#5779); this PR is DPA4's turn on the same infrastructure. Four layers, built and reviewed incrementally: 1. **DPA4 graph-native port** — dense `call` becomes a thin adapter over one graph-native math owner (`_call_graph_impl`); the model-level graph seam drives DPA4 through `call_graph`. Parity vs the dense route at fp64 1e-12; parity vs the pt reference at ~6e-15. 2. **Multi-rank (with-comm) graph `.pt2`** — per-block `border_op` ghost-feature exchange on the graph lower; graph-kind `.pt2` embeds the nested with-comm artifact. Dense (nlist) lower stays comm-less and fails fast on multi-rank; bridging models also fail fast (a rank cannot observe a ghost owner's full outgoing-edge set). 3. **Native spin (magnetic moments)** — per-local-atom spin `(nf, nloc, 3)` conditions the descriptor and is a second autograd leaf giving `force_mag = -dE/dspin`. Rides the graph route only; frozen to a graph-kind `.pt2` (spin at positional ABI index 10), evaluated in Python (`DeepEval`) and C++ (`DeepSpinPTExpt` graph route), run in LAMMPS (`pair_style deepspin`). Single-rank. `DPA4NativeSpinModel` in dpmodel + pt_expt. 4. **Charge-spin FiLM + SFPG bridging** — the last dense-only conditioning inputs, now wired onto the graph route (the edge-native trunk already implemented both; this threads `charge_spin` and flips the capability gates). ## Validation - dpmodel ↔ pt reference DPA4 parity: 387 tests. - native-spin autograd `force_mag` vs finite difference: 5.5e-10 (atol 1e-6); pt weight-copied parity 1e-12. - charge-spin graph-vs-dense parity 1e-12; SFPG bridging frozen-sphere invariance <1e-10 (ablation reopens the leak >1e-6). - C++ (CUDA, Tesla T4): `deeppot_universal` DPA4 dense + graph rows 38 passed / 8 skipped (NoPbc profile); native-spin `DeepSpinPTExpt` graph gtest 16 passed (double/float × PBC/NoPbc). - LAMMPS (CUDA): DPA4 graph energy/force vs live DeepEval ~1e-13; native-spin `deepspin` energy/force/force_mag; 2-rank multi-rank fail-fast aborts cleanly. - Full Python native-spin + charge-spin + bridging battery green on CUDA (device-conditional tolerances where the graph-route `index_add` atomics introduce 1-2 fp64 ULP run-to-run nondeterminism). ## Known limitations - **Single-rank only for spin / charge-spin FiLM / bridging.** Multi-rank spin-graph (ghost-spin exchange) is a follow-up; the graph-kind spin `.pt2` carries `has_comm_artifact=false` and C++ fails fast on multi-rank. Bridging is intrinsically single-rank (per-node freeze fold over the full outgoing-edge set). - **Charge-spin FiLM + native spin combined** is rejected at build (`add_chg_spin_ebd` + native spin) — a small follow-up. - **dpmodel is energy-only** for the native-spin wrapper; force/force_mag come from pt_expt autograd. - **Whole-model conversion of pt-serialized `sezm_native_spin` checkpoints is not claimed** — the deserialize alias covers the type string with dpmodel structure only. - **Sel-free carry-all divergence when sel binds** — the graph route uses the full physical edge set, so it diverges from a sel-truncated dense route only when sel actually binds (the same deliberate divergence as dpa1/dpa2). - **The legacy pt-side `edge_vec` `.pt2` lower schema** (DPA4's original inference rail) is quarantined with `DeprecationWarning`, not removed — its removal is a follow-up gated on repointing the pt-SeZM freeze (and, for its spin part, on graph-spin support, now landed here). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added native per-atom spin conditioning for DPA4/SeZM on the NeighborGraph route, including magnetic outputs (e.g., `force_mag`) and spin masking. * Introduced native-spin `.pt2` export/inference support on the graph route (with appropriate single-rank/MPI behavior). * Extended graph-route input/output handling for spin and charge+spin routing. * **Bug Fixes** * Improved energy+spin loss for flat per-atom label shapes via consistent reshaping. * **Documentation** * Expanded guidance on graph vs dense capabilities, native-spin limitations, and legacy `edge_vec` deprecation. * **Tests** * Added/extended coverage for graph routing, spin outputs, export parity, and multi-rank fail-fast/empty-rank scenarios. <!-- 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 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.
Summary
Adds the graph-native forward path for
dpa1(attn_layer=0)(the factorizable, mixed-types case), built on theNeighborGraphfoundation from #5581. Geometry enters the descriptor only through per-edgeedge_vec; the neighbor-axis reduction becomes asegment_sumover edge centers. Forpt_exptthis becomes the default forward (force/virial via a single autograd backward throughedge_vec).What it adds
edge_env_mat(per-edge env-mat 4-vector),DescrptBlockSeAtten._call_graph+DescrptDPA1.call_graph, modelcall_lower_graph(energy),neighbor_graph_from_ijs+ an optional ASE O(N) carry-all builder.edge_energy_deriv(autogradgrad(E, edge_vec)→edge_force_virial) +forward_common_lower_graph(energy + force + virial + atom_virial).DescrptDPA1.callbecomes a thin adapter (from_dense_quartet → call_graph) preserving the 5-tuple ABI; a shape-static converter keeps itjax.jit/torch.export-traceable.Default behavior
dpa1(attn_layer=0, concat tebd, no exclude_types)models to the carry-all graph (it has the autograd force/virial path).sel.exclude_types, linear/ZBL) fall back to the dense path unchanged.neighbor_graph_method="legacy"forces dense;"dense"/"ase"force the graph.Parity (graph vs legacy dense lower, fp64 CPU)
atom_virial matches the canonical TF==pt-legacy full-to-src convention. dpa1 descriptor + model consistency suites green across dp/jax/pt_expt.
Known limitations
make_fx(forward + grad) traces; full.pt2AOTI export is a follow-up (PR-B). The carry-all builders (build_neighbor_graph/from_ijs) still usenonzero(eager-only); their static variants land with the export PR.Also folds in three follow-up fixes to the #5581 foundation from @OutisLi's review (dangling spec refs → design discussion,
edge_force_virialjax int-sum short-circuit,Arraytyping).Summary by CodeRabbit
neighbor_graph_methodrouting for energy/force/virial, with carry-all neighbor graphs and graph-output fitting/post-processing.(i,j,S)conversion, and per-edge environment-matrix computation), exported as part of the public API.