feat: add use_loc_mapping #4772
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughA new boolean parameter Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant DescrptDPA3
participant DescrptBlockRepflows
participant get_graph_index
User->>DescrptDPA3: call forward/call(coord_ext, atype_ext, nlist, mapping, ...)
DescrptDPA3->>DescrptBlockRepflows: call forward/call(..., use_loc_mapping)
DescrptBlockRepflows->>get_graph_index: get_graph_index(nlist, ..., use_loc_mapping)
get_graph_index-->>DescrptBlockRepflows: returns graph indices (with or without local mapping)
DescrptBlockRepflows-->>DescrptDPA3: returns descriptor output
DescrptDPA3-->>User: returns descriptor output
Suggested reviewers
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
deepmd/pt/model/descriptor/dpa3.py (1)
478-478: Remove unused variable assignment.The variable
nallis assigned but never used in this method.- nall = extended_coord.view(nframes, -1).shape[1] // 3🧰 Tools
🪛 Ruff (0.11.9)
478-478: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
deepmd/pt/model/descriptor/repflows.py (1)
542-542: Rename unused loop variable.The loop variable
idxis not used within the loop body.Apply this diff to follow Python conventions:
- for idx, ll in enumerate(self.layers): + for _, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.11.9)
542-542: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
deepmd/pt/model/descriptor/dpa3.py(2 hunks)deepmd/pt/model/descriptor/repflow_layer.py(1 hunks)deepmd/pt/model/descriptor/repflows.py(5 hunks)deepmd/pt/model/network/utils.py(2 hunks)deepmd/utils/argcheck.py(1 hunks)source/tests/pt/model/test_loc_mapping.py(1 hunks)source/tests/universal/dpmodel/descriptor/test_descriptor.py(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
deepmd/pt/model/descriptor/repflows.py (1)
deepmd/pt/model/network/utils.py (1)
get_graph_index(49-138)
🪛 Ruff (0.11.9)
deepmd/pt/model/descriptor/repflows.py
542-542: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
deepmd/pt/model/descriptor/dpa3.py
156-156: Undefined name use_loc_mapping
(F821)
164-164: Undefined name use_loc_mapping
(F821)
478-478: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🔇 Additional comments (15)
deepmd/pt/model/descriptor/repflow_layer.py (1)
687-687: LGTM! Documentation improvement aligns with the new feature.The comment update accurately reflects that
node_ebd_extcan have different shapes depending on whether local mapping is used in non-parallel mode, which aligns with the newuse_loc_mappingfunctionality.source/tests/universal/dpmodel/descriptor/test_descriptor.py (3)
489-489: LGTM! Proper addition of the new parameter.The
use_loc_mappingparameter is correctly added with a sensible default value ofTrue.
535-535: LGTM! Parameter properly passed to configuration.The parameter is correctly included in the input dictionary for descriptor configuration.
560-560: LGTM! Comprehensive test parameterization.The parameterization with
(True, False)ensures that both local mapping modes are tested, providing good test coverage for the new feature.deepmd/pt/model/network/utils.py (2)
54-54: LGTM! Well-designed parameter addition.The
use_loc_mappingparameter is properly added with a sensible default value ofTrueand correct type annotation.
104-106: LGTM! Correct conditional logic implementation.The frame shift calculation properly switches between
nall(global mapping) andnloc(local mapping) based on theuse_loc_mappingflag. This enables the core functionality for local vs global atom indexing.deepmd/pt/model/descriptor/dpa3.py (2)
474-474: LGTM! Clear parallel mode detection.The
parrallel_modevariable provides a clear way to detect parallel inference mode based on the presence ofcomm_dict.
480-483: LGTM! Correct conditional embedding logic.The conditional logic properly switches between local and global atom type embeddings based on parallel mode and the
use_loc_mappingflag. When not in parallel mode and local mapping is enabled, it uses only the local subset of atom types.deepmd/pt/model/descriptor/repflows.py (5)
212-212: LGTM! Well-designed backward-compatible parameter.The addition of
use_loc_mappingwith a default value ofTruemaintains backward compatibility while providing control over local atom index mapping behavior.
420-422: Good defensive programming with parallel mode detection.The addition of
parallel_modeflag and the assertion formappingin non-parallel mode ensures proper data flow and catches potential issues early.
474-476: Proper handling of extended atom type embeddings.The unconditional assertion and extraction of local atom embeddings from the extended embeddings is correct and aligns with the local mapping design.
495-510: Well-implemented neighbor list remapping logic.The conditional remapping of neighbor indices from global to local using
torch.gatheris correctly implemented. The logic properly:
- Checks for non-parallel mode and
use_loc_mappingflag- Reshapes for gathering and restores original shape
- Passes the
use_loc_mappingflag toget_graph_indexfor consistent behavior
537-553: Correct implementation of node embedding extension logic.The conditional logic properly handles node embedding extension based on
use_loc_mapping:
- When
use_loc_mappingis True in non-parallel mode,node_ebd_extuses local embeddings directly- Otherwise, it gathers from global embeddings using the mapping tensor
This aligns perfectly with the neighbor list remapping logic above.
🧰 Tools
🪛 Ruff (0.11.9)
542-542: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
source/tests/pt/model/test_loc_mapping.py (2)
34-146: Excellent test coverage for the local mapping feature.The
test_consistencymethod thoroughly validates that the descriptor outputs remain consistent regardless of theuse_loc_mappingsetting. The test:
- Covers multiple parameter combinations through itertools.product
- Uses duplicated data to ensure identical inputs
- Properly handles precision tolerances for GPU testing
- Sets consistent random seeds for reproducibility
148-258: Well-designed test for dynamic selection mode.The
test_consistency_noselmethod extends the testing to cover dynamic selection (use_dynamic_sel=True) with reduction factor, ensuring the local mapping feature works correctly in this mode as well.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Chun Cai <amoycaic@gmail.com>
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
deepmd/pt/model/descriptor/dpa3.py (3)
475-475: Fix the typo in variable name.The variable name has a spelling error.
Apply this diff to fix the typo:
- parrallel_mode = comm_dict is not None + parallel_mode = comm_dict is not NoneAlso update the usage in the conditional below:
- if not parrallel_mode and self.use_loc_mapping: + if not parallel_mode and self.use_loc_mapping:
481-484: LGTM! Conditional embedding logic is correct.The conditional logic properly implements the local mapping feature by computing embeddings only for local atoms when not in parallel mode and
use_loc_mappingis enabled.Consider adding the
use_loc_mappingparameter to the class docstring for better documentation:use_tebd_bias : bool, Optional Whether to use bias in the type embedding layer. + use_loc_mapping : bool, Optional + Whether to use local atom index mapping in non-parallel inference. type_map : list[str], Optional A list of strings. Give the name to each type of atoms.
479-479: Remove unused variable.The variable
nallis computed but never used in the function.Apply this diff to remove the unused variable:
nframes, nloc, nnei = nlist.shape - nall = extended_coord.view(nframes, -1).shape[1] // 3🧰 Tools
🪛 Ruff (0.11.9)
479-479: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/model/descriptor/dpa3.py(3 hunks)deepmd/utils/argcheck.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/utils/argcheck.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pt/model/descriptor/dpa3.py
479-479: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🔇 Additional comments (2)
deepmd/pt/model/descriptor/dpa3.py (2)
111-111: LGTM! Parameter addition is correct.The
use_loc_mappingparameter is properly added with appropriate type annotation and default value for backward compatibility.
157-157: LGTM! Parameter propagation is correct.The
use_loc_mappingparameter is properly passed to theDescrptBlockRepflowsconstructor and stored as an instance attribute for use in the forward method.Also applies to: 165-165
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pt/model/descriptor/repflows.py (1)
538-552: Well-designed conditional node embedding logic.The conditional construction of
node_ebd_extproperly handles different execution modes:
- Non-parallel with
use_loc_mapping=False: Usestorch.gatherfor global mapping (preserves original behavior)- Non-parallel with
use_loc_mapping=True: Usesnode_ebddirectly (new optimization)- Parallel mode: Uses existing communication logic
However, address the unused loop variable:
- for idx, ll in enumerate(self.layers): + for _idx, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.11.9)
543-543: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/pt/model/descriptor/dpa3.py(4 hunks)deepmd/pt/model/descriptor/repflow_layer.py(1 hunks)deepmd/pt/model/descriptor/repflows.py(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/descriptor/repflow_layer.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/pt/model/descriptor/repflows.py
543-543: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
deepmd/pt/model/descriptor/dpa3.py
481-481: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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 (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 (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
🔇 Additional comments (9)
deepmd/pt/model/descriptor/dpa3.py (4)
92-93: LGTM: Clear parameter documentation.The documentation for
use_loc_mappingis well-written and clearly explains its purpose for controlling local atom index mapping in non-parallel inference.
113-113: Good fix: Constructor parameter properly added.This correctly addresses the previous review comment about the missing
use_loc_mappingparameter in the constructor. The default value ofTruemaintains backward compatibility.
159-159: LGTM: Proper parameter handling.The
use_loc_mappingparameter is correctly passed to theDescrptBlockRepflowsconstructor and stored as an instance attribute following standard patterns.Also applies to: 167-167
477-477: Well-implemented conditional logic.The parallel mode detection and conditional type embedding computation are well-designed:
- Clean parallel mode detection using
comm_dict is not None- Logical optimization: when not in parallel mode and
use_loc_mapping=True, compute embeddings only on local atoms (extended_atype[:, :nloc]) rather than all extended atoms- Maintains backward compatibility when
use_loc_mapping=FalseAlso applies to: 483-486
deepmd/pt/model/descriptor/repflows.py (5)
212-212: LGTM: Consistent parameter implementation.The
use_loc_mappingparameter is properly added to the constructor and stored as an instance variable, following the same pattern as indpa3.py.Also applies to: 243-243
421-422: Good refactoring: Improved readability.The parallel mode detection is consistent with
dpa3.py, and refactoring fromcomm_dict is Nonetonot parallel_modeimproves code readability and maintainability.
475-477: Good defensive programming.The assertions ensure that
extended_atype_embdis available and properly shaped, which helps catch issues early in the pipeline.
496-503: Core optimization implementation looks correct.The neighbor list remapping logic properly converts global indices to local indices when
use_loc_mapping=Truein non-parallel mode. The use oftorch.gatherwith proper reshaping is the correct approach for this index transformation.
507-511: Proper parameter propagation.The
use_loc_mappingparameter is correctly passed toget_graph_index, ensuring consistent behavior throughout the neighbor graph construction pipeline.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/dpmodel/descriptor/repflows.py (1)
541-541: Remove trailing whitespace.- ) - + ) +
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/dpmodel/descriptor/dpa3.py(5 hunks)deepmd/dpmodel/descriptor/repflows.py(5 hunks)deepmd/dpmodel/utils/network.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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 (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: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (cpu, cpu)
🔇 Additional comments (3)
deepmd/dpmodel/utils/network.py (1)
1004-1026: LGTM! Theuse_loc_mappingparameter is correctly implemented.The changes properly implement the conditional frame shift calculation based on whether local atom index mapping is used. The default value of
Truemaintains backward compatibility.Also applies to: 1066-1066
deepmd/dpmodel/descriptor/dpa3.py (1)
274-275: LGTM! Theuse_loc_mappingparameter is properly integrated into the DPA3 descriptor.The implementation correctly handles the conditional selection of type embeddings based on whether local atom index mapping is used. The parameter is properly documented and propagated to child components.
Also applies to: 295-295, 341-341, 350-350, 549-558
deepmd/dpmodel/descriptor/repflows.py (1)
148-149: Theuse_loc_mappingimplementation in repflows is well-designed.The neighbor list remapping and node embedding handling are correctly implemented. The logic properly differentiates between local and global atom indexing modes, which is essential for non-parallel inference scenarios.
Also applies to: 201-201, 235-235, 534-540, 545-545, 576-580
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
source/tests/pd/model/test_dpa3.py (1)
73-76: Test modification foruse_loc_mappingfeature looks correct.The change to use duplicated input frames is consistent with testing the new
use_loc_mappingfeature. The implementation correctly duplicates the first frame across all input arrays (coord_ext, atype_ext, nlist, mapping) using proper array concatenation.However, consider that this modification reduces test coverage by only testing with identical duplicated frames rather than varied input data.
Consider adding additional test cases that use varied frames to ensure the descriptor handles diverse input scenarios correctly:
+ # Test with varied frames as well + coord_ext_varied = np.concatenate([self.coord_ext[:1], self.coord_ext[1:2] if self.coord_ext.shape[0] > 1 else self.coord_ext[:1]], axis=0) + # Similar for other arrays...
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pd/model/test_dpa3.py(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
source/tests/pd/model/test_dpa3.py (1)
deepmd/pd/model/descriptor/dpa3.py (3)
DescrptDPA3(65-573)deserialize(405-435)serialize(368-402)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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 (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Analyze (python)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (1)
source/tests/pd/model/test_dpa3.py (1)
114-117: Consistent usage of duplicated input arrays.The test correctly uses the new duplicated input arrays throughout all descriptor calls (dd0, dd1 deserialized, and dd2 dp implementation). This maintains consistency and ensures all three implementations are tested with the same input data.
Also applies to: 122-125, 136-136
for more information, see https://pre-commit.ci
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## devel #4772 +/- ##
==========================================
- Coverage 84.79% 84.78% -0.01%
==========================================
Files 698 698
Lines 67734 67775 +41
Branches 3540 3542 +2
==========================================
+ Hits 57432 57465 +33
- Misses 9171 9174 +3
- Partials 1131 1136 +5 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
A consistent test also need to be added in source/tests/consistent/descriptor/test_dpa3.py (and skip pd), see https://github.com/deepmodeling/deepmd-kit/pull/4760/files#diff-2175c5f9ca14ee18e75a5fe46c7ecd66d02fae4612599a7763a2e29acdf12943 for example.
Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| optim_update, | ||
| use_exp_switch, | ||
| use_dynamic_sel, | ||
| use_loc_mapping, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/pt/model/descriptor/repflows.py (1)
546-546: Address static analysis hint about unused loop variable.The loop variable
idxis not used within the loop body. Consider renaming it to_idxto indicate it's intentionally unused, or remove the enumeration if the index isn't needed.- for idx, ll in enumerate(self.layers): + for _, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.11.9)
546-546: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
deepmd/pt/model/descriptor/dpa3.py (1)
483-483: Remove unused variable assignment.The
nallvariable is assigned but never used in the current implementation.Apply this diff to remove the unused variable:
- nall = extended_coord.view(nframes, -1).shape[1] // 3🧰 Tools
🪛 Ruff (0.11.9)
483-483: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
deepmd/dpmodel/descriptor/dpa3.py(8 hunks)deepmd/dpmodel/descriptor/repflows.py(6 hunks)deepmd/pd/model/descriptor/dpa3.py(7 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pt/model/descriptor/dpa3.py(7 hunks)deepmd/pt/model/descriptor/repflows.py(7 hunks)deepmd/utils/argcheck.py(2 hunks)source/tests/consistent/descriptor/test_dpa3.py(11 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/utils/argcheck.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/descriptor/dpa3.py
- deepmd/dpmodel/descriptor/repflows.py
🧰 Additional context used
🧬 Code Graph Analysis (2)
deepmd/pd/model/descriptor/dpa3.py (1)
deepmd/utils/version.py (1)
check_version_compatibility(2-27)
deepmd/pt/model/descriptor/dpa3.py (1)
deepmd/utils/version.py (1)
check_version_compatibility(2-27)
🪛 Ruff (0.11.9)
deepmd/pt/model/descriptor/dpa3.py
483-483: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pt/model/descriptor/repflows.py
546-546: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (3, 3.9)
- 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 (5, 3.9)
- 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 (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-macosx_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-manylinux_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
🔇 Additional comments (20)
deepmd/pd/model/descriptor/repflows.py (2)
115-117: Good documentation of the new parameter.The docstring clearly explains the purpose of
use_loc_mappingand explicitly states it's not supported yet in Paddle, which sets proper user expectations.
167-167: Appropriate preparatory implementation for interface consistency.The parameter is added with a safe default (
False) and immediately asserted to prevent usage until the feature is fully implemented in Paddle. This maintains API consistency with the PyTorch implementation while preventing incorrect usage.Also applies to: 203-204
source/tests/consistent/descriptor/test_dpa3.py (2)
70-70: Comprehensive test parameterization.Good addition of
use_loc_mappingto the parameterized test inputs to ensure both True and False values are tested across different configurations.
174-174: Correct skip logic for unsupported feature.Properly skips Paddle tests when
use_loc_mappingis enabled, which is consistent with the Paddle implementation that explicitly disables this feature.deepmd/pd/model/descriptor/dpa3.py (3)
92-94: Clear documentation with implementation status.The docstring appropriately documents the new parameter and clearly states it's not supported yet in Paddle, maintaining transparency about feature availability.
159-159: Proper parameter integration.The parameter is correctly passed to
DescrptBlockRepflowsand stored as an instance attribute, maintaining consistency with the overall architecture.Also applies to: 168-168
379-379: Excellent serialization version handling.The version increment from 1 to 2, inclusion of
use_loc_mappingin serialization, and updated compatibility check properly handle the new parameter while maintaining backward compatibility.Also applies to: 390-390, 415-415
deepmd/pt/model/descriptor/repflows.py (5)
159-161: Well-documented new parameter.The docstring clearly explains the purpose and behavior of
use_loc_mapping, providing good context for when and how it should be used.
215-215: Verify default value choice.The default value is
Truein PyTorch butFalsein Paddle. Since this is a memory optimization feature that's fully implemented in PyTorch, theTruedefault makes sense to enable the optimization by default.Please confirm that the different default values between PyTorch (
True) and Paddle (False) are intentional, given that Paddle doesn't support the feature yet.
424-426: Clear parallel mode detection.Good separation of logic between parallel and non-parallel modes, which helps maintain existing functionality while adding the new optimization.
499-506: Core memory optimization implementation.This is the key change that enables the memory optimization by converting neighbor lists from global (
nall) to local (nloc) indices whenuse_loc_mappingis enabled. This allows keeping embeddings at local size instead of expanding them.
541-555: Efficient embedding assignment strategy.The conditional logic correctly chooses between the traditional approach (expanding embeddings to
nallsize) and the optimized approach (keeping embeddings atnlocsize) based on theuse_loc_mappingflag.🧰 Tools
🪛 Ruff (0.11.9)
546-546: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
deepmd/pt/model/descriptor/dpa3.py (8)
92-94: LGTM! Clear documentation for the new parameter.The documentation clearly explains the purpose and behavior of
use_loc_mapping, addressing when it's applied and its effect on local indexing and mapping during descriptor computation.
114-114: LGTM! Constructor parameter properly added.The
use_loc_mappingparameter has been correctly added to the constructor with a sensible default value ofTrue. This addresses the previous review comment about the missing constructor parameter.
160-160: LGTM! Parameter properly passed and stored.The
use_loc_mappingparameter is correctly passed to theDescrptBlockRepflowsconstructor and stored as an instance attribute for later use.Also applies to: 168-168
374-374: LGTM! Version increment for serialization compatibility.The version has been appropriately incremented from 1 to 2 to reflect the addition of the new
use_loc_mappingparameter to the serialized data structure.
385-385: LGTM! Parameter included in serialization.The
use_loc_mappingparameter is correctly included in the serialized data to ensure the configuration is preserved when saving/loading models.
410-410: LGTM! Version compatibility properly updated.The version compatibility check has been correctly updated to accept both version 1 and 2, ensuring backward compatibility while supporting the new serialization format.
479-479: LGTM! Clear parallel mode detection.The
parallel_modevariable provides clear detection of whether the system is running in parallel inference mode by checking ifcomm_dictis not None.
485-488: LGTM! Memory optimization logic correctly implemented.The conditional logic properly implements the memory optimization strategy described in the PR objectives. When not in parallel mode and
use_loc_mappingis enabled, only local atoms (extended_atype[:, :nloc]) are processed for type embedding, reducing memory usage from O(nall) to O(nloc).
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
source/tests/consistent/descriptor/test_dpa3.py (1)
145-145: Consider removing unused variable unpacking.The
optim_updatevariable is unpacked but not used in this method. While this follows a consistent pattern across the class, consider whether all parameters need to be unpacked in every method.This issue was previously flagged by static analysis tools and remains unresolved.
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 145-145: Unused local variable
Variable optim_update is not used.
🧹 Nitpick comments (2)
source/tests/consistent/descriptor/test_dpa3.py (2)
63-63: Parameter restriction reduces test coverage.The
update_angleparameter has been restricted from(True, False)to(True,), which reduces test coverage. Ensure this restriction is intentional and doesn't miss important test scenarios.
72-72: Parameter restriction reduces test coverage.Similar to
update_angle, then_multi_edge_messageparameter has been restricted from multiple values to(1,). Consider whether this reduction in test coverage is appropriate.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/consistent/descriptor/test_dpa3.py(11 hunks)
🧰 Additional context used
🪛 GitHub Check: CodeQL
source/tests/consistent/descriptor/test_dpa3.py
[notice] 145-145: Unused local variable
Variable optim_update is not used.
⏰ Context from checks skipped due to timeout of 90000ms (26)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- 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-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
🔇 Additional comments (5)
source/tests/consistent/descriptor/test_dpa3.py (5)
70-70: New parameter integration looks good.The
use_loc_mappingparameter is properly added to the parameterized test configuration. Testing with bothTrueandFalsevalues ensures comprehensive coverage of the new feature.
88-88: Parameter unpacking correctly updated.The
use_loc_mappingparameter is properly added to the parameter unpacking tuple, maintaining consistency with the parameterized decorator.
129-129: Parameter correctly added to data dictionary.The
use_loc_mappingparameter is properly included in the data dictionary that configures the descriptor, ensuring the test parameter is passed through to the implementation.
164-164: Proper handling of unsupported feature in PD backend.The skip condition correctly excludes PD tests when
use_loc_mappingis enabled, acknowledging that this feature is not yet supported in the PaddlePaddle backend.Also applies to: 174-174
189-189: Consistent parameter handling across all methods.The
use_loc_mappingparameter is consistently unpacked in all relevant methods, maintaining the established pattern for parameter handling throughout the test class.Also applies to: 208-208, 269-269, 351-351, 376-376
2 participants
node_ebd_extcontains embedding on expanded atoms, which might be large for a large cut-off. Current implementation do matmul first, then gather tensor by the neighbors. This introduces savingsub_node_ext_updateof size nf * nall * ndim in each repflow layer, where nall might be multiple times larger than nloc.This PR do gathering first, then compute matmul. After this PR, the peak memory size is of O(nlayer * nf * nloc), unrelated to
nall.Summary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Style