feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank - #5939
feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank#5939wanghan-iapcm wants to merge 26 commits into
Conversation
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
|
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:
📝 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 💡 2⚔️ Resolve merge conflicts 💡
🛠️ Fix failing CI checks 💡
🧪 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}
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.41% 79.20% -0.22%
==========================================
Files 1072 1072
Lines 124893 125017 +124
Branches 4531 4536 +5
==========================================
- Hits 99187 99015 -172
- Misses 24085 24379 +294
- Partials 1621 1623 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
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__.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The native-spin + ZBL multi-rank magnetic-force contract is not yet validated: the PR's own transition-zone self-comm parity case fails reproducibly on the current HEAD.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The existing graph-freeze regressions still pin the obsolete single-rank composition contract and need to be updated together with this gate change.
Dismissed because this review was based on a stale locally installed C++ operator library. After explicitly loading the current build artifact from this checkout, the border-backward behavior is correct and the exact native-spin + ZBL self-comm test passes. The reported force_mag mismatch is therefore not a defect in this PR. Sorry for the false positive.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The new capability interface remains incomplete for composite descriptors: a supported hybrid containing DPA4 is misclassified as dense-comm capable, causing freeze to trace an implementation that always raises.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One maintainability issue remains in the updated DPA4 capability documentation.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The current SFPG communication sequence adds an avoidable MPI exchange and global synchronization on every force evaluation; please simplify this hot path before merging.
| @@ -193,16 +202,31 @@ def compute_edge_src_gate( | |||
| log_eta = torch.zeros( | |||
| n_nodes, dtype=edge_w.dtype, device=edge_w.device | |||
| ).scatter_add(0, src, log_safe) | |||
There was a problem hiding this comment.
[P2] Build eta on the owner instead of reverse-accumulating source partials
The graph contract defines src = neighbor and dst = center, and LAMMPS requests a full neighbor list, so the owner of node j already has every edge k -> j. Since w depends only on the symmetric distance (and pair exclusion is symmetric), eta_j = prod_k w(r_jk) can be computed by reducing log_safe and is_zero over dst. Then only one forward border_op is needed to broadcast the completed eta to ghost source rows.
Reducing over src here makes the partials rank-incomplete by construction and forces _gate_partial_exchange to run border_op_backward followed by border_op. Both are MPI exchanges, and each current kernel ends with an MPI_Barrier, so this adds one avoidable round trip/global synchronization to every bridged DPA4 force evaluation and also requires the extra backward-op autograd machinery. Please use destination reduction plus one broadcast on the comm path (while preserving the current non-comm behavior if arbitrary non-reciprocal graphs must remain supported), or identify a supported multi-rank topology where an owner does not hold its complete destination-centered neighborhood.
There was a problem hiding this comment.
I would like to push back on this one, because the equivalence rests on the edge set being reciprocal and our supported topologies do not guarantee that.
Your reasoning holds exactly when for every edge j -> k the mirror k -> j also exists with the same length. That is true for a LAMMPS full neighbour list built purely by cutoff. It is not true for the nlist-derived edge schemas this same code serves: format_nlist caps each centre at sel, and when sel binds it truncates by distance per centre, so a dense-environment atom can drop a neighbour that the sparser neighbour keeps. Measured on a deliberately inhomogeneous system (one dense cluster plus a sparse tail, rcut=4.0, sel=6, ghosts mapped back to owner images):
sel = 6 edges = 90
NON-RECIPROCAL ordered pairs: 12 [(0,5), (6,2), (3,6), (0,9), (7,2), (3,9)]
So on that path prod_{src=j} w and prod_{dst=j} w are different numbers, not two spellings of one. That leaves only bad options: reduce over dst everywhere and the gate is silently redefined for sel-capped inputs (pt forward_with_edges is exactly this path, and it is what test_sezm_parallel_bridging_parity.py compares against), or reduce over dst only when comm_dict is set and the comm path stops computing the same quantity as the non-comm path — which is precisely the kind of split this branch is trying to remove, and it would make the folded-vs-parallel parity test tautological rather than meaningful.
The SFPG definition is source-keyed (eta_j = prod_{e: src_e = j} w_e), so I would rather keep one reduction that is correct for every supported input than a faster one that is correct only under an invariant nothing asserts.
Your cost point stands on its own though, and I do not want to lose it: one extra border_op_backward round trip with an MPI_Barrier per bridged force evaluation is real. The clean version of your optimisation is to establish reciprocity as an explicit, asserted property of the graph (which the LAMMPS-fed graph genuinely has) and then take the single-broadcast path under that guarantee, keeping the general reduction otherwise. That is a larger change than this PR should carry, so I would rather file it as a follow-up than bolt a conditional fast path on here. Happy to open that issue if you agree with the framing — and if you think the reciprocity invariant already holds by construction on every path that reaches this function, say so and I will re-examine, since that would change the conclusion.
There was a problem hiding this comment.
Your non-reciprocal counterexample is valid for the sel-capped, nlist-derived PT compatibility path, so that path must retain the source-keyed two-phase completion. It does not apply to the production pt_expt graph-lower path, however. That path uses the carry-all, sel-free NeighborGraph (all edges inside rcut), and DPA4's dense call() rejects comm_dict before _run_graph; consequently the pt_expt communication hook is reached through call_graph, not through a sel-truncated dense graph.
For that graph-native route, the owner has the complete destination-centred neighbourhood. Reducing [log_eta, zero_count] over dst and performing one forward border_op is therefore the appropriate fast path, while the generic/nlist-derived PT route can keep its current source reduction plus reverse-accumulate/broadcast.
Please either (1) add this as a clean, explicitly carry-all/reciprocal graph fast path, without changing the generic source-keyed semantics, or (2) open a focused follow-up issue that records the graph contract, routing boundary, and required benchmark/parity validation. Graph lower is the primary route going forward, so I do not want the two-exchange/two-barrier implementation to become its only long-term path. I will leave this thread unresolved until the fast path lands or the tracking issue is linked.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The graph with-comm exporter is energy-contract machinery and should not be attached to every model generated by the generic model factory. Please keep this capability in an energy-specific abstraction before merging.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The new spin+ZBL LAMMPS coverage should not duplicate the existing native-spin harness wholesale. Please keep the variant-specific tests independent while consolidating their shared reference, LAMMPS, and MPI mechanics before merging.
| expected_v = None | ||
|
|
||
|
|
||
| def _cell_from_lammps_box(lmp_box: np.ndarray) -> np.ndarray: |
There was a problem hiding this comment.
[P2] Extract the shared native-spin LAMMPS harness
This new 566-line module shares 416 identical lines with test_lammps_dpa4_spin_graph_pt2.py (about 71% whole-file similarity). The duplicated surface includes _cell_from_lammps_box, live DeepPot reference generation, setup/teardown, LAMMPS construction, force_mag gathering, the single-rank energy/force/virial checks, and the MPI subprocess/parser/parity machinery. The actual variant-specific inputs are primarily the archive path, the close-pair geometry, and a small number of ZBL assertions.
Keeping two copies of the whole harness means any future change to the LAMMPS invocation, reference conversion, tolerances, timeout handling, or MPI output schema must be updated in parallel, and makes another stale-contract divergence likely. Please retain separate plain-native-spin and spin+ZBL test cases, but extract these shared mechanics into one helper/harness and parameterize only the archive, geometry, and variant-specific assertions.
There was a problem hiding this comment.
Agreed on the substance — I measured it and it matches your estimate: 421 of the 566 lines are byte-identical to test_lammps_dpa4_spin_graph_pt2.py, and the genuinely variant-specific parts are the archive, the geometry and a handful of ZBL assertions. Two copies of the MPI subprocess/parser/parity machinery will drift.
What I do not want to do is perform that extraction blind, and right now I cannot execute either file to check it. The .pt2 fixtures are not present in this working tree, generating them needs AOTI compiles that hit a vectorised-scatter codegen assertion in the local torch build (workable only with cpp.simdlen=1), the multi-rank spin parity comparisons are additionally unreliable on this box at O(1e-4) for environmental reasons unrelated to this branch, and the GPU box I would normally use to adjudicate is unreachable at the moment. Restructuring 400+ lines of MPI harness — including the subprocess output schema and teardown — with only --collect-only as evidence is a worse trade than the duplication, especially since it would also rewrite a pre-existing file that is not otherwise touched by this PR.
So my preference is a dedicated follow-up PR that does the extraction and is validated on a machine that can actually run both files, rather than an unvalidated refactor here. If you would rather it land in this PR and are content to let the C++/LAMMPS CI jobs be the validation, say so and I will do it — I just want that trade made explicitly rather than by accident.
There was a problem hiding this comment.
Please perform the extraction in this PR. I do not want to merge a newly added test module with 416 of its 566 lines duplicated and rely on an unspecified follow-up to remove that debt.
The repository already has source/lmp/tests/lammps_test_utils.py, which owns shared LAMMPS setup and MPI-runner mechanics (make_spin_lammps, run_mpi_pair_runner). Please either extend that existing seam or add a narrowly scoped shared DPA4-native-spin harness, while keeping the plain native-spin and spin+ZBL archives, geometries, and variant-specific assertions in their respective tests. The common live-reference setup, teardown, LAMMPS construction, magnetic-force gathering, subprocess/output parsing, and common parity assertions should have one owner.
It is acceptable for the C++/LAMMPS CI jobs to provide the end-to-end runtime validation; local collection/static checks can validate the refactor's imports and parametrization. Lack of a locally reliable MPI environment is not a reason to commit two copies of the harness. Please keep this thread unresolved until the shared structure lands.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please update the DPA4/SeZM factory documentation to match the capabilities implemented by this PR.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please remove the stale single-rank limitation from the spin+ZBL fixture documentation now that this fixture pins the with-comm multi-rank path.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please make the PT SeZM capability documentation consistent with the bridging multi-rank behavior implemented by this PR.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please align the DPA4 freeze-route documentation with the public pt_expt freeze behavior that always selects graph for graph-capable models.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please update the shared native-spin marker documentation to match the graph with-comm contract implemented here.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Please make the new standard + DPA4 bridging route a fully supported, normalized equivalent of type: dpa4, rather than a direct-builder-only partial path.
A DescrptHybrid inherited the concrete `dense_lower_supports_comm() == True` default while aggregating `has_message_passing_across_ranks()` over its children. For hybrid([dpa4, se_e2_a]) that combination makes `_needs_with_comm_artifact(model, "nlist")` true, so an auto/nlist freeze builds a dense with-comm trace that dies in DPA4's NotImplementedError as soon as comm_dict is passed. Aggregate it (and supports_edge_parallel) with ALL, matching the veto semantics already used for the atomic-model composition.
Two executable assertions still pinned the pre-deepmodeling#5906 single-rank contract (`has_comm_artifact is False`) for the ZBL and native-spin+ZBL graph freezes. They are skipped only when CI=true, so a local full-suite run compiled the archive and then failed on the obsolete expectation. Flip both to the multi-rank contract and also assert the nested forward_lower_with_comm.pt2 entry is present.
make_model() produces DOS, dipole, polar, property and energy models, but
forward_lower_graph_exportable_with_comm hard-codes do_grad_r("energy"),
do_grad_c("energy") and _translate_energy_keys. Defining it on the common
factory made every non-energy model advertise an API that would raise on
the missing energy output; the serializer's energy gate only hid it.
Give it the same ownership as its non-comm twin: define it on EnergyModel
and alias it in LinearEnergyModel. NativeSpinEnergyModel is built on
EnergyModel, so the spin route keeps it through the MRO.
Six sites still described the pre-deepmodeling#5906 contract or a CLI behaviour that does not exist: - dpa4.uses_graph_lower / sezm.has_message_passing_across_ranks claimed bridging still fails multi-rank fast, contradicting the capability split this branch introduces. - NativeSpinModelKind cited "native-spin lowers are single-rank only" as its motivating example; only the DENSE lower is excluded now. - get_sezm_model advertised bridging and spin as unsupported although the factory deliberately constructs both; replace with the precise matrix (deepspin scheme, lora, use_compile, preset_out_bias remain unsupported). - doc/model/dpa4.md offered plain DPA4 a choice of dense or graph freeze via a '--lower-kind auto' default. There is no 'auto' choice (nlist and graph only, default nlist) and freeze() overrides any non-graph request to graph for graph-capable models. - gen_dpa4_spin_zbl.py still headed its contract 'Single-rank only'.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/dpmodel/model/native_spin_model.py`:
- Line 37: Update the docstring near the shared cross-backend marker to remove
the duplicated word, changing “cross-backend family test test” to a single
“test” while preserving the rest of the marker wording.
🪄 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: 2213b926-20f0-40ce-a83d-12c1285fa49b
📒 Files selected for processing (12)
deepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/hybrid.pydeepmd/dpmodel/model/native_spin_model.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt_expt/model/dp_linear_model.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/model/make_model.pydoc/model/dpa4.mdsource/tests/common/dpmodel/test_atomic_model_capabilities.pysource/tests/infer/gen_dpa4_spin_zbl.pysource/tests/pt_expt/model/test_zbl_bridging.py
💤 Files with no reviewable changes (1)
- deepmd/pt_expt/model/make_model.py
🚧 Files skipped from review as they are similar to previous changes (6)
- deepmd/pt_expt/model/ener_model.py
- doc/model/dpa4.md
- deepmd/dpmodel/descriptor/dpa4.py
- deepmd/pt/model/descriptor/sezm.py
- source/tests/infer/gen_dpa4_spin_zbl.py
- source/tests/pt_expt/model/test_zbl_bridging.py
Left over from removing the stale parenthetical in 4021124.
Conflict in deepmd/pt_expt/utils/serialization.py: upstream deepmodeling#5913 relocated the graph edge-dtype helper to deepmd/pt_expt/model/graph_lower.py as a public graph_edge_dtype(), while this branch had rewritten the same helper in place to answer from the atomic-model capability instead of reaching through a single .descriptor. Resolution keeps BOTH: upstream's location (serialization.py imports it from graph_lower) with this branch's implementation. The reach-in version is wrong for compositions -- a LinearEnergyAtomicModel has no .descriptor, so a bridged/ZBL model would silently report float64 regardless of its children -- which is exactly what the capability aggregation fixes. Dropped the now-unused torch import from graph_lower.py.
get_standard_model honored `bridging_method` and returned a LinearEnergyModel -- the type asked for was not the type returned. It also made a second owner of the bridging build, and the two owners had already drifted: get_sezm_model promotes descriptor.exclude_types to model-level pair_exclude_types and this one never did, changing a 0.9 A Ni-O dimer by 79.97 eV (max |dF| 318.48 eV/A). Reject `bridging_method` in the standard builder instead. Rejecting rather than ignoring keeps the original fail-loud property: silently dropping the term yields a physically different model than the config requests. DPA4 (`type: "dpa4"`/`"sezm"`) is now the single bridging owner in this backend. The route was unreachable from a validated input.json anyway -- argcheck declares `bridging_method` only under the dpa4/sezm variant -- so no supported configuration changes behavior. Follow-ups: deepmodeling#5947 (drop the exclusion promotion), deepmodeling#5948 (express bridging as an explicit linear_ener composition, after which the restriction is moot).
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The PT freeze comment still describes bridged SeZM as single-rank.
deepmd/pt/entrypoints/freeze_pt2.py:951-956 says that bridging models report supports_edge_parallel() == False, that Source Freeze Propagation is not rank-decomposable, and that they fall back to single-rank. That is now the opposite of the contract implemented by this PR: bridged SeZM reports True, _gate_partial_exchange() completes the SFPG partials across ranks, and the surrounding predicate therefore embeds forward_lower_with_comm.pt2 for the edge_vec route. The new PT capability tests also pin that behavior.
Please update this production comment together with the capability change. Leaving the explanation inverted next to the artifact-selection predicate makes the multi-rank contract particularly easy to misread during future maintenance.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The DPA4 documentation makes an unsupported performance claim about the new SFPG communication path.
| *outgoing*-edge set, which no single rank observes for ghost owners; the | ||
| with-comm artifact completes the gate's per-node `[log eta, zero count]` | ||
| partials across ranks with one reverse-accumulate plus one forward-broadcast | ||
| border exchange of an `(N, 2)` tensor per forward pass — negligible next to |
There was a problem hiding this comment.
[P2] Do not label the SFPG synchronization negligible without measurements
The (N, 2) payload is narrow, but its cost is not determined by bandwidth alone. This path invokes border_op_backward followed by border_op, and both C++ kernels end in an MPI_Barrier; because the exported force graph differentiates through both operations, their registered transpose callbacks add the corresponding communication pair to the force evaluation as well. On small systems or at higher rank counts, the added round trips/global synchronizations can dominate despite the small tensor. This is also the hot path covered by the still-open graph fast-path thread. Please remove “negligible”, or support it with benchmarks across representative system sizes and MPI rank counts.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The refactored DPA4/SeZM factory now has an internally inconsistent return-type contract.
| atom_exclude_types=data.get("atom_exclude_types", []), | ||
| pair_exclude_types=pair_exclude_types, | ||
|
|
||
| def _compose_bridging(model: Any, data: dict, bridging_method: str) -> Any: |
There was a problem hiding this comment.
[P2] Keep the factory return contract truthful after composition
The new composition branch makes the public annotation on get_sezm_model() incorrect: the plain path returns DPA4EnergyModel, but the bridging path returns LinearEnergyModel, which is not an EnergyModel. Typing this helper and get_standard_model() as Any only hides that mismatch and discards useful structural information from the factory API. The PT twin already uses the honest common contract, BaseModel. Please type get_sezm_model() and get_standard_model() against BaseModel, and give this helper either its concrete LinearEnergyModel result or the same BaseModel contract, instead of erasing all three paths with Any.
3 participants
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