fix(dpmodel): remap virtual type embeddings - #5856
Conversation
|
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:
📝 WalkthroughWalkthroughThe PR adds shared padding-aware type-embedding helpers. It applies them across dpmodel, legacy PyTorch, and experimental descriptors. Regression tests cover negative virtual types across embedding, strip, graph, and fused paths. ChangesVirtual type embedding handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
Map negative virtual atom types to the explicit padding embedding row before array-api gathers and type-pair indexing. Apply the shared boundary across DPA1 through DPA4, SeZM, and SeTTebd descriptor paths. Cover NumPy and array_api_strict embedding gathers, DPA1 concat and strip modes, SeZM, and an independent SeTTebd strip pair-index regression comparing -1 with the explicit padding type. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
1f8ca16 to
1e74ddd
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5856 +/- ##
==========================================
- Coverage 79.53% 79.28% -0.25%
==========================================
Files 1075 1075
Lines 126134 126178 +44
Branches 4592 4592
==========================================
- Hits 100315 100036 -279
- Misses 24164 24488 +324
+ Partials 1655 1654 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix itself is correct and it prevents a real failure, but the test suite almost entirely fails to demonstrate that, so it does not currently work as the regression gate it claims to be.
I ran the PR's own tests against the base commit (6c3b985cb): 7 of the 8 cases pass unpatched. Only test_se_t_tebd_strip_strict_virtual_type_matches_explicit_padding fails pre-fix. Details in the first two inline comments.
On the stated rationale - "negative gather indices either wrap or fail depending on the array backend" - it is half right, and the half that is right is the half that is untested. Through xp.take every backend wraps to the last row (numpy, array_api_strict, array_api_compat's torch, jax), and since that row is the padding row for every caller here, those substitutions are behavioural no-ops. Through xp_take_along_axis, deepmd/dpmodel/array_api.py routes torch arrays to raw torch.gather, which raises RuntimeError: index -1 is out of bounds. So the stripped pair-index path is the one genuinely broken thing: silently wrong on numpy, a hard crash on torch.
Eight inline comments. Smaller notes not worth their own thread:
-
The padding row historically means "empty neighbour slot", reached through the positive index
ntypes(seePairExcludeMask.build_type_exclude_mask), not "virtual atom". For strip tables it is not zero in any case, sincecal_g_stripputs a biased MLP after it. -
SeZMTypeEmbedding.call()'s docstring still says "Valid type range is [0, ntypes-1]" directly above the new code that makes negative types a supported input. -
The new
if self.padding:conditional has no test for itsFalsebranch. -
The fused CUDA path (
deepmd/kernels/cuda/dpa1/graph_descriptor.py) passes rawatypeintotorch.ops.deepmd.dpa1_graph_descriptorand does raw pointer arithmetic on it. Reachability is narrow - opt-in env var, and the C++ inference path filterstype < 0upstream - but it is unguarded. -
take_type_embeddingderives the padding index fromshape[0]with no check that the table is actually padded. Every current caller passespadding=True, so this is a latent contract issue only.
Apply the padding-row convention to strip pair indices across dpmodel, pt_expt, PyTorch, and Paddle implementations. Keep real-type-only exclusion and normalization lookups clamped separately, document the sentinel invariant, and add focused DPA1/DPA2/DPA3 and Torch regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/dpmodel/descriptor/dpa1.py (1)
1666-1670: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the two-side pair-index formula into a shared helper.
center_type * ntypes_with_padding + nei_typeis computed here for the dense strip path and again, separately, in_graph_edge_gg_stripfor the graph strip path (same file, later in the class). A past review identified this exact duplicated, unremapped formula as the root cause of the original virtual-type bug, and it had also drifted out of sync across thept_expt/legacy PyTorch/Paddle backends. Extracting the core index computationcenter*ntypes_with_padding + neiinto one small helper function, called from both the tiled dense-path site and the per-edge graph-path site, reduces the chance that a future edit re-diverges the two implementations. Each call site can still handle its own tiling/broadcast shape around the shared core formula.🤖 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 1666 - 1670, Extract the shared two-side pair-index calculation into a small helper using center and neighbor types with ntypes_with_padding, returning center*ntypes_with_padding + nei. Update the dense strip computation near remap_atype_to_padding and the graph strip implementation in _graph_edge_gg_strip to call this helper, while preserving each path’s existing tiling or broadcast shape handling.
🤖 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/pd/model/descriptor/se_atten.py`:
- Around line 585-591: Environment-statistics lookup currently receives negative
local center types before padded embedding remapping, causing invalid or
unintended mean/stddev indexing. In
deepmd/pd/model/descriptor/se_atten.py:585-591,
deepmd/pd/model/descriptor/se_t_tebd.py:932-936,
deepmd/pt/model/descriptor/se_atten.py:695-702, and
deepmd/pt/model/descriptor/se_t_tebd.py:1024-1028, create and use a sanitized
center-type tensor for prod_env_mat statistics lookup; retain
remap_atype_to_padding for embedding and type-pair tables. In
deepmd/pt_expt/descriptor/dpa1.py:173-176,
deepmd/pt_expt/descriptor/dpa2.py:420-426, and
deepmd/pt_expt/descriptor/se_t_tebd.py:243-245, sanitize types before _env_mat
or env_mat.call receives them. Add virtual-center regression coverage for dense,
compressed, and graph routes.
---
Nitpick comments:
In `@deepmd/dpmodel/descriptor/dpa1.py`:
- Around line 1666-1670: Extract the shared two-side pair-index calculation into
a small helper using center and neighbor types with ntypes_with_padding,
returning center*ntypes_with_padding + nei. Update the dense strip computation
near remap_atype_to_padding and the graph strip implementation in
_graph_edge_gg_strip to call this helper, while preserving each path’s existing
tiling or broadcast shape handling.
🪄 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: 3ff1b7ac-7180-4d69-8432-ed28455e1e4d
📒 Files selected for processing (16)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/dpa2.pydeepmd/dpmodel/descriptor/dpa3.pydeepmd/dpmodel/descriptor/dpa4_nn/embedding.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/utils/type_embed.pydeepmd/pd/model/descriptor/se_atten.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pt/model/descriptor/se_atten.pydeepmd/pt/model/descriptor/se_t_tebd.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/se_t_tebd.pysource/tests/common/dpmodel/test_type_embedding_virtual.pysource/tests/pt/model/test_virtual_type_embedding.pysource/tests/pt_expt/descriptor/test_dpa1.py
🚧 Files skipped from review as they are similar to previous changes (4)
- deepmd/dpmodel/descriptor/dpa2.py
- deepmd/dpmodel/descriptor/dpa3.py
- deepmd/dpmodel/utils/type_embed.py
- deepmd/dpmodel/descriptor/se_t_tebd.py
Clamp virtual center types before real-type-only environment-statistics lookups in PyTorch, Paddle, and pt_expt paths while retaining padding-row remapping for embedding tables. Add dense, fused-prologue, and graph regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
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 `@source/tests/pt_expt/descriptor/test_dpa1.py`:
- Around line 131-145: Update the test around _env_mat to exercise the fused
Triton branch by patching triton_infer_level to return 1 and stubbing
_env_mat_triton. Capture its invocation and assert that the center-type argument
is clamped to type 0 when the input contains -1, while preserving the existing
comparison behavior as appropriate.
🪄 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: d711a1bc-0f86-46bc-8223-854d3550a575
📒 Files selected for processing (10)
deepmd/pd/model/descriptor/se_atten.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pt/model/descriptor/se_atten.pydeepmd/pt/model/descriptor/se_t_tebd.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/se_t_tebd.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/pt/model/test_virtual_type_embedding.pysource/tests/pt_expt/descriptor/test_dpa1.py
🚧 Files skipped from review as they are similar to previous changes (3)
- deepmd/pt_expt/descriptor/dpa2.py
- deepmd/pt_expt/descriptor/se_t_tebd.py
- deepmd/pt_expt/descriptor/dpa1.py
Coding-Agent: Codex\nCodex-Version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
This is a thorough round and I am happy with it. I checked every thread against HEAD 609d77f2d rather than the replies, and rather than reading the new tests I ran them against pre-fix source with the helper restored, so the signal is the descriptor change itself and not a missing symbol.
What that showed:
| test | pre-fix |
|---|---|
test_dpa1_strict_virtual_type_matches_explicit_padding[strip-False] |
fails |
test_se_t_tebd_strip_strict_virtual_type_matches_explicit_padding |
fails |
test_dpa2_virtual_neighbor_matches_explicit_padding |
fails |
source/tests/pt/model/test_virtual_type_embedding.py (4 tests) |
4/4 fail |
pt_expt test_strip_virtual_neighbor…, test_strip_pair_index_remaps_virtual_types, test_fused_env_prologue_clamps_virtual_center_statistics |
3/3 fail |
So the DPA1 case now fails at exactly the ("strip", type_one_side=False) combination I asked for, the pair-index duplication is genuinely closed across pt_expt, legacy PyTorch and Paddle with failing regressions to prove it, and _strip_pair_index is exercised directly so raw torch.gather semantics are covered rather than assumed.
Two resolutions I want to call out as better than what I suggested. Splitting call_graph into safe_real_atype for apply_pair_exclusion and the davg/dstd normalization, while keeping raw types only for the padded lookups, is the right shape -- it makes the two conventions explicit at the point of use instead of leaving a reader to infer which applies. And the Notes section on remap_atype_to_padding naming davg, dstd and spin masks as excluded is precisely the guard against the next contributor copying the wrong convention. I also appreciate that take_type_embedding's note says plainly that the table implementation stays responsible for keeping the reserved row neutral -- that is the honest statement of the SeZM situation rather than a claim that it is structurally zero.
One leftover, inline and non-blocking -- it does not affect correctness, only what the suite can catch. Approving; please take it or leave it.
A virtual -1 sentinel wraps to the final padding row under NumPy and array_api_strict, so the old DPA3 test passed even without the remap. Use -2 as the negative sentinel so a raw take lands on a real row and only the explicit remap-to-padding satisfies the test. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
for more information, see https://pre-commit.ci
Resolve the DPA1 pt_expt merge conflicts by combining the virtual-type clamping fix from this branch with master's DPA1 l=2 moment-basis work (deepmodeling#5911): clamp center types before real-type-only mean/stddev/radial-stddev lookups in _env_mat and _call_graph_compress_reference, while keeping the master moment_basis construction. Coding-Agent: opencode opencode-Version: 1.18.11 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
3 participants
Closes #5665.
Summary
Sentinel convention
Only tables explicitly constructed with a final padding row may use ntypes as the remapped virtual type. Real-type-only tables such as davg, dstd, exclusion inputs, and spin masks must instead receive masked or clamped real-type indices. TypeEmbedNet reconstructs a literal zero padding row; SeZMTypeEmbedding reserves and initializes its stored final row to zero.
Validation
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Bug Fixes
Tests