feat(dpmodel): add backend-independent trainer abstraction (#5603)

## Summary
- Add backend-independent training abstractions under
`deepmd.dpmodel.train` for task/rank normalization, display scheduling,
learning-curve output, checkpoint cadence, lifecycle hooks, and shared
train entrypoint orchestration.
- Factor common training-data helpers so single-task training is handled
as a one-task collection and multi-task data
construction/summary/probability handling is shared where possible.
- Add a backend-independent finetune rule builder in
`deepmd.utils.finetune`, and reduce the PT, PT-exportable, Paddle, and
JAX backend finetune modules to backend-specific checkpoint loading plus
shared rule generation.
- Migrate JAX train entrypoint/trainer onto the shared pipeline and add
JAX finetune plus multi-task support on top of the new abstractions.
- Migrate `pt_expt` train entrypoint/trainer behavior further onto the
shared pipeline, including single-task-as-multi-task normalization, data
summaries, checkpoint retention, stat-file parent creation, relative
latest checkpoint symlinks, and checkpoint parent creation.
- Address PR review comments around task-key validation, learning-curve
metric ordering, lifecycle cleanup, `print_summary` fallback behavior,
broken `__len__` handling, JAX finetune branch/alias validation,
numeric-looking JAX task keys, HDF5 stat paths, and `pt_expt` checkpoint
symlinks.
- Move the new dpmodel trainer/entrypoint tests from
`source/tests/test_dpmodel_*.py` into `source/tests/common/dpmodel/`.

Refs #5229, #5230, #5231

## Tests
- `ruff format .`
- `ruff check .`
- `git diff --check`
- `PYTHONPATH=/home/jzzeng/codes/deepmd-kit pytest
source/tests/common/dpmodel/test_train_abstract_trainer.py
source/tests/common/dpmodel/test_train_entrypoint.py
source/tests/common/dpmodel/test_train_data.py
source/tests/common/dpmodel/test_training_utils.py
source/tests/common/test_finetune_utils.py
source/tests/jax/test_training.py
source/tests/pt_expt/test_entrypoint.py
source/tests/pt_expt/test_multitask.py::TestMultiTaskSeA::test_multitask_finetune
source/tests/pt_expt/test_multitask.py::TestMultiTaskSeA::test_multitask_finetune_from_single_task
source/tests/pt_expt/test_multitask.py::TestMultiTaskSeA::test_multitask_finetune_no_change_model_params
-q` (`53 passed, 2 subtests passed`)
- `PYTHONPATH=/home/jzzeng/codes/deepmd-kit timeout 180 srun
--gres=gpu:1 dp --jax train input.json --skip-neighbor-stat --finetune
pretrain.jax --use-pretrain-script` on a temporary 1-step water finetune
smoke; completed on NVIDIA GeForce RTX 5090 and saved `ft-model-1.jax`.
- `PYTHONPATH=/home/jzzeng/codes/deepmd-kit timeout 180 srun
--gres=gpu:1 dp --pt-expt train input.json --skip-neighbor-stat` on a
temporary 2-step water smoke; completed on NVIDIA GeForce RTX 5090,
saved `ckpts/pt-model-2.pt`, created `stats/stat.hdf5`, and verified
`ckpts/pt-model.pt -> pt-model-2.pt` with old step checkpoint pruned by
`max_ckpt_keep=1`.

## Notes
- Paddle-specific runtime tests were not run locally because `paddle` is
not installed in this environment.
- Plain PyTorch backend test collection is blocked in this environment
by external `deepmd_gnn`/CUDA initialization, not by the shared finetune
rule builder changes.

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

* **New Features**
* Introduced a unified, backend-independent training framework with
consistent single-task and multi-task handling, learning-curve output,
and structured training/validation steps.
* Added a common training entrypoint abstraction that standardizes
config normalization, neighbor-stat updates, and lifecycle teardown.
* Implemented full-validation with best-checkpoint tracking, top-K
selection, and `val.log` reporting (including backend-specific
checkpoint suffixes).

* **Bug Fixes**
* Improved checkpoint save/restore and retention (including “latest”
link updates and older checkpoint cleanup).
* Improved task-weighting logic to better handle datasets with/without
sizing information.
* Fixed multi-task neighbor-stat updates and JAX full-validation error
propagation across ranks.

* **Tests**
* Expanded unit and smoke tests for training orchestration, finetuning,
validation, and checkpoint reconciliation.

* **Documentation**
* Updated validation-configuration help text to reflect broader backend
support.
<!-- 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>
Co-authored-by: njzjz-bot (driven by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5))[bot] <48687836+njzjz-bot@users.noreply.github.com>

Author: 3 people

deepmd/dpmodel/train/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-independent training abstractions."""
+
+from .data import (
+    TrainingTaskConfig,
+    iter_training_task_configs,
+    make_task_maps,
+    print_data_summaries,
+)
+from .entrypoint import (
+    AbstractTrainEntrypoint,
+    TrainEntrypointOptions,
+)
+from .trainer import (
+    DEFAULT_TASK_KEY,
+    AbstractTrainer,
+    LearningCurveWriter,
+    RankContext,
+    TrainerConfig,
+    TrainingTask,
+    TrainingTaskCollection,
+    TrainStepResult,
+)
+
+__all__ = [
+    "DEFAULT_TASK_KEY",
+    "AbstractTrainEntrypoint",
+    "AbstractTrainer",
+    "LearningCurveWriter",
+    "RankContext",
+    "TrainEntrypointOptions",
+    "TrainStepResult",
+    "TrainerConfig",
+    "TrainingTask",
+    "TrainingTaskCollection",
+    "TrainingTaskConfig",
+    "iter_training_task_configs",
+    "make_task_maps",
+    "print_data_summaries",
+]

deepmd/dpmodel/train/data.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Shared training-data helpers for backend entrypoints."""
+
+from __future__ import (
+    annotations,
+)
+
+import inspect
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from .trainer import (
+    DEFAULT_TASK_KEY,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+        Iterator,
+        Mapping,
+    )
+
+
+@dataclass(frozen=True)
+class TrainingTaskConfig:
+    """Normalized config view for one training task."""
+
+    key: str
+    model_params: Mapping[str, Any]
+    training_data_params: Mapping[str, Any]
+    validation_data_params: Mapping[str, Any] | None
+    stat_file: str | None
+    valid_numb_batch: int
+
+
+def iter_training_task_configs(
+    config: Mapping[str, Any],
+) -> Iterator[TrainingTaskConfig]:
+    """Yield task configs, treating single-task input as one ``Default`` task."""
+    model_params = config["model"]
+    training_params = config["training"]
+    if "model_dict" not in model_params:
+        validation_data_params = training_params.get("validation_data")
+        yield TrainingTaskConfig(
+            key=DEFAULT_TASK_KEY,
+            model_params=model_params,
+            training_data_params=training_params["training_data"],
+            validation_data_params=validation_data_params,
+            stat_file=training_params.get("stat_file"),
+            valid_numb_batch=_valid_numb_batch(validation_data_params),
+        )
+        return
+
+    data_dict = training_params["data_dict"]
+    for task_key, task_model_params in model_params["model_dict"].items():
+        task_data_params = data_dict[task_key]
+        validation_data_params = task_data_params.get("validation_data")
+        yield TrainingTaskConfig(
+            key=task_key,
+            model_params=task_model_params,
+            training_data_params=task_data_params["training_data"],
+            validation_data_params=validation_data_params,
+            stat_file=task_data_params.get("stat_file"),
+            valid_numb_batch=_valid_numb_batch(validation_data_params),
+        )
+
+
+def make_task_maps(
+    config: Mapping[str, Any],
+    factory: Callable[[TrainingTaskConfig], tuple[Any, Any | None, Any | None]],
+) -> tuple[dict[str, Any], dict[str, Any | None], dict[str, Any | None]]:
+    """Build training, validation, and stat maps from normalized task configs."""
+    training_data: dict[str, Any] = {}
+    validation_data: dict[str, Any | None] = {}
+    stat_data: dict[str, Any | None] = {}
+    for task_config in iter_training_task_configs(config):
+        train_item, valid_item, stat_item = factory(task_config)
+        training_data[task_config.key] = train_item
+        validation_data[task_config.key] = valid_item
+        stat_data[task_config.key] = stat_item
+    return training_data, validation_data, stat_data
+
+
+def print_data_summaries(
+    training_data: Mapping[str, Any],
+    validation_data: Mapping[str, Any | None],
+    *,
+    probabilities: Mapping[str, float] | None = None,
+) -> None:
+    """Print train/validation data summaries for one or more tasks."""
+    multi_task = len(training_data) > 1
+    for task_key, data in training_data.items():
+        name = f"training data({task_key})" if multi_task else "training"
+        _print_summary(data, name, _task_probability(probabilities, task_key))
+        valid_data = validation_data.get(task_key)
+        if valid_data is not None:
+            name = f"validation data({task_key})" if multi_task else "validation"
+            _print_summary(valid_data, name, None)
+
+
+def _valid_numb_batch(validation_data_params: Mapping[str, Any] | None) -> int:
+    if validation_data_params is None:
+        return 1
+    return max(int(validation_data_params.get("numb_btch", 1)), 1)
+
+
+def _task_probability(
+    probabilities: Mapping[str, float] | None,
+    task_key: str,
+) -> list[float] | None:
+    if probabilities is None or task_key not in probabilities:
+        return None
+    return [float(probabilities[task_key])]
+
+
+def _print_summary(data: Any, name: str, prob: list[float] | None) -> None:
+    printer = data.print_summary
+    try:
+        signature = inspect.signature(printer)
+    except (TypeError, ValueError):
+        printer(name, prob)
+        return
+    try:
+        signature.bind(name, prob)
+    except TypeError as exc:
+        try:
+            signature.bind(name)
+        except TypeError:
+            raise exc from None
+        printer(name)
+    else:
+        printer(name, prob)

deepmd/dpmodel/train/entrypoint.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-independent training entrypoint pipeline."""
+
+from __future__ import (
+    annotations,
+)
+
+import json
+import logging
+from abc import (
+    ABC,
+    abstractmethod,
+)
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Any,
+)
+
+from deepmd.common import (
+    j_loader,
+)
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainEntrypointOptions:
+    """Common command options for backend train entrypoints."""
+
+    input_file: str
+    output: str = "out.json"
+    init_model: str | None = None
+    restart: str | None = None
+    init_frz_model: str | None = None
+    finetune: str | None = None
+    model_branch: str = ""
+    use_pretrain_script: bool = False
+    skip_neighbor_stat: bool = False
+
+
+class AbstractTrainEntrypoint(ABC):
+    """Shared pipeline for backend train entrypoints.
+
+    Backend subclasses keep ownership of backend-specific feature handling,
+    neighbor-stat updates, distributed setup, data construction, and trainer
+    construction.  This pipeline only coordinates the common command flow.
+    """
+
+    def run(self, options: TrainEntrypointOptions) -> None:
+        """Run the training entrypoint."""
+        log.info("Configuration path: %s", options.input_file)
+        options = self.prepare_options(options)
+        config = self.load_config(options.input_file)
+        self.validate_options(config, options)
+
+        config = self.preprocess_config(config, options)
+        multi_task = self.is_multi_task(config)
+        config = self.update_input(config)
+        config = self.normalize_config(config, multi_task=multi_task)
+
+        neighbor_stat = None
+        if not options.skip_neighbor_stat:
+            config, neighbor_stat = self.update_neighbor_stat(
+                config,
+                options,
+                multi_task=multi_task,
+            )
+
+        self.dump_config(config, options.output)
+        self.print_summary()
+
+        try:
+            self.setup_run(options, config)
+            self.run_training(config, options, neighbor_stat)
+        finally:
+            self.teardown_run(options, config)
+
+    def prepare_options(
+        self,
+        options: TrainEntrypointOptions,
+    ) -> TrainEntrypointOptions:
+        """Normalize command options before reading or preprocessing config."""
+        return options
+
+    def load_config(self, input_file: str) -> dict[str, Any]:
+        """Load the JSON/YAML training config."""
+        return j_loader(input_file)
+
+    def validate_options(
+        self,
+        config: dict[str, Any],
+        options: TrainEntrypointOptions,
+    ) -> None:
+        """Validate backend feature support before mutating the config."""
+        return None
+
+    def preprocess_config(
+        self,
+        config: dict[str, Any],
+        options: TrainEntrypointOptions,
+    ) -> dict[str, Any]:
+        """Apply backend-specific config preprocessing before argcheck."""
+        return config
+
+    def is_multi_task(self, config: dict[str, Any]) -> bool:
+        """Return whether the config is in multi-task layout."""
+        return "model_dict" in config.get("model", {})
+
+    def update_input(self, config: dict[str, Any]) -> dict[str, Any]:
+        """Apply DeePMD input-version compatibility conversion."""
+        return update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
+
+    def normalize_config(
+        self,
+        config: dict[str, Any],
+        *,
+        multi_task: bool,
+    ) -> dict[str, Any]:
+        """Run DeePMD argcheck normalization."""
+        return normalize(config, multi_task=multi_task)
+
+    def update_neighbor_stat(
+        self,
+        config: dict[str, Any],
+        options: TrainEntrypointOptions,
+        *,
+        multi_task: bool,
+    ) -> tuple[dict[str, Any], Any]:
+        """Update descriptor selections from neighbor statistics."""
+        return config, None
+
+    def dump_config(self, config: dict[str, Any], output: str) -> None:
+        """Dump the normalized config used for training."""
+        with open(output, "w") as fp:
+            json.dump(config, fp, indent=4)
+
+    def print_summary(self) -> None:
+        """Print backend summary information."""
+        return None
+
+    def setup_run(
+        self,
+        options: TrainEntrypointOptions,
+        config: dict[str, Any],
+    ) -> None:
+        """Set up backend runtime state before trainer execution."""
+        return None
+
+    def teardown_run(
+        self,
+        options: TrainEntrypointOptions,
+        config: dict[str, Any],
+    ) -> None:
+        """Tear down backend runtime state after trainer execution."""
+        return None
+
+    @abstractmethod
+    def run_training(
+        self,
+        config: dict[str, Any],
+        options: TrainEntrypointOptions,
+        neighbor_stat: Any,
+    ) -> None:
+        """Build backend data/trainer objects and run training."""
+        raise NotImplementedError