diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 95fe69422b..de6bc5d45f 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -255,11 +255,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 @@ -557,6 +560,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..fef0da2d03 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,49 @@ 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, + ) + deserialized_descriptor = DescrptDPA3.deserialize(descriptor.serialize()) + + 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_consistency( self, ) -> None: