feat(pt_expt): dpa1(attn_layer=0) graph-native NeighborGraph forward#5583
Draft
wanghan-iapcm wants to merge 44 commits into
Draft
feat(pt_expt): dpa1(attn_layer=0) graph-native NeighborGraph forward#5583wanghan-iapcm wants to merge 44 commits into
wanghan-iapcm wants to merge 44 commits into
Conversation
added 26 commits
June 25, 2026 17:26
…_graph The dense path masks excluded type pairs; the graph path does not yet, so raise NotImplementedError instead of silently diverging.
…pter (attn_layer=0)
…k kernel; sw to dense adapter
…ph_method (Option B)
…gy_deriv) + parity
…d/pt_expt/model/)
serialize roundtrip + dpmodel->pt_expt interop on the attn_layer=0 graph path are already covered by test_dpa1.py::test_consistency (lines 86-113), which routes through the graph forward via the Task-3 dense-call adapter.
…all back to dense Task 3's adapter routed ALL attn_layer==0 through the graph, but the graph only supports tebd_input_mode='concat', no exclude_types, and needs mapping for ghosts. strip-mode / exclude / mapping-None-with-ghosts attn_layer=0 models raised/IndexError'd. uses_graph_lower() now encodes full eligibility and ineligible configs fall back to the legacy dense body unchanged. Fixes test_compressed_forward (attn_layer=0 strip).
…ping-ghosts) dense fallback
…pt graph mask key; legacy opt-out in Option-B test - _resolve_graph_method/_call_common_graph use getattr(atomic_model,'descriptor',None) so Linear/ZBL models (no descriptor) fall back to dense instead of AttributeError - pt_expt _call_common_graph override adds the all-ones mask key for dense parity - test_dpa1_graph_model_energy dense refs use neighbor_graph_method='legacy' to opt out of the now-default carry-all graph (decision deepmodeling#17 default-flip)
…nse default dpmodel/jax compute force/virial analytically inside call_common (energy_derv_r); the energy-only graph lower drops it -> KeyError when force is requested. Only pt_expt has the autograd graph force/virial path, so only pt_expt defaults eligible models to the graph. dpmodel base _resolve_graph_method no longer auto-routes; pt_expt overrides it to re-enable AUTO.
…nse call jit/export-traceable (decision deepmodeling#16)
…x int-sum, Array typing) - swap dangling memory/spec_unified_edge_nlist.md refs -> public design discussion (#4) so the references resolve - edge_force_virial: short-circuit n_out=int(node_capacity) when supplied so the static jax/export path never calls int() on a traced sum(n_node) - derivatives.py: move Array import under TYPE_CHECKING (+ from __future__ import annotations) for subpackage uniformity
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 Run ID: 📒 Files selected for processing (7)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe PR adds graph-based neighbor construction, DPA1 graph lowering, graph-routed energy execution in the array API and PyTorch model paths, and tests that compare the graph flow against dense references. ChangesGraph-native execution
Estimated code review effort🎯 5 (Critical) | ⏱️ ~90+ minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
added 2 commits
June 25, 2026 18:00
…rop for-loops + subTest)
![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: 7
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_from_ijs.py (1)
55-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd one torch-backed ASE smoke test alongside these parity checks.
These tests only cover NumPy inputs, but
build_neighbor_graph_aseis also used from the torchpt_exptgraph path. A minimal CPU torch case here would lock in that contract and catch host-conversion regressions in the ASE builder much earlier.🤖 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 `@source/tests/common/dpmodel/test_from_ijs.py` around lines 55 - 86, Add a small CPU torch smoke test for build_neighbor_graph_ase alongside the existing NumPy parity checks, since the function is also exercised through the pt_expt path. Reuse the same test module and symbols (build_neighbor_graph_ase, build_neighbor_graph) but pass torch tensors for coord, atype, and box, then assert the neighbor set matches the in-tree reference to catch host-conversion regressions. Keep the new case minimal and focused on verifying torch-backed inputs work end to end.
🤖 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/make_model.py`:
- Around line 406-423: The graph path in make_model.py is treating padded
virtual atoms as real by passing the full flattened atype into
call_lower_graph() and setting mask to all ones. Update the lower-graph route in
the same model methods to either fall back to the legacy/dense path whenever any
local type is negative, or propagate a real-atom mask through the graph lowering
so atom_energy and energy are only exposed for valid atoms. Be sure the fix is
applied consistently in the call_lower_graph handling and the related graph-path
code referenced by the same symbols.
- Around line 318-347: Graph routing in model prediction is missing
`charge_spin`, so `self._call_common_graph()` can silently change behavior for
charge-spin models while the legacy `model_call_from_call_lower()` path still
receives `charge_spin=cs`. Update the routing in `make_model.py` around
`graph_method` / `neighbor_graph_method` to detect when `cs` is present and
force those cases through the legacy dense-nlist path instead of calling
`_call_common_graph()` or `call_lower_graph()` until the graph lower layer can
accept `charge_spin`. Apply the same fallback in the other affected prediction
block referenced by the comment so both paths stay behaviorally consistent.
In `@deepmd/dpmodel/utils/neighbor_graph/ase_builder.py`:
- Around line 86-90: The ASE neighbor-graph path in ase_builder.py is passing
backend tensors straight into np.asarray() for coord and box, which breaks when
make_model.py provides raw PyTorch tensors. Update the ASE builder flow around
the coord_np/box_np construction to explicitly detach tensors, move them to CPU,
and then convert to NumPy before reshaping, so neighbor_graph_method="ase" works
for CUDA and autograd-backed inputs.
In `@deepmd/dpmodel/utils/neighbor_graph/env.py`:
- Around line 82-87: The `EnvMat.call` neighbor-graph path is overriding all
near-zero lengths with `1`, which changes `denom` and breaks parity with the
dense formula for real overlapping atoms when `protection > 0`. Update the
`safe_len` handling in `env.py` so only masked padding edges use an arbitrary
finite fallback, while real zero/near-zero edges still preserve `length +
protection` semantics in the `t0` computation.
In `@deepmd/dpmodel/utils/neighbor_graph/from_ijs.py`:
- Around line 97-100: Normalize box to the same backend/device as coord before
any reshaping in from_ijs so it cannot remain on a mismatched CPU/GPU device. In
the box-handling branch, coerce box with the same dev/xp path used for i, j, and
nframe_id before calling xp.reshape, then continue with xp.take and S conversion
as before. Keep the fix localized to the from_ijs flow and preserve the existing
backend inference from coord.
In `@deepmd/pt_expt/model/make_model.py`:
- Around line 283-350: The graph-native default in forward_common_lower_graph is
unsafe for padded or virtual atoms because it flattens the full atype tensor and
then uses an all-ones mask, which can let negative-type slots contribute
energy/forces. Update this path to either fall back to the legacy lower whenever
any local atom type is negative, or pass the real-atom mask through the graph
lower so only valid atoms are carried by call_lower_graph and the related
builder logic in the referenced lower-graph code paths.
- Around line 283-293: The graph auto-routing currently selects the lower-graph
path even for models that require charge-spin input, but
`forward_common_lower_graph()` and `_call_common_graph()` do not accept
`charge_spin`, so those models lose part of their inputs. Update
`_resolve_graph_method()` to keep charge-spin-capable models on the dense/legacy
path unless the graph implementation is extended to consume `charge_spin`, and
make sure the routing logic checks for that input before choosing the graph
path. Use the existing method names `forward_common_lower_graph`,
`_call_common_graph`, and `_resolve_graph_method` to keep the fix scoped to the
model dispatch flow.
---
Nitpick comments:
In `@source/tests/common/dpmodel/test_from_ijs.py`:
- Around line 55-86: Add a small CPU torch smoke test for
build_neighbor_graph_ase alongside the existing NumPy parity checks, since the
function is also exercised through the pt_expt path. Reuse the same test module
and symbols (build_neighbor_graph_ase, build_neighbor_graph) but pass torch
tensors for coord, atype, and box, then assert the neighbor set matches the
in-tree reference to catch host-conversion regressions. Keep the new case
minimal and focused on verifying torch-backed inputs work end to end.
🪄 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: db2db2ec-2d28-4dd2-ae92-b2b72ddaf4b2
📒 Files selected for processing (21)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/ase_builder.pydeepmd/dpmodel/utils/neighbor_graph/builder.pydeepmd/dpmodel/utils/neighbor_graph/derivatives.pydeepmd/dpmodel/utils/neighbor_graph/env.pydeepmd/dpmodel/utils/neighbor_graph/from_ijs.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/pt_expt/model/edge_transform_output.pydeepmd/pt_expt/model/make_model.pysource/tests/common/dpmodel/test_call_lower_graph.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_call_graph_descriptor.pysource/tests/common/dpmodel/test_dpa1_graph_model_energy.pysource/tests/common/dpmodel/test_edge_env_mat.pysource/tests/common/dpmodel/test_from_ijs.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_edge_energy_deriv.py
Comment on lines
+283
to
+293
| def forward_common_lower_graph( | ||
| self, | ||
| atype: torch.Tensor, | ||
| n_node: torch.Tensor, | ||
| edge_index: torch.Tensor, | ||
| edge_vec: torch.Tensor, | ||
| edge_mask: torch.Tensor, | ||
| do_atomic_virial: bool = False, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| ) -> dict[str, torch.Tensor]: |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Default graph routing ignores charge_spin.
forward_common_lower_graph() / _call_common_graph() have no charge_spin input, but _resolve_graph_method() now auto-selects the graph path for eligible descriptors. That means charge-spin models silently lose part of their input unless callers force "legacy". Please keep those models on the dense path until the graph lower can consume charge_spin.
Also applies to: 376-442
🤖 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/make_model.py` around lines 283 - 293, The graph
auto-routing currently selects the lower-graph path even for models that require
charge-spin input, but `forward_common_lower_graph()` and `_call_common_graph()`
do not accept `charge_spin`, so those models lose part of their inputs. Update
`_resolve_graph_method()` to keep charge-spin-capable models on the dense/legacy
path unless the graph implementation is extended to consume `charge_spin`, and
make sure the routing logic checks for that input before choosing the graph
path. Use the existing method names `forward_common_lower_graph`,
`_call_common_graph`, and `_resolve_graph_method` to keep the fix scoped to the
model dispatch flow.
…all_graph_adapter/_call_dense
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/descriptor/dpa1.py (1)
669-716: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHonor the optional
type_embeddingdefault before calling_call_graph.Line 669 advertises
type_embeddingas optional, butNoneis forwarded into_call_graphand later passed toxp.asarray. Either make it required or populate it fromself.type_embedding.call()before use.Proposed fix
xp = array_api_compat.array_namespace(graph.edge_vec) dev = array_api_compat.device(graph.edge_vec) + if type_embedding is None: + type_embedding = self.type_embedding.call() + atype = xp.asarray(atype, device=dev) + type_embedding = xp.asarray(type_embedding, device=dev) grrg_node, rot_mat_node = self.se_atten._call_graph( graph, atype, type_embedding=type_embedding ) @@ # descriptor-level concat_output_tebd if self.concat_output_tebd: - tebd = xp.asarray(type_embedding, device=dev) - atype_local = xp.asarray(atype, device=dev) - atype_embd = xp.take(tebd, atype_local, axis=0) # (nf*nloc, tebd_dim) + atype_embd = xp.take(type_embedding, atype, axis=0) # (nf*nloc, tebd_dim)🤖 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/dpmodel/descriptor/dpa1.py` around lines 669 - 716, The optional type_embedding handling in the descriptor graph path is inconsistent because dpa1.py’s graph-native forward method forwards None into DescrptBlockSeAtten._call_graph and later into xp.asarray. Update the forward flow so type_embedding is resolved before use: either make it required in the public signature or, if absent, fetch it from self.type_embedding.call() before calling _call_graph and before the concat_output_tebd branch. Keep the fix localized around forward and the self.se_atten._call_graph path.
🤖 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/descriptor/dpa1.py`:
- Around line 625-627: The dense-path handling of sw in dpa1.py applies the 4-D
nlist_mask before sw has the trailing singleton dimension, which can
mis-broadcast the raw switch values in the graph adapter. Update the sw flow in
the descriptor logic around the dense-path masking so sw is reshaped to (nf,
nloc, nnei, 1) before calling xp.where with nlist_mask, using the same
sw/nlist_mask path in dpa1’s descriptor code.
---
Outside diff comments:
In `@deepmd/dpmodel/descriptor/dpa1.py`:
- Around line 669-716: The optional type_embedding handling in the descriptor
graph path is inconsistent because dpa1.py’s graph-native forward method
forwards None into DescrptBlockSeAtten._call_graph and later into xp.asarray.
Update the forward flow so type_embedding is resolved before use: either make it
required in the public signature or, if absent, fetch it from
self.type_embedding.call() before calling _call_graph and before the
concat_output_tebd branch. Keep the fix localized around forward and the
self.se_atten._call_graph path.
🪄 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: 89804960-df1b-4a87-8104-06fe0bb20ecc
📒 Files selected for processing (3)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/utils/neighbor_graph/builder.pysource/tests/pt_expt/model/test_dpa1_graph_lower.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/dpmodel/utils/neighbor_graph/builder.py
Comment on lines
+625
to
+627
| nlist_mask = (nlist != -1)[:, :, :, None] | ||
| sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) | ||
| sw = xp.reshape(sw, (nf, nloc, nnei, 1)) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reshape sw before applying the 4-D mask.
The dense path adds the trailing singleton before masking; doing the mask first can mis-broadcast raw (nf, nloc, nnei) switch values in the graph adapter.
Proposed fix
- nlist_mask = (nlist != -1)[:, :, :, None]
- sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))
sw = xp.reshape(sw, (nf, nloc, nnei, 1))
+ nlist_mask = (nlist != -1)[:, :, :, None]
+ sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| nlist_mask = (nlist != -1)[:, :, :, None] | |
| sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) | |
| sw = xp.reshape(sw, (nf, nloc, nnei, 1)) | |
| sw = xp.reshape(sw, (nf, nloc, nnei, 1)) | |
| nlist_mask = (nlist != -1)[:, :, :, None] | |
| sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) |
🤖 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/dpmodel/descriptor/dpa1.py` around lines 625 - 627, The dense-path
handling of sw in dpa1.py applies the 4-D nlist_mask before sw has the trailing
singleton dimension, which can mis-broadcast the raw switch values in the graph
adapter. Update the sw flow in the descriptor logic around the dense-path
masking so sw is reshaped to (nf, nloc, nnei, 1) before calling xp.where with
nlist_mask, using the same sw/nlist_mask path in dpa1’s descriptor code.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5583 +/- ##
==========================================
+ Coverage 82.30% 82.33% +0.03%
==========================================
Files 887 891 +4
Lines 100161 100758 +597
Branches 4043 4058 +15
==========================================
+ Hits 82433 82955 +522
- Misses 16269 16345 +76
+ Partials 1459 1458 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…g keys, not just energy) call_lower_graph now reduces via the general fit_output_to_model_output (any fitting output, not hard-coded 'energy'); pt_expt adds fit_output_to_model_output_graph mirroring take_deriv but differentiating w.r.t. edge_vec per scalar component via edge_energy_deriv, so force/virial/ atom_virial generalize to any r_differentiable output. Fixes the KeyError on non-energy models (dos/dipole/polar/property) with attn_layer=0 descriptors.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/model/make_model.py (1)
585-589: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winReject multi-rank graph metadata instead of ignoring it.
n_local/comm_dictare accepted but ignored, while the reduction uses an all-local mask. A caller passing multi-rank graph data would double-count halo nodes as owned atoms. Since PR-A is single-rank only, fail fast when these are provided.Proposed guard
def call_lower_graph( self, atype: Array, n_node: Array, edge_index: Array, edge_vec: Array, edge_mask: Array, n_local: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, comm_dict: dict | None = None, ) -> dict[str, Array]: @@ """ + if n_local is not None or comm_dict is not None: + raise NotImplementedError( + "call_lower_graph currently supports single-rank graphs only" + ) xp = array_api_compat.array_namespace(edge_vec)Also applies to: 646-653
🤖 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/dpmodel/model/make_model.py` around lines 585 - 589, The model reduction path in make_model should not silently accept multi-rank graph metadata: in the relevant method(s) that take n_local and comm_dict, add a fail-fast guard that rejects non-None values instead of proceeding with the all-local mask. Update the same handling in both affected spots so callers of the model-building/reduction flow get an explicit error when passing multi-rank graph data, rather than having halo nodes counted as owned atoms.
🧹 Nitpick comments (1)
source/tests/pt_expt/model/test_dos_graph.py (1)
127-161: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert graph returns dense keys and cover eval mode.
The current intersection hides missing graph outputs, including derivative keys, and default
Module.training=Trueskips thecreate_graph=Falseinference path.Proposed test tightening
`@pytest.mark.parametrize`( "make_model", [_make_dos, _make_dipole, _make_polar, _make_property], ) # one builder per fitting kind - def test_graph_matches_dense(self, make_model) -> None: + `@pytest.mark.parametrize`("training", [True, False]) + def test_graph_matches_dense(self, make_model, training) -> None: @@ ds = _make_descriptor() m = make_model(ds) + m.train(training) graph = m.call_common(self.coord, self.atype, None, do_atomic_virial=True) @@ - shared = [ - k - for k in graph - if k in dense and graph[k] is not None and dense[k] is not None - ] - # at least the reduced + per-atom output must be present and shared - assert len(shared) >= 2 - for k in shared: + dense_keys = {k for k, v in dense.items() if v is not None} + graph_keys = {k for k, v in graph.items() if v is not None} + assert dense_keys <= graph_keys + for k in sorted(dense_keys): torch.testing.assert_close( graph[k].to(torch.float64), dense[k].to(torch.float64), **tol )🤖 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 `@source/tests/pt_expt/model/test_dos_graph.py` around lines 127 - 161, Tighten test_graph_matches_dense so it verifies the graph path returns the same expected keys as the dense legacy path instead of only comparing the intersection, which can hide missing outputs like derivative entries. Use the existing m.call_common flow and compare against the dense result by asserting the key sets (or expected shared keys) match for each make_model builder, and add an eval-mode case on the model instance so the inference path with create_graph=False is exercised instead of relying on the default Module.training=True behavior.
🤖 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/make_model.py`:
- Around line 555-566: The `make_model` path is converting a traced graph value
to a Python int via `int(n_node[0])`, which breaks JAX/JIT-style tracing in the
`dpmodel` backend. Update the `make_model` logic (and the matching reshape path
in the same class) to derive `nloc` from the static atom count using
`atype.shape[0] // nf` instead of reading `n_node[0]`, so the reshape to `(nf,
nloc)` remains compatible with graph tracing.
In `@deepmd/pt_expt/model/edge_transform_output.py`:
- Around line 109-116: The masked intensive reduction in
edge_transform_output.py is still summing padded atoms in the numerator and uses
a denominator shape that is not safe for higher-rank tensors. Update the masked
branch in the reduction logic around the kk_redu handling to apply mask before
summing vv, then divide by a denominator reshaped to match the non-atom
dimensions so it broadcasts correctly for outputs like (nf, 3, 3). Use the
existing tensor/atom_axis logic in this block and keep the unmasked and
extensive paths unchanged.
---
Outside diff comments:
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 585-589: The model reduction path in make_model should not
silently accept multi-rank graph metadata: in the relevant method(s) that take
n_local and comm_dict, add a fail-fast guard that rejects non-None values
instead of proceeding with the all-local mask. Update the same handling in both
affected spots so callers of the model-building/reduction flow get an explicit
error when passing multi-rank graph data, rather than having halo nodes counted
as owned atoms.
---
Nitpick comments:
In `@source/tests/pt_expt/model/test_dos_graph.py`:
- Around line 127-161: Tighten test_graph_matches_dense so it verifies the graph
path returns the same expected keys as the dense legacy path instead of only
comparing the intersection, which can hide missing outputs like derivative
entries. Use the existing m.call_common flow and compare against the dense
result by asserting the key sets (or expected shared keys) match for each
make_model builder, and add an eval-mode case on the model instance so the
inference path with create_graph=False is exercised instead of relying on the
default Module.training=True 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: 5a42be10-6fff-44b0-b303-ddf2472280ba
📒 Files selected for processing (5)
deepmd/dpmodel/model/make_model.pydeepmd/pt_expt/model/edge_transform_output.pydeepmd/pt_expt/model/make_model.pysource/tests/common/dpmodel/test_call_lower_graph.pysource/tests/pt_expt/model/test_dos_graph.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/common/dpmodel/test_call_lower_graph.py
Comment on lines
+109
to
+116
| if mask is not None: | ||
| model_ret[kk_redu] = torch.sum( | ||
| vv.to(redu_prec), dim=atom_axis | ||
| ) / torch.sum(mask, dim=-1, keepdim=True) | ||
| else: | ||
| model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis) | ||
| else: | ||
| model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Apply the mask to the intensive numerator and reshape it rank-safely.
Line 110 sums all atoms even when mask marks padded atoms invalid, then divides by the real-atom count. Also, the (nf, 1) denominator is not broadcast-safe for multi-rank shapes like (nf, 3, 3).
Proposed fix
if vdef.intensive:
if mask is not None:
+ v_red = vv.to(redu_prec)
+ mask_view = mask.to(dtype=redu_prec, device=v_red.device).reshape(
+ mask.shape + (1,) * len(shap)
+ )
model_ret[kk_redu] = torch.sum(
- vv.to(redu_prec), dim=atom_axis
- ) / torch.sum(mask, dim=-1, keepdim=True)
+ v_red * mask_view, dim=atom_axis
+ ) / torch.sum(mask_view, dim=atom_axis)
else:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if mask is not None: | |
| model_ret[kk_redu] = torch.sum( | |
| vv.to(redu_prec), dim=atom_axis | |
| ) / torch.sum(mask, dim=-1, keepdim=True) | |
| else: | |
| model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis) | |
| else: | |
| model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis) | |
| if mask is not None: | |
| v_red = vv.to(redu_prec) | |
| mask_view = mask.to(dtype=redu_prec, device=v_red.device).reshape( | |
| mask.shape + (1,) * len(shap) | |
| ) | |
| model_ret[kk_redu] = torch.sum( | |
| v_red * mask_view, dim=atom_axis | |
| ) / torch.sum(mask_view, dim=atom_axis) | |
| else: | |
| model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis) | |
| else: | |
| model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis) |
🤖 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/edge_transform_output.py` around lines 109 - 116, The
masked intensive reduction in edge_transform_output.py is still summing padded
atoms in the numerator and uses a denominator shape that is not safe for
higher-rank tensors. Update the masked branch in the reduction logic around the
kk_redu handling to apply mask before summing vv, then divide by a denominator
reshaped to match the non-atom dimensions so it broadcasts correctly for outputs
like (nf, 3, 3). Use the existing tensor/atom_axis logic in this block and keep
the unmasked and extensive paths unchanged.
…ual-atom fallthrough, env protection parity, ASE/from_ijs device)
The dpa1 graph dense->graph adapter (from_dense_quartet -> call_graph) derived the per-edge neighbor type from the local owner atype[mapping[ neighbor]]. That equals the neighbor's true extended type atype_ext[ neighbor] for every PHYSICAL input (a ghost is a periodic image of its owner), but diverges for an arbitrary external quartet whose mapping is inconsistent -- e.g. the naively-permuted mapping in the universal descriptor fixture. The dense path reads atype_ext[neighbor] directly, so test_exclude_types (which routes the graph-eligible attn_layer=0 dpa1 through the adapter) failed graph-vs-dense parity. Thread an optional per-edge nei_type override into call_graph/_call_graph; the adapter computes it from atype_ext[neighbor] (matching the converter's compact=False edge ordering). Geometry-native (carry-all) builders pass None and keep using atype[src], correct for all real inputs. Adds a non-vacuous regression test (corrupts ghost types so ghost type != owner type and asserts a ghost is actually a neighbor).
…o ghost-free The test_exclude_types failure on deepmodeling#5583 was NOT a graph bug -- it was an invalid test fixture. extend_coord_with_ghosts tiles the local atype and aidx (nlist.py:369-370), so atype_ext[k] == atype[mapping[k]] is a hard single-rank invariant: a ghost is a periodic image of its owner and shares its type. The ghost-free graph (src = mapping[neighbor], nei_type = atype[src]) is therefore correct. The shared TestCaseSingleFrameWithNlist fixtures permuted `mapping` as `mapping[:, perm]` WITHOUT remapping the local-index VALUES through inv_perm (which they DO apply to nlist), desyncing atype_ext from mapping and even breaking local self-mapping -- a quartet extend_coord_with_ghosts can never emit. The ghost-free graph (correct) then diverged from dense on this impossible input. - Fix the 4 duplicated fixture copies: mapping permutation now mirrors the nlist one, `inv_perm[mapping[:, perm]]` (universal/common/cases/cases.py -- the copy the universal CI test imports -- plus common/dpmodel, common/test_mixins, pd/model/test_env_mat). - Revert the per-edge nei_type channel (commit 99c707a) back to the ghost-free atype[src] gather; it was treating a symptom of the bad fixture. - Replace the now-invalid inconsistent-mapping regression test with test_single_rank_extension_keeps_type_invariant, which pins atype_ext[k]==atype[mapping[k]] on extend_coord_with_ghosts output. Multi-rank (real halo ghosts with independent types) needs the extended atype table + comm fold and is the documented PR-B follow-on.
…tomic_graph; drop _graph_descriptor_fitting
… + exclude fallback
for more information, see https://pre-commit.ci
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
1 participant
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.
Summary
Adds the graph-native forward path for
dpa1(attn_layer=0)(the factorizable, mixed-types case), built on theNeighborGraphfoundation from #5581. Geometry enters the descriptor only through per-edgeedge_vec; the neighbor-axis reduction becomes asegment_sumover edge centers. Forpt_exptthis becomes the default forward (force/virial via a single autograd backward throughedge_vec).What it adds
edge_env_mat(per-edge env-mat 4-vector),DescrptBlockSeAtten._call_graph+DescrptDPA1.call_graph, modelcall_lower_graph(energy),neighbor_graph_from_ijs+ an optional ASE O(N) carry-all builder.edge_energy_deriv(autogradgrad(E, edge_vec)→edge_force_virial) +forward_common_lower_graph(energy + force + virial + atom_virial).DescrptDPA1.callbecomes a thin adapter (from_dense_quartet → call_graph) preserving the 5-tuple ABI; a shape-static converter keeps itjax.jit/torch.export-traceable.Default behavior
dpa1(attn_layer=0, concat tebd, no exclude_types)models to the carry-all graph (it has the autograd force/virial path).sel.exclude_types, linear/ZBL) fall back to the dense path unchanged.neighbor_graph_method="legacy"forces dense;"dense"/"ase"force the graph.Parity (graph vs legacy dense lower, fp64 CPU)
atom_virial matches the canonical TF==pt-legacy full-to-src convention. dpa1 descriptor + model consistency suites green across dp/jax/pt_expt.
Known limitations
make_fx(forward + grad) traces; full.pt2AOTI export is a follow-up (PR-B). The carry-all builders (build_neighbor_graph/from_ijs) still usenonzero(eager-only); their static variants land with the export PR.Also folds in three follow-up fixes to the #5581 foundation from @OutisLi's review (dangling spec refs → design discussion,
edge_force_virialjax int-sum short-circuit,Arraytyping).Summary by CodeRabbit
neighbor_graph_methodrouting for energy/force/virial workflows, with carry-all graphs and new graph-output fitting/post-processing utilities.(i,j,S)conversion, and edge environment matrices) and exported them publicly.edge_force_virialbehavior when node capacity is provided, and improved masking/ABI parity for invalid/virtual atoms.