feat(dpmodel,pt_expt,cc): graph-native DPA-2 — first message-passing model on the NeighborGraph lower (PR-G-dpa2)#5779
Open
wanghan-iapcm wants to merge 80 commits into
Open
Conversation
added 25 commits
July 12, 2026 14:18
…erLayer.call_graph parity Add parametrized test case 'no_chnnl_2_no_gg1' to test_repformer_layer_call_graph_parity that exercises two previously untested branches: 1. update_chnnl_2=False: Tests the short-circuit path where g2/h2 updates are skipped and returned unchanged (lines 2397-2401 in repformers.py call_graph). This is the critical production path for the LAST layer of every DPA2 model. 2. cal_gg1=False (gg1=None): Tests the graph path when no gg1-dependent updates are needed (all of update_g1_has_drrd, update_g1_has_conv, update_g1_has_attn, update_g2_has_g1g1 are False). With g1_out_mlp=False, g1_mlp is seeded with the identity [g1] for xp.concat. Test passes all 30 cases (29 baseline + 1 new).
gen_dpa2.py section B builds a graph-eligible dpa2 (use_three_body=False; three-body is graph-ineligible), exports the graph-form .pt2 (which embeds the nested with-comm artifact automatically for message-passing models) plus an independent dense-nlist .pt2 of the same weights, cross-checks atomic energies/forces/total virial between the two at gen time (NaN-safe), and writes deeppot_dpa2_graph.expected from the graph artifact eval. Unlike the dpa1 precedent the .expected is not copied from the nlist oracle: for a message-passing model the graph path's full-to-source per-edge atomic-virial decomposition legitimately differs from the dense per-atom decomposition (only the sum agrees), and cross-path e/f noise (~1e-10) sits at the universal gtest's 1e-10 double tolerance. The dpa2_graph_ptexpt VariantDeepPotCase row mirrors dpa1_graph_ptexpt (same tolerances and capability flags).
- Add nall_real == 0 guard at the top of graph with-comm arm (lines 922-929) to throw a clear error instead of silently crashing with Dim violation - Fix stale comment (lines 1050-1051) to reflect that ghost forces fold back via both plain and with-comm graph routes
…h-comm route The with-comm GRAPH artifact derives n_local from the nlocal comm tensor IN-GRAPH (owned-node energy mask); after move_to_device_pass that derivation is a device kernel, so feeding it a CPU tensor makes the kernel read a host pointer as device memory -> CUDA illegal memory access on every multi-rank run (first live exercise of the route, T4). The other six comm tensors are consumed only by the opaque border_op whose host code dereferences their data_ptr, so they stay on CPU -- same placement the dense with-comm route uses for all eight. Verified on Tesla T4: minimal AOTI repro crashes with CPU nlocal and matches the plain graph artifact bit-for-bit with device nlocal.
… DEVICE
_trace_and_compile_graph built its synthetic NeighborGraph trace inputs on
the module-level DEVICE constant instead of the device the model's own
parameters/buffers actually live on. In real training these always match
(the trainer moves the model to DEVICE before compiling), but any caller
that intentionally holds the model on a different device (e.g. a test
pinning the model to CPU while running on a CUDA host, mirroring the .pt2
export's model.to("cpu") pattern) hits a device split between the traced
graph tensors and the model params.
DPA2.call_graph's tebd gather (dpmodel/descriptor/dpa2.py:1159) surfaced
this first: xp.take(tebd_table, atype_local, axis=0) requires tebd_table
(a model param, left un-asarray'd to avoid detaching its gradient -- the
documented dpa1 lesson) and atype_local (moved to device(graph.edge_vec))
to share a device. Under test_compiled_training_graph_smoke the model is
pinned to CPU but the trace sample was still built on the global DEVICE
(cuda:0 on a GPU box), so tebd_table stayed CPU while atype_local was CUDA.
DPA1's call_graph has the identical gather pattern and would fail the same
way, but no existing dpa1 test calls _trace_and_compile_graph directly with
a device-pinned model, so the bug was latent there since PR-B (deepmodeling#5604).
Fix: add _model_trace_device(), reading the device off the model's own
parameters/buffers (falling back to DEVICE only if the model exposes
neither), and use it to build the trace sample instead of the global
DEVICE. No change to the gather itself -- gradient flow through
type_embedding is untouched.
…back test dpa2 became graph-eligible (uses_graph_lower=True), so neighbor_list=None now routes to the carry-all graph while explicit DefaultNeighborList() forces the dense route. The smooth repformer attention diverges intentionally between these paths (sel-independent on graph, sel-padded on dense), amplified through message passing. Observed divergence at this fixture: - energy rel ~1e-4, abs ~9.5e-4 - force abs ~5e-3 - virial abs ~1.3e-2 Set atol=2e-2, rtol=1e-3 to accommodate dpa2's larger divergence compared to dpa1_smooth/se_atten_v2.
… into feat-graph-dpa2
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
OutisLi
pushed a commit
to OutisLi/dpmd-public
that referenced
this pull request
Jul 16, 2026
…adapter (deepmodeling#5785) ## Problem `DescrptDPA1.call` routes graph-eligible configs through `_call_graph_adapter` (decision deepmodeling#14: dense call = thin adapter). That adapter is bit-exact vs `_call_dense` **only in the trivial-statistics regime** (`davg == 0`): the dense se_atten body leaks a phantom padding-neighbor `-davg/dstd` residual (`EnvMat.call` subtracts `davg` after the padding rows' geometry is weight-zeroed; with empty `exclude_types` nothing re-masks it, at **any** `attn_layer`) that the graph path deliberately omits. Measured on the consistency fixture: dense vs adapter agree at 0.0 for `davg=0`, diverge by ~20 (abs) for nonzero `davg`, at both `attn_layer=0` and `2`. The gate also made `mapping` — an argument that only enables ghost folding on graph routes — silently change the numerics of a dense call: `call(..., mapping)` and `call(...)` returned different descriptors for the same physical system. ## Why it never surfaced The pt/pd consistency tests invoke `dd2.call()` **without** `mapping`, so the routing gate (`mapping is not None or nall == nloc`) is off and the tests exercise the dense body — a code path production never takes: `DPAtomicModel.forward_atomic` always forwards `mapping`. Consequence: dpmodel/jax/tf2 inference of a **trained** dpa1 with `set_davg_zero=False` (the constructor default) silently used the adapter and differed from the pt/tf dense backends. Cross-backend suites stayed green only because fresh models sit in the trivial-stat regime. This is the dpa1 twin of the dpa2 routing fix in deepmodeling#5779 (wanghan-iapcm/deepmd-kit@bc0d70ade, where the same class of divergence was caught by CI because the dpa2 consistency test does pass `mapping`). ## Fix - `DescrptDPA1.call` always runs `_call_dense` — it is the cross-backend consistency reference. The graph-native route is reached exclusively through `call_graph` (pt_expt `forward_atomic_graph` and the graph `.pt2`); nothing in shipping code consumed the `call` flip. - `_call_graph_adapter` is retained as the bit-exact-regime reference; its docstring now states the actual regime (including the slot-0 statistics simplification, exact for real slot-uniform stat tables). ## Tests - pt/pd `test_dpa1.py::test_consistency` now pass `mapping` (the production invocation). Red on the old routing (71%-scale mismatch), green with the fix. - New regression: `call(mapping) == call(None) == _call_dense` at zero tolerance with nonzero injected `davg` and ghosts present. - Adapter parity tests (`test_dpa1_call_graph_descriptor.py`, `test_dpa1_graph_attention_parity.py`) exercise `_call_graph_adapter` directly instead of relying on the removed `call` routing — adapter coverage is unchanged. - Local sweep: dpa1/se_atten_v2 graph+consistency+universal battery — 94 + 304 + 852 passed. ## Notes - No conflict expected with deepmodeling#5779 (different files); the two PRs are independent. - pt_expt training/export graph routes are untouched (`uses_graph_lower` still gates them; it no longer influences the dense `call`). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Ensured descriptor calculations consistently use the dense path, regardless of mapping or graph-related conditions. * Preserved graph-based processing through its dedicated pathway. * Prevented mapping settings from changing dense numerical results. * **Tests** * Expanded parity and consistency coverage across dense and graph-based calculations. * Added regression checks for mapping, ghost atoms, exclusions, and single-rank scenarios. * Improved validation of bit-exact results and production-path behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
added 2 commits
July 16, 2026 17:20
The CI gen_dpa2 graph export (repformer attention ON) failed with GuardOnDataDependentSymNode: the pair count entering the slot-occupancy softmax is a data-dependent (unbacked) symbol, and (a) the Python 'if n == 0' early-out guarded Eq(u1, 0) at trace time, (b) the gradient-carrying cumsum over that axis guarded 'numel <= 1' in its backward. Drop the early-out (every op is empty-safe) and prepend two zero pads to the cumsum operand so the backward guard is statically false for any entry count (bitwise-identical prefix sums). Regression: test_graph_lower_symbolic_trace_with_attention mirrors the existing symbolic-trace test with update_g1/g2_has_attn=True (the CI export config); verified to reproduce both guards without the fix. Note: test_graph_with_comm_export dpa2 cases fail locally with a CppVecKernel atomic_add AssertionError in inductor codegen -- verified pre-existing on the parent commit 1a01380 (the known local-venv AVX2 inductor issue; CI was green there).
- test_dpa1_triton.py: import the dpa1 module via a single 'from' style (py/import-and-import-from). - test_graph_with_comm_export.py: replace the documentary 'del sendlist_indices' with an assert on the live buffer address -- a real use that both keeps the ctypes-referenced buffer alive through the compiled calls and satisfies py/unnecessary-delete.
OutisLi
requested changes
Jul 16, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
The soft slot-occupancy update fixes the forward jump, but the new denominator is only C0 and introduces a reproducible first-derivative discontinuity in LocalAtten. See the inline P1 finding.
added 3 commits
July 16, 2026 21:48
…nbacked_bindings)
The pad-and-slice workaround for the cumsum backward guard traded one
export failure for another: AOTI lowering of the [2:] slice on the
unbacked pair axis raised InductorError: KeyError: 'unbacked_bindings'
in CI's gen_dpa2 (repformer attention ON).
Rework: the cumsum-derived per-rank suffix T_k now feeds ONLY the cut
feasibility comparison (no backward node is ever built for it, so its
backward numel<=1 guard never traces), and the differentiable water
mass T_{k*} is rebuilt as segment_sum(sw * (rank > kstar)) -- an
index_add over the unsaturated entries with mathematically identical
value and gradient (d lam / d sw_j = unsaturated_j / (cap - kstar)).
No pad, no slice, no gradient-carrying scan on the unbacked axis; the
kstar == 0 special case collapses into the same expression.
Verified: all 59 segment/repformer regressions (water-filling bisection
reference, f32 torch/jax backward, cutoff-crossing continuity) and the
attention symbolic-trace regression pass; the full dpa2+attention graph
.pt2 AOTI export (primary + nested with-comm artifact) completes
locally with the venv's simdlen=1 vectorizer workaround.
test_graph_with_comm_export.py sat in tests/pt_expt/utils/ while its dense counterpart test_export_with_comm.py lives in tests/pt_expt/model/ alongside test_graph_export.py; the graph/dense pairs now mirror each other: model/test_export_with_comm.py <-> model/test_graph_export_with_comm.py. Audit of the other test files added by this PR (test_dpa2_call_graph.py, test_repformer_graph_ops.py in common/dpmodel; test_dpa2_graph_lower.py in pt_expt/model) confirms they already match their dpa1 siblings' placement. No content changes; imports are absolute.
…ow-phantom surface) OutisLi round 6: the relu bracket made the round-5 denominator only C0 -- at e == P the slope w.r.t. the logit jumped theta*P -> P, a force discontinuity on a surface reachable with finite q.k (measured through LocalAtten: d(out)/d(logit) -0.00916 vs -0.15365), and the min-based water-filling had active-set kinks of the same class. Two C1 replacements, keeping every earlier property: - Bracket: B = theta*e + (1-theta)*(e-P)*chi with chi = (e/P)**(1/theta) below P (else 1). The damped tail meets the in-design branch at e == P with matching value AND slope; B = e*[theta + (1-theta)(a-1)a**(1/theta)] >= theta*(1-1/e0)*e stays strictly positive with weights bounded by ~1.6/theta, and theta == 1 yields BITWISE e with no special case (the (1-theta)=0 factor kills the finite tail exactly). - Occupancy: theta = h(w*u) with the plateau saturator h(t)=t(2-t) capped at 1; h'(1)=0 makes both the saturation boundary and every active-set transition C1 automatically. The water equation is quadratic per cut; the cancellation-free root (cap-k)/(S+sqrt(S^2-Q(cap-k))) is built from masked segment_sums on the gradient path (per-rank suffix cumsums stay comparison-only -- the torch.export constraint from the previous fix). Deliberate limit, documented: C1 is the ceiling -- exact in-design linearity (bitwise dense parity) and smoothness beyond C1 at e == P are mutually exclusive; second derivatives still jump there. Tests: left/right FD slope-match regressions at the below-phantom surface (primitive, LocalAtten with the reviewer's exact construction pinning his -0.153650933331561 in-design slope, Atten2Map with a negative-definite pairing so the crossing pair stays visible), FD derivative-step-scaling across a water-filling active-set change, bisection reference updated to the h-form; all round 2-5 regressions (bitwise dense, telescoping, positivity, cutoff continuity, f32 torch/jax backward) retained. Full local sweep: 845 dpmodel, 21 graph lower (incl. attention symbolic trace), 272 dpa2, and the 4 AOTI with-comm exports (simdlen=1 env workaround).
OutisLi
requested changes
Jul 16, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
The new bracket removes the original analytic C1 kink, but its inactive tail branch now produces NaN first derivatives for finite, supported float32 attention inputs. This affects graph forces and training gradients and needs to be fixed before merge.
OutisLi
requested changes
Jul 16, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
The new quadratic occupancy solver has a separate float32 forward-correctness failure: cancellation in its suffix moments and discriminant can select the wrong active set by a large margin for valid cutoff weights.
added 2 commits
July 16, 2026 23:56
… tail OutisLi round 7: the C1 tail's safe-where substituted ph_e on BOTH sides of the ratio, so the inactive branch still evaluated log(ph_e / ph_e). For in-design HIGH logits (count >= 0, no negative phantom needed) ph_e = exp(phantom_logit - segment_max) goes subnormal in float32 and the unselected branch's backward overflows 1/ph_e (torch, logits ~68) or ph_e**2 (jax, logits ~24); 0 * inf then poisons the selected gradient with NaN while the forward stays exactly dense. chi is now computed entirely in log space: log(e/P) is the already available shifted - ph_arg, so no log() and no division by ph_e enter the autodiff graph at all; the log-ratio is clamped at -1e4 (exp == 0 in both precisions) so masked entries' -inf shifted logits cannot reach the division by theta either. Regressions (verified to fail on the parent commit): float32 torch and jax backward at the reviewer's high-logit thresholds against the analytic softmax jacobian, and his exact float32 LocalAtten.call_graph repro (sel == n_real == 2, raw logits [68, 69]) with graph gradients checked finite AND equal to the dense route's autograd.
Conflict: source/api_cc/src/common.cc -- both sides fixed the same daparam-stride aparam buffer sizing (deepmodeling#5790 upstream, e95f2ba here); code identical, comment text taken from upstream so the hunk drops out of this PR's diff. deepmd/dpmodel/descriptor/dpa1.py auto-merged with master's deepmodeling#5785 (dense call is unconditionally _call_dense; verified post-merge).
OutisLi
reviewed
Jul 16, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
One non-blocking correction to the explanatory algebra in the new C1 bracket implementation.
… positivity proof OutisLi round 8, two items: - [P1] The active-set selection formed suffix moments as total - prefix and the discriminant as S^2 - Q*room -- ill-conditioned subtractions that collapse in float32 once cutoff weights span ordinary decades (sw = [1, 0.1, 0.01, 5e-7], capacity 3: the valid saturated cut was rejected and sum(theta) came out 1.52 instead of 3.0 -- a ~23% forward error through LocalAtten). The selection scan now runs in float64 internally: it feeds ONLY comparisons (kstar), so the upcast costs no gradient and no export surface, while the differentiable S/Q at the chosen cut stay in the compute dtype (positive-only masked sums, no cancellation). At the disc == 0 tangency the compute-dtype root degrades gracefully (u -> cap_rem / S equals the exact double root). - The positivity comment factored the tail incorrectly: (e - P) a**(1/theta) = e (a - 1) a**(1/theta - 1), so the minimized factor is (a - 1) a**(1/theta - 1) with minimum -theta (1 - theta)**(1/theta - 1) at a = 1 - theta, giving B/e >= theta (1 - (1 - theta)**(1/theta)) >= theta (1 - 1/e0). Same final bound, corrected derivation (documentation only). Regressions (first three verified to fail on the parent commit): the reviewer's float32 repro against a float64 bisection reference, a float32-vs-float64 segment_softmax forward at sel=3 with his weights, a LocalAtten.call_graph float32 forward at the same shape, and a broader log-spaced float32 weight sweep against bisection.
OutisLi
approved these changes
Jul 17, 2026
![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: 6
🤖 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/pt_expt/model/ener_model.py`:
- Around line 673-679: Update the traced module input signature in the Returns
docstring to match the make_fx forward ABI: place n_local in slot 3 and include
destination_order, destination_row_ptr, source_order, and source_row_ptr in
their actual order, removing the incorrect trailing n_local placement.
In `@deepmd/pt_expt/model/make_model.py`:
- Around line 580-585: Remove the trailing duplicate n_local documentation block
after charge_spin in the affected function docstring, keeping the existing
n_local entry at its correct signature position and leaving the surrounding
parameter documentation unchanged.
In `@deepmd/pt/utils/compile_compat.py`:
- Around line 190-196: Update the getter loop around model.get_dim_fparam and
model.get_dim_aparam so attribute lookup occurs inside the existing try block
via getattr, allowing missing attributes to be handled by the best-effort
exception path without interrupting execution.
In `@source/api_cc/include/DeepPotPTExpt.h`:
- Around line 635-637: Update the positional AOTI input-order documentation in
DeepPotPTExpt to include n_local, both edge permutations, and both CSR pointers,
matching the declared signature and run_model_graph ABI while preserving the
existing ordering of the other inputs.
In `@source/op/pt/comm.cc`:
- Around line 108-114: Move the CUDA synchronization in the communication flow
so it executes after the `index_select`/`index_add_` pack operations have
populated the MPI send buffers, on both CUDA-aware paths corresponding to the
current fences. Ensure CUDA-aware `MPI_Send` cannot access device buffers before
packing completes, while preserving the existing non-CUDA-aware ordering.
In `@source/tests/infer/gen_dpa2.py`:
- Around line 239-248: Update the LSAN early-return branch in the DPA2
generation flow to delete all Section B and Section C graph output artifacts
before returning. Reuse the existing output path and cleanup mechanisms,
ensuring stale archives cannot remain available to skip_if_artifact_missing
while preserving the current LSAN skip behavior.
🪄 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
Run ID: c6505564-49df-4943-ae33-07f92bb68b4e
📒 Files selected for processing (37)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/dpa2.pydeepmd/dpmodel/descriptor/repformers.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/model/edge_transform_output.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/utils/neighbor_graph/segment.pydeepmd/pt/utils/compile_compat.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/repformers.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/edge_transform_output.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/graph_lower.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/serialization.pydoc/model/dpa2.mdsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/tests/test_deeppot_dpa1_graph_aparam_ptexpt.ccsource/install/test_cc_local.shsource/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.pysource/lmp/tests/test_lammps_dpa2_graph_pt2.pysource/op/pt/comm.ccsource/tests/common/dpmodel/test_dpa2_call_graph.pysource/tests/common/dpmodel/test_neighbor_graph.pysource/tests/common/dpmodel/test_repformer_graph_ops.pysource/tests/common/dpmodel/test_segment_softmax.pysource/tests/infer/deeppot_dpa3_spin.yamlsource/tests/infer/deeppot_dpa3_spin_mpi.yamlsource/tests/infer/gen_dpa1.pysource/tests/infer/gen_dpa2.py
🚧 Files skipped from review as they are similar to previous changes (10)
- deepmd/pt_expt/descriptor/repformers.py
- source/install/test_cc_local.sh
- deepmd/dpmodel/model/edge_transform_output.py
- source/tests/common/dpmodel/test_neighbor_graph.py
- doc/model/dpa2.md
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- source/api_cc/include/commonPT.h
- deepmd/dpmodel/model/make_model.py
- deepmd/pt_expt/model/edge_transform_output.py
- deepmd/dpmodel/descriptor/repformers.py
added 2 commits
July 18, 2026 13:25
…tale fixtures, accessor guard, ABI docs) - comm.cc: the pre-loop gpuDeviceSynchronize only orders work launched BEFORE it; the per-swap index_select pack (forward) and the chained reverse-communication index_add_ (backward) launch afterwards, so CUDA-aware MPI_Send could read stale device buffers. Fence inside the loop before each cross-rank send on both paths (gated on cuda_aware and a CUDA buffer; self-send memcpy is already stream-ordered). - gen_dpa2.py: the LSAN early-return skipped REgeneration but left any graph artifacts from a previous non-LSAN run of a reused workspace in place, so skip_if_artifact_missing would execute them under LSAN and hit the crash the skip exists to avoid; delete the four Section B/C outputs before returning. - compile_compat.py: resolve get_dim_fparam/get_dim_aparam via getattr INSIDE the try -- building the accessor tuple eagerly raised AttributeError before the best-effort guard could catch a model without the accessors. Both branches pinned by a new unit test. - DeepPotPTExpt.h: the with-comm run_model doc omitted n_local and the four CSR permutation/pointer inputs from the positional AOTI order; now matches the 13-input graph ABI. - make_model.py: merge the duplicated (out-of-order) n_local docstring entry into its signature-position entry, keeping the detailed text. - ener_model.py: the with-comm traced-forward Returns docstring listed the inputs in the wrong order (n_local appended at the end, CSR tensors missing); now mirrors the actual make_fx signature. comm.cc compile-checked via the deepmd_op_pt target; runtime MPI behavior is CI/LAMMPS-validated (multi-rank CUDA-aware path has no local coverage).
… into feat-graph-dpa2
wanghan-iapcm
enabled auto-merge
July 18, 2026 05:40
Member
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 18, 2026
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
6 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.
DPA-2 becomes graph-native full-stack — the first message-passing model on the NeighborGraph lower (PR-G-dpa2 of the NeighborGraph series; follows dpa1's route from #5581/#5583/#5604/#5714/#5715/#5717/#5733). The dense path coexists untouched; the graph route is opt-in for inference (
--lower-kind graph) and the default for pt_expt training/eager on graph-eligible configs, exactly like dpa1.What's new
dpmodel (backend-agnostic math)
DescrptBlockRepformers.call_graph: every repformer op ported to the flat edge list — conv/grrg/drrd/g1g1/LocalAttenviasegment_*overdst; the densennei×nneigated attention (Atten2Map/Atten2MultiHeadApply/Atten2EquiVarApply) viacenter_edge_pairs(ordered, include_self)+ per-headsegment_softmax(extended to trailing feature dims).DescrptDPA2.call_graph+uses_graph_lower+ dense-call adapter. The multi-level nlist (build_multiple_neighbor_list) becomes per-block edge masks; in the shape-static adapter layout the per-block slice replicates the densenlist[:, :, :ns]truncation, making the adapter bit-exact vs the dense path at ANY sel (verified to 5e-16 at deliberately binding sel with attention enabled) — this is what keeps the cross-backend consistency suites green after the routing flip. Ineligible configs (use_three_body, compression, spin) keep the legacy dense route.fit_output_to_model_output_graphconsumesn_local, excluding halo rows from the differentiated energy (prevents cross-rank double counting; force stays full-N, halo partials reverse-commed by LAMMPS).pt_expt / export
_exchange_ghosts_graph— identity on ghost-free single-rank graphs (srcIS the owner),deepmd_export::border_opoverwrite of halo rows on extended-region multi-rank graphs..pt2archives embed a with-comm AOTInductor artifact (model/extra/forward_lower_with_comm.pt2), traced with the 8-tensor comm ABI shared with the dense flow.C++ / LAMMPS
DeepPotPTExpt::run_model_graph_with_comm+ dispatch replacing the PR-G fail-fast: MP message-passing graph models run multi-rank on the extended-region graph + with-comm artifact; non-MP (dpa1) multi-rank unchanged; old graph archives without the artifact get a clear re-freeze error.gen_dpa2.pysection B),dpa2_graph_ptexptuniversal gtest row, andtest_lammps_dpa2_graph_pt2.py(single-rank vs reference, per-atom virial,mpirun -n 2≡-n 1, graph-vs-nlist cross-artifact, bounded empty-rank).Validation
dpa2_graphC++ gtest variant passed; python GPU suites green.nlocal/nghostcomm scalars (the graph route consumes them in on-device owned-mask index math);_trace_and_compile_graphhardcoded the global device for trace samples (latent since feat(pt_expt): graph-native dpa1 — .pt2 export + compiled training + C++ inference single & multi-rank (NeighborGraph PR-B) #5604).Known limitations / deliberate divergences (documented in
doc/model/dpa2.md+ code Notes)KNOWN_GRAPH_DENSE_DIVERGENT).border_opcollectives, so the job stalls until MPI timeout — phantom-node support is a follow-up. The LAMMPS test bounds this with a hard timeout and asserts it never silently succeeds.se_attenatattn_layer=0leaks a deterministic-davg/dstdpadding residual whenset_davg_zero=Falseandexclude_types == [](PairExcludeMaskall-ones short-circuit is the only padding mask on that path). The graph path masks padding correctly, so graph and dense deliberately differ in that regime (pinned by test + docstring). A dense-side fix will be proposed separately.Follow-ups (separate issues/PRs)
se_attenpadding-residual fix (above).BUILD_PT_EXPTC++ gtest suite into CI (pre-existing gap — the universal gtest rows, including the new one, never run in CI today).TestDeepPotPTExptWithCommLoadFailure.multi_rank_compute_throwsis vacuous (setsnswapbut notnprocs; pre-existing since fix(pt_expt): fail-fast on .pt2 GNN inference without LAMMPS atom-map #5450).PR-G-dpa3 (repflows + angle channel +
charge_spin, reusing this PR's MP machinery) comes next.Summary by CodeRabbit
call_graph) for eligible energy models, including multi-rank “with-comm” ghost exchange.comm_dict) and persistent graph-lower enable/disable controls.n_localfor multi-rank reductions and exported graph artifacts (including with-comm variants).aparamshapes.