[pull] master from deepmodeling:master - #170
Merged
Merged
Conversation
Avoid wasting time on broken test that pending the workflow for a long time like this: https://github.com/deepmodeling/deepmd-kit/actions/runs/22224311896 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added test timeout configuration to prevent indefinitely hanging test runs, improving test suite reliability and execution stability. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: link89 <3314130+link89@users.noreply.github.com>
## Summary This PR adds PyTorch support for DPA4, the DeePMD-kit implementation of SeZM (Smooth Equivariant Zone-bridging Model). It introduces the DPA4/SeZM model, descriptor, fitting network, training integration, export path, documentation, examples, and tests. ## Main Changes - Add the DPA4/SeZM PyTorch model stack: - `model.type: "dpa4"` / `"sezm"` - `descriptor.type: "dpa4"` / `"sezm"` - `fitting_net.type: "dpa4_ener"` / `"sezm_ener"` - Implement the SO(3)-equivariant descriptor with edge-local SO(2) convolutions, angular schedules, smooth radial envelopes, attention/focus streams, and environment-seeded initial features. - Add zone-bridging support for short-range analytical repulsion, including ZBL coupling and descriptor-side short-range clamping. - Add DPA4 training support for: - conservative energy/force training through `loss.type: "ener"` - experimental direct-force denoising through `loss.type: "dens"` - spin models in the PyTorch backend - shared-fitting multitask case FiLM conditioning - LoRA fine-tuning and merged checkpoint export - Add the DPA4 `.pt2` freeze/export path using AOTInductor for checkpoints that cannot be represented by the regular TorchScript freeze path. - Add CLI, argcheck, validation, data-system, and inference integration needed to route DPA4 configs and exported models correctly. - Add water examples for standard DPA4, ZBL bridging, spin, DeNS, multitask/shared-fitting, LoRA fine-tuning, and LAMMPS inference. - Add official model documentation at `doc/model/dpa4.md`. ## Tests This PR adds coverage for: - DPA4/SeZM model and descriptor construction - DPA4 aliases in model, descriptor, and fitting configuration - SO(3)/SO(2) equivariance behavior - conservative energy/force paths - `torch.compile` eager/compiled consistency - DPA4 `.pt2` export and DeepPot inference - spin model behavior - ZBL zone bridging - DeNS loss and direct-force mode - LoRA adapter injection, freezing, merging, and compile compatibility - optional Triton kernel dispatch and numerical consistency - supporting utility changes in neighbor-list, LMDB data, and distributed checks Relevant test files include: - `source/tests/pt/model/test_descriptor_sezm.py` - `source/tests/pt/model/test_descriptor_sezm_s2_equivariance.py` - `source/tests/pt/model/test_descriptor_sezm_triton.py` - `source/tests/pt/model/test_sezm_model.py` - `source/tests/pt/model/test_sezm_spin_model.py` - `source/tests/pt/model/test_sezm_export.py` - `source/tests/pt/test_training.py` - `source/tests/pt/test_train_utils.py` - `source/tests/common/dpmodel/test_dist_check.py` - `source/tests/common/dpmodel/test_lmdb_data.py` ## Notes DPA4 is currently implemented for the PyTorch backend. Model compression is not supported, and DPA4 checkpoints use the `.pt2` export path instead of the regular TorchScript freeze path. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added SeZM model family and DeNS denoising loss for training; new optimized ".pt2" export path with embedded metadata. * **Improvements** * LoRA fine-tuning workflow (apply/merge/strip) for lightweight adapters. * On-demand minimum pairwise-distance computation during data reads. * Better JAX neighbor-list handling and optional GPU/Triton-accelerated descriptor kernels for faster inference/training. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/deepmodeling/deepmd-kit/pull/5448?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
](https://app.coderabbit.ai/change-stack/deepmodeling/deepmd-kit/pull/5448?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack){kind=link}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )