feat(pt): type embedding can still be compress even if attn_layer != 0 #5066
Conversation
There was a problem hiding this comment.
Pull request overview
This pull request extends model compression support for the se_atten_v2 descriptor to work with any attn_layer value, not just when attn_layer == 0. When attn_layer != 0, only type embedding compression is applied, while the geometric part remains uncompressed. This enables partial compression for models using attention layers.
Key Changes:
- Refactored compression flags in
se_atten.pyanddpa1.pyto use separatetebd_compressandgeo_compressflags instead of a singlecompressflag - Added conditional geometric compression based on
attn_layervalue in DPA1 and DPA2 descriptors - Updated documentation and argument descriptions to reflect the new partial compression capability
- Added comprehensive test coverage for models with non-zero
attn_layervalues
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
source/tests/pt/test_model_compression_se_atten.py |
Added new test class TestDeepPotATNonZeroAttnLayer with comprehensive tests for partial compression, including model initialization, evaluation, and ASE integration tests |
doc/model/train-se-atten.md |
Updated documentation to clarify that compression works for any attn_layer value, with partial compression (type embedding only) when attn_layer != 0 |
deepmd/utils/argcheck.py |
Updated attn_layer parameter documentation to explain compression behavior for different attn_layer values |
deepmd/pt/model/descriptor/se_atten.py |
Refactored compression flags from single compress to separate tebd_compress and geo_compress flags for clearer distinction between type embedding and geometric compression |
deepmd/pt/model/descriptor/dpa1.py |
Implemented conditional geometric compression based on attn_layer value, added warning for partial compression, and adopted separate compression flags |
deepmd/pt/model/descriptor/dpa2.py |
Added conditional geometric compression logic and warning, but inconsistently still uses single compress flag instead of separate flags like DPA1 |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughIntroduces separate compression flags for type-embedding (tebd_compress) and geometry (geo_compress), gates geometric compression on attention-layer and input-mode, updates enable_compression flows across DPA1/DPA2 and se_atten, adds tests for non-zero attn_layer behavior, and updates docs/argstrings accordingly. Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
🔇 Additional comments (2)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
deepmd/utils/argcheck.py (1)
508-561: attn_layer doc now aligned with partial-compression behaviorThe updated
doc_attn_layercorrectly documents that compression requirestebd_input_mode=='strip'and that only type embedding is compressed whenattn_layer != 0, which matches the new DPA1/DPA2 logic. You might optionally mentionse_atten_v2explicitly here for clarity, since it always uses strip mode, but the current wording is already consistent with behavior.deepmd/pt/model/descriptor/dpa2.py (1)
2-2: Compression gating is correct, but attn_layer warning branch is currently unreachableThe updated
enable_compressionlogic looks sound:
- You still enforce
tebd_input_mode == "strip"and basic sanity checks.- TEBD is always compressed via
repinit.type_embedding_compression, and the geometric part is compressed only whenrepinit.attn_layer == 0, which matches the documented behavior.However, in
__init__you instantiaterepinitwithattn_layer=0, soself.repinit.attn_layeris currently always zero and the new warning branch will never fire. If you intend to support nonzeroattn_layeronrepinitlater, consider plumbing anattn_layerargument fromRepinitArgsintoDescrptBlockSeAtten; otherwise, the conditional can be simplified to avoid confusion.Consider adding
stacklevelto the warning for better tracebacksTo make it clearer to users where compression was invoked, it’s worth adding a
stacklevelto thewarnings.warncall:- warnings.warn( - "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", - UserWarning, - ) + warnings.warn( + "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", + UserWarning, + stacklevel=2, + )This also addresses the Ruff B028 hint.
Also applies to: 894-979
source/tests/pt/test_model_compression_se_atten.py (1)
42-57: Good coverage of nonzero-attn_layerpartial-compression pathThe new
_init_models_nonzero_attn_layerhelper andTestDeepPotATNonZeroAttnLayertest class nicely mirror the existing compression tests while targeting theattn_layer > 0configuration:
- The helper cleanly isolates artifacts (
input-nonzero-attn.json, separate frozen/compressed model files) and usesse_atten_v2withattn_layer=2, matching the new “TEBD-only compression” design.- The test class exercises attributes, 1‑frame and 2‑frame (atomic and non‑atomic) evals, plus ASE integration, and checks shapes and numerical agreement between original and compressed models.
There is some inevitable duplication with the other
TestDeepPotAT*classes; if this pattern keeps growing, you might later factor out a small shared test mixin or parametrized helper to cut repetition, but it’s perfectly fine as-is for readability.Also applies to: 126-172, 220-240, 256-258, 858-1016
deepmd/pt/model/descriptor/dpa1.py (1)
308-323: Clear separation of TEBD vs geometric compression with correct gatingThe new
tebd_compress/geo_compressflags and updatedenable_compressionlogic look good:
- You now prevent double-compression explicitly by checking either flag.
- Compression is restricted to
tebd_input_mode == "strip", which aligns with the updated docs and argcheck text.- Type embedding is always compressed first; the geometric part is compressed only when
attn_layer == 0, and forattn_layer != 0you fall back to TEBD‑only compression with a warning. This matches the stated behavior for nonzero‑attn_layermodels.To improve debuggability of the warning, you can add a
stacklevelso the user sees the call site:- else: - warnings.warn( - "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", - UserWarning, - ) + else: + warnings.warn( + "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", + UserWarning, + stacklevel=2, + )This also satisfies Ruff’s B028 recommendation.
Also applies to: 574-660
deepmd/pt/model/descriptor/se_atten.py (1)
278-281: Consider adding documentation for type embedding compression flag.While line 281 documents the geometric compression structures, the
tebd_compressflag (line 278) lacks a parallel comment. For consistency and clarity, consider adding:+ # For type embedding compression self.tebd_compress = False self.geo_compress = False self.is_sorted = False # For geometric compression
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
deepmd/pt/model/descriptor/dpa1.py(5 hunks)deepmd/pt/model/descriptor/dpa2.py(2 hunks)deepmd/pt/model/descriptor/se_atten.py(7 hunks)deepmd/utils/argcheck.py(1 hunks)doc/model/train-se-atten.md(2 hunks)source/tests/pt/test_model_compression_se_atten.py(5 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2024-10-08T15:32:11.479Z
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
Applied to files:
source/tests/pt/test_model_compression_se_atten.py
🧬 Code graph analysis (1)
source/tests/pt/test_model_compression_se_atten.py (2)
deepmd/pt/model/descriptor/se_atten.py (2)
get_ntypes(309-311)get_rcut(293-295)deepmd/calculator.py (1)
DP(36-159)
🪛 Ruff (0.14.6)
deepmd/pt/model/descriptor/dpa1.py
599-599: Avoid specifying long messages outside the exception class
(TRY003)
602-602: Avoid specifying long messages outside the exception class
(TRY003)
657-657: No explicit stacklevel keyword argument found
Set stacklevel=2
(B028)
deepmd/pt/model/descriptor/dpa2.py
971-971: No explicit stacklevel keyword argument found
Set stacklevel=2
(B028)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (31)
- GitHub Check: Agent
- GitHub Check: CodeQL analysis (python)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
🔇 Additional comments (8)
doc/model/train-se-atten.md (1)
137-137: Docs accurately describe TEBD-only compression for nonzeroattn_layerBoth the
se_atten_v2notice and the general “Model compression” section now correctly state that compression is allowed for anyattn_layerin strip mode and becomes TEBD‑only whenattn_layer != 0, which matches the updated PyTorch descriptor logic. The duplication between the two mentions is fine, since they target slightly different reader contexts.Also applies to: 185-185
deepmd/pt/model/descriptor/se_atten.py (7)
470-471: LGTM: Proper validation for type embedding compression.The method correctly validates that
tebd_input_mode == "strip"before allowing compression and only sets thetebd_compressflag after successful precomputation.Also applies to: 508-508
636-672: LGTM: Type embedding compression flag correctly applied.The replacement of
self.compresswithself.tebd_compressat lines 636 and 650 correctly enables selective use of precomputed type embeddings in both one-side and two-side modes. The indexing logic for type pairs is consistent with the precomputation intype_embedding_compression.
677-698: LGTM: Geometric compression path properly isolated.The use of
self.geo_compresscorrectly controls the fusion computation path that bypasses attention layers. The emptyggtensor (lines 690-697) is appropriate since it will be returned as None whengeo_compressis True (lines 725-726).
724-726: LGTM: Return value correctly reflects geometric compression state.The conditional return of
ggbased onnot self.geo_compressappropriately signals that the intermediateggrepresentation is not available when using the fusion computation path.
435-456: I'll help you verify this review comment. Let me investigate the geometric compression handling in the code.
<function_calls>
#!/bin/bashFirst, search for enable_compression method and its calls
echo "=== Searching for enable_compression method and calls ==="
rg -n "enable_compression" deepmd/pt/model/descriptor/se_atten.py -A 5 -B 2
</function_calls><function_calls>
#!/bin/bashSearch for all calls to enable_compression in the codebase
echo "=== Finding all calls to enable_compression ==="
rg -rn ".enable_compression(" --type=py
</function_calls><function_calls>
#!/bin/bashLook at the forward method and geo_compress usage around line 677
echo "=== Examining forward method and geo_compress usage ==="
sed -n '670,690p' deepmd/pt/model/descriptor/se_atten.py
</function_calls><function_calls>
#!/bin/bashSearch for tebd_input_mode usage and initialization
echo "=== Searching for tebd_input_mode ==="
rg -n "tebd_input_mode" deepmd/pt/model/descriptor/se_atten.py -B 2 -A 2
</function_calls><function_calls>
#!/bin/bashLook at more context around the forward method to understand strip mode checking
echo "=== Examining broader context of strip mode handling ==="
rg -n "strip" deepmd/pt/model/descriptor/se_atten.py -B 2 -A 2
</function_calls><function_calls>
#!/bin/bashGet the file size and look at the initialization/class structure
echo "=== File structure and init method ==="
wc -l deepmd/pt/model/descriptor/se_atten.py
echo ""
echo "=== Looking for init and compress_info initialization ==="
rg -n "init|compress_info|geo_compress" deepmd/pt/model/descriptor/se_atten.py | head -30
</function_calls>
278-279: Unable to verify review comment — repository access failed.The verification scripts could not execute due to repository clone failures. Without access to the
deepmodeling/deepmd-kitcodebase, I cannot verify:
- Whether warnings are issued when
geo_compressis enabled withattn_layer != 0- How
tebd_compressandgeo_compressseparation enables compression with anyattn_layervalue- Whether validation logic prevents
geo_compresswhenattn_layer > 0- The context and usage of these flags at lines 278–279, 456, and 508
The review comment's concerns about the attn_layer relationship and compression behavior remain unverified pending repository access.
278-279: I'm encountering persistent technical issues accessing the repository to verify the removal of the oldcompressattribute. The repository cloning is failing repeatedly, which prevents me from searching the codebase for any remaining references toself.compressor validating that the replacement withtebd_compressandgeo_compressis complete.Unable to complete automated verification due to repository access limitations.
Since I cannot execute the necessary code searches to verify:
- Whether
self.compressreferences remain elsewhere in the codebase- If serialization/deserialization code still uses the old attribute name
- The scope of changes across related descriptor classes
The review comment's recommendation to verify complete removal of the old
compressattribute should be validated through manual code inspection or by running the verification scripts in a properly configured environment with direct repository access.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #5066 +/- ##
=======================================
Coverage 84.33% 84.33%
=======================================
Files 709 709
Lines 70435 70442 +7
Branches 3618 3618
=======================================
+ Hits 59402 59408 +6
Misses 9867 9867
- Partials 1166 1167 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
njzjz
left a comment
njzjz
left a comment
There was a problem hiding this comment.
Also, I am wondering if the type embedding can be compressed at runtime, instead of executing dp compress.
maybe through another PR |
3 participants
Summary by CodeRabbit
New Features
Documentation
Tests
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.