[pull] master from deepmodeling:master - #308
Merged
Merged
Conversation
…n+ZBL multi-rank (#5939) Closes #5906. The DPA4/SeZM Source Freeze Propagation Gate computes each node's `eta_j = prod over outgoing edges of w(r_e)`; under MPI domain decomposition a rank only holds edges with owned destinations, so the src-keyed per-node partials are rank-incomplete and bridged models were single-rank only. This PR completes the gate across ranks and, as a prerequisite, promotes the export-time questions to atomic-model capabilities so compositions answer by aggregation. ## Phase 1 — capability aggregation (issue Task 4) - Split the conflated descriptor capability: `has_message_passing_across_ranks` (needs the per-block exchange; unconditionally true for SeZM) vs the new `supports_edge_parallel` (can run under domain decomposition). - Six capabilities on `BaseAtomicModel` with concrete defaults, descriptor delegation on `DPAtomicModel`, and any/all aggregation on `LinearEnergyAtomicModel`: `has_message_passing_across_ranks` (any), `supports_edge_parallel` (all), `dense_lower_supports_comm` (all), `uses_compact_edge_pairs` (any), `graph_edge_dtype` (float32 iff all children), `supports_graph_export` (all). - `forward_lower_graph_exportable_with_comm` hoisted from `EnergyModel` into `make_model` (one owner, next to the non-comm twin) so `LinearEnergyModel` compositions can export it. - The four `serialization.py` helpers now consult the atomic model — no `isinstance`-on-concrete-model checks, no `.descriptor` walks. Regression fixed: a linear composition of two DPA2 children now gets its with-comm artifact (previously denied by wrapper type). - Composition-safe reach-ins outside serialization: `.pt`-checkpoint eval no longer crashes on compositions (`ntypes` via the model API), `enable_compile` degrades gracefully, and pt_expt `get_standard_model` honors `bridging_method` like its dpmodel twin (`_compose_bridging` is the single composition owner). ## Phase 2 — SFPG cross-rank completion (issue Tasks 2 and 3) No new communication machinery: the fix is one extra invocation of the existing `deepmd_export::border_op_backward` + `border_op` pair (they are exact transposes, `R = B^T`) on an `(N, 2)` `[log_eta, zero_count]` tensor before the gate is applied — reverse-accumulate ghost partials into owners, then broadcast the completed values back. Zero C++ changes. - `border_op_backward` gains autograd (its gradient is `border_op`'s forward), so gate gradients cross ranks. - dpmodel: `compute_edge_src_gate` packs the partials through an optional `node_partial_exchange` hook; the dpmodel `_gate_partial_exchange` raises (single-process reference), the pt_expt subclass implements it on the border-op pair. - pt backend wired the same way. The red run of the new pt parity test demonstrated the issue's claim and more: pt's bridged parallel path did not just compute a silently wrong gate — it crashed outright (the ZBL injection indexed per-local types with extended ghost `src` indices); fixed by reading extended types. - Gates flipped: `supports_edge_parallel` is now `True` for bridged SeZM in both backends; bridged (and spin+ZBL) graph freezes embed the nested `forward_lower_with_comm.pt2`. ## Verification Anti-vacuous discipline throughout: every parity test places a sub-`r_outer` pair ACROSS the periodic/rank boundary (without it every cross-rank gate contribution is `log w = 0`), covers both bridging channels (hard-freeze `zero_count` at 0.4 Å, transition-zone `log_eta` at ~1 Å), and carries an identity-exchange ablation that must diverge. - Eager self-comm parity vs the folded reference at rtol/atol 1e-12 (energy, force, and force_mag for the spin variant), pt and pt_expt. - make_fx traces both border ops symbolically (21-input with-comm ABI unchanged); freeze embeds the nested artifact for ZBL and spin+ZBL compositions. - LAMMPS end-to-end on a Tesla T4: 2-rank vs 1-rank close-pair parity for ZBL (`pair_style deepmd`) and spin+ZBL (`pair_style deepspin`, incl. magnetic forces) — the spin+ZBL variant gets its first LAMMPS file. All 24 `*Dpa4Zbl*` C++ gtests pass (CPU + T4). - Variant-alignment coverage: ZBL empty-rank fail-fast twin, the first test of the DeepSpin owned-empty phantom path, charge-spin through `pair_style deepspin`, and default-CLI `dp freeze` resolution (nlist→graph auto-override + with-comm artifact) for both compositions. ## Known limitations 1. pt eager multi-rank bridging has no true-MPI pt test (no pt `.pth` LAMMPS ZBL fixtures exist); its parity rung is self-comm. 2. `graph_edge_dtype` composition rule (float32 iff ALL children) is conservative; fp64 is the universal ABI. 3. `supports_graph_export` keeps the hardcoded `"cuda"` probe inside pt_expt DPA1 (capability promoted; probe internals unchanged). 4. The `NativeSpinModelKind` marker-base check in `_needs_with_comm_artifact` remains (a cross-backend family test, not a concrete-type reach-through). 5. Model-deviation coverage stays absent for all dpa4 variants (pre-existing; Python `model_devi` has no spin support at all). 6. DeepPot vs DeepSpin empty-rank designs deliberately differ (fail-fast vs phantom-pad, PR #5485); both are now pinned per variant, not unified. 7. Found while testing, left for a separate fix: `source/api_c/include/deepmd.hpp` uses `&vec[0]` on possibly-empty vectors (~33 sites) — undefined behavior that SIGABRTs under `_GLIBCXX_ASSERTIONS` before the empty-rank guard's message can fire (benign on non-hardened builds). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added multi-rank inference for bridged DPA4/SeZM models, including native-spin and ZBL configurations. * Improved graph export detection, metadata, edge precision, and communication-aware export. * Added atomic-output-only inference for statistics workflows. * Standard model loading now preserves bridging configurations. * **Bug Fixes** * Improved handling of atom types, ghost atoms, empty MPI ranks, charge-spin inputs, and cross-rank calculations. * **Documentation** * Updated DPA4 and native-spin documentation for expanded multi-rank and graph export support. * **Tests** * Added regression coverage for MPI parity, graph exports, bridging, charge-spin behavior, and capability reporting. <!-- 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 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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )