feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank - #5939
Open
wanghan-iapcm wants to merge 19 commits into
Open
feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank#5939wanghan-iapcm wants to merge 19 commits into
wanghan-iapcm wants to merge 19 commits into
Conversation
added 14 commits
July 30, 2026 08:55
has_message_passing_across_ranks conflated the two; the SFPG bridging veto moves to the new supports_edge_parallel() (issue deepmodeling#5906 Task 4 groundwork).
…ties Six capabilities with concrete defaults on BaseAtomicModel, descriptor delegation on DPAtomicModel, and any/all aggregation on LinearEnergyAtomicModel (issue deepmodeling#5906 Task 4).
LinearEnergyModel compositions need it once the with-comm gate opens for them (issue deepmodeling#5906 Task 4); one owner, next to the non-comm twin. _translate_energy_keys moves along (make_model cannot import from ener_model without a cycle); importers re-pointed.
_needs_with_comm_artifact / compact-edge-pairs / edge-dtype / graph-export answer via the atomic model; two-DPA2 linear compositions now get their with-comm artifact (issue deepmodeling#5906 Task 4).
.pt-checkpoint eval and enable_compile no longer reach for a single descriptor; get_standard_model honors bridging_method like its dpmodel twin, with _compose_bridging as the one composition owner (issue deepmodeling#5906 Task 4 audit findings).
R = B^T as linear maps over node rows, so the reverse-accumulate's vector-Jacobian product is the forward broadcast (issue deepmodeling#5906: gate gradients cross ranks).
compute_edge_src_gate packs (N,2) [log_eta, zero_count] through the optional node_partial_exchange hook; zero_count becomes float so both partials ride one border exchange. dpmodel's _gate_partial_exchange raises (single-process reference); pt_expt overrides it (issue deepmodeling#5906).
_gate_partial_exchange = reverse-accumulate (border_op_backward) then forward-broadcast (border_op) of the (N,2) [log_eta, zero_count] partials. Eager self-comm parity vs the folded reference at 1e-12 with a cross-boundary close pair in both bridging channels (hard-freeze and transition zone); identity-exchange ablation diverges (issue deepmodeling#5906).
supports_edge_parallel flips to True (the SFPG exchange completes the gate partials across ranks); the graph freeze embeds the nested forward_lower_with_comm.pt2 for bridged compositions and gen_dpa4_zbl asserts it (issue deepmodeling#5906 Task 2).
Closes the deepmodeling#5906 pt gap: the gate's per-node partials are completed via border_op_backward + border_op before the gate is applied, and the model-level ZBL injection reads EXTENDED types (the parallel path used to crash on ghost src indices). Bridged SeZM reports supports_edge_parallel=True; close-pair parity pinned on both bridging channels with an identity-exchange ablation.
Replaces the fails-fast test: the with-comm artifact now completes the SFPG partials across ranks, so a 2-rank run with the 0.9 A pair straddling the processors 2 1 1 boundary must (and does) match the 1-rank reference (issue deepmodeling#5906 Task 2 E2E).
…ling#5906 Task 3) The spin wrapper needed NO production change: the gate exchange sits below it. gen_dpa4_spin_zbl asserts the with-comm artifact; python self-comm parity (energy/force/force_mag, both bridging channels) at 1e-12; first-ever LAMMPS file for the spin+ZBL variant, incl. 2-rank close-pair parity via pair_style deepspin.
Empty-rank behavior (zbl fail-fast twin; first test of the DeepSpin phantom path -- owned-empty ranks with ghosts phantom-pad and match), charge_spin via pair_deepspin, and dp-freeze default-CLI resolution for zbl and native-spin compositions (issue deepmodeling#5906 Task 12b).
The SFPG per-node partials are completed across ranks by one reverse-accumulate + forward-broadcast border exchange of an (N,2) tensor per forward pass (issue deepmodeling#5906); also sweeps the stale pre-deepmodeling#5884 native-spin single-rank claims.
for more information, see https://pre-commit.ci
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR adds capability queries, cross-rank SFPG partial exchange for bridged DPA4/SeZM inference, capability-driven graph export, bridging-aware model composition, atomic-output statistics, training cleanup, and expanded PyTorch and LAMMPS validation. ChangesSFPG and export pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
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
🤖 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 `@doc/model/dpa4.md`:
- Around line 595-597: Update the native-spin support section heading in the
graph route documentation to use a positive statement, since the listed
multi-rank inference, charge-spin FiLM conditioning, and ZBL zone bridging
combinations are supported with spin.scheme: native.
In `@source/tests/infer/gen_dpa4_spin_zbl.py`:
- Around line 357-369: Update the _check_metadata docstring to state that it
verifies the presence of the with-comm artifact and the metadata flag, matching
the assertions for has_comm_artifact and forward_lower_with_comm.pt2. Remove the
contradictory wording that says it checks their absence.
In `@source/tests/pt/model/test_sezm_parallel_bridging_parity.py`:
- Around line 152-162: Update the parallel forward call to pass
sysm["edge_index"] as the graph connectivity argument where it currently repeats
sysm["edge_scatter_index"], matching the reference call while preserving
edge_scatter_index for its intended argument.
🪄 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: b6eb425d-2677-4159-a996-819c2ef72fc7
📒 Files selected for processing (40)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/atomic_model/linear_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/dpa2.pydeepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/edge_cache.pydeepmd/dpmodel/descriptor/make_base_descriptor.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/edge_cache.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/dpa4.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/native_spin_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/comm.pydeepmd/pt_expt/utils/serialization.pydoc/model/dpa4.mdsource/lmp/tests/test_lammps_dpa4_chg_spin_deepspin_pt2.pysource/lmp/tests/test_lammps_dpa4_spin_graph_pt2.pysource/lmp/tests/test_lammps_dpa4_spin_zbl_pt2.pysource/lmp/tests/test_lammps_dpa4_zbl_pt2.pysource/tests/common/dpmodel/test_atomic_model_capabilities.pysource/tests/common/dpmodel/test_descrpt_dpa4.pysource/tests/infer/gen_dpa4_spin_zbl.pysource/tests/infer/gen_dpa4_zbl.pysource/tests/pt/model/test_sezm_parallel.pysource/tests/pt/model/test_sezm_parallel_bridging_parity.pysource/tests/pt_expt/model/test_dpa4_zbl_parallel.pysource/tests/pt_expt/model/test_export_with_comm.pysource/tests/pt_expt/model/test_get_model_bridging.pysource/tests/pt_expt/model/test_zbl_bridging.pysource/tests/pt_expt/test_dp_freeze.pysource/tests/pt_expt/utils/test_border_op_backward.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.py
👮 Files not reviewed due to content moderation or server errors (11)
- deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py
- deepmd/pt/model/descriptor/sezm.py
- deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
- deepmd/pt/model/model/sezm_model.py
- deepmd/pt_expt/descriptor/dpa4.py
- deepmd/pt_expt/infer/deep_eval.py
- deepmd/pt_expt/model/ener_model.py
- deepmd/pt_expt/model/get_model.py
- deepmd/pt_expt/model/make_model.py
- deepmd/pt_expt/model/native_spin_model.py
- deepmd/pt_expt/utils/comm.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5939 +/- ##
==========================================
- Coverage 79.37% 79.15% -0.22%
==========================================
Files 1070 1070
Lines 124791 124897 +106
Branches 4531 4531
==========================================
- Hits 99050 98860 -190
- Misses 24119 24418 +299
+ Partials 1622 1619 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
added 3 commits
July 31, 2026 09:17
Resolutions: keep the capability-based _needs_with_comm_artifact and the _compose_bridging owner (renamed to upstream's InnerPotentialAtomicModel, deepmodeling#5910); drop this branch's pt extended-atype ZBL fix in favor of upstream's equivalent descriptor_atype at the relocated injection site.
CodeRabbit review on deepmodeling#5939.
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 (1)
deepmd/pt_expt/train/training.py (1)
2147-2176: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winKeep legacy epoch lengths consistent with pt semantics.
pt_exptdividescompute_total_numb_batch(...)byworld_sizefor legacy data, whileptonly uses the sharded LMDB dataloader length and otherwise returns the full weighted batch count. If legacyDeepmdDataSystemremains replicated per rank, this makesnum_stepsper epoch differ between backends for the samenumb_epoch/num_epoch_dictconfiguration, changing effective epoch length and the derived LR schedule. Align legacypt_exptwithptunless this parallel epoch definition is intentional.🤖 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/train/training.py` around lines 2147 - 2176, The _epoch_length method currently divides every dataset’s weighted batch count by world_size, unlike pt semantics for legacy DeepmdDataSystem data. Preserve the world_size division only for sharded LMDB data, and return the full compute_total_numb_batch result for replicated legacy data so epoch lengths and LR schedules remain consistent across backends.
🧹 Nitpick comments (1)
deepmd/pt_expt/model/get_model.py (1)
246-249: 📐 Maintainability & Code Quality | 🔵 TrivialDuplicated inner-clamp radii injection.
The 2-line inner-clamp radii injection (
inner_clamp_r_inner/inner_clamp_r_outerfrombridging_r_inner/bridging_r_outer) is repeated verbatim inget_sezm_model(lines 127-128) and here. Consider extracting a small helper (e.g._inject_bridging_radii(data)) shared by both builders to avoid future divergence.🤖 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/model/get_model.py` around lines 246 - 249, Extract the duplicated descriptor-radii assignment into a shared helper such as _inject_bridging_radii(data), preserving the existing default values and descriptor initialization. Replace the inline logic in both get_sezm_model and the shown bridging_enabled path with calls to that helper.
🤖 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/train/training.py`:
- Around line 2147-2176: The _epoch_length method currently divides every
dataset’s weighted batch count by world_size, unlike pt semantics for legacy
DeepmdDataSystem data. Preserve the world_size division only for sharded LMDB
data, and return the full compute_total_numb_batch result for replicated legacy
data so epoch lengths and LR schedules remain consistent across backends.
---
Nitpick comments:
In `@deepmd/pt_expt/model/get_model.py`:
- Around line 246-249: Extract the duplicated descriptor-radii assignment into a
shared helper such as _inject_bridging_radii(data), preserving the existing
default values and descriptor initialization. Replace the inline logic in both
get_sezm_model and the shown bridging_enabled path with calls to that helper.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 50f94545-ef5a-4f70-939d-8fa0ac908f8a
📒 Files selected for processing (9)
deepmd/pt/model/model/sezm_model.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/train/training.pysource/lmp/tests/test_lammps_dpa4_zbl_pt2.pysource/tests/common/dpmodel/test_atomic_model_capabilities.pysource/tests/infer/gen_dpa4_spin_zbl.pysource/tests/infer/gen_dpa4_zbl.pysource/tests/pt_expt/model/test_get_model_bridging.pysource/tests/pt_expt/model/test_zbl_bridging.py
🚧 Files skipped from review as they are similar to previous changes (3)
- source/tests/infer/gen_dpa4_zbl.py
- source/tests/pt_expt/model/test_get_model_bridging.py
- source/lmp/tests/test_lammps_dpa4_zbl_pt2.py
Two GitHub Advanced Security findings on the new file: - 'unittest' was imported both as a module and via 'from unittest import mock'; use a plain 'import unittest.mock' instead. - the 'deepmd.pt.utils.env' import was dead: this test pins CPU explicitly, and env is pulled in transitively by get_model plus the source/tests/pt package __init__.
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
2 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.
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)
has_message_passing_across_ranks(needs the per-block exchange; unconditionally true for SeZM) vs the newsupports_edge_parallel(can run under domain decomposition).BaseAtomicModelwith concrete defaults, descriptor delegation onDPAtomicModel, and any/all aggregation onLinearEnergyAtomicModel: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_commhoisted fromEnergyModelintomake_model(one owner, next to the non-comm twin) soLinearEnergyModelcompositions can export it.serialization.pyhelpers now consult the atomic model — noisinstance-on-concrete-model checks, no.descriptorwalks. Regression fixed: a linear composition of two DPA2 children now gets its with-comm artifact (previously denied by wrapper type)..pt-checkpoint eval no longer crashes on compositions (ntypesvia the model API),enable_compiledegrades gracefully, and pt_exptget_standard_modelhonorsbridging_methodlike its dpmodel twin (_compose_bridgingis 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_oppair (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_backwardgains autograd (its gradient isborder_op's forward), so gate gradients cross ranks.compute_edge_src_gatepacks the partials through an optionalnode_partial_exchangehook; the dpmodel_gate_partial_exchangeraises (single-process reference), the pt_expt subclass implements it on the border-op pair.srcindices); fixed by reading extended types.supports_edge_parallelis nowTruefor bridged SeZM in both backends; bridged (and spin+ZBL) graph freezes embed the nestedforward_lower_with_comm.pt2.Verification
Anti-vacuous discipline throughout: every parity test places a sub-
r_outerpair ACROSS the periodic/rank boundary (without it every cross-rank gate contribution islog w = 0), covers both bridging channels (hard-freezezero_countat 0.4 Å, transition-zonelog_etaat ~1 Å), and carries an identity-exchange ablation that must diverge.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).pair_style deepspin, and default-CLIdp freezeresolution (nlist→graph auto-override + with-comm artifact) for both compositions.Known limitations
.pthLAMMPS ZBL fixtures exist); its parity rung is self-comm.graph_edge_dtypecomposition rule (float32 iff ALL children) is conservative; fp64 is the universal ABI.supports_graph_exportkeeps the hardcoded"cuda"probe inside pt_expt DPA1 (capability promoted; probe internals unchanged).NativeSpinModelKindmarker-base check in_needs_with_comm_artifactremains (a cross-backend family test, not a concrete-type reach-through).model_devihas no spin support at all).source/api_c/include/deepmd.hppuses&vec[0]on possibly-empty vectors (~33 sites) — undefined behavior that SIGABRTs under_GLIBCXX_ASSERTIONSbefore the empty-rank guard's message can fire (benign on non-hardened builds).Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests