DPA4/SeZM zone-bridging (SFPG) gate is incomplete under domain decomposition; blocks ZBL and spin+ZBL multi-rank

[wanghan-iapcm]

Summary

The DPA4/SeZM zone-bridging gate (Source Freeze Propagation Gate, SFPG) is computed from a per-node reduction over a node's outgoing edge set. Under MPI domain decomposition no single rank observes that full set, so the gate is incomplete on every rank. This blocks multi-rank inference for ZBL-bridged DPA4 models, and — separately — for the combination of native spin + ZBL.

Native-spin multi-rank on the NeighborGraph route is already working (PR #5884); bridging is the one remaining graph-route feature that cannot go multi-rank without a new communication primitive.

Why the gate is rank-incomplete

compute_edge_src_gate computes, per node j:

w_e   = bridging_switch(r_e)              in [0, 1]
eta_j = prod over { e : src_e == j } w_e  in [0, 1]
gate_e = eta_{src_e}

• dpmodel/pt_expt: deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py:127 (reduction at lines 209 and 218, keyed by src)
• pt: deepmd/pt/model/descriptor/sezm_nn/edge_cache.py:113 (same structure)

The reduction key is src — the neighbor side of the edge. But a rank's edge list only contains edges whose dst (center) is an owned atom, which is the LAMMPS neighbor-list contract. An owned boundary atom j is a neighbor of centers owned by other ranks, and those edges exist only on those ranks. So eta_j is a partial product on every rank, for owned atoms as well as ghosts.

This is structurally different from every exchange the descriptor already does. deepmd_export::border_op broadcasts an owner's already-complete value onto its ghost copies; here no rank ever holds a complete value to broadcast. The partial products must first be combined across ranks.

When it produces wrong numbers

Edges at r >= r_outer contribute w = 1 (log w = 0), so omitting them is harmless. The error appears only when a close contact inside the bridging window straddles a rank boundary: atom j within r_outer of a center i owned by another rank. Both ranks see that pair, but with opposite orientation — j's rank holds (src=i, dst=j), while eta_j needs (src=j, dst=i), which lives only on i's rank.

The result is that eta_j fails to freeze (or freezes too weakly) for exactly the cross-boundary close pairs that ZBL bridging exists to handle. Any test that does not place a sub-r_outer pair across a subdomain boundary passes vacuously.

Current state per backend

pt_expt — safe, fails fast. DescrptDPA4.has_message_passing_across_ranks() returns self.bridging_switch is None (deepmd/dpmodel/descriptor/dpa4.py:2280-2294), so a bridged model freezes with has_comm_artifact=false (asserted in source/tests/infer/gen_dpa4_zbl.py:155) and the C++ inference path rejects multi-rank. Wrong answers are not reachable; multi-rank is simply unsupported.

pt — no guard. DescrptSeZM.forward accepts comm_dict (deepmd/pt/model/descriptor/sezm.py:1472) and then builds the edge cache with bridging_switch=self.bridging_switch unconditionally (sezm.py:1526). There is no bridging-vs-parallel rejection on that path, and comm_dict reaches the descriptor through the ordinary model plumbing (deepmd/pt/model/model/make_model.py:402). A multi-rank run of a bridged .pth would therefore compute a wrong gate at subdomain boundaries with no diagnostic.

Caveat on the pt claim: this is established by reading the code, not by a failing test. It is reachable by construction, but has not been demonstrated end-to-end.

Proposed fix: reverse-accumulate, then forward-broadcast

Three phases per forward pass:

1. Local partials. Scatter-sum log w and the zero-count into a per-node buffer spanning owned and ghost rows, so a ghost row accumulates the contributions of edges where that ghost is the src.
2. Reverse-accumulate. Send ghost rows to their owning rank and += them into the owner's row — the transpose of border_op's forward, walking the same LAMMPS swap schedule backwards. Log-products are additive and each edge lives on exactly one rank, so after this every owner row holds the complete global value with nothing double-counted.
3. Forward-broadcast. A plain border_op copies the completed owner values back onto ghost copies.

Then eta = exp(log_eta) (with the zero-count hard-freeze rule) is correct everywhere and gate_e = take(eta, src) is unchanged.

The two per-node scalars (log_eta, zero_count) pack into one (N, 2) tensor, so the cost is one reverse plus one forward exchange of a small tensor per step — negligible next to the per-block feature exchanges.

Much of the primitive already exists. deepmd_export::border_op_backward is registered as a standalone op with CPU and CUDA implementations (source/op/pt/comm.cc:650, 659-668), and its kernel is exactly the ghost-to-owner index_add_ reverse walk (comm.cc:322-386), including the ghost-row zeroing (comm.cc:400-402) that phase 3 then refills. The remaining work is to expose it as a first-class forward op that make_fx can trace into the .pt2, pair it with the transposed autograd (its backward is border_op's forward, so gradients from the gate cross ranks correctly), add the fake/meta registration and the Python fallback, and thread comm_dict into the edge-cache build — today comm_dict only reaches the interaction blocks, while the SFPG runs before block 0.

Rejected alternative: widening the halo so each rank sees every relevant node's outgoing edges. That needs a 2×rcut communication skin plus ghost-ghost edges, which LAMMPS neighbor lists do not give pair styles.

Scope

Task 2 — ZBL bridging multi-rank. The primitive above, wired at the edge-cache build in both dpmodel/pt_expt and pt (one owner per backend, same seam), lifting the has_message_passing_across_ranks exclusion, plus C++/LAMMPS enablement.

Task 3 — native spin + ZBL multi-rank. Expected to follow from Task 2 with wiring and verification only: native-spin multi-rank on the graph route is already verified, so the bridging gate is the only missing piece of the combination.

Task 4 — export-time questions must be atomic-model capabilities, not reach-through. (prerequisite of Task 2 — same seam, not a separate concern)

The general statement: the export/serialization layer asks questions the atomic-model interface does not expose, so it substitutes structural knowledge of a concrete model — either the wrapper's TYPE (isinstance) or its INTERNALS (.descriptor). Both assume "one atomic model == one descriptor, of a class I recognise", which every composition falsifies, and both fail SILENTLY: isinstance returning False and getattr(..., None) returning None each land on a plausible default instead of an error.

Model-time questions already do this correctly — uses_graph_lower(), supports_native_spin(), has_chg_spin_ebd(), get_dim_chg_spin() are capabilities on BaseAtomicModel that LinearEnergyAtomicModel aggregates (all / any / max as the semantics require). Export-time questions were never migrated.

Four capabilities need promoting to the atomic model with composition aggregation, after which every isinstance and every .descriptor lookup below becomes unreachable rather than merely fixed:

#	Question	Asked today at	How
1	needs cross-rank message passing	serialization.py:193 (isinstance(LinearEnergyAtomicModel)) and :200 (desc.has_message_passing_across_ranks())	already a capability, but on the DESCRIPTOR; aggregate: ANY child needs it, EVERY child supports edge-parallel
2	uses compact edge pairs (torch>=2.6 unbacked-SymInt guard)	serialization.py:248 (desc.uses_compact_edge_pairs())	a real capability asked of the wrong object
3	graph edge dtype	serialization.py:1027 — reads desc.se_atten.mean.dtype, THREE levels deep into a DPA1-specific internal	no capability exists; add one
4	supports graph export	serialization.py:1047 — reads desc.geo_compress and the PRIVATE desc._fused_eligible	no public capability exists at all

Item 1 is the one that blocks Task 2: _needs_with_comm_artifact denies a composition through BOTH the isinstance check and the desc is None -> return False fallthrough immediately after it, so deleting only the former changes nothing, and fixing the SFPG gate alone would stay invisible at export. Reported and reproduced by @OutisLi in the #5884 review (at 2f4e63dd8) with a two-DPA2 linear_ener model that is graph-eligible, whose children both report has_message_passing_across_ranks() == True, and which still gets no with-comm artifact.

Items 2-4 are latent rather than live: for DPA4 every default they fall back on happens to be correct (DPA4 has no geo_compress, and uses_compact_edge_pairs() is False), so bridged DPA4 works today. A bridged geo_compress DPA1/DPA2 would silently get fp64 edges and no torch-version guard — a broken artifact, not an error.

Fixing these one function at a time invites the next helper to reach through .descriptor again, which is how this accumulated; the interface completion is the actual fix.

Regressions: a two-DPA2 linear_ener asserting the with-comm artifact is requested (unblocked by item 1 alone), and a bridged composition asserting it is not, until Task 2 lands.

Also stale and worth correcting in the same pass: the comment at serialization.py:248 says models without a single descriptor "take the dense route anyway". Since 881e2087a, graph-capable models — bridged ones included — are auto-resolved ONTO the graph route, so the comment now documents the opposite of the behaviour.

Deferred out of #5884 rather than fixed there.

Interim hardening (independent of the above). pt's parallel path should get the same rejection pt_expt already has — a comm_dict is not None and self.bridging_switch is not None guard in DescrptSeZM.forward — so the silent-wrong-answer window closes now rather than at fix time.

Verification

The regression test both backends need is a cross-boundary close pair: a fixture with a sub-r_outer contact positioned so that a processors 2 1 1 decomposition splits it, compared against the 1-rank result. Without that geometry the test is vacuous.

The bisection ladder that resolved the native-spin multi-rank bug in #5884 applies here too, each level exonerated before descending: eager self-comm parity (single-rank comm_dict self-send, no MPI runtime) → make_fx trace → locally compiled AOTI artifact → C++ single-rank forced through the with-comm branch → per-rank tensor dumps in a real 2-rank run.

Out of scope for #5884

PR #5884 delivers the graph-native DPA4 port, multi-rank for non-bridged models, native spin (single and multi-rank), and charge-spin/bridging on the graph route for single-rank. The bridged multi-rank cases above are tracked here and deliberately left to a follow-up.

[wanghan-iapcm]

PR #5939 implements this issue end-to-end:

• Task 4 (prerequisite) — export-time questions are atomic-model capabilities with LinearEnergyAtomicModel any/all aggregation; the four serialization.py helpers consume them (no isinstance / .descriptor reach-throughs). The two-DPA2 linear regression now gets its with-comm artifact.
• Task 2 (ZBL bridging multi-rank) — the SFPG per-node [log_eta, zero_count] partials are completed across ranks by one reverse-accumulate (border_op_backward, which gains autograd: its gradient is border_op's forward, the two are exact transposes) + one forward-broadcast (border_op) before the gate is applied. No new C++ — one extra call site of the existing kernels. Verified with cross-boundary close-pair geometry at every rung: eager self-comm 1e-12, identity-exchange ablation diverges, LAMMPS 2-rank vs 1-rank parity on T4.
• Task 3 (native-spin + ZBL) — no spin-side production change was needed (the exchange sits below the spin wrapper); verified with force_mag parity, freeze-embed asserts, and the variant's first-ever LAMMPS file incl. a 2-rank close-pair run through pair_style deepspin.
• Interim hardening — superseded by the single-PR structure; instead of a temporary guard, the pt fix itself lands, and its red parity run demonstrated the "pt has no guard" claim (plus a crash: the pt ZBL injection indexed per-local types with extended ghost src indices).

One finding outside this issue's scope, left for a separate fix: source/api_c/include/deepmd.hpp uses &vec[0] on possibly-empty vectors (~33 sites) — UB that SIGABRTs under _GLIBCXX_ASSERTIONS before the empty-rank guard's message can fire.