Skip to content

feat(pt): implement shared_dict functionality for linear models#4933

Closed
Copilot wants to merge 3 commits into
masterfrom
copilot/fix-4412
Closed

feat(pt): implement shared_dict functionality for linear models#4933
Copilot wants to merge 3 commits into
masterfrom
copilot/fix-4412

Conversation

Copilot AI commented Aug 28, 2025

Copy link
Copy Markdown
Contributor

This PR implements shared_dict parameter support for LinearEnergyModel to enable parameter sharing between submodels, following the same pattern used in multi-task models.

Problem

Users requested the ability to share parameters between submodels in linear models, similar to how multi-task models work. The original linear model implementation only supported independent submodels without parameter sharing capabilities.

Solution

Added support for the shared_dict parameter to LinearEnergyModel by:

  1. Parameter preprocessing: Integrated preprocess_shared_params() to convert shared references to actual parameter dictionaries
  2. Parameter sharing logic: Implemented _share_params() method adapted from ModelWrapper.share_params() to enable actual parameter sharing between submodels
  3. Backward compatibility: All existing linear model usage patterns continue to work unchanged

Usage Example

config = {
    "type": "linear_ener",
    "shared_dict": {
        "type_map_all": ["O", "H"],
        "dpa1_descriptor_1": {
            "type": "dpa1",
            "rcut": 6.00,
            "rcut_smth": 0.50,
            "sel": 138,
            "neuron": [25, 50, 100],
            "axis_neuron": 16,
            "seed": 1,
        },
    },
    "models": [
        {
            "type_map": "type_map_all",
            "descriptor": "dpa1_descriptor_1",
            "fitting_net": {
                "neuron": [240, 240, 240],
                "resnet_dt": True,
                "seed": 1,
            },
        },
        {
            "type_map": "type_map_all", 
            "descriptor": "dpa1_descriptor_1",
            "fitting_net": {
                "type": "dipole",
                "neuron": [100, 100, 100],
                "resnet_dt": True,
                "seed": 1,
            },
        },
    ],
    "weights": "mean",
}

model = LinearEnergyModel(**config)

The implementation reuses existing multi-task infrastructure to minimize code changes while providing the requested functionality.

Fixes #4412.


✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
@njzjz

njzjz commented Aug 28, 2025

Copy link
Copy Markdown
Member

@copilot continue your work.

Copilot AI commented Aug 28, 2025

Copy link
Copy Markdown
Contributor Author

@copilot continue your work.

Continuing work on the shared_dict implementation for linear models. I've completed the core implementation and testing:

✅ LinearEnergyModel now supports shared_dict parameter
✅ Parameter sharing logic integrated using existing multi-task patterns
✅ Tests added and passing for both shared_dict and backward compatibility
✅ Issue #4412 example configuration confirmed working
✅ Existing tests still pass (no regressions)
✅ Code properly linted and formatted

The implementation allows parameter sharing between submodels in linear models following the same pattern as multi-task models, with minimal changes to the codebase.

Complete implementation of shared_dict parameter support for LinearEnergyModel to enable parameter sharing between submodels, following the existing multi-task pattern.

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
Copilot AI changed the title [WIP] [Feature Request] Use shared_dict for linear model feat(pt): implement shared_dict functionality for linear models Aug 28, 2025
@njzjz njzjz closed this Mar 2, 2026
@njzjz njzjz deleted the copilot/fix-4412 branch March 2, 2026 13:20
njzjz added a commit to njzjz/deepmd-kit that referenced this pull request Jul 1, 2026
## Summary
- Add `shared_dict` support to the PyTorch `linear_ener` model config,
reusing the multi-task shared-parameter preprocessor.
- Share descriptor/fitting parameters among linear sub-models with the
same shared-link semantics as multi-task training.
- Preserve shared references through `update_sel`, including configs
where `type_map` only lives in `shared_dict`.
- Add unit coverage for shared descriptor/type_map references and
top-level `type_map` plus shared descriptors.

Closes deepmodeling#4412.
Supersedes/continues deepmodeling#4933.

## Tests
- `uvx ruff check deepmd/pt/model/model/dp_linear_model.py
deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py
deepmd/utils/argcheck.py
source/tests/pt/model/test_linear_model_shared_dict.py`
- `uvx ruff format deepmd/pt/model/model/dp_linear_model.py
deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py
deepmd/utils/argcheck.py
source/tests/pt/model/test_linear_model_shared_dict.py`
- `python3 -m py_compile deepmd/pt/model/model/dp_linear_model.py
deepmd/pt/model/model/__init__.py deepmd/pt/utils/multi_task.py
deepmd/utils/argcheck.py
source/tests/pt/model/test_linear_model_shared_dict.py`

Note: local `pytest -q
source/tests/pt/model/test_linear_model_shared_dict.py` could not run in
this worktree because the compiled `deepmd.lib` extension is unavailable
in the local environment.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added `shared_dict` support for `linear_ener` models to enable shared
parameters across atomic sub-models, including descriptor and
fitting-network components (with hybrid descriptor support).
* Enhanced `linear_ener` neighbor-selection updates to correctly
round-trip `shared_dict` configurations while preserving shared
descriptor references.
* **Bug Fixes**
* Strengthened validation for shared setups, including consistent
`type_map` handling and safer enforcement of compatible sharing
definitions and link ordering.
* **Tests**
* Added/expanded coverage for `shared_dict` behavior, `shared_links`
population, hybrid descriptors, successful round-trips, and expected
failure cases.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Use shared_dict for linear model

2 participants