pd: support dpa1

[HydrogenSulfate]

Summary of this PR:

1. upload DPA-1 related code
2. merge much develop code
3. add all eager composite operators except softmax_grad, p_norm_grad, split_grad, and concat_grad to the composite operator blacklist(https://github.com/deepmodeling/deepmd-kit/pull/4414/files#diff-e678abb052b278f8a479f8d13b839a9ec0effd9923478a850bc13758f918e1e9R134-R148) to significantly improve model execution speed (reducing the time taken from 100% more than PyTorch to about 10% to 15% more).

related PR: lanpa/tensorboardX#728

Training curve:

Accuracy test(left: paddle, right: torch):

Ralated optimization of Paddle framework:

• [Setup] Update nvidia-cudnn-cu12 version from 9.0.0.312 PaddlePaddle/Paddle#69349
• [Fix] Fix astype copy when target dtype is same as source Tensor PaddlePaddle/Paddle#69333
• [Prim/Comp] Fix place setting to avoid redundant H2D/D2H copy for backward decomposition PaddlePaddle/Paddle#69479
• [Prim] Optimize the use of the reshape operator in eager composite backward PaddlePaddle/Paddle#69515
• [Comp] Fix place setting to avoid redundant H2D/D2H copy in forward decomposition PaddlePaddle/Paddle#69487
• [API] Optimize clip_grad_norm PaddlePaddle/Paddle#69661
• [Kernel] Optimize p_norm gpu PaddlePaddle/Paddle#69660
• [Kernel] Optimize isfinite kernel PaddlePaddle/Paddle#69596
• [API] Optimize paddle.where and paddle.where_ in eager mode PaddlePaddle/Paddle#69556

Summary by CodeRabbit

Release Notes

• New Features

  • Introduced several new classes for molecular descriptors, including DescrptDPA1, DescrptBlockSeAtten, and LayerNorm, enhancing the modeling capabilities for molecular simulations.
  • Added new JSON configuration files for model parameters and multitask models related to water simulations.
  • Implemented new test classes for validating the functionality of the DPAtomicModel and various descriptor classes.
  • Added new test classes for evaluating denoising models, including TestDenoiseModelDPA1 and TestDenoiseModelDPA2.
  • Enhanced the ModelWrapper class to clarify the handling of model parameters and state management.

• Bug Fixes

  • Improved internal logic for handling model state saving and loading, ensuring consistency in outputs.

• Documentation

  • Enhanced type hints and return annotations across various classes and methods for better clarity.

• Tests

  • Expanded the testing framework with new test cases for denoising models and descriptor functionalities, ensuring robust validation of features.
  • Activated previously skipped tests for energy models, improving test coverage.
  • Enhanced multitask training tests with new configuration handling and test classes.

[HydrogenSulfate]

@iProzd PR of pd: DPA-1 and pd: DPA-2 is ready for review.

[coderabbitai]

Actionable comments posted: 2

🧹 Outside diff range and nitpick comments (3)

source/tests/pd/model/water/multitask_sharefit.json (2)

11-15: Document the rationale behind selection parameters

The selection parameters [46, 92, 4] seem arbitrary. Consider adding a comment explaining why these specific values were chosen and their significance in the context of water molecule simulations.

---

56-61: Consider adjusting the minimum learning rate

The stop_lr value of 3.51e-08 is extremely small. Consider using a larger value (e.g., 1e-6) as training might not benefit from such small learning rates and could unnecessarily extend training time.

source/tests/pd/test_multitask.py (1)

71-76: Document the conditions for parameter sharing

The conditions for parameter sharing in fitting networks exclude certain parameters (bias_atom_e and case_embd). Consider adding a comment explaining why these specific parameters are excluded from sharing.

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

📥 Commits

Reviewing files that changed from the base of the PR and between 95ac41d and 63ccdbd.

📒 Files selected for processing (2)

• source/tests/pd/model/water/multitask_sharefit.json (1 hunks)
• source/tests/pd/test_multitask.py (6 hunks)

🧰 Additional context used

🪛 GitHub Check: CodeQL

source/tests/pd/test_multitask.py

[notice] 48-48: Unused global variable
The global variable 'multitask_sharefit_template' is not used.

🔇 Additional comments (2)

source/tests/pd/model/water/multitask_sharefit.json (1)

64-82: LGTM! Well-structured loss weighting strategy

The loss configuration implements a good strategy by:

1. Initially prioritizing force accuracy (pref_f: 1000 → 1)
2. Gradually increasing energy contribution (pref_e: 0.02 → 1)
3. Maintaining identical configurations for both models

source/tests/pd/test_multitask.py (1)

43-48: ⚠️ Potential issue

Remove or utilize the unused template variable

The multitask_sharefit_template variable is loaded but never used in the code. Either:

1. Remove it if it's not needed, or
2. Update the tests to utilize this shared fitting configuration

🧰 Tools

🪛 GitHub Check: CodeQL

[notice] 48-48: Unused global variable
The global variable 'multitask_sharefit_template' is not used.

[njzjz]

@coderabbitai resolve