breaking(pt/dp): tune new sub-structures for DPA2 #4089
Conversation
WalkthroughWalkthroughThe changes enhance the descriptor classes and functions within the DeepMD framework by integrating a three-body interaction mechanism and modifying existing parameters for improved configurability. New parameters are added across several classes and functions, enabling more complex configurations and interactions in the model. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Descriptor
participant RepinitArgs
participant RepformerArgs
participant RepformerLayer
User->>Descriptor: Initialize with parameters
Descriptor->>RepinitArgs: Create instance with three-body parameters
RepinitArgs->>Descriptor: Return initialized instance
Descriptor->>RepformerArgs: Create instance with new configuration options
RepformerArgs->>Descriptor: Return initialized instance
Descriptor->>RepformerLayer: Call forward method with new parameters
RepformerLayer->>Descriptor: Process and return output
Recent review detailsConfiguration used: CodeRabbit UI Files selected for processing (3)
Files skipped from review as they are similar to previous changes (1)
Additional context usedRuff
Additional comments not posted (8)
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
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 (
|
|
One question: should we set default values for these new params (e.g. Weighted average RMSE comparison on 27 datasets: 
Eamples will be added after this comment is resovled. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Outside diff range, codebase verification and nitpick comments (6)
source/tests/consistent/descriptor/test_dpa2.py (1)
40-40: Update all references to the renamed function.The old function name
descrpt_se_atten_argsis still present in several files. Please ensure that all occurrences are updated todescrpt_dpa2_argsto maintain consistency across the codebase.
source/tests/consistent/descriptor/test_se_atten_v2.pysource/tests/consistent/descriptor/test_dpa1.pydeepmd/utils/argcheck.pyAnalysis chain
Verify the usage of the new function name across the test suite.
The function name has been changed from
descrpt_se_atten_argstodescrpt_dpa2_args. Please ensure that all references to this function are updated accordingly.Run the following script to verify the function usage:
Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify all function calls to `descrpt_dpa2_args` match the new name. # Test: Search for the old function name. Expect: No occurrences. rg --type python -A 5 $'descrpt_se_atten_args'Length of output: 76
Script:
#!/bin/bash # Description: Verify all function calls to `descrpt_dpa2_args` match the new name. # Test: Search for the old function name. Expect: No occurrences. rg -t py -A 5 'descrpt_se_atten_args'Length of output: 2504
deepmd/pt/model/descriptor/dpa2.py (4)
Line range hint
447-471: Review of three-body handling inchange_type_map.The method now includes logic to handle the three-body representation when new types are introduced. This is a critical addition for maintaining state consistency. Ensure that the logic correctly updates the three-body representation and does not introduce any inconsistencies, especially when handling new types.
Consider adding unit tests to verify that the three-body representation is correctly updated when new types are introduced. This will help catch any potential issues early.
Propose adding unit tests for
change_type_map:def test_change_type_map_with_three_body(): # Setup initial state descriptor = DescrptDPA2(ntypes=2, repinit={'use_three_body': True}, repformer={}) original_three_body_state = descriptor.repinit_three_body.state_dict() # Change type map and check if three-body state is updated correctly new_type_map = ['H', 'He', 'Li'] descriptor.change_type_map(new_type_map) assert descriptor.repinit_three_body.state_dict() != original_three_body_state
Line range hint
539-614: Review of three-body serialization inserialize.The method now includes logic to serialize the three-body representation. This is a critical addition for ensuring that the state can be correctly saved and restored. Ensure that the serialization logic is robust and correctly handles all properties of the three-body representation, including any custom objects or configurations that might be part of the three-body setup.
Consider adding tests to verify the serialization and deserialization processes, ensuring that no data is lost or incorrectly handled during these operations.
Propose adding tests for serialization:
def test_three_body_serialization(): # Setup initial state with three-body interactions descriptor = DescrptDPA2(ntypes=2, repinit={'use_three_body': True}, repformer={}) serialized_data = descriptor.serialize() # Deserialize and verify that no data is lost new_descriptor = DescrptDPA2.deserialize(serialized_data) assert new_descriptor.repinit_three_body.state_dict() == descriptor.repinit_three_body.state_dict()
Line range hint
625-681: Review of three-body deserialization indeserialize.The method now includes logic to deserialize the three-body representation. This is a critical addition for ensuring that the state can be correctly restored from saved data. Ensure that the deserialization logic is robust and correctly handles all properties of the three-body representation, including any custom objects or configurations that might be part of the three-body setup.
Consider adding tests to verify the deserialization process, ensuring that no data is lost or incorrectly handled during this operation.
Propose adding tests for deserialization:
def test_three_body_deserialization(): # Setup initial state with three-body interactions and serialize descriptor = DescrptDPA2(ntypes=2, repinit={'use_three_body': True}, repformer={}) serialized_data = descriptor.serialize() # Deserialize and verify that no data is lost new_descriptor = DescrptDPA2.deserialize(serialized_data) assert new_descriptor.repinit_three_body.state_dict() == descriptor.repinit_three_body.state_dict()Tools
Ruff
685-685: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
Line range hint
737-773: Review of three-body handling inforward.The method now incorporates the output from the three-body representation into the overall output. This is a critical addition for ensuring that the descriptor correctly processes and includes three-body interactions in its output. Ensure that the logic correctly integrates the three-body representation and does not introduce any inconsistencies or errors in the output.
Consider adding integration tests to verify that the three-body interactions are correctly processed and included in the output. This will help catch any potential issues early.
Propose adding integration tests for
forward:def test_forward_with_three_body(): # Setup initial state with three-body interactions descriptor = DescrptDPA2(ntypes=2, repinit={'use_three_body': True}, repformer={}) # Mock input data extended_coord = torch.randn(10, 30, 3) extended_atype = torch.randint(0, 2, (10, 30)) nlist = torch.randint(0, 30, (10, 10, 5)) # Call forward and verify output includes three-body interactions output = descriptor.forward(extended_coord, extended_atype, nlist) assert output is not None # Add more specific checks based on expected output structuredeepmd/pt/model/descriptor/repformer_layer.py (1)
Line range hint
928-981: Enhanced neighbor count handling in_cal_hg.The modification to conditionally apply a square root transformation to the neighbor count based on
use_sqrt_nneiis a significant change. This feature can affect the model's performance and should be thoroughly tested to ensure it behaves as expected under various configurations.Consider adding unit tests to verify the behavior of the square root transformation on the neighbor count calculation to ensure its correctness and performance implications.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| if data["repinit"].use_three_body: | ||
| # deserialize repinit_three_body | ||
| statistic_repinit_three_body = repinit_three_body_variable.pop("@variables") | ||
| env_mat = repinit_three_body_variable.pop("env_mat") |
Check warning
Code scanning / CodeQL
Variable defined multiple times Warning
| @@ -162,6 +177,7 @@ | |||
| repinit_tebd_input_mode, | |||
| repinit_set_davg_zero, | |||
| repinit_type_one_side, | |||
| repinit_use_three_body, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -176,6 +192,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -176,6 +192,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
| repformer_g1_out_conv, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -176,6 +192,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
| repformer_g1_out_conv, | |||
| repformer_g1_out_mlp, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -234,6 +258,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
| repformer_g1_out_conv, | |||
| repformer_g1_out_mlp, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -285,6 +312,7 @@ | |||
| repinit_tebd_input_mode, | |||
| repinit_set_davg_zero, | |||
| repinit_type_one_side, | |||
| repinit_use_three_body, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -299,6 +327,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -299,6 +327,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
| repformer_g1_out_conv, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
| @@ -299,6 +327,9 @@ | |||
| repformer_set_davg_zero, | |||
| repformer_trainable_ln, | |||
| repformer_ln_eps, | |||
| repformer_use_sqrt_nnei, | |||
| repformer_g1_out_conv, | |||
| repformer_g1_out_mlp, | |||
Check notice
Code scanning / CodeQL
Unused local variable Note test
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4089 +/- ##
==========================================
- Coverage 83.05% 83.01% -0.05%
==========================================
Files 532 532
Lines 51971 52164 +193
Branches 3030 3030
==========================================
+ Hits 43167 43306 +139
- Misses 7857 7910 +53
- Partials 947 948 +1 ☔ View full report in Codecov by Sentry. |
I don't think it is a problem before the stable version is released. |
There was a problem hiding this comment.
Actionable comments posted: 4
Outside diff range, codebase verification and nitpick comments (3)
deepmd/pt/model/descriptor/dpa2.py (3)
Line range hint
433-470: Review: Handling of new types inchange_type_map.The method
change_type_maphas been extended to handle the three-body representation. This includes updating statistics for new types, which is a critical operation. It's important to ensure that these updates are performed correctly and that they do not introduce inconsistencies in the model's state, especially when new types are introduced.Consider adding tests to verify the correct handling of new types and the consistency of the model's state after such updates.
Line range hint
525-600: Review: Serialization ofrepinit_three_body.The serialization process now includes
repinit_three_body, which is essential for maintaining the state of the model across different sessions. It's crucial to ensure that the serialization and deserialization processes are symmetrical and that no data is lost or corrupted during these operations.Consider adding tests to verify the integrity of the serialization and deserialization processes, especially for the new three-body representation.
Line range hint
723-759: Review: Integration of three-body interactions inforwardmethod.The
forwardmethod has been updated to incorporate the output from the three-body representation. This is a significant change that affects the core functionality of the model. It's crucial to ensure that this integration is done correctly and that it does not introduce any performance regressions or bugs.Consider adding performance benchmarks and functional tests to verify that the integration of three-body interactions does not negatively impact the model's performance or accuracy.
iProzd
changed the title
There was a problem hiding this comment.
Actionable comments posted: 6
Outside diff range, codebase verification and nitpick comments (1)
deepmd/pt/model/descriptor/repformer_layer.py (1)
Line range hint
892-917: Update method forg1based ong1_out_conv.The method
_update_g1_convuses theg1_out_convparameter to determine the flow of data through different layers. This method's implementation is crucial for the correct application of the model's convolutional features based on the configuration.Ensure comprehensive testing of this method under different configurations to verify that both branches (conv and non-conv) perform as expected.
There was a problem hiding this comment.
Actionable comments posted: 1
Outside diff range, codebase verification and nitpick comments (1)
deepmd/pt/model/descriptor/dpa2.py (1)
Line range hint
525-600: Address variable redefinition and approve serialization logic.The serialization and deserialization logic correctly handles the new
repinit_three_bodyattribute. However, there is a static analysis warning about variable redefinition at line 655 that should be addressed to prevent potential bugs.Consider removing the redundant assignment to
env_matto resolve the warning:- env_mat = repinit_three_body_variable.pop("env_mat")Also applies to: 606-668
There was a problem hiding this comment.
Actionable comments posted: 0
Outside diff range, codebase verification and nitpick comments (1)
source/tests/pt/model/models/dpa2.json (1)
41-41: Discrepancy in Default Value foruse_sqrt_nneiThe
use_sqrt_nneiparameter is set tofalsein the JSON configuration but defaults toTruein the codebase. This discrepancy could lead to inconsistent behavior depending on the configuration source. Consider aligning the JSON configuration with the code default if the intention is to use square root normalization by default. Further testing or documentation review is recommended to confirm the intended default behavior.
- Files to Review:
source/tests/pt/model/models/dpa2.jsondeepmd/pt/model/descriptor/repformer_layer.pydeepmd/pt/model/descriptor/repformers.pydeepmd/dpmodel/descriptor/dpa2.pyAnalysis chain
New parameter:
use_sqrt_nneiThe addition of the
use_sqrt_nneiparameter allows for more granular control over the model's behavior. It is set tofalseby default, which is conservative and avoids breaking changes. However, consider the discussion in the PR about potentially setting this toTrueto enhance user experience by default.Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the impact of setting `use_sqrt_nnei` to `True` on model performance. # Test: Search for any performance metrics related to `use_sqrt_nnei`. Expect: Performance metrics showing the impact of changing this setting. rg --type json -A 5 $'use_sqrt_nnei'Length of output: 402
Script:
#!/bin/bash # Search for references to `use_sqrt_nnei` in the codebase to understand its impact. rg 'use_sqrt_nnei' -A 5Length of output: 25943
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Introduced support for three-body interactions in the descriptor model, enhancing predictive capabilities. - Added new parameters for greater configurability in the `RepinitArgs`, `RepformerArgs`, and related classes. - **Bug Fixes** - Improved handling of serialization and deserialization processes to include new parameters. - **Tests** - Expanded test cases to validate new configurations and parameters related to three-body interactions and model flexibility. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
Summary by CodeRabbit
New Features
RepinitArgs,RepformerArgs, and related classes.Bug Fixes
Tests