Repository navigation
SFPG: add a carry-all/reciprocal graph fast path (dst reduction + one broadcast) for the pt_expt graph lower #5957
Description
Follow-up from review of #5939 (thread on deepmd/pt/model/descriptor/sezm_nn/edge_cache.py).
Current implementation
The Source Freeze Propagation Gate folds each node's full outgoing-edge set, eta_j = prod_{e: src_e = j} w(r_e). It reduces [log_eta, zero_count] over src. Under domain decomposition a rank only holds edges whose destination it owns, so those per-node partials are rank-incomplete and are completed by a two-phase exchange in _gate_partial_exchange:
border_op_backward— reverse-accumulate ghost rows into their owners;border_op— forward-broadcast the completed owner values back onto ghosts.
Why source-keyed reduction is required in general
Destination reduction equals source reduction only when the edge set is reciprocal. That does not hold for the sel-capped, nlist-derived edge schema (the PT compatibility path): format_nlist truncates each centre at sel by distance, so a dense-environment atom can drop a neighbour that the sparser neighbour keeps. Measured on an inhomogeneous system (rcut=4.0, sel=6, ghosts mapped 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 the generic source-keyed semantics must stay for that route.
The fast path this issue tracks
The pt_expt graph-lower route is different, and it is the primary route going forward:
- it uses the carry-all, sel-free
NeighborGraph(every edge insidercut), so the edge set IS reciprocal; - DPA4's dense
call()rejectscomm_dictbefore_run_graph, so the communication hook is reached throughcall_graph, never through a sel-truncated dense graph; - therefore the owner holds its complete destination-centred neighbourhood.
On that route the partials can be reduced over dst — already complete for owned nodes — and completed with one forward border_op instead of a reverse-accumulate plus a broadcast. That halves the exchanges, and with them the MPI_Barrier count. It matters more than the payload size suggests: both kernels end in a barrier, and because the exported force graph differentiates through them, their transposes add the matching pair to the force evaluation too.
Requirements
- Add it as an explicit, clearly-labelled carry-all/reciprocal fast path. Do not change the generic source-keyed semantics used by the nlist-derived route.
- State the reciprocity contract and the routing boundary at the seam (which graphs qualify, and what selects the path), so the assumption is pinned rather than implied.
- A negative test that the sel-capped/nlist-derived route still takes the two-phase completion, so the fast path cannot silently capture it.
- Parity: fast path vs current two-phase completion, bit-for-bit on a carry-all graph, including a bridging-zone pair straddling a rank boundary (the geometry used in
test_dpa4_zbl_parallel.py/test_lammps_dpa4_zbl_pt2.py). - Benchmark across representative system sizes and MPI rank counts. The docs currently make no performance claim about this exchange (feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank #5939 removed the unsupported "negligible" wording); numbers from this work are what should replace it.
Context
- feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank #5939 — introduced the SFPG cross-rank completion; that PR keeps the general two-phase form on both routes.
- DPA4/SeZM zone-bridging (SFPG) gate is incomplete under domain decomposition; blocks ZBL and spin+ZBL multi-rank #5906 — the original rank-incompleteness issue.
Activity