Add support for multi-task DPA3 models - #1246
Conversation
pfebrer
left a comment
pfebrer
left a comment
There was a problem hiding this comment.
Thanks for taking care of this! Here are a few comments
| Using a pretrained model | ||
| ------------------------ | ||
|
|
||
| Set ``dpa3_model`` to a deepmd-kit model file to fine-tune from pretrained | ||
| weights instead of training from scratch: | ||
|
|
||
| .. code-block:: yaml | ||
|
|
||
| model: | ||
| dpa3_model: path/to/deepmd-model.pt | ||
|
|
||
| Energy biases and standard deviations are extracted from the loaded model and | ||
| handed to metatrain's composition model and scaler automatically. | ||
|
|
There was a problem hiding this comment.
This part needs to be removed
| # Remove the "dpa3_model" key from the hypers dictionary if it exists | ||
| if "dpa3_model" not in checkpoint["model_data"]["model_hypers"]: | ||
| checkpoint["model_data"]["model_hypers"]["dpa3_model"] = None |
There was a problem hiding this comment.
The comment is wrong, and I think this is not needed since dpa3_model will always be there.
| if "dpa3_model_branch" not in checkpoint["model_data"]["model_hypers"]: | ||
| checkpoint["model_data"]["model_hypers"]["dpa3_model_branch"] = None |
There was a problem hiding this comment.
You can assume dpa3_model_branch is not there, the previous version didn't have it.
| "The loaded model must be a torch.nn.Module or a " | ||
| "collections.OrderedDict." |
There was a problem hiding this comment.
Let's say something more user friendly like We could not load a DPA3 model from the provided path: "the path here". The file must contain a torch module or an ordered dict.
| if not torch.all(std == std[0]): | ||
| raise NotImplementedError( | ||
| "Loaded DPA3 models with non-uniform per-type 'out_std' are " | ||
| "not supported." | ||
| ) |
There was a problem hiding this comment.
What do the DPA3 checkpoints contain usually?
| def test_multi_task_branch_loading(tmp_path): | ||
| """One branch of a multi-task checkpoint is loaded as a standalone model.""" | ||
| path = tmp_path / "multi_task.pt" | ||
| base = _small_pretrained_checkpoint(path, branches=["Alpha", "Beta"]) | ||
|
|
||
| pretrained = DPA3(_hypers_for(path, "Beta"), _make_dataset_info()) | ||
|
|
||
| # The structure comes from the branch configuration, not from the hypers | ||
| base_state = base.model.state_dict() | ||
| loaded_state = pretrained.model.state_dict() | ||
| assert set(loaded_state) == set(base_state) | ||
| for key, value in base_state.items(): | ||
| torch.testing.assert_close(loaded_state[key], value) | ||
|
|
||
|
|
||
| def test_single_task_loading(tmp_path): | ||
| """A single-task checkpoint is loaded without asking for a branch.""" | ||
| path = tmp_path / "single_task.pt" | ||
| base = _small_pretrained_checkpoint(path) | ||
|
|
||
| pretrained = DPA3(_hypers_for(path), _make_dataset_info()) | ||
|
|
||
| assert pretrained.loaded_dpa3 is True | ||
| torch.testing.assert_close( | ||
| pretrained.model.state_dict(), base.model.state_dict(), rtol=0, atol=0 | ||
| ) |
There was a problem hiding this comment.
Do you think we could use these two as a way of initializing a model for the normal testing in test_basic.py, like we do for MACE? (i.e. by parametrizing the model_hypers fixture). If it gets too expensive, maybe just the multi branch case, although since we are testing only energies in this case it will probably be cheap.
2 participants
This PR adds support for multi-task DPA3 models. Currently publicly available DPA3 foundation models are all a multi-task model, but only the single-task model is supported in
metatrain.The structure of the selected
branchof the pretrained model is get fromloaded["model"]["_extra_state"]["model_params"]["model_dict"][branch]. After the model is loaded,self.model.get_model_def_script()is saved toself._deepmd_cfgfor reloading from it next timeContributor (creator of pull-request) checklist
Maintainer/Reviewer checklist
📚 Documentation preview 📚: https://metatrain--1246.org.readthedocs.build/en/1246/