feat(jax): add local training entrypoint

Port the JAX training entrypoint from the parallel branch onto current master,
but keep it local-only by removing distributed, sharding, and Hessian hooks.
Use the current dpmodel compute_or_load_stat data-stat path and add regression
coverage for the cleanup constraints.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

Author: njzjz-bot

deepmd/backend/jax.py

@@ -62,7 +62,11 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
         Callable[[Namespace], None]
             The entry point hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.entrypoints.main import (
+            main,
+        )
+
+        return main
 
     @property
     def deep_eval(self) -> type["DeepEvalBackend"]:

deepmd/jax/entrypoints/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later

deepmd/jax/entrypoints/main.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DeePMD-Kit entry point module."""
+
+import argparse
+from pathlib import (
+    Path,
+)
+from typing import (
+    Optional,
+    Union,
+)
+
+from deepmd.backend.suffix import (
+    format_model_suffix,
+)
+from deepmd.jax.entrypoints.freeze import (
+    freeze,
+)
+from deepmd.jax.entrypoints.train import (
+    train,
+)
+from deepmd.loggers.loggers import (
+    set_log_handles,
+)
+from deepmd.main import (
+    parse_args,
+)
+
+__all__ = ["main"]
+
+
+def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
+    """DeePMD-Kit entry point.
+
+    Parameters
+    ----------
+    args : list[str] or argparse.Namespace, optional
+        list of command line arguments, used to avoid calling from the subprocess,
+        as it is quite slow to import tensorflow; if Namespace is given, it will
+        be used directly
+
+    Raises
+    ------
+    RuntimeError
+        if no command was input
+    """
+    if not isinstance(args, argparse.Namespace):
+        args = parse_args(args=args)
+
+    dict_args = vars(args)
+    set_log_handles(
+        args.log_level,
+        Path(args.log_path) if args.log_path else None,
+        mpi_log=None,
+    )
+
+    if args.command == "train":
+        train(**dict_args)
+    elif args.command == "freeze":
+        dict_args["output"] = format_model_suffix(
+            dict_args["output"], preferred_backend=args.backend, strict_prefer=True
+        )
+        freeze(**dict_args)
+    elif args.command is None:
+        pass
+    else:
+        raise RuntimeError(f"unknown command {args.command}")

deepmd/jax/entrypoints/train.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DeePMD training entrypoint script.
+
+Can handle local training.
+"""
+
+import json
+import logging
+import time
+from typing import (
+    Any,
+    Optional,
+)
+
+
+from deepmd.common import (
+    j_loader,
+)
+from deepmd.jax.env import (
+    jax,
+    jax_export,
+)
+from deepmd.jax.train.trainer import (
+    DPTrainer,
+)
+from deepmd.utils import random as dp_random
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
+from deepmd.utils.data_system import (
+    get_data,
+)
+from deepmd.utils.summary import SummaryPrinter as BaseSummaryPrinter
+
+__all__ = ["train"]
+
+log = logging.getLogger(__name__)
+
+
+class SummaryPrinter(BaseSummaryPrinter):
+    """Summary printer for JAX."""
+
+    def is_built_with_cuda(self) -> bool:
+        """Check if the backend is built with CUDA."""
+        return jax_export.default_export_platform() == "cuda"
+
+    def is_built_with_rocm(self) -> bool:
+        """Check if the backend is built with ROCm."""
+        return jax_export.default_export_platform() == "rocm"
+
+    def get_compute_device(self) -> str:
+        """Get Compute device."""
+        return jax.default_backend()
+
+    def get_ngpus(self) -> int:
+        """Get the number of GPUs."""
+        return jax.device_count()
+
+    def get_backend_info(self) -> dict:
+        """Get backend information."""
+        return {
+            "Backend": "JAX",
+            "JAX ver": jax.__version__,
+        }
+
+    def get_device_name(self) -> str:
+        """Get the name of the device."""
+        devices = jax.devices()
+        if devices:
+            return devices[0].device_kind
+        else:
+            return "Unknown"
+
+
+def train(
+    *,
+    INPUT: str,
+    init_model: Optional[str],
+    restart: Optional[str],
+    output: str,
+    init_frz_model: str,
+    mpi_log: str,
+    log_level: int,
+    log_path: Optional[str],
+    skip_neighbor_stat: bool = False,
+    finetune: Optional[str] = None,
+    use_pretrain_script: bool = False,
+    **kwargs: Any,
+) -> None:
+    """Run DeePMD model training.
+
+    Parameters
+    ----------
+    INPUT : str
+        json/yaml control file
+    init_model : Optional[str]
+        path prefix of checkpoint files or None
+    restart : Optional[str]
+        path prefix of checkpoint files or None
+    output : str
+        path for dump file with arguments
+    init_frz_model : str
+        path to frozen model or None
+    mpi_log : str
+        mpi logging mode
+    log_level : int
+        logging level defined by int 0-3
+    log_path : Optional[str]
+        logging file path or None if logs are to be output only to stdout
+    skip_neighbor_stat : bool, default=False
+        skip checking neighbor statistics
+    finetune : Optional[str]
+        path to pretrained model or None
+    use_pretrain_script : bool
+        Whether to use model script in pretrained model when doing init-model or init-frz-model.
+        Note that this option is true and unchangeable for fine-tuning.
+    **kwargs
+        additional arguments
+
+    Raises
+    ------
+    RuntimeError
+        if the training command fails.
+    """
+    # load json database
+    jdata = j_loader(INPUT)
+
+    origin_type_map = None
+
+    jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
+
+    jdata = normalize(jdata)
+    jdata = update_sel(jdata)
+
+    with open(output, "w") as fp:
+        json.dump(jdata, fp, indent=4)
+    SummaryPrinter()()
+
+    # make necessary checks
+    assert "training" in jdata
+
+    # init the model
+
+    model = DPTrainer(
+        jdata,
+        init_model=init_model,
+        restart=restart,
+    )
+    rcut = model.model.get_rcut()
+    type_map = model.model.get_type_map()
+    if len(type_map) == 0:
+        ipt_type_map = None
+    else:
+        ipt_type_map = type_map
+
+    # init random seed of data systems
+    seed = jdata["training"].get("seed", None)
+    if seed is not None:
+        seed += jax.process_index()
+        seed = seed % (2**32)
+    dp_random.seed(seed)
+
+    # init data
+    train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, None)
+    train_data.add_data_requirements(model.data_requirements)
+    train_data.print_summary("training")
+    if jdata["training"].get("validation_data", None) is not None:
+        valid_data = get_data(
+            jdata["training"]["validation_data"],
+            rcut,
+            train_data.type_map,
+            None,
+        )
+        valid_data.add_data_requirements(model.data_requirements)
+        valid_data.print_summary("validation")
+    else:
+        valid_data = None
+
+    # get training info
+    stop_batch = jdata["training"]["numb_steps"]
+    origin_type_map = jdata["model"].get("origin_type_map", None)
+    if (
+        origin_type_map is not None and not origin_type_map
+    ):  # get the type_map from data if not provided
+        origin_type_map = get_data(
+            jdata["training"]["training_data"], rcut, None, None
+        ).get_type_map()
+
+    # train the model with the provided systems in a cyclic way
+    start_time = time.time()
+    model.train(train_data, valid_data)
+    end_time = time.time()
+    log.info("finished training")
+    log.info(f"wall time: {(end_time - start_time):.3f} s")
+
+
+def update_sel(jdata: dict) -> dict:
+    """Update descriptor selections from neighbor statistics when available."""
+    log.info(
+        "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)"
+    )
+    jdata_cpy = jdata.copy()
+    type_map = jdata["model"].get("type_map")
+    train_data = get_data(
+        jdata["training"]["training_data"],
+        0,  # not used
+        type_map,
+        None,  # not used
+    )
+    # TODO: OOM, need debug
+    # jdata_cpy["model"], min_nbor_dist = BaseModel.update_sel(
+    #     train_data, type_map, jdata["model"]
+    # )
+    return jdata_cpy

deepmd/jax/train/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later