fix(paddle): freeze DPA3 descriptor parameters

deepmd/pd/model/descriptor/dpa3.py

@@ -215,19 +215,22 @@ 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(
                 100,
                 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(

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: