feat: support DPA4 and DPA4C model formats

doc/run/param.rst

@@ -8,3 +8,68 @@ dpgen run param parameters
 .. dargs::
    :module: dpgen.generator.arginfo
    :func: run_jdata_arginfo
+
+DPA4 and DPA4C
+---------------
+
+DPA4 and the PyTorch-exportable backend require DeePMD-kit 3.2 or later. DPA4
+uses the regular PyTorch backend for both training and export:
+
+.. code-block:: json
+
+   {
+     "train_backend": "pytorch",
+     "model_format": "pt2",
+     "default_training_param": {
+       "model": {
+         "type": "dpa4",
+         "use_compile": true,
+         "enable_tf32": true
+       }
+     }
+   }
+
+DPA4C uses the PyTorch-exportable backend for both training and graph export:
+
+.. code-block:: json
+
+   {
+     "train_backend": "pytorch-exportable",
+     "model_format": "pt2",
+     "dp_compress": true,
+     "default_training_param": {
+       "model": {
+         "descriptor": {"type": "dpa4c"}
+       },
+       "training": {
+         "enable_compile": true,
+         "enable_tf32": true
+       }
+     }
+   }
+
+The default ``train_backend`` remains ``tensorflow``. ``pt-expt`` is accepted
+as an alias of ``pytorch-exportable``. PyTorch-exportable model deviation with
+LAMMPS defaults to ``pt2``. Training checkpoints keep the ``.pt`` suffix
+independently of the frozen model format.
+
+Freeze and export use the same backend as training. Regular PyTorch and
+PyTorch-exportable checkpoints are backend-specific and cannot be converted by
+switching the ``dp`` backend flag after training.
+
+The acceleration controls belong to different sections of the DeePMD training
+template: DPA4 uses ``model.use_compile`` and ``model.enable_tf32``; DPA4C uses
+``training.enable_compile`` and ``training.enable_tf32``. DP-GEN validates the
+backend and these locations but does not inject either policy. A template cannot
+mix DPA4 and DPA4C branches because they require different training backends.
+
+AOTInductor ``.pt2`` archives are specific to the target CPU or GPU, GPU
+compute capability, and libtorch version. DP-GEN therefore finishes the
+training submission with the checkpoint, then runs ``freeze`` and optional
+``compress`` in a separate submission using ``model_devi_machine`` and
+``model_devi_resources``. Those resources must select the same hardware and
+software target used by all subsequent model-deviation jobs. The resulting
+models are linked into the model-deviation stage automatically.
+
+The dense PyTorch-exportable ``pte`` format remains available for non-LAMMPS
+workflows but is not supported by LAMMPS model deviation.

dpgen/generator/arginfo.py

@@ -94,12 +94,37 @@ def training_args_dp() -> list[Argument]:
     list[dargs.Argument]
         List of training arguments.
     """
-    doc_train_backend = (
-        "The backend of the training. Currently only support tensorflow and pytorch."
+    doc_train_backend = textwrap.dedent(
+        """\
+        The DeePMD-kit training backend. Supported values are ``tensorflow``,
+        ``pytorch``, ``pytorch-exportable`` (or its ``pt-expt`` alias), and ``jax``.
+        The PyTorch-exportable backend and DPA4 ``pt2`` export require DeePMD-kit
+        3.2 or later.
+        """
+    )
+    doc_model_format = textwrap.dedent(
+        """\
+        The frozen model format. Defaults are ``pb`` for TensorFlow, ``pth`` for
+        PyTorch, ``pt2`` for PyTorch-exportable model deviation with LAMMPS,
+        and ``savedmodel`` for JAX. PyTorch ``pt2`` is the DPA4 export;
+        PyTorch-exportable ``pt2`` is the graph export used by DPA4C. The
+        PyTorch-exportable ``pte`` format is not supported by LAMMPS. Freeze
+        and export use the training backend; cross-backend checkpoint conversion
+        is not supported.
+        """
     )
     doc_training_iter0_model_path = "The model used to init the first iter training. Number of element should be equal to numb_models."
     doc_training_init_model = "Iteration > 0, the model parameters will be initilized from the model trained at the previous iteration. Iteration == 0, the model parameters will be initialized from training_iter0_model_path."
-    doc_default_training_param = "Training parameters for deepmd-kit in 00.train. You can find instructions from `DeePMD-kit documentation <https://docs.deepmodeling.com/projects/deepmd/>`_."
+    doc_default_training_param = textwrap.dedent(
+        """\
+        Training parameters for DeePMD-kit in 00.train. DPA4 uses
+        ``model.use_compile`` and ``model.enable_tf32`` with the PyTorch backend.
+        DPA4C uses ``training.enable_compile`` and ``training.enable_tf32`` with
+        the PyTorch-exportable backend. DP-GEN validates these locations but does
+        not inject numerical-policy settings. See the `DeePMD-kit documentation
+        <https://docs.deepmodeling.com/projects/deepmd/>`_.
+        """
+    )
     doc_dp_train_skip_neighbor_stat = "Append --skip-neighbor-stat flag to dp train."
     doc_dp_compress = "Use dp compress to compress the model."
     doc_training_reuse_iter = "The minimal index of iteration that continues training models from old models of last iteration."
@@ -139,6 +164,12 @@ def training_args_dp() -> list[Argument]:
             default="tensorflow",
             doc=doc_train_backend,
         ),
+        Argument(
+            "model_format",
+            str,
+            optional=True,
+            doc=doc_model_format,
+        ),
         Argument(
             "training_iter0_model_path",
             list[str],

dpgen/generator/lib/calypso_run_opt.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 import os
 import time
 
@@ -110,9 +111,9 @@ def read_stress_fmax():
     return fmax, pstress
 
 
-def run_opt(fmax, stress):
+def run_opt(fmax, stress, model):
     """Using the ASE&DP to Optimize Configures."""
-    calc = DP(model="../graph.000.pb")  # init the model before iteration
+    calc = DP(model=model)  # init the model before iteration
     os.system("mv OUTCAR OUTCAR-last")
 
     print("Start to Optimize Structures by DP----------")
@@ -164,8 +165,11 @@ def run_opt(fmax, stress):
 
 
 def run():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="../graph.000.pb")
+    args = parser.parse_args()
     fmax, stress = read_stress_fmax()
-    run_opt(fmax, stress)
+    run_opt(fmax, stress, args.model)
 
 
 if __name__ == "__main__":

dpgen/generator/lib/lammps.py

@@ -59,7 +59,7 @@ def make_lammps_input(
         while power < nbeads:
             power *= 10
         ret += "variable        ibead           uloop %d pad\n" % (power - 1)  # noqa: UP031
-    if nbeads is not None:
+    if nbeads is not None or jdata.get("model_format") == "pt2":
         ret += "atom_modify        map yes\n"
     ret += "variable        THERMO_FREQ     equal %d\n" % trj_freq  # noqa: UP031
     ret += "variable        DUMP_FREQ       equal %d\n" % trj_freq  # noqa: UP031

dpgen/generator/lib/run_calypso.py

@@ -29,8 +29,27 @@
 calypso_model_devi_name = "model_devi_results"
 
 
+def _find_models(path, model_suffix=".pb"):
+    """Return model-deviation artifacts for the resolved deployment format."""
+    return glob.glob(os.path.join(path, f"graph*{model_suffix}"))
+
+
+def _make_calypso_opt_command(deepmdkit_python, model_name):
+    """Return the CALYPSO optimization command for the resolved model."""
+    return (
+        f"{deepmdkit_python} calypso_run_opt.py --model ../{model_name} "
+        "1>> model_devi.log 2>> model_devi.log"
+    )
+
+
 def gen_structures(
-    iter_index, jdata, mdata, caly_run_path, current_idx, length_of_caly_runopt_list
+    iter_index,
+    jdata,
+    mdata,
+    caly_run_path,
+    current_idx,
+    length_of_caly_runopt_list,
+    model_suffix=".pb",
 ):
     # run calypso
     # vsc means generate elemental, binary and ternary at the same time
@@ -50,13 +69,11 @@ def gen_structures(
     calypso_path = mdata.get("model_devi_calypso_path")
     # calypso_input_path = jdata.get('calypso_input_path')
 
-    all_models = glob.glob(os.path.join(calypso_run_opt_path, "graph*pb"))
+    all_models = _find_models(calypso_run_opt_path, model_suffix)
     model_names = [os.path.basename(ii) for ii in all_models]
 
     deepmdkit_python = mdata.get("model_devi_deepmdkit_python")
-    command = (
-        f"{deepmdkit_python} calypso_run_opt.py  1>> model_devi.log 2>> model_devi.log"
-    )
+    command = _make_calypso_opt_command(deepmdkit_python, sorted(model_names)[0])
     # command = "%s calypso_run_opt.py %s 1>> model_devi.log 2>> model_devi.log" % (deepmdkit_python,os.path.abspath(calypso_run_opt_path))
     # command += "  ||  %s check_outcar.py %s " % (deepmdkit_python,os.path.abspath(calypso_run_opt_path))
     command += f"  ||  {deepmdkit_python} check_outcar.py  "
@@ -335,7 +352,7 @@ def gen_structures(
     os.chdir(cwd)
 
 
-def gen_main(iter_index, jdata, mdata, caly_run_opt_list, gen_idx):
+def gen_main(iter_index, jdata, mdata, caly_run_opt_list, gen_idx, model_suffix=".pb"):
     iter_name = make_iter_name(iter_index)
     work_path = os.path.join(iter_name, model_devi_name)
 
@@ -353,7 +370,13 @@ def gen_main(iter_index, jdata, mdata, caly_run_opt_list, gen_idx):
     for iidx, temp_path in enumerate(caly_run_opt_list):
         if iidx >= indice:
             gen_structures(
-                iter_index, jdata, mdata, temp_path, iidx, len(caly_run_opt_list)
+                iter_index,
+                jdata,
+                mdata,
+                temp_path,
+                iidx,
+                len(caly_run_opt_list),
+                model_suffix=model_suffix,
             )
 
 
@@ -448,7 +471,7 @@ def analysis(iter_index, jdata, calypso_model_devi_path):
     os.chdir(cwd)
 
 
-def run_calypso_model_devi(iter_index, jdata, mdata):
+def run_calypso_model_devi(iter_index, jdata, mdata, model_suffix=".pb"):
     dlog.info("start running CALYPSO")
 
     iter_name = make_iter_name(iter_index)
@@ -483,7 +506,14 @@ def run_calypso_model_devi(iter_index, jdata, mdata):
         if lines[-1].strip().strip("\n").split()[0] == "1":
             # Gen Structures
             gen_index = lines[-1].strip().strip("\n").split()[1]
-            gen_main(iter_index, jdata, mdata, caly_run_opt_list, gen_index)
+            gen_main(
+                iter_index,
+                jdata,
+                mdata,
+                caly_run_opt_list,
+                gen_index,
+                model_suffix=model_suffix,
+            )
 
         elif lines[-1].strip().strip("\n") == "2":
             # Analysis & to deepmd/raw
@@ -492,7 +522,7 @@ def run_calypso_model_devi(iter_index, jdata, mdata):
         elif lines[-1].strip().strip("\n") == "3":
             # Model Devi
             _calypso_run_opt_path = os.path.abspath(caly_run_opt_list[0])
-            all_models = glob.glob(os.path.join(_calypso_run_opt_path, "graph*pb"))
+            all_models = _find_models(_calypso_run_opt_path, model_suffix)
             cwd = os.getcwd()
             os.chdir(calypso_model_devi_path)
             args = " ".join(