fix(pt): share direct descriptor parameters across tasks

Author: OutisLi

deepmd/pt/model/descriptor/descriptor.py

@@ -168,8 +168,10 @@ def share_params(
                 # must share, even if not do stat
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev
-            # self.load_state_dict(base_class.state_dict()) # this does not work, because it only inits the model
-            # the following will successfully link all the params except buffers
+            # Direct parameters are not part of the child-module registry.
+            for item in self._parameters:
+                self._parameters[item] = base_class._parameters[item]
+            # Child modules carry their own parameters and buffers.
             for item in self._modules:
                 self._modules[item] = base_class._modules[item]
         else:

source/tests/pt/model/test_dpa1.py

@@ -368,6 +368,19 @@ def test_degree_gain_zero_recovers_lmax_one(self) -> None:
         (gradient,) = torch.autograd.grad(output.sum(), raw_gain)
         self.assertGreater(torch.linalg.vector_norm(gradient).item(), 0.0)
 
+    def test_level_zero_shares_degree_gain_parameter(self) -> None:
+        base = self._build_descriptor(lmax=4)
+        branch = self._build_descriptor(lmax=4)
+        base_gain = base.se_atten.adam_degree_gain_raw
+        branch_gain = branch.se_atten.adam_degree_gain_raw
+        assert base_gain is not None
+        assert branch_gain is not None
+        self.assertIsNot(branch_gain, base_gain)
+
+        branch.share_params(base, shared_level=0, resume=True)
+
+        self.assertIs(branch.se_atten.adam_degree_gain_raw, base_gain)
+
     def test_lmax_two_resolves_quadrupole_collision(self) -> None:
         square = self._square_directions()
         tetrahedral = self._tetrahedral_directions()