feat(pt_expt): add training infrastructure (#5270)

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

* **New Features**
* New "PyTorch-Exportable" backend with end-to-end training CLI, model
factory, loss wrapper, model wrapper, checkpointing, neighbor/stat
utilities, and optional torch.compile acceleration with automatic
re-tracing.
* **Bug Fixes**
* Added runtime validation for required parameter fields (fparam/aparam)
and ensured loss/reporting arrays are device-consistent.
* **Refactor**
* Reworked stat input, batch normalization, and model-stat collection
for clearer, backend-agnostic data flow.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: wanghan-iapcm

deepmd/backend/pt_expt.py

@@ -32,7 +32,7 @@
 class PyTorchExportableBackend(Backend):
     """PyTorch exportable backend."""
 
-    name = "PyTorch Exportable"
+    name = "PyTorch-Exportable"
     """The formal name of the backend."""
     features: ClassVar[Backend.Feature] = (
         Backend.Feature.ENTRY_POINT
@@ -63,7 +63,7 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
         Callable[[Namespace], None]
             The entry point hook of the backend.
         """
-        from deepmd.pt.entrypoints.main import main as deepmd_main
+        from deepmd.pt_expt.entrypoints.main import main as deepmd_main
 
         return deepmd_main
 

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -394,19 +394,19 @@ 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()
-            ):
+            # For systems where fparam is missing (find_fparam == 0),
+            # fill with default fparam if available and mark as found.
+            if self.has_default_fparam():
                 default_fparam = self.get_default_fparam()
                 if default_fparam is not None:
                     default_fparam_np = np.array(default_fparam)
                     for sample in sampled:
-                        nframe = sample["atype"].shape[0]
-                        sample["fparam"] = np.tile(
-                            default_fparam_np.reshape(1, -1), (nframe, 1)
-                        )
+                        if "find_fparam" in sample and not sample["find_fparam"]:
+                            nframe = sample["atype"].shape[0]
+                            sample["fparam"] = np.tile(
+                                default_fparam_np.reshape(1, -1), (nframe, 1)
+                            )
+                            sample["find_fparam"] = np.bool_(True)
             return sampled
 
         return wrapped_sampler

deepmd/dpmodel/descriptor/dpa1.py

@@ -242,6 +242,8 @@ class DescrptDPA1(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2208.08236.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -662,7 +664,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, sel = UpdateSel().update_one_sel(
+        min_nbor_dist, sel = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
         )
         local_jdata_cpy["sel"] = sel[0]

deepmd/dpmodel/descriptor/dpa2.py

@@ -441,6 +441,8 @@ class DescrptDPA2(NativeOP, BaseDescriptor):
        Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -1114,7 +1116,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repinit_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/dpa3.py

@@ -337,6 +337,8 @@ class DescrptDPA3(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2506.01686 (2025).
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -729,7 +731,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repflow_e_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -51,7 +51,7 @@ class BD(ABC, PluginVariant, make_plugin_registry("descriptor")):
         def __new__(cls, *args: Any, **kwargs: Any) -> Any:
             if cls is BD:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
-            return super().__new__(cls)
+            return object.__new__(cls)
 
         @abstractmethod
         def get_rcut(self) -> float:

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -149,6 +149,8 @@ class DescrptSeA(NativeOP, BaseDescriptor):
        Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -582,7 +584,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_r.py

@@ -128,6 +128,8 @@ class DescrptSeR(NativeOP, BaseDescriptor):
        Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -505,7 +507,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_t.py

@@ -116,6 +116,8 @@ class DescrptSeT(NativeOP, BaseDescriptor):
             Not used in this descriptor, only to be compat with input.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -505,7 +507,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+        min_nbor_dist, local_jdata_cpy["sel"] = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
         )
         return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -141,6 +141,8 @@ class DescrptSeTTebd(NativeOP, BaseDescriptor):
 
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -500,7 +502,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, sel = UpdateSel().update_one_sel(
+        min_nbor_dist, sel = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
         )
         local_jdata_cpy["sel"] = sel[0]