feat(pt): Implement type embedding compression for se_atten #5057
Conversation
|
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. 📝 WalkthroughWalkthroughAdds type-embedding precomputation and caching to the SE-Attention descriptor ( Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as enable_compression()
participant DPA as DPA1/DPA2
participant SE as DescrptBlockSeAtten
rect rgba(200,230,255,0.35)
Caller->>DPA: enable_compression()
DPA->>SE: se_atten.enable_compression(...)
SE->>SE: type_embedding_compression(TypeEmbedNet)
SE-->>SE: precompute embeddings -> store in type_embd_data
SE-->>DPA: return
DPA->>Caller: set compress flag
end
rect rgba(220,255,220,0.25)
participant Forward as forward(mode="strip")
Caller->>Forward: forward(input, mode="strip")
alt type_embd_data present
Forward->>SE: use cached type_embd_data for embedding lookup
else
Forward->>SE: compute embeddings on-the-fly via TypeEmbedNet
end
SE-->>Forward: descriptor outputs
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 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 (1)
🚧 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)
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 |
There was a problem hiding this comment.
Pull Request Overview
This PR implements type embedding compression for the se_atten descriptor to optimize inference performance. The implementation precomputes embedding network outputs for all type combinations and uses direct indexing during forward passes instead of computing embeddings on-the-fly.
Key changes:
- Added
type_embedding_compression()method to precompute type embeddings for all combinations (one-side: ntypes+1, two-side: (ntypes+1)²) - Refactored forward pass to use precomputed embeddings when available, with cleaner indexing operations
- Integrated type embedding compression into the DPA1 descriptor's
enable_compression()workflow
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| deepmd/pt/model/descriptor/se_atten.py | Added type embedding compression method with precomputation logic and updated forward pass to conditionally use precomputed embeddings with simplified indexing |
| deepmd/pt/model/descriptor/dpa1.py | Integrated type embedding compression call into the existing compression enablement flow |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
The speed comparison be shown below through the same simulation using LAMMPS to a 1000 atoms 3C-SiC system. This PR leads to roughly 20% speedup. Considering only 2 types of atoms in total, the improvement is quite obvious. atten_v2 with no attention layer before compression: after compression with no type embedding compression： after compression with type embedding compression： |
![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 (1)
deepmd/pt/model/descriptor/se_atten.py (1)
455-499: LGTM! Validation checks address previous review feedback.The implementation correctly precomputes type embeddings for all combinations:
- One-side: (ntypes+1) neighbor embeddings
- Two-side: (ntypes+1)² pair embeddings using index formula
center_type * nt + neighbor_typeThe use of
.detach()is appropriate for cached precomputed values.The static analysis TRY003 warnings at lines 468 and 470-472 are style suggestions about error message length. The current messages are clear and descriptive, so addressing this is optional. If you prefer to follow the style guide strictly, you could define custom exception classes.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/model/descriptor/se_atten.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/model/descriptor/se_atten.py (1)
deepmd/pt/model/network/network.py (2)
TypeEmbedNet(268-358)get_full_embedding(317-332)
🪛 Ruff (0.14.5)
deepmd/pt/model/descriptor/se_atten.py
468-468: Avoid specifying long messages outside the exception class
(TRY003)
470-472: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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). (22)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test C++ (true)
- 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: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
🔇 Additional comments (3)
deepmd/pt/model/descriptor/se_atten.py (3)
30-32: LGTM!The import is correctly added to support the type hint in the new
type_embedding_compressionmethod.
278-288: LGTM!The attribute initialization and comments clearly separate geometric compression from type embedding compression concerns.
627-663: LGTM! Index formula correctly matches precomputation.The forward path properly leverages cached embeddings when available and falls back to on-the-fly computation otherwise:
- One-side (627-633): Directly indexes cached embeddings by neighbor type
- Two-side (634-663): Uses index formula
center_type * ntypes_with_padding + neighbor_typewhich matches the precomputation layoutThe fallback logic ensures backward compatibility with non-compressed models.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #5057 +/- ##
==========================================
+ Coverage 84.24% 84.26% +0.02%
==========================================
Files 709 709
Lines 70236 70305 +69
Branches 3623 3619 -4
==========================================
+ Hits 59169 59244 +75
+ Misses 9900 9893 -7
- Partials 1167 1168 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
…se_atten - Added functionality to enable type embedding compression in the class, specifically for two-side mode. - Introduced a new method to precompute type embedding network outputs for all type pair combinations, optimizing inference by using precomputed values. - Updated the method to utilize precomputed embeddings when compression is enabled, improving performance during inference. This enhancement allows for more efficient handling of type embeddings in the descriptor model.
…andling in se_atten - Streamlined the index calculation for neighbor types by removing unnecessary reshaping and expansion. - Improved clarity in the handling of type embeddings, ensuring that the logic for two-side embeddings remains intact while enhancing readability. - Maintained functionality for both compressed and uncompressed type embeddings, optimizing the inference process. These changes contribute to cleaner code and maintain the performance benefits introduced in previous commits.
… se_atten - Simplified the type embedding compression process by consolidating methods and removing unnecessary conditions. - Enhanced clarity in the handling of one-side and two-side type embeddings, ensuring consistent functionality across both modes. - Updated comments for better understanding of the compression logic and its implications on performance. These changes contribute to cleaner code and improved maintainability of the descriptor model.
- Added runtime checks to ensure type embedding compression only operates in "strip" mode. - Introduced validation for the initialization of `filter_layers_strip` to prevent runtime errors. - Updated comments to clarify the expected dimensions of type embeddings, enhancing code readability. These changes improve error handling and maintain the integrity of the type embedding compression logic.
OutisLi
force-pushed
the
pr/atten_compress
branch
from
7f776fa to
4f72994
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
deepmd/pt/model/descriptor/se_atten.py (2)
639-663: Core logic is correct; consider improving shape comment consistency.The index calculation correctly matches the precomputation formula, and the conditional logic properly uses cached embeddings when available.
However, the shape comments at lines 645, 650, and 655 use
*as a dimension separator, which is inconsistent with standard Python tuple notation used elsewhere (e.g., line 642). Consider updating for clarity:- # ((ntypes+1)^2) * (ntypes+1)^2 * nt + # (ntypes+1, ntypes+1, nt) type_embedding_nei = torch.tile( type_embedding.view(1, ntypes_with_padding, nt), [ntypes_with_padding, 1, 1], ) - # (ntypes+1)^2 * ((ntypes+1)^2) * nt + # (ntypes+1, ntypes+1, nt) type_embedding_center = torch.tile( type_embedding.view(ntypes_with_padding, 1, nt), [1, ntypes_with_padding, 1], ) - # ((ntypes+1)^2 * (ntypes+1)^2) * (nt+nt) + # ((ntypes+1)^2, 2*nt) two_side_type_embedding = torch.cat( [type_embedding_nei, type_embedding_center], -1 ).reshape(-1, nt * 2)
467-472: Static analysis: Consider shorter exception messages (optional).Ruff flags TRY003 warnings for the exception messages. While the current messages are clear and helpful, you could optionally shorten them or define custom exception classes if you want to follow strict style guidelines. Given the clarity of the current implementation, this is purely a nitpick.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/model/descriptor/dpa1.py(1 hunks)deepmd/pt/model/descriptor/se_atten.py(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/descriptor/dpa1.py
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/model/descriptor/se_atten.py (1)
deepmd/pt/model/network/network.py (2)
TypeEmbedNet(268-358)get_full_embedding(317-332)
🪛 Ruff (0.14.5)
deepmd/pt/model/descriptor/se_atten.py
468-468: Avoid specifying long messages outside the exception class
(TRY003)
470-472: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- 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: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.9)
🔇 Additional comments (4)
deepmd/pt/model/descriptor/se_atten.py (4)
30-32: LGTM!The TypeEmbedNet import is necessary for the new type_embedding_compression method and is correctly placed.
278-288: LGTM!Clear separation between geometric compression state and type embedding compression cache. The None initialization is appropriate for optional cached data.
467-472: Validation checks correctly implemented.The validation checks match the requirements from past review comments, ensuring the method is only called in strip mode with initialized filter_layers_strip.
627-633: LGTM!The conditional logic correctly uses cached embeddings when available and falls back to on-the-fly computation otherwise, ensuring backward compatibility.
OutisLi
force-pushed
the
pr/atten_compress
branch
from
93bb015 to
4f72994
Compare
- Replaced direct assignment of `type_embd_data` with a call to `register_buffer` for better memory management. - Improved clarity by using a temporary variable `embd_tensor` for storing the output of the embedding network before registration. - Maintained functionality for both one-side and two-side type embeddings, ensuring consistent behavior across modes. These changes enhance the maintainability and performance of the descriptor model.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/model/descriptor/se_atten.py (1)
278-289: Critical bug confirmed:register_bufferwill raiseKeyErrorat runtimePyTorch's
register_bufferraisesKeyErrorif an attribute with that name already exists on the Module and is not already a registered buffer. The current code at line 288 setstype_embd_dataas a plain attribute, and then lines 483 and 499 attempt to register it as a buffer, which will fail with the error:KeyError("attribute 'type_embd_data' already exists").The review's suggested fix is correct and idiomatic: register the buffer in
__init__and update it later via direct assignment, as this keeps the registered buffer object and its registration. Simply change line 288 to register the buffer withNoneinitially, then assign the computed tensor directly intype_embedding_compression.Apply the suggested diff to fix the immediate runtime error.
🧹 Nitpick comments (1)
deepmd/pt/model/descriptor/se_atten.py (1)
455-472: Strip‑mode type‑embedding compression logic looks consistent; only minor comment nitsFunctionally, the new
type_embedding_compressionand its use in the strip‑mode forward path look sound:
- One‑side case: precomputes
filter_layers_strip.networks[0](full_embd)with shape(ntypes+1, ng)and reuses it via direct indexing in forward — matching the previous on‑the‑fly behavior.- Two‑side case: constructs
(neighbor, center)pairs so that the flattened layout matchesidx = center_type * ntypes_with_padding + neighbor_typeused in forward; the reuse viatt_full[idx]is consistent with the original computation.- Forward’s
if self.type_embd_data is not Nonebranches correctly fall back to the old path when precomputation is disabled, so behavior is preserved when compression isn’t used.The only issues here are cosmetic: a few shape comments around Lines 645‑656 are slightly misleading (e.g., repeated
((ntypes+1)^2) * ...text rather than clearly stating the 3‑D tensor shapes), but they don’t affect correctness. You may optionally clean those up for future readers.Also applies to: 627-633, 640-643, 645-661, 663-663
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/model/descriptor/se_atten.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/model/descriptor/se_atten.py (1)
deepmd/pt/model/network/network.py (3)
TypeEmbedNet(268-358)Tensor(36-37)get_full_embedding(317-332)
🪛 Ruff (0.14.5)
deepmd/pt/model/descriptor/se_atten.py
468-468: Avoid specifying long messages outside the exception class
(TRY003)
470-472: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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)
- GitHub Check: Build C++ (clang, clang)
- 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++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test C++ (true)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test C++ (false)
🔇 Additional comments (1)
deepmd/pt/model/descriptor/se_atten.py (1)
30-32: TypeEmbedNet import matches new compression APIImporting
TypeEmbedNetfromdeepmd.pt.model.network.networkaligns with the newtype_embedding_compressionsignature and the provided implementation ofTypeEmbedNet.get_full_embedding; no issues here.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/pt/model/descriptor/se_atten.py (2)
455-504: type_embedding_compression implementation is sound; consider minor robustness and style tweaksThe precomputation logic for both one-side and two-side cases, plus registering
type_embd_dataas a buffer, aligns with the forward-path indexing and fixes prior device-movement issues—looks correct overall.Non-blocking refinements you might consider:
- Add light shape checks on
full_embd(e.g.,nt == self.ntypes + 1andt_dim == self.tebd_dim) to fail fast with clearer errors if a mismatchedTypeEmbedNetis accidentally wired in.- Ruff’s TRY003 hints here: if you want to satisfy it, you could slightly shorten the
RuntimeErrormessages or factor them into dedicated exception types/constants instead of long inline strings.All of these are optional; behavior as written is fine.
631-667: Strip-mode forward reuse of cached embeddings is consistent; comment shapes could be clearerThe new checks on
self.type_embd_datain both one-side and two-side strip paths correctly match the buffer layout fromtype_embedding_compression, and the index construction (center_type * ntypes_with_padding + neighbor_type) is consistent with the(ntypes+1)²table size.For readability only, the shape comments in the two-side fallback (around lines 649–655) currently read like chained products (e.g.
((ntypes+1)^2) * (ntypes+1)^2 * nt); rewriting them in a more conventional tensor-shape style such as(ntypes_with_padding, ntypes_with_padding, nt)would make the intent clearer to future readers.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/model/descriptor/se_atten.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/model/descriptor/se_atten.py (1)
deepmd/pt/model/network/network.py (3)
TypeEmbedNet(268-358)Tensor(36-37)get_full_embedding(317-332)
🪛 Ruff (0.14.5)
deepmd/pt/model/descriptor/se_atten.py
468-468: Avoid specifying long messages outside the exception class
(TRY003)
470-472: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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). (28)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Analyze (c-cpp)
- 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)
🔇 Additional comments (2)
deepmd/pt/model/descriptor/se_atten.py (2)
30-32: TypeEmbedNet import usage is appropriateImporting
TypeEmbedNethere cleanly supports thetype_embedding_compressionsignature without introducing unused or circular dependencies in the shown context. No issues.
278-288: Compression state separation looks goodGrouping
compress_*and the newtype_embd_datafield together makes it clear which state belongs to geometric vs type-embedding compression, and theOptional[torch.Tensor]initialization aligns with the later buffer registration logic. Nothing blocking here.
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
Could you add a simple ut for se_atten compression? See source/tests/pt/test_model_compression_se_a.py
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
source/tests/pt/test_model_compression_se_atten.py (1)
209-808: Consider adding consistent cleanup across test classes.Only
TestDeepPotAPBCExcludeTypeshas atearDownClassmethod. Consider adding cleanup to other test classes or adding a module-leveltearDownModulefunction to remove all test artifacts consistently.Example module-level cleanup:
def tearDownModule() -> None: # Clean up files created by _init_models _file_delete(INPUT) _file_delete(FROZEN_MODEL) _file_delete(COMPRESSED_MODEL) # Clean up files created by _init_models_skip_neighbor_stat _file_delete(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) _file_delete(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) # Clean up files created by _init_models_exclude_types _file_delete(INPUT_ET) _file_delete(FROZEN_MODEL_ET) _file_delete(COMPRESSED_MODEL_ET) # Clean up other artifacts _file_delete("out.json") _file_delete("compress.json") _file_delete("checkpoint") _file_delete("lcurve.out") _file_delete("model.ckpt") _file_delete("model-compression/checkpoint") _file_delete("model-compression")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_model_compression_se_atten.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2024-09-19T04:25:12.408Z
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-09-19T04:25:12.408Z
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 (1)
deepmd/calculator.py (1)
DP(36-159)
⏰ 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)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (6)
source/tests/pt/test_model_compression_se_atten.py (6)
1-25: LGTM!The imports and precision setup follow standard patterns for deepmd-kit test files.
34-81: LGTM!The model initialization correctly configures se_atten with compression-compatible settings and documents the critical
attn_layer=0requirement.
84-132: LGTM!The exclude_types variant properly tests compression with type exclusion configured.
135-189: LGTM!The skip-neighbor-stat variant correctly uses
--skip-neighbor-statduring training and provides the input config during compression via-tflag.
192-207: LGTM!Module-level setup appropriately initializes all model variants needed for the test suite.
810-811: LGTM!Standard unittest main block.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
source/tests/pt/test_model_compression_se_atten.py (1)
27-31: Fix cleanup helper to handle non‑empty directories safely
os.rmdironly removes empty directories; intearDownClassthis helper is used on paths like"checkpoint"and"model-compression", which are expected to contain files. This can raiseOSErrorand make tests fail during cleanup instead of during assertions.Consider a more robust implementation:
-import os +import os +import shutil @@ -def _file_delete(file) -> None: - if os.path.isdir(file): - os.rmdir(file) - elif os.path.isfile(file): - os.remove(file) +def _file_delete(path: str) -> None: + if not os.path.exists(path): + return + if os.path.isdir(path): + shutil.rmtree(path, ignore_errors=True) + else: + try: + os.remove(path) + except OSError: + # Best-effort cleanup; don't fail tests just because deletion failed + pass
🧹 Nitpick comments (2)
source/tests/pt/test_model_compression_se_atten.py (2)
34-190: Consider factoring shared model‑init logic to reduce duplication
_init_models,_init_models_exclude_types, and_init_models_skip_neighbor_statduplicate most of the descriptor andtype_embeddingsetup and training/compression pipeline, differing only in a few keys and CLI flags. A small helper that accepts overrides (e.g.,exclude_typesandskip_neighbor_stat) would make future descriptor schema changes less error‑prone.
192-207: Optional: centralize cleanup for all generated models and artifacts
setUpModulecreates multiple models (FROZEN_MODEL,COMPRESSED_MODEL,*_SKIP_NEIGHBOR_STAT,*_ET), but onlyTestDeepPotAPBCExcludeTypes.tearDownClasscleans up a subset of artifacts. Not strictly wrong, but you could:
- Add a
tearDownModule()that calls_file_deleteon all generated models and common files, or- Extend the existing teardown to also cover the non‑ET variants,
to keep the test data directory clean across repeated runs and reduce potential crosstalk with other tests sharing
tests_path.Also applies to: 579-591
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_model_compression_se_atten.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2024-09-19T04:25:12.408Z
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-09-19T04:25:12.408Z
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 (1)
deepmd/calculator.py (1)
DP(36-159)
⏰ 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). (30)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
🔇 Additional comments (1)
source/tests/pt/test_model_compression_se_atten.py (1)
519-546: Import verified — code change is correctThe
from deepmd.calculator import DPimport is valid and confirmed. The DP class is properly defined and exported indeepmd/calculator.py(line 33:__all__ = ["DP"]), and this import pattern is already used throughout the codebase in multiple PT tests and utilities (source/tests/infer/test_models.py,deepmd/pt/utils/ase_calc.py, etc.), confirming it's the standard approach for ASE/PT integration.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
source/tests/pt/test_model_compression_se_atten.py (2)
520-547: Consider expanding ASE calculator test coverage.The
test_asemethod validates the compressed model via the ASE interface, but it only exists inTestDeepPotALargeBoxNoPBC. Extending this test to other scenarios (PBC, exclude_types, skip-neighbor-stat) would strengthen verification that compression preserves correctness across all integration paths.
210-795: Optional: Consider reducing test code duplication.The five test classes share significant boilerplate: identical coordinate arrays,
atypelists, and near-identical test method implementations. Extracting common test logic to a base class or using parameterized tests (e.g.,@pytest.mark.parametrizeorunittest'ssubTest) could reduce maintenance burden and improve readability.Example structure:
class BaseCompressionTest(unittest.TestCase): """Base class with shared test logic""" def _compare_eval(self, atomic=False, nframes=1): # Common comparison logic pass def test_1frame(self): self._compare_eval(atomic=False, nframes=1)This is a nice-to-have improvement and can be deferred if test clarity is prioritized over DRY principles.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_model_compression_se_atten.py(1 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 (1)
deepmd/calculator.py (1)
DP(36-159)
⏰ 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)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
source/tests/pt/test_model_compression_se_atten.py (1)
85-104: Inconsistenttebd_dimconfiguration across test scenarios.The first model initialization (
_init_models) explicitly sets"tebd_dim": 8(line 49), but this configuration omits it, allowing the descriptor to use its default value. Since this PR introduces type embedding compression andtebd_dimcontrols the type embedding dimension, inconsistent values across test scenarios may lead to incomplete test coverage or unexpected behavior.Please clarify whether omitting
tebd_dimhere is intentional (to test default behavior) or an oversight. If testing the default is not the goal, consider adding it for consistency:jdata["model"]["descriptor"] = { "type": "se_atten", "exclude_types": [[0, 1]], "sel": 120, "rcut_smth": 0.50, "rcut": 6.00, "neuron": [25, 50, 100], + "tebd_dim": 8, "axis_neuron": 16, "attn": 64,The same concern applies to
_init_models_skip_neighbor_stat(lines 130-148).
- Replaced the optional `type_embd_data` with a registered buffer to improve memory management. - Adjusted conditional checks for type embedding usage to rely on the `compress` attribute, enhancing clarity and functionality. - Removed the associated test file for model compression of se_atten, streamlining the test suite. These changes enhance the maintainability and performance of the descriptor model while ensuring consistent behavior in type embedding compression.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pt/model/descriptor/se_atten.py (1)
469-474: Consider extracting error messages to reduce line length.The validation logic is correct and addresses previous review concerns. However, the error messages are verbose. Consider defining them as constants or extracting to a separate line for better maintainability.
Based on static analysis hints.
Apply this diff:
def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: """Enable type embedding compression for strip mode. Precomputes embedding network outputs for all type combinations: - One-side: (ntypes+1) combinations (neighbor types only) - Two-side: (ntypes+1)² combinations (neighbor x center type pairs) Parameters ---------- type_embedding_net : TypeEmbedNet The type embedding network that provides get_full_embedding() method """ + strip_mode_error = "Type embedding compression only works in strip mode" + filter_init_error = ( + "filter_layers_strip must be initialized for type embedding compression" + ) if self.tebd_input_mode != "strip": - raise RuntimeError("Type embedding compression only works in strip mode") + raise RuntimeError(strip_mode_error) if self.filter_layers_strip is None: - raise RuntimeError( - "filter_layers_strip must be initialized for type embedding compression" - ) + raise RuntimeError(filter_init_error)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/model/descriptor/se_atten.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/model/descriptor/se_atten.py (1)
deepmd/pt/model/network/network.py (2)
TypeEmbedNet(268-358)get_full_embedding(317-332)
🪛 Ruff (0.14.5)
deepmd/pt/model/descriptor/se_atten.py
470-470: Avoid specifying long messages outside the exception class
(TRY003)
472-474: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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)
- 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++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (4)
deepmd/pt/model/descriptor/se_atten.py (4)
30-32: LGTM!The import of
TypeEmbedNetis necessary for the new type embedding compression feature, specifically to access theget_full_embedding()method.
287-290: LGTM!Using
register_buffer()fortype_embd_dataensures it automatically moves with the model during device transfers (e.g.,.to('cuda')). This correctly addresses the device handling concern from previous reviews.
476-506: LGTM!The precomputation logic is correct and efficient:
- One-side path: Precomputes embeddings for all (ntypes+1) neighbor types
- Two-side path: Precomputes embeddings for all (ntypes+1)² neighbor×center pairs using memory-efficient broadcasting
The indexing formula (
idx = center_type * nt + neighbor_typeon line 646) correctly matches the flattened 2D structure created here. Usingtorch.no_grad()and.detach()prevents gradient tracking for cached data.
633-669: LGTM!The forward method correctly integrates the cached type embeddings:
- One-side path (633-639): Uses cached
type_embd_datawhen compression is enabled, otherwise computes on-the-fly- Two-side path (641-669): Uses cached data indexed by
center_type * nt + neighbor_type, with fallback that correctly replicates the precomputation structureThe branching logic preserves existing behavior when compression is disabled while enabling the performance optimization when enabled.
This can be merged first, I will commit another request later. I am having some trouble with the test. |
4 participants
Summary by CodeRabbit
New Features
Behavioral Improvements
Tests
✏️ Tip: You can customize this high-level summary in your review settings.