Skip to content

Commit f143171

Browse files
authored
feat(dpmodel): add descriptor compression (#5592)
## Summary - add dpmodel compression entrypoints and wire dp --dp/--jax compress - persist and restore compressed dpmodel/JAX descriptor state, including HLO export metadata - implement dpmodel compression for se_e2_a, se_e2_r, and se_atten/dpa1 (with se_atten_v2 state sync) - add common dpmodel and JAX compression coverage ## Tests - python -m pytest source/tests/common/dpmodel/test_model_compression.py -q - python -m pytest source/tests/jax/test_model_compression.py -q - python -m pytest source/tests/common/test_argument_parser.py -q - python -m pytest source/tests/pt_expt/descriptor/test_se_r.py source/tests/pt_expt/descriptor/test_dpa1.py source/tests/pt_expt/descriptor/test_se_atten_v2.py -q - python -m pytest source/tests/pt_expt/model/test_model_compression.py -q - ruff check . - ruff format --check . <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added model compression support for DPModel and JAX workflows. * Introduced a new `compress` command in the CLI for creating compressed model files. * Expanded support for compressed descriptors, including saving and restoring compressed models. * **Documentation** * Updated CLI help text with new examples and supported model file formats. * **Tests** * Added end-to-end and descriptor-level coverage for compression behavior and model reloads. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 1550599 commit f143171

17 files changed

Lines changed: 2021 additions & 57 deletions

File tree

deepmd/backend/dpmodel.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,10 @@ class DPModelBackend(Backend):
3434
name = "DPModel"
3535
"""The formal name of the backend."""
3636
features: ClassVar[Backend.Feature] = (
37-
Backend.Feature.DEEP_EVAL | Backend.Feature.NEIGHBOR_STAT | Backend.Feature.IO
37+
Backend.Feature.ENTRY_POINT
38+
| Backend.Feature.DEEP_EVAL
39+
| Backend.Feature.NEIGHBOR_STAT
40+
| Backend.Feature.IO
3841
)
3942
"""The features of the backend."""
4043
suffixes: ClassVar[list[str]] = [".dp", ".yaml", ".yml"]
@@ -59,7 +62,9 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
5962
Callable[[Namespace], None]
6063
The entry point hook of the backend.
6164
"""
62-
raise NotImplementedError(f"Unsupported backend: {self.name}")
65+
from deepmd.dpmodel.entrypoints.main import main as deepmd_main
66+
67+
return deepmd_main
6368

6469
@property
6570
def deep_eval(self) -> type["DeepEvalBackend"]:

0 commit comments

Comments
 (0)