feat: support DPA4 and DPA4C deployment

docs/input.md

@@ -59,6 +59,76 @@ The `"type" : "dp"` tell the traning method is {dargs:argument}`"dp" <train>`, i
 The `"config"` key defines the training configs, see {ref}`the full documentation<train[dp]/config>`.
 The {dargs:argument}`"template_script" <train[dp]/template_script>` provides the template training script in `json` format.
 
+For DPA4, use the regular PyTorch training backend and deploy `.pt2` models in
+LAMMPS exploration:
+
+```json
+"train": {"type": "dp", "config": {"impl": "pytorch"}},
+"explore": {
+    "type": "lmp",
+    "config": {"model_devi_backend": "pytorch", "model_format": "pt2"}
+}
+```
+
+Put DPA4 training acceleration controls in the DeePMD training template under
+`model`, not in the DPGEN2 workflow file:
+
+```json
+"model": {
+	"type": "dpa4",
+	"use_compile": true,
+	"enable_tf32": true
+}
+```
+
+For DPA4C, both training and deployment use the PyTorch Exportable backend.
+Compression is optional:
+
+```json
+"train": {"type": "dp", "config": {"impl": "pytorch-exportable"}},
+"explore": {
+    "type": "lmp",
+    "config": {
+        "model_devi_backend": "pytorch-exportable",
+        "model_format": "pt2",
+        "dp_compress": true
+    }
+}
+```
+
+Put DPA4C training acceleration controls in the DeePMD training template file
+referenced by `train.template_script`, under `training`:
+
+```json
+"model": {"descriptor": {"type": "dpa4c"}},
+"training": {
+	"training_data": {
+		"systems": [],
+		"batch_size": "auto:512"
+	},
+	"numb_steps": 1000000,
+	"enable_compile": true,
+	"enable_tf32": true
+}
+```
+
+Do not copy the DPA4 paths `model.use_compile` or `model.enable_tf32` into a
+DPA4C template. Conversely, DPA4 does not use the DPA4C paths
+`training.enable_compile` or `training.enable_tf32`. DPGEN2 validates these
+backend-specific placements before creating the workflow, but it does not inject
+or change performance and numerical-policy settings. Therefore setting only
+`train.config.impl` is not sufficient to enable compilation: a DPA4 template
+must contain `model.use_compile: true`, while a DPA4C template must contain
+`training.enable_compile: true`. Run a bounded smoke test and inspect the
+generated `task.*/input.json` before launching a long campaign.
+
+TensorFlow remains the default when `impl` is omitted.
+For PyTorch and PyTorch Exportable checkpoints, `model_devi_backend` must match
+the training `impl`; checkpoints cannot be frozen across these backends. PT2
+export runs with `run_explore_config`, which must select hardware, the libtorch
+version, and other runtime libraries compatible with the LAMMPS exploration
+environment.
+
 
 ### Exploration
 

dpgen2/constants.py

@@ -4,7 +4,8 @@
 train_log_name = "train.log"
 model_name_pattern = "model.%03d.pb"
 pytorch_model_name_pattern = "model.%03d.pth"
-model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth)"
+pt2_model_name_pattern = "model.%03d.pt2"
+model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth|\.pt2)"
 lmp_index_pattern = "%06d"
 lmp_task_pattern = "task." + lmp_index_pattern
 lmp_conf_name = "conf.lmp"

dpgen2/entrypoint/args.py

@@ -45,6 +45,7 @@ def make_link(content, ref_key):
 
 
 def dp_dist_train_args():
+    doc_numb_models = "Number of student models trained for model deviation"
     doc_config = "Configuration of training"
     doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
     dock_student_model_path = "The path of student model"
@@ -63,6 +64,7 @@ def dp_dist_train_args():
         Argument(
             "template_script", [List[str], str], optional=False, doc=doc_template_script
         ),
+        Argument("numb_models", int, optional=True, default=1, doc=doc_numb_models),
         Argument("student_model_path", str, optional=True, doc=dock_student_model_path),
         Argument(
             "student_model_uri",

dpgen2/entrypoint/submit.py

@@ -107,6 +107,7 @@
     RunRelax,
     RunRelaxHDF5,
     SelectConfs,
+    validate_model_backend,
 )
 from dpgen2.op.caly_evo_step_merge import (
     CalyEvoStepMerge,
@@ -462,6 +463,92 @@ def get_systems_from_data(data, data_prefix=None):
     return data
 
 
+def _normalize_model_backend(backend: str) -> str:
+    return {"pt-expt": "pytorch-exportable"}.get(backend, backend)
+
+
+def _iter_model_sections(template_script: dict):
+    model = template_script.get("model", {})
+    yield "model", model
+    for name, branch in (model.get("model_dict") or {}).items():
+        yield f"model.model_dict.{name}", branch
+
+
+def _model_family(template_script: dict) -> Optional[str]:
+    families = set()
+    for _, model in _iter_model_sections(template_script):
+        model_type = model.get("type")
+        descriptor = model.get("descriptor", {})
+        descriptor_type = (
+            descriptor.get("type") if isinstance(descriptor, dict) else None
+        )
+        model_type = model_type.lower() if isinstance(model_type, str) else model_type
+        descriptor_type = (
+            descriptor_type.lower()
+            if isinstance(descriptor_type, str)
+            else descriptor_type
+        )
+        if descriptor_type == "dpa4c":
+            families.add("dpa4c")
+        if model_type == "dpa4" or descriptor_type in {"dpa4", "sezm"}:
+            families.add("dpa4")
+    if len(families) > 1:
+        raise RuntimeError(
+            "A training template cannot mix DPA4 and DPA4C branches because "
+            "they require different DeePMD backends"
+        )
+    return next(iter(families), None)
+
+
+def validate_dpa_training_template(
+    train_backend: str,
+    explore_config: dict,
+    template_script: dict,
+) -> None:
+    """Validate DPA4/DPA4C backend and compile-option placement."""
+    family = _model_family(template_script)
+    if family is None:
+        return
+
+    train_backend = _normalize_model_backend(train_backend)
+    expected_backend = "pytorch" if family == "dpa4" else "pytorch-exportable"
+    if train_backend != expected_backend:
+        raise RuntimeError(
+            f"{family.upper()} training requires impl='{expected_backend}', "
+            f"not '{train_backend}'"
+        )
+
+    normalized_explore = RunLmp.normalize_config(explore_config)
+    if normalized_explore["model_format"] != "pt2":
+        raise RuntimeError(
+            f"{family.upper()} LAMMPS exploration requires model_format='pt2'"
+        )
+
+    misplaced = []
+    if family == "dpa4c":
+        for scope, model in _iter_model_sections(template_script):
+            for key in ("use_compile", "enable_tf32"):
+                if key in model:
+                    misplaced.append(f"{scope}.{key}")
+        if misplaced:
+            raise RuntimeError(
+                "DPA4C uses training.enable_compile and training.enable_tf32; "
+                f"remove misplaced {', '.join(misplaced)}"
+            )
+    else:
+        training = template_script.get("training", {})
+        misplaced = [
+            f"training.{key}"
+            for key in ("enable_compile", "enable_tf32")
+            if key in training
+        ]
+        if misplaced:
+            raise RuntimeError(
+                "DPA4 uses model.use_compile and model.enable_tf32; "
+                f"remove misplaced {', '.join(misplaced)}"
+            )
+
+
 def workflow_concurrent_learning(
     config: Dict,
 ) -> Step:
@@ -472,6 +559,19 @@ def workflow_concurrent_learning(
     train_style = config["train"]["type"]
     explore_style = config["explore"]["type"]
     fp_style = config["fp"]["type"]
+    template_script_ = config["train"]["template_script"]
+    if isinstance(template_script_, list):
+        template_script = [json.loads(Path(ii).read_text()) for ii in template_script_]
+    else:
+        template_script = json.loads(Path(template_script_).read_text())
+    if train_style in ["dp", "dp-dist"] and explore_style == "lmp":
+        train_backend = train_config.get("impl", "tensorflow")
+        validate_model_backend(train_backend, explore_config)
+        templates = (
+            template_script if isinstance(template_script, list) else [template_script]
+        )
+        for template in templates:
+            validate_dpa_training_template(train_backend, explore_config, template)
     prep_train_config = config["step_configs"]["prep_train_config"]
     run_train_config = config["step_configs"]["run_train_config"]
     prep_explore_config = config["step_configs"]["prep_explore_config"]
@@ -493,12 +593,17 @@ def workflow_concurrent_learning(
                 "not match numb_models={numb_models}"
             )
     elif train_style == "dp-dist":
+        numb_models = config["train"]["numb_models"]
+        if "student_model_path" in config["train"] and numb_models != 1:
+            raise RuntimeError(
+                "student_model_path initializes one model; omit it for multiple "
+                "from-scratch students or set numb_models=1"
+            )
         init_models_paths = (
             [config["train"]["student_model_path"]]
             if "student_model_path" in config["train"]
             else None
         )
-        config["train"]["numb_models"] = 1
     else:
         raise RuntimeError(f"unknown params, train_style: {train_style}")
 
@@ -554,11 +659,6 @@ def workflow_concurrent_learning(
 
     type_map = config["inputs"]["type_map"]
     numb_models = config["train"]["numb_models"]
-    template_script_ = config["train"]["template_script"]
-    if isinstance(template_script_, list):
-        template_script = [json.loads(Path(ii).read_text()) for ii in template_script_]
-    else:
-        template_script = json.loads(Path(template_script_).read_text())
 
     if (
         "teacher_model_path" in explore_config

dpgen2/op/__init__.py

@@ -35,8 +35,10 @@
     RunDPTrain,
 )
 from .run_lmp import (
+    PrepareDPModels,
     RunLmp,
     RunLmpHDF5,
+    validate_model_backend,
 )
 from .run_relax import (
     RunRelax,

dpgen2/op/run_dp_train.py

@@ -61,7 +61,10 @@ def _make_train_command(
     # find checkpoint
     if impl == "tensorflow" and os.path.isfile("checkpoint"):
         checkpoint = "model.ckpt"
-    elif impl == "pytorch" and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0:
+    elif (
+        impl in ["pytorch", "pytorch-exportable"]
+        and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0
+    ):
         checkpoint = "model.ckpt-%s.pt" % max(
             [int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")]
         )
@@ -184,10 +187,14 @@ def execute(
         finetune_mode = ip["optional_parameter"]["finetune_mode"]
         config = ip["config"] if ip["config"] is not None else {}
         impl = ip["config"].get("impl", "tensorflow")
+        if impl == "pt-expt":
+            impl = "pytorch-exportable"
         dp_command = ip["config"].get("command", "dp").split()
-        assert impl in ["tensorflow", "pytorch"]
+        assert impl in ["tensorflow", "pytorch", "pytorch-exportable"]
         if impl == "pytorch":
             dp_command.append("--pt")
+        elif impl == "pytorch-exportable":
+            dp_command.append("--pt-expt")
         finetune_args = config.get("finetune_args", "")
         train_args = config.get("train_args", "")
         config = RunDPTrain.normalize_config(config)
@@ -318,7 +325,7 @@ def clean_before_quit():
                 shutil.copy2("input_v2_compat.json", train_script_name)
 
             # freeze model
-            if impl == "pytorch":
+            if impl in ["pytorch", "pytorch-exportable"]:
                 model_file = "model.ckpt.pt"
             else:
                 ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"])
@@ -339,10 +346,10 @@ def clean_before_quit():
                     )
                     raise FatalError("dp freeze failed")
                 model_file = "frozen_model.pb"
-            fplog.write("#=================== freeze std out ===================\n")
-            fplog.write(out)
-            fplog.write("#=================== freeze std err ===================\n")
-            fplog.write(err)
+                fplog.write("#=================== freeze std out ===================\n")
+                fplog.write(out)
+                fplog.write("#=================== freeze std err ===================\n")
+                fplog.write(err)
 
             clean_before_quit()
 
@@ -503,7 +510,7 @@ def decide_init_model(
     @staticmethod
     def training_args():
         doc_command = "The command for DP, 'dp' for default"
-        doc_impl = "The implementation/backend of DP. It can be 'tensorflow' or 'pytorch'. 'tensorflow' for default."
+        doc_impl = "The implementation/backend of DP. It can be 'tensorflow', 'pytorch', or 'pytorch-exportable' (alias 'pt-expt'). 'tensorflow' for default."
         doc_init_model_policy = "The policy of init-model training. It can be\n\n\
     - 'no': No init-model training. Traing from scratch.\n\n\
     - 'yes': Do init-model training.\n\n\