Commit f143171
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
- descriptor
- entrypoints
- utils
- jax
- entrypoints
- utils
- source/tests
- common/dpmodel
- jax
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
34 | 34 | | |
35 | 35 | | |
36 | 36 | | |
37 | | - | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
38 | 41 | | |
39 | 42 | | |
40 | 43 | | |
| |||
59 | 62 | | |
60 | 63 | | |
61 | 64 | | |
62 | | - | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
63 | 68 | | |
64 | 69 | | |
65 | 70 | | |
| |||
0 commit comments