feat(pt_expt): multi-task training support (#5397)

## Summary

- Add multi-task training to the pt_expt backend: multiple models share
a descriptor but have separate fitting nets/losses, trained on
heterogeneous datasets
- Support shared descriptor stats merging (probability-weighted
`merge_env_stat`), shared fitting via `share_params`, and `case_embd`
per-task embedding
- Add DDP (distributed) training support with rank-0 gating, stat
broadcast, and `find_unused_parameters`
- Add multi-task finetune (multi-task→multi-task and
single-task→multi-task) with extended type_map
- Enable `torch.compile` for multi-task with `backend="inductor"`,
`dynamic=True`, and symbolic tracing
(`make_fx(tracing_mode="symbolic")`); includes `silu_backward`
decomposition for second-order gradient compatibility
- Add `silut` activation variant and `DescrptBlockRepformers` accessors
(ported from #5393)

## Compile training benchmark (V100-SXM2-16GB, batch_size=1)

### Speed

| Model | Uncompiled (ms/step) | Compiled (ms/step) | Speedup |
|-------|---------------------|-------------------|---------|
| DPA1 (se_atten) | 21.5 | 10.2 | **2.11x** |
| DPA2 | 66.3 | 23.5 | **2.82x** |
| DPA3 | 161.9 | 50.6 | **3.20x** |

### Convergence (1000 steps)

Training data does not contain virial labels; only energy and force RMSE
are reported.

**DPA1** (se_atten_compressible, float64, rcut=6, sel=120)

| step | Uncompiled (e_rmse / f_rmse) | Compiled (e_rmse / f_rmse) |
|---|---|---|
| 1 | 4.24e-01 / 1.31 | 4.35e-01 / 1.31 |
| 500 | 3.74e-02 / 0.446 | 7.96e-03 / 0.437 |
| 1000 | 9.22e-03 / 0.285 | 8.65e-03 / 0.284 |

**DPA2** (input_torch_small, float32)

| step | Uncompiled (e_rmse / f_rmse) | Compiled (e_rmse / f_rmse) |
|---|---|---|
| 1 | 1.63e-01 / 0.983 | 1.59e-01 / 0.870 |
| 500 | 9.12e-02 / 0.562 | 6.50e-02 / 0.506 |
| 1000 | 4.78e-02 / 0.547 | 4.34e-02 / 0.430 |

**DPA3** (input_torch, float32, static sel)

| step | Uncompiled (e_rmse / f_rmse) | Compiled (e_rmse / f_rmse) |
|---|---|---|
| 1 | 3.88e-02 / 0.784 | 3.13e-02 / 0.893 |
| 500 | 2.25e-02 / 0.829 | 3.01e-03 / 0.879 |
| 1000 | 1.47e-03 / 0.400 | 1.78e-03 / 0.429 |

All models converge to comparable accuracy. Variation between compiled
and uncompiled is within normal run-to-run noise from random batch
ordering.

## Known limitations

- No `num_epoch_dict`: only `numb_steps` + `model_prob`; epoch-based
scheduling deferred
- Only `EnergyFittingNet` has `share_params`; other fitting types (DOS,
dipole, polar, property) need the same override
- No cross-backend (PT vs pt_expt) multi-task consistency test
- `share_fitting` + single-task finetune is incompatible (no
`dim_case_embd` in pretrained)
- DPA3 `use_dynamic_sel: true` cannot compile (symbolic tracer fails on
data-dependent `int()` in `get_graph_index`)

## Test plan

- [ ] 78 multi-task tests (`test_multitask.py`): training, freeze,
finetune, compile, shared fitting, DPA1/DPA2/DPA3/SeA
- [ ] 21 single-task training tests (`test_training.py`): compile
correctness for se_e2_a, DPA1 (with and without attention), DPA2, DPA3;
dynamic shapes, silu compile
- [ ] DDP training tests (`test_training_ddp.py`): single-task +
multi-task with compile
- [ ] Descriptor stat merge tests (`test_descrpt_stat_merge.py`)
- [ ] Fitting stat tests (`test_fitting_stat.py`)
- [ ] share_params tests for all descriptor types
- [ ] Activation tests (`test_activation.py`): silut export + compile
compatibility

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/dpmodel/descriptor/dpa1.py

@@ -405,6 +405,10 @@ def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
 
+    def get_numb_attn_layer(self) -> int:
+        """Returns the number of se_atten attention layers."""
+        return self.se_atten.attn_layer
+
     def share_params(
         self, base_class: "DescrptDPA1", shared_level: int, resume: bool = False
     ) -> NoReturn:

deepmd/dpmodel/descriptor/repformers.py

@@ -345,6 +345,14 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_rcut_smth(self) -> float:
+        """Returns the radius where the neighbor information starts to smoothly decay to 0."""
+        return self.rcut_smth
+
+    def get_env_protection(self) -> float:
+        """Returns the protection of building environment matrix."""
+        return self.env_protection
+
     def get_nsel(self) -> int:
         """Returns the number of selected atoms in the cut-off radius."""
         return sum(self.sel)

deepmd/dpmodel/fitting/general_fitting.py

@@ -255,6 +255,7 @@ def compute_input_stats(
         stat_file_path : Optional[DPPath]
             The path to the stat file.
         """
+        self._param_stats: dict[str, list[StatItem]] = {}
         if self.numb_fparam == 0 and self.numb_aparam == 0:
             # skip data statistics
             return
@@ -296,6 +297,7 @@ def compute_input_stats(
                     self._save_param_stats_to_file(
                         stat_file_path, "fparam", fparam_stats
                     )
+            self._param_stats["fparam"] = fparam_stats
             fparam_avg = np.array(
                 [s.compute_avg() for s in fparam_stats], dtype=np.float64
             )
@@ -362,6 +364,7 @@ def compute_input_stats(
                     self._save_param_stats_to_file(
                         stat_file_path, "aparam", aparam_stats
                     )
+            self._param_stats["aparam"] = aparam_stats
             aparam_avg = np.array(
                 [s.compute_avg() for s in aparam_stats], dtype=np.float64
             )
@@ -407,6 +410,10 @@ def _load_param_stats_from_file(
             for ii in range(numb)
         ]
 
+    def get_param_stats(self) -> dict[str, list[StatItem]]:
+        """Get the stored fparam/aparam statistics (populated by compute_input_stats)."""
+        return getattr(self, "_param_stats", {})
+
     @abstractmethod
     def _net_out_dim(self) -> int:
         """Set the FittingNet output dim."""
@@ -666,11 +673,13 @@ def _call_common(
         # check fparam dim, concate to input descriptor
         if self.numb_fparam > 0:
             assert fparam is not None, "fparam should not be None"
-            if fparam.shape[-1] != self.numb_fparam:
+            try:
+                fparam = xp.reshape(fparam, (nf, self.numb_fparam))
+            except (ValueError, RuntimeError) as e:
                 raise ValueError(
-                    f"get an input fparam of dim {fparam.shape[-1]}, "
-                    f"which is not consistent with {self.numb_fparam}."
-                )
+                    f"input fparam: cannot reshape {fparam.shape} "
+                    f"into ({nf}, {self.numb_fparam})."
+                ) from e
             fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
             fparam = xp.tile(
                 xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
@@ -687,12 +696,13 @@ def _call_common(
         # check aparam dim, concate to input descriptor
         if self.numb_aparam > 0 and not self.use_aparam_as_mask:
             assert aparam is not None, "aparam should not be None"
-            if aparam.shape[-1] != self.numb_aparam:
+            try:
+                aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
+            except (ValueError, RuntimeError) as e:
                 raise ValueError(
-                    f"get an input aparam of dim {aparam.shape[-1]}, "
-                    f"which is not consistent with {self.numb_aparam}."
-                )
-            aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
+                    f"input aparam: cannot reshape {aparam.shape} "
+                    f"into ({nf}, {nloc}, {self.numb_aparam})."
+                ) from e
             aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]
             xx = xp.concat(
                 [xx, aparam],
@@ -735,7 +745,8 @@ def _call_common(
                 )
             for type_i in range(self.ntypes):
                 mask = xp.tile(
-                    xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out)
+                    xp.reshape((atype == type_i), (nf, nloc, 1)),
+                    (1, 1, net_dim_out),
                 )
                 atom_property = self.nets[(type_i,)](xx)
                 if self.remove_vaccum_contribution is not None and not (

deepmd/dpmodel/utils/env_mat.py

@@ -68,7 +68,11 @@ def _make_env_mat(
     xp = array_api_compat.array_namespace(nlist)
     nf, nloc, nnei = nlist.shape
     # nf x nall x 3
-    coord = xp.reshape(coord, (nf, -1, 3))
+    # Callers may pass either (nf, nall*3) or (nf, nall, 3); normalise
+    # both to (nf, nall, 3) using shape-based inference so the concrete nf
+    # value is not baked into the reshape.
+    if coord.ndim == 2:
+        coord = xp.reshape(coord, (-1, coord.shape[1] // 3, 3))
     mask = nlist >= 0
     nlist = nlist * xp.astype(mask, nlist.dtype)
     # nf x (nloc x nnei) x 3
@@ -77,7 +81,7 @@ def _make_env_mat(
     # nf x nloc x nnei x 3
     coord_r = xp.reshape(coord_r, (nf, nloc, nnei, 3))
     # nf x nloc x 1 x 3
-    coord_l = xp.reshape(xp_take_first_n(coord, 1, nloc), (nf, -1, 1, 3))
+    coord_l = xp.reshape(xp_take_first_n(coord, 1, nloc), (nf, nloc, 1, 3))
     # nf x nloc x nnei x 3
     diff = coord_r - coord_l
     # nf x nloc x nnei

deepmd/dpmodel/utils/env_mat_stat.py

@@ -40,6 +40,75 @@
     )
 
 
+def merge_env_stat(
+    base_obj: Union["Descriptor", "DescriptorBlock"],
+    link_obj: Union["Descriptor", "DescriptorBlock"],
+    model_prob: float = 1.0,
+) -> None:
+    """Merge descriptor env mat stats from link_obj into base_obj.
+
+    Uses probability-weighted merging: merged = base_stats + link_stats * model_prob,
+    where model_prob = link_prob / base_prob.
+    Mutates base_obj.stats for chaining (3+ models).
+
+    Parameters
+    ----------
+    base_obj : Descriptor or DescriptorBlock
+        The base descriptor whose stats will be updated.
+    link_obj : Descriptor or DescriptorBlock
+        The linked descriptor whose stats will be merged in.
+    model_prob : float
+        The probability weight ratio (link_prob / base_prob).
+    """
+    if (
+        getattr(base_obj, "stats", None) is None
+        or getattr(link_obj, "stats", None) is None
+    ):
+        return
+    if getattr(base_obj, "set_stddev_constant", False) and getattr(
+        base_obj, "set_davg_zero", False
+    ):
+        return
+
+    # Weighted merge of StatItem objects
+    base_stats = base_obj.stats
+    link_stats = link_obj.stats
+    merged_stats = {}
+    for kk in base_stats:
+        merged_stats[kk] = base_stats[kk] + link_stats[kk] * model_prob
+
+    # Compute mean/stddev from merged stats
+    base_env = EnvMatStatSe(base_obj)
+    base_env.stats = merged_stats
+    mean, stddev = base_env()
+
+    # Update base_obj stats for chaining
+    base_obj.stats = merged_stats
+
+    # Update buffers in-place: davg/dstd (simple) or mean/stddev (blocks)
+    # mean/stddev are numpy arrays; convert to match the buffer's backend
+    if hasattr(base_obj, "davg"):
+        xp = array_api_compat.array_namespace(base_obj.dstd)
+        device = array_api_compat.device(base_obj.dstd)
+        if not getattr(base_obj, "set_davg_zero", False):
+            base_obj.davg[...] = xp.asarray(
+                mean, dtype=base_obj.davg.dtype, device=device
+            )
+        base_obj.dstd[...] = xp.asarray(
+            stddev, dtype=base_obj.dstd.dtype, device=device
+        )
+    elif hasattr(base_obj, "mean"):
+        xp = array_api_compat.array_namespace(base_obj.stddev)
+        device = array_api_compat.device(base_obj.stddev)
+        if not getattr(base_obj, "set_davg_zero", False):
+            base_obj.mean[...] = xp.asarray(
+                mean, dtype=base_obj.mean.dtype, device=device
+            )
+        base_obj.stddev[...] = xp.asarray(
+            stddev, dtype=base_obj.stddev.dtype, device=device
+        )
+
+
 class EnvMatStat(BaseEnvMatStat):
     def compute_stat(self, env_mat: dict[str, Array]) -> dict[str, StatItem]:
         """Compute the statistics of the environment matrix for a single system.

deepmd/pt/model/descriptor/dpa1.py

@@ -383,6 +383,10 @@ def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
 
+    def get_numb_attn_layer(self) -> int:
+        """Returns the number of se_atten attention layers."""
+        return self.se_atten.attn_layer
+
     def share_params(
         self, base_class: Any, shared_level: int, resume: bool = False
     ) -> None:

deepmd/pt/model/task/fitting.py

@@ -779,10 +779,10 @@ def _forward_common(
             assert fparam is not None, "fparam should not be None"
             assert self.fparam_avg is not None
             assert self.fparam_inv_std is not None
-            if fparam.shape[-1] != self.numb_fparam:
+            if fparam.numel() != nf * self.numb_fparam:
                 raise ValueError(
-                    "get an input fparam of dim {fparam.shape[-1]}, ",
-                    "which is not consistent with {self.numb_fparam}.",
+                    f"input fparam: cannot reshape {list(fparam.shape)} "
+                    f"into ({nf}, {self.numb_fparam})."
                 )
             fparam = fparam.view([nf, self.numb_fparam])
             nb, _ = fparam.shape
@@ -804,10 +804,10 @@ def _forward_common(
             assert aparam is not None, "aparam should not be None"
             assert self.aparam_avg is not None
             assert self.aparam_inv_std is not None
-            if aparam.shape[-1] != self.numb_aparam:
+            if aparam.numel() % (nf * self.numb_aparam) != 0:
                 raise ValueError(
-                    f"get an input aparam of dim {aparam.shape[-1]}, ",
-                    f"which is not consistent with {self.numb_aparam}.",
+                    f"input aparam: cannot reshape {list(aparam.shape)} "
+                    f"into ({nf}, nloc, {self.numb_aparam})."
                 )
             aparam = aparam.view([nf, -1, self.numb_aparam])
             nb, nloc, _ = aparam.shape

deepmd/pt_expt/descriptor/dpa1.py

@@ -9,6 +9,9 @@
     cast_precision,
 )
 from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP
+from deepmd.dpmodel.utils.env_mat_stat import (
+    merge_env_stat,
+)
 from deepmd.pt_expt.common import (
     torch_module,
 )
@@ -26,6 +29,31 @@
 class DescrptDPA1(DescrptDPA1DP):
     _update_sel_cls = UpdateSel
 
+    def share_params(
+        self,
+        base_class: Any,
+        shared_level: int,
+        model_prob: float = 1.0,
+        resume: bool = False,
+    ) -> None:
+        """Share parameters with base_class for multi-task training.
+
+        Level 0: share type_embedding and se_atten (all modules and buffers).
+        Level 1: share type_embedding only.
+        """
+        assert self.__class__ == base_class.__class__, (
+            "Only descriptors of the same type can share params!"
+        )
+        if shared_level == 0:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+            if not resume:
+                merge_env_stat(base_class.se_atten, self.se_atten, model_prob)
+            self._modules["se_atten"] = base_class._modules["se_atten"]
+        elif shared_level == 1:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+        else:
+            raise NotImplementedError
+
     def enable_compression(
         self,
         min_nbor_dist: float,

deepmd/pt_expt/descriptor/dpa2.py

@@ -14,6 +14,9 @@
     build_multiple_neighbor_list,
     get_multiple_nlist_key,
 )
+from deepmd.dpmodel.utils.env_mat_stat import (
+    merge_env_stat,
+)
 from deepmd.pt_expt.common import (
     torch_module,
 )
@@ -30,6 +33,47 @@
 class DescrptDPA2(DescrptDPA2DP):
     _update_sel_cls = UpdateSel
 
+    def share_params(
+        self,
+        base_class: "DescrptDPA2",
+        shared_level: int,
+        model_prob: float = 1.0,
+        resume: bool = False,
+    ) -> None:
+        """Share parameters with base_class for multi-task training.
+
+        Level 0: share type_embedding, repinit, repinit_three_body,
+                 g1_shape_tranform, and repformers.
+        Level 1: share type_embedding only.
+        """
+        assert self.__class__ == base_class.__class__, (
+            "Only descriptors of the same type can share params!"
+        )
+        if shared_level == 0:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+            if not resume:
+                merge_env_stat(base_class.repinit, self.repinit, model_prob)
+                if self.use_three_body and "repinit_three_body" in base_class._modules:
+                    merge_env_stat(
+                        base_class.repinit_three_body,
+                        self.repinit_three_body,
+                        model_prob,
+                    )
+                merge_env_stat(base_class.repformers, self.repformers, model_prob)
+            self._modules["repinit"] = base_class._modules["repinit"]
+            if self.use_three_body and "repinit_three_body" in base_class._modules:
+                self._modules["repinit_three_body"] = base_class._modules[
+                    "repinit_three_body"
+                ]
+            self._modules["g1_shape_tranform"] = base_class._modules[
+                "g1_shape_tranform"
+            ]
+            self._modules["repformers"] = base_class._modules["repformers"]
+        elif shared_level == 1:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+        else:
+            raise NotImplementedError
+
     def enable_compression(
         self,
         min_nbor_dist: float,

deepmd/pt_expt/descriptor/dpa3.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3DP
+from deepmd.dpmodel.utils.env_mat_stat import (
+    merge_env_stat,
+)
 from deepmd.pt_expt.common import (
     torch_module,
 )
@@ -16,3 +19,28 @@
 @torch_module
 class DescrptDPA3(DescrptDPA3DP):
     _update_sel_cls = UpdateSel
+
+    def share_params(
+        self,
+        base_class: "DescrptDPA3",
+        shared_level: int,
+        model_prob: float = 1.0,
+        resume: bool = False,
+    ) -> None:
+        """Share parameters with base_class for multi-task training.
+
+        Level 0: share type_embedding and repflows.
+        Level 1: share type_embedding only.
+        """
+        assert self.__class__ == base_class.__class__, (
+            "Only descriptors of the same type can share params!"
+        )
+        if shared_level == 0:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+            if not resume:
+                merge_env_stat(base_class.repflows, self.repflows, model_prob)
+            self._modules["repflows"] = base_class._modules["repflows"]
+        elif shared_level == 1:
+            self._modules["type_embedding"] = base_class._modules["type_embedding"]
+        else:
+            raise NotImplementedError