perf(pt_expt): share compiled forward_lower across multi-task shared-fitting tasks (deepmodeling#5457)

make dataset embedding and energy bias as input not buffer for compile,
this allows multitask training share compiled model thus resolve OOM and
NCCL timeout issue. Since the empty_cache and del are removed, no GC
complaints.

Regression Test
<img width="3600" height="2100" alt="lcurve"
src="https://github.com/user-attachments/assets/c043bf6c-53bb-441f-ac98-0d021b68ec1b"
/>


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

* **New Features**
* Multi-task training groups models by structure and caches/reuses
compiled computation graphs, with per-task buffer handling to support
shared fitting nets.

* **Bug Fixes**
* Checkpoint loading now skips extraneous per-task buffer entries so
only original model parameters are restored.
* Training aggregation coerces tensor-like loss/metric values to floats
for accurate reporting.

* **Tests**
* Added regression test ensuring compiled and eager outputs match per
task for shared-fitting, different-descriptor setups.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/deepmodeling/deepmd-kit/pull/5457?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: anyangml

deepmd/pt_expt/descriptor/dpa2.py

@@ -69,6 +69,8 @@ def share_params(
                 "g1_shape_tranform"
             ]
             self._modules["repformers"] = base_class._modules["repformers"]
+            if "tebd_transform" in base_class._modules:
+                self._modules["tebd_transform"] = base_class._modules["tebd_transform"]
         elif shared_level == 1:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
         else:

deepmd/pt_expt/descriptor/dpa3.py

@@ -40,6 +40,10 @@ def share_params(
             if not resume:
                 merge_env_stat(base_class.repflows, self.repflows, model_prob)
             self._modules["repflows"] = base_class._modules["repflows"]
+            if self.add_chg_spin_ebd:
+                for key in ("chg_embedding", "spin_embedding", "mix_cs_mlp"):
+                    if key in base_class._modules:
+                        self._modules[key] = base_class._modules[key]
         elif shared_level == 1:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
         else:

deepmd/pt_expt/infer/deep_eval.py

@@ -368,10 +368,15 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None:
         # eager inference).  Drop the latter and unwrap the former.
         cleaned: dict[str, Any] = {}
         compiled_marker = ".compiled_forward_lower."
+        # Per-task buffer copies registered on _CompiledModel (bias_atom_e,
+        # case_embd) — real values live on the original model's fitting net.
+        task_buf_marker = "._task_"
         wrapper_infix = ".original_model."
         for key, value in state_dict.items():
             if compiled_marker in key:
                 continue
+            if task_buf_marker in key:
+                continue
             if wrapper_infix in key:
                 key = key.replace(wrapper_infix, ".", 1)
             cleaned[key] = value