feat(tf2): add training workflow (#5735)

## Summary
- add TensorFlow 2 training/freeze/compress entrypoints and trainer
support
- wire multi-task, validation, finetune/pretrain, TensorBoard,
checkpointing, and bias adjustment paths
- keep TF2 training atomic virial disabled and add performance
optimizations, including DP_JIT-controlled lower-forward XLA compilation

## Validation
- `python -m pytest source/tests/tf2/test_training.py -q`
- `ruff check .`
- `ruff format --check deepmd/tf2 source/tests/tf2`
- `DP_JIT=1 srun --gres=gpu:1 dp --tf2 train input.json
--skip-neighbor-stat` on `examples/water/se_e2_a` temp copy: XLA lower
compiled; first step ~140s, steady windows ~0.0305 s/step through step
900 before timeout
- `DP_JIT=1 srun --gres=gpu:1 dp --tf2 train input.json
--skip-neighbor-stat` on fixed-selection dpa3 temp input: XLA lower
compiled; first step ~227s, steady windows ~0.119 s/step through step
500 before timeout

## Notes
- `DP_JIT` is opt-in for training lower forward only; the whole
train/eval step is intentionally not XLA compiled because neighbor/outer
training logic is too broad and previously hit unsupported XLA ops.
- dpa3 lower JIT has high first-compile cost and high memory pressure,
so benchmark numbers above separate warm steady-state from first-step
compile cost.


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

* **New Features**
* Expanded TensorFlow 2 backend with new train, freeze, and checkpoint
compression entry points, including full validation support.
* Added multi-task model-branch selection for freeze/compress via
`--head` / `--model-branch`.
* Introduced optional TF2 JIT acceleration and TF2-specific automatic
batch sizing.
* **Bug Fixes**
  * Improved compression dtype handling for more consistent outputs.
* Enabled additional TF2 runtime neighbor-statistics and corrected TF2
backend wiring so more workflows complete end-to-end.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: njzjz

.github/workflows/test_python.yml

@@ -60,7 +60,7 @@ jobs:
           restore-keys: |
             test2-durations-combined-${{ matrix.python }}-${{ github.sha }}
             test2-durations-combined-${{ matrix.python }}
-      - run: pytest --cov=deepmd source/tests  --splits 12 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration
+      - run: pytest --cov=deepmd source/tests --ignore=source/tests/tf2 --splits 12 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration
         env:
           NUM_WORKERS: 0
           DP_CI_IMPORT_PADDLE_BEFORE_TF: 1
@@ -81,6 +81,10 @@ jobs:
             return "$status"
           }
 
+          run_pytest_allow_no_tests --cov=deepmd --cov-append \
+            source/tests/tf2 \
+            --splits 12 \
+            --group ${{ matrix.group }}
           run_pytest_allow_no_tests --cov=deepmd --cov-append \
             source/tests/consistent/io/test_io.py \
             source/jax2tf_tests \

deepmd/backend/tf2.py

@@ -33,7 +33,12 @@ class TensorFlow2Backend(Backend):
     """TensorFlow 2 eager backend."""
 
     name = "TensorFlow2"
-    features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL | Backend.Feature.IO
+    features: ClassVar[Backend.Feature] = (
+        Backend.Feature.ENTRY_POINT
+        | Backend.Feature.DEEP_EVAL
+        | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.IO
+    )
     suffixes: ClassVar[list[str]] = [".savedmodeltf"]
 
     @classmethod
@@ -45,7 +50,11 @@ def is_available(self) -> bool:
 
     @property
     def entry_point_hook(self) -> Callable[["Namespace"], None]:
-        raise NotImplementedError("Training entry point is not implemented for TF2")
+        from deepmd.tf2.entrypoints.main import (
+            main,
+        )
+
+        return main
 
     @property
     def deep_eval(self) -> type["DeepEvalBackend"]:
@@ -57,7 +66,11 @@ def deep_eval(self) -> type["DeepEvalBackend"]:
 
     @property
     def neighbor_stat(self) -> type["NeighborStat"]:
-        raise NotImplementedError("Neighbor statistics are not implemented for TF2")
+        from deepmd.dpmodel.utils.neighbor_stat import (
+            NeighborStat,
+        )
+
+        return NeighborStat
 
     @property
     def serialize_hook(self) -> Callable[[str], dict]:

deepmd/dpmodel/common.py

@@ -90,6 +90,18 @@ def get_xp_precision(
         raise ValueError(f"unsupported precision {precision} for {xp}")
 
 
+def to_numpy_dtype(dtype: Any) -> np.dtype:
+    """Normalize backend dtype objects to a NumPy dtype."""
+    dtype = getattr(dtype, "as_numpy_dtype", dtype)
+    try:
+        return np.dtype(dtype)
+    except TypeError:
+        dtype_name = getattr(dtype, "name", None)
+        if dtype_name is not None:
+            return np.dtype(dtype_name)
+        raise
+
+
 class NativeOP(ABC):
     """The unit operation of a native model."""
 

deepmd/dpmodel/descriptor/dpa1.py

@@ -28,6 +28,7 @@
     cast_precision,
     get_xp_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -1448,7 +1449,7 @@ def enable_compression(
     ) -> None:
         """Store tabulated geometric embedding-net data."""
         net = "filter_net"
-        dtype = self.mean.dtype
+        dtype = to_numpy_dtype(self.mean.dtype)
         self.compress_info = [
             np.asarray(
                 [

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -22,6 +22,7 @@
 from deepmd.dpmodel.common import (
     cast_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -431,7 +432,7 @@ def _store_compress_data(
         """Store tabulated embedding-net data in the descriptor state."""
         compress_data = []
         compress_info = []
-        dtype = self.davg.dtype
+        dtype = to_numpy_dtype(self.davg.dtype)
         ndim = 1 if self.type_one_side else 2
         for embedding_idx in range(self.ntypes**ndim):
             if self.type_one_side:

deepmd/dpmodel/descriptor/se_r.py

@@ -22,6 +22,7 @@
     cast_precision,
     get_xp_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -410,7 +411,7 @@ def _store_compress_data(
         """Store tabulated embedding-net data in the descriptor state."""
         compress_data = []
         compress_info = []
-        dtype = self.davg.dtype
+        dtype = to_numpy_dtype(self.davg.dtype)
         for embedding_idx in range(self.ntypes):
             net = "filter_-1_net_" + str(embedding_idx)
             if net not in table_data:

deepmd/main.py

@@ -327,7 +327,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--output",
         type=str,
         default="frozen_model",
-        help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; Paddle backend: suffix is .json and .pdiparams",
+        help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: suffix is .savedmodeltf; PyTorch backend: suffix is .pth; Paddle backend: suffix is .json and .pdiparams",
     )
     parser_frz.add_argument(
         "-n",
@@ -605,14 +605,14 @@ def main_parser() -> argparse.ArgumentParser:
         "--input",
         default="frozen_model",
         type=str,
-        help="The original frozen model, which will be compressed by the code. Filename (prefix) of the input model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
+        help="The original frozen model or checkpoint, which will be compressed by the code. Filename (prefix) of the input model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: .tf2 checkpoint directory or checkpoint prefix; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
     )
     parser_compress.add_argument(
         "-o",
         "--output",
         default="frozen_model_compressed",
         type=str,
-        help="The compressed model. Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
+        help="The compressed model. Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: suffix is .savedmodeltf; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
     )
     parser_compress.add_argument(
         "-s",
@@ -659,6 +659,13 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="The training script of the input frozen model",
     )
+    parser_compress.add_argument(
+        "--head",
+        "--model-branch",
+        default=None,
+        type=str,
+        help="Task head (alias: model branch) to compress if in multi-task mode.",
+    )
 
     # * print docs script **************************************************************
     parsers_doc = subparsers.add_parser(

deepmd/tf2/atomic_model/dp_atomic_model.py

@@ -11,6 +11,7 @@
 )
 from deepmd.tf2.env import (
     stop_gradient,
+    tf,
     xp,
 )
 from deepmd.tf2.fitting.base_fitting import (
@@ -41,6 +42,13 @@ class tf2_atomic_model(dpmodel_atomic_model):
         base_fitting_cls = BaseFitting
         """The base fitting class."""
 
+        @tf.autograph.experimental.do_not_convert
+        def make_atom_mask(
+            self,
+            atype: xp.ndarray,
+        ) -> xp.ndarray:
+            return atype >= 0
+
         def forward_common_atomic(
             self,
             extended_coord: xp.ndarray,

deepmd/tf2/entrypoints/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Entry points for the TensorFlow 2 backend."""
+
+from deepmd.tf2.entrypoints.compress import (
+    enable_compression,
+)
+from deepmd.tf2.entrypoints.freeze import (
+    freeze,
+)
+from deepmd.tf2.entrypoints.train import (
+    train,
+)
+
+__all__ = ["enable_compression", "freeze", "train"]

deepmd/tf2/entrypoints/compress.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compress TensorFlow 2 checkpoints by tabulating embedding networks."""
+
+from __future__ import (
+    annotations,
+)
+
+import logging
+from typing import (
+    Any,
+)
+
+from deepmd.backend.suffix import (
+    format_model_suffix,
+)
+from deepmd.dpmodel.entrypoints.compress_common import (
+    enable_model_compression,
+    resolve_min_nbor_dist,
+)
+from deepmd.dpmodel.utils.update_sel import (
+    UpdateSel,
+)
+from deepmd.tf2.entrypoints.freeze import (
+    select_model_branch,
+)
+from deepmd.tf2.model.base_model import (
+    BaseModel,
+)
+from deepmd.tf2.utils.serialization import (
+    deserialize_to_file,
+    serialize_from_file,
+)
+
+log = logging.getLogger(__name__)
+
+
+def enable_compression(
+    input_file: str,
+    output: str,
+    stride: float = 0.01,
+    extrapolate: int = 5,
+    check_frequency: int = -1,
+    training_script: str | None = None,
+    head: str | None = None,
+    **kwargs: Any,
+) -> None:
+    """Compress a TF2 training checkpoint and export a SavedModel."""
+    del kwargs
+    output = format_model_suffix(
+        output,
+        preferred_backend="tf2",
+        strict_prefer=True,
+    )
+    data = serialize_from_file(input_file)
+    data = select_model_branch(data, head=head)
+    model = BaseModel.deserialize(data["model"])
+    min_nbor_dist = resolve_min_nbor_dist(
+        model,
+        [data],
+        training_script,
+        UpdateSel,
+    )
+    enable_model_compression(
+        model,
+        min_nbor_dist,
+        stride,
+        extrapolate,
+        check_frequency,
+    )
+
+    compressed_data = data.copy()
+    compressed_data["model"] = model.serialize()
+    compressed_data["min_nbor_dist"] = float(min_nbor_dist)
+    deserialize_to_file(output, compressed_data)
+    log.info("Compressed TF2 model saved to %s", output)