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

@@ -14,7 +14,18 @@ def run_mdata_arginfo() -> Argument:
     Argument
         arginfo
     """
-    return general_mdata_arginfo("run_mdata", ("train", "model_devi", "fp"))
+    arginfo = general_mdata_arginfo("run_mdata", ("train", "model_devi", "fp"))
+    train = arginfo.sub_fields["train"]
+    train.sub_fields["export_command"] = Argument(
+        "export_command",
+        str,
+        optional=True,
+        doc=(
+            "DeePMD command used to freeze and compress pt2 models on the "
+            "model-deviation machine. Defaults to train.command."
+        ),
+    )
+    return arginfo
 
 
 # basics
@@ -94,12 +105,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 +175,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_check_outcar.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 import os
 
 import numpy as np
@@ -82,11 +83,11 @@ def Write_Outcar(element, ele, volume, lat, pos, ene, force, stress, pstress):
     f.write(f"enthalpy is  TOTEN    = {enthalpy:20.6f} {enthalpy:20.6f}\n")
 
 
-def check():
+def check(model):
     from ase.io import read
     from deepmd.calculator import DP
 
-    calc = DP(model="../graph.000.pb")  # init the model before iteration
+    calc = DP(model=model)  # init the model before iteration
 
     to_be_opti = read("POSCAR")
     to_be_opti.calc = calc
@@ -118,4 +119,6 @@ def check():
 
 cwd = os.getcwd()
 if not os.path.exists(os.path.join(cwd, "OUTCAR")):
-    check()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="../graph.000.pb")
+    check(parser.parse_args().model)

dpgen/generator/lib/calypso_run_opt.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 import os
 import time
 
@@ -110,9 +111,23 @@ def read_stress_fmax():
     return fmax, pstress
 
 
-def run_opt(fmax, stress):
-    """Using the ASE&DP to Optimize Configures."""
-    calc = DP(model="../graph.000.pb")  # init the model before iteration
+def run_opt(fmax, stress, model):
+    """Optimize a CALYPSO structure with ASE and a DeePMD model.
+
+    Parameters
+    ----------
+    fmax : float
+        Maximum force convergence threshold for the ASE optimizer.
+    stress : float
+        Target external pressure in kbar.
+    model : str
+        Path to the frozen DeePMD model artifact.
+
+    Returns
+    -------
+    None
+    """
+    calc = DP(model=model)  # init the model before iteration
     os.system("mv OUTCAR OUTCAR-last")
 
     print("Start to Optimize Structures by DP----------")
@@ -164,8 +179,17 @@ def run_opt(fmax, stress):
 
 
 def run():
+    """Run the CALYPSO optimization command-line entry point.
+
+    Returns
+    -------
+    None
+    """
+    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
@@ -149,15 +149,15 @@ def make_lammps_input(
             else:
                 ret += f"pair_style      deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file model_devi${{ibead}}.out {keywords}\n"
 
-    # Add pair_coeff lines
+    # Use DP-GEN's LAMMPS type order when it is available. Direct callers that
+    # omit type_map retain the historical bare pair_coeff behavior.
+    type_map_str = " ".join(jdata.get("type_map", []))
+    type_map_args = f" {type_map_str}" if type_map_str else ""
     if d3_enabled:
-        # D3 requires type maps (element symbols)
-        type_map = jdata.get("type_map", [])
-        type_map_str = " ".join(type_map)
-        ret += "pair_coeff      * * deepmd\n"
+        ret += f"pair_coeff      * * deepmd{type_map_args}\n"
         ret += f"pair_coeff      * * dispersion/d3 {type_map_str}\n"
     else:
-        ret += "pair_coeff      * *\n"
+        ret += f"pair_coeff      * *{type_map_args}\n"
     ret += "\n"
     ret += "thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz\n"
     ret += "thermo          ${THERMO_FREQ}\n"