diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 95fe69422b..01e4b741ec 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -215,6 +215,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.tebd_dim, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) # 100 is a conservative upper bound self.spin_embedding = TypeEmbedNet( @@ -222,12 +223,14 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.tebd_dim, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) self.mix_cs_mlp = MLPLayer( 2 * self.tebd_dim, self.tebd_dim, precision=precision, seed=child_seed(seed, 5), + trainable=trainable, ) else: self.chg_embedding = None @@ -255,11 +258,14 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") ) - # set trainable - for param in self.parameters(): - param.requires_grad = trainable + self._apply_trainable() self.compress = False + def _apply_trainable(self) -> None: + """Apply the descriptor-level trainable setting to every parameter.""" + for param in self.parameters(): + param.stop_gradient = not self.trainable + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut @@ -354,6 +360,11 @@ def share_params( assert self.__class__ == base_class.__class__, ( "Only descriptors of the same type can share params!" ) + if self.trainable != base_class.trainable: + raise ValueError( + "DPA3 descriptors must use the same trainable setting before " + "sharing parameters." + ) # For DPA3 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in type_embedding, repflow @@ -384,6 +395,9 @@ def change_type_map( remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.type_embedding.change_type_map(type_map=type_map) + # Type remapping replaces the first embedding matrix with a newly + # created parameter, so restore the descriptor-level frozen state. + self._apply_trainable() self.exclude_types = map_pair_exclude_types(self.exclude_types, remap_index) self.ntypes = len(type_map) repflow = self.repflows @@ -557,6 +571,9 @@ def t_cvt(xx: Any) -> paddle.Tensor: obj.repflows.layers = paddle.nn.LayerList( [RepFlowLayer.deserialize(layer) for layer in repflow_layers] ) + # Deserialization replaces several sublayers after construction, so apply + # the descriptor-level setting again to their newly registered parameters. + obj._apply_trainable() return obj def forward( diff --git a/source/tests/pd/model/test_dpa3.py b/source/tests/pd/model/test_dpa3.py index 7b3c8ba18f..a42069b0c3 100644 --- a/source/tests/pd/model/test_dpa3.py +++ b/source/tests/pd/model/test_dpa3.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools import unittest +from unittest.mock import ( + patch, +) import numpy as np import paddle @@ -36,6 +39,125 @@ class TestDescrptDPA3(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) + def test_non_trainable_parameters_stop_gradient(self) -> None: + """A non-trainable descriptor must disable Paddle autograd parameters.""" + # Paddle 3.0 does not expose the newer ``requires_grad`` compatibility + # alias. Ignore writes to that alias so this test retains minimum-version + # semantics when the CI environment uses a newer Paddle release. + requires_grad_alias = property( + lambda parameter: not parameter.stop_gradient, + lambda _parameter, _value: None, + ) + with patch.object( + paddle.Tensor, "requires_grad", requires_grad_alias, create=True + ): + descriptor = DescrptDPA3( + self.nt, + repflow=RepFlowArgs( + n_dim=4, + e_dim=4, + a_dim=4, + nlayers=1, + e_sel=2, + a_sel=1, + axis_neuron=2, + ), + trainable=False, + add_chg_spin_ebd=True, + seed=GLOBAL_SEED, + ) + serialized_descriptor = descriptor.serialize() + deserialized_descriptor = DescrptDPA3.deserialize(serialized_descriptor) + + # Optional charge/spin layers must serialize their effective frozen + # state so cross-backend deserializers do not re-enable gradients. + self.assertFalse(serialized_descriptor["chg_embedding"]["trainable"]) + self.assertFalse(serialized_descriptor["spin_embedding"]["trainable"]) + self.assertFalse(serialized_descriptor["mix_cs_mlp"]["trainable"]) + + for stage, checked_descriptor in ( + ("constructed", descriptor), + ("deserialized", deserialized_descriptor), + ): + with self.subTest(stage=stage): + parameters = list(checked_descriptor.named_parameters()) + self.assertTrue(parameters) + parameters_with_grad = [ + name + for name, parameter in parameters + if not parameter.stop_gradient + ] + self.assertEqual([], parameters_with_grad) + + def test_non_trainable_change_type_map_stays_frozen(self) -> None: + """Type remapping must not make a frozen embedding trainable again.""" + descriptor = DescrptDPA3( + self.nt, + repflow=RepFlowArgs( + n_dim=4, + e_dim=4, + a_dim=4, + nlayers=1, + e_sel=2, + a_sel=1, + axis_neuron=2, + ), + trainable=False, + type_map=["O", "H"], + seed=GLOBAL_SEED, + ) + + descriptor.change_type_map(["H", "O"]) + + parameters_with_grad = [ + name + for name, parameter in descriptor.named_parameters() + if not parameter.stop_gradient + ] + self.assertEqual([], parameters_with_grad) + + def test_share_params_rejects_mismatched_trainable(self) -> None: + """Sharing must not let one descriptor rewrite another's gradients.""" + repflow = RepFlowArgs( + n_dim=4, + e_dim=4, + a_dim=4, + nlayers=1, + e_sel=2, + a_sel=1, + axis_neuron=2, + ) + for shared_level in (0, 1): + with self.subTest(shared_level=shared_level): + trainable_descriptor = DescrptDPA3( + self.nt, + repflow=repflow, + trainable=True, + seed=GLOBAL_SEED, + ) + frozen_descriptor = DescrptDPA3( + self.nt, + repflow=repflow, + trainable=False, + seed=GLOBAL_SEED, + ) + + with self.assertRaisesRegex(ValueError, "same trainable setting"): + trainable_descriptor.share_params( + frozen_descriptor, shared_level=shared_level + ) + + # Validation happens before either the embedding or repflow + # layers can be aliased to the other descriptor. + self.assertIsNot( + trainable_descriptor.type_embedding, + frozen_descriptor.type_embedding, + ) + self.assertIsNot( + trainable_descriptor.repflows, + frozen_descriptor.repflows, + ) + def test_consistency( self, ) -> None: