fix(stat): Caculate correct fitting stat when using default fparam and using share fitting. (#5038)

In this PR:
1. Support writing fitting stat to `stat_file` and loading fitting stat
from `stat_file`
2. Ensure the fitting stat calculate is correct when using
`default_fparam`
3. Support sharing fitting stat when using `share_fitting` in multitask
mode.
4. Print the process of calculating fitting stat to the board via
`log.info`.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Default frame parameters auto-fill missing samples and are exposed via
a new accessor.
* Compute/load/save per-parameter statistics to/from optional stat
files.
* Multitask training adds probability-weighted parameter sharing and
propagates default fparam into data requirements.
  * Statistic items support scalar scaling for aggregation.

* **Refactor**
* Parameter-sharing and statistic-propagation flows reorganized for
consistent buffering and persistence.

* **Tests**
* New tests and test data for stats computation, persistence, and
multitask sharing.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenqqian Zhang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Chengqian-Zhang

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -397,13 +397,14 @@ def wrapped_sampler():
             return sampled
 
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
-        self.compute_fitting_input_stat(wrapped_sampler)
+        self.compute_fitting_input_stat(wrapped_sampler, stat_file_path)
         if compute_or_load_out_stat:
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
     def compute_fitting_input_stat(
         self,
         sample_merged: Union[Callable[[], list[dict]], list[dict]],
+        stat_file_path: Optional[DPPath] = None,
     ) -> None:
         """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.
 
@@ -416,9 +417,13 @@ def compute_fitting_input_stat(
             - Callable[[], list[dict]]: A lazy function that returns data samples in the above format
                 only when needed. Since the sampling process can be slow and memory-intensive,
                 the lazy function helps by only sampling once.
+        stat_file_path : Optional[DPPath]
+            The dictionary of paths to the statistics files.
         """
         self.fitting_net.compute_input_stats(
-            sample_merged, protection=self.data_stat_protect
+            sample_merged,
+            protection=self.data_stat_protect,
+            stat_file_path=stat_file_path,
         )
 
     def get_dim_fparam(self) -> int:

deepmd/pd/model/task/fitting.py

@@ -40,6 +40,9 @@
     get_index_between_two_maps,
     map_atom_exclude_types,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 dtype = env.GLOBAL_PD_FLOAT_PRECISION
 device = env.DEVICE
@@ -76,6 +79,7 @@ def compute_input_stats(
         self,
         merged: Union[Callable[[], list[dict]], list[dict]],
         protection: float = 1e-2,
+        stat_file_path: Optional[DPPath] = None,
     ) -> None:
         """
         Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.
@@ -91,6 +95,8 @@ def compute_input_stats(
                 the lazy function helps by only sampling once.
         protection : float
             Divided-by-zero protection
+        stat_file_path : Optional[DPPath]
+            The path to the stat file.
         """
         if self.numb_fparam == 0 and self.numb_aparam == 0:
             # skip data statistics

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -326,16 +326,26 @@ def wrapped_sampler() -> list[dict]:
                 atom_exclude_types = self.atom_excl.get_exclude_types()
                 for sample in sampled:
                     sample["atom_exclude_types"] = list(atom_exclude_types)
+            if (
+                "find_fparam" not in sampled[0]
+                and "fparam" not in sampled[0]
+                and self.has_default_fparam()
+            ):
+                default_fparam = self.get_default_fparam()
+                for sample in sampled:
+                    nframe = sample["atype"].shape[0]
+                    sample["fparam"] = default_fparam.repeat(nframe, 1)
             return sampled
 
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
-        self.compute_fitting_input_stat(wrapped_sampler)
+        self.compute_fitting_input_stat(wrapped_sampler, stat_file_path)
         if compute_or_load_out_stat:
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
     def compute_fitting_input_stat(
         self,
         sample_merged: Union[Callable[[], list[dict]], list[dict]],
+        stat_file_path: Optional[DPPath] = None,
     ) -> None:
         """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.
 
@@ -348,9 +358,13 @@ def compute_fitting_input_stat(
             - Callable[[], list[dict]]: A lazy function that returns data samples in the above format
                 only when needed. Since the sampling process can be slow and memory-intensive,
                 the lazy function helps by only sampling once.
+        stat_file_path : Optional[DPPath]
+            The dictionary of paths to the statistics files.
         """
         self.fitting_net.compute_input_stats(
-            sample_merged, protection=self.data_stat_protect
+            sample_merged,
+            protection=self.data_stat_protect,
+            stat_file_path=stat_file_path,
         )
 
     def get_dim_fparam(self) -> int:
@@ -361,6 +375,9 @@ def has_default_fparam(self) -> bool:
         """Check if the model has default frame parameters."""
         return self.fitting_net.has_default_fparam()
 
+    def get_default_fparam(self) -> Optional[torch.Tensor]:
+        return self.fitting_net.get_default_fparam()
+
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.fitting_net.get_dim_aparam()

deepmd/pt/model/model/make_model.py

@@ -532,6 +532,9 @@ def has_default_fparam(self) -> bool:
             """Check if the model has default frame parameters."""
             return self.atomic_model.has_default_fparam()
 
+        def get_default_fparam(self) -> Optional[torch.Tensor]:
+            return self.atomic_model.get_default_fparam()
+
         @torch.jit.export
         def get_dim_aparam(self) -> int:
             """Get the number (dimension) of atomic parameters of this atomic model."""