diff --git a/doc/run/param.rst b/doc/run/param.rst
index 74772253e..251557340 100644
--- a/doc/run/param.rst
+++ b/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.
diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py
index 0f8abb865..24cb61362 100644
--- a/dpgen/generator/arginfo.py
+++ b/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 `_."
+ 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
+ `_.
+ """
+ )
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],
diff --git a/dpgen/generator/lib/calypso_run_opt.py b/dpgen/generator/lib/calypso_run_opt.py
index 1bc765af1..cca0200a1 100644
--- a/dpgen/generator/lib/calypso_run_opt.py
+++ b/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__":
diff --git a/dpgen/generator/lib/lammps.py b/dpgen/generator/lib/lammps.py
index b4bf603d7..80c37c72a 100644
--- a/dpgen/generator/lib/lammps.py
+++ b/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"
diff --git a/dpgen/generator/lib/run_calypso.py b/dpgen/generator/lib/run_calypso.py
index dbff1f0cb..77b0cb40f 100644
--- a/dpgen/generator/lib/run_calypso.py
+++ b/dpgen/generator/lib/run_calypso.py
@@ -29,9 +29,77 @@
calypso_model_devi_name = "model_devi_results"
+def _find_models(path, model_suffix=".pb"):
+ """Find model-deviation artifacts for a deployment format.
+
+ Parameters
+ ----------
+ path : str or os.PathLike
+ Directory containing committee model artifacts.
+ model_suffix : str, optional
+ Resolved model filename suffix.
+
+ Returns
+ -------
+ list[str]
+ Paths matching the resolved committee model suffix.
+ """
+ return glob.glob(os.path.join(path, f"graph*{model_suffix}"))
+
+
+def _make_calypso_opt_command(deepmdkit_python, model_name):
+ """Build the CALYPSO optimization command for a resolved model.
+
+ Parameters
+ ----------
+ deepmdkit_python : str
+ Python executable used by the DeePMD environment.
+ model_name : str
+ Filename of the model artifact forwarded to CALYPSO.
+
+ Returns
+ -------
+ str
+ Shell command that invokes the optimization script.
+ """
+ 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",
):
+ """Generate and optimize one CALYPSO structure batch.
+
+ Parameters
+ ----------
+ iter_index : int
+ DP-GEN iteration index.
+ jdata : dict
+ DP-GEN workflow parameters.
+ mdata : dict
+ Machine and resource parameters.
+ caly_run_path : str
+ CALYPSO generation and optimization working directory.
+ current_idx : int
+ Index of the current CALYPSO generation.
+ length_of_caly_runopt_list : int
+ Number of CALYPSO generation directories.
+ model_suffix : str, optional
+ Resolved deployment-model suffix.
+
+ Returns
+ -------
+ None
+ """
# run calypso
# vsc means generate elemental, binary and ternary at the same time
vsc = jdata.get("vsc", False) # take CALYPSO as confs generator
@@ -50,13 +118,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 +401,28 @@ 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"):
+ """Run CALYPSO generation from the selected generation index.
+
+ Parameters
+ ----------
+ iter_index : int
+ DP-GEN iteration index.
+ jdata : dict
+ DP-GEN workflow parameters.
+ mdata : dict
+ Machine and resource parameters.
+ caly_run_opt_list : list[str]
+ Ordered CALYPSO generation directories.
+ gen_idx : int
+ Generation index from which to resume.
+ model_suffix : str, optional
+ Resolved deployment-model suffix.
+
+ Returns
+ -------
+ None
+ """
iter_name = make_iter_name(iter_index)
work_path = os.path.join(iter_name, model_devi_name)
@@ -353,7 +440,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 +541,24 @@ 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"):
+ """Run the CALYPSO model-deviation workflow.
+
+ Parameters
+ ----------
+ iter_index : int
+ DP-GEN iteration index.
+ jdata : dict
+ DP-GEN workflow parameters.
+ mdata : dict
+ Machine and resource parameters.
+ model_suffix : str, optional
+ Resolved deployment-model suffix.
+
+ Returns
+ -------
+ None
+ """
dlog.info("start running CALYPSO")
iter_name = make_iter_name(iter_index)
@@ -483,7 +593,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 +609,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(
diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py
index 0ce513a96..d26963450 100644
--- a/dpgen/generator/run.py
+++ b/dpgen/generator/run.py
@@ -125,21 +125,235 @@
run_opt_file = os.path.join(ROOT_PATH, "generator/lib/calypso_run_opt.py")
-def _get_model_suffix(jdata) -> str:
- """Return the model suffix based on the backend."""
+_BACKEND_ALIASES = {"pt-expt": "pytorch-exportable"}
+_BACKEND_CONFIG = {
+ "tensorflow": {
+ "flag": "",
+ "checkpoint_suffix": ".index",
+ "default_model_format": "pb",
+ "model_formats": {"pb"},
+ },
+ "pytorch": {
+ "flag": "--pt",
+ "checkpoint_suffix": ".pt",
+ "default_model_format": "pth",
+ "model_formats": {"pth", "pt2"},
+ },
+ "pytorch-exportable": {
+ "flag": "--pt-expt",
+ "checkpoint_suffix": ".pt",
+ "default_model_format": "pte",
+ "model_formats": {"pte", "pt2"},
+ },
+ "jax": {
+ "flag": "--jax",
+ "checkpoint_suffix": ".jax",
+ "default_model_format": "savedmodel",
+ "model_formats": {"savedmodel"},
+ },
+}
+
+
+def _get_backend(jdata, key, default) -> tuple[str, dict]:
+ """Return and validate a DeePMD backend."""
mlp_engine = jdata.get("mlp_engine", "dp")
- if mlp_engine == "dp":
- suffix_map = {"tensorflow": ".pb", "pytorch": ".pth", "jax": ".savedmodel"}
- backend = jdata.get("train_backend", "tensorflow")
- if backend in suffix_map:
- suffix = suffix_map[backend]
- else:
+ if mlp_engine != "dp":
+ raise ValueError(f"Unsupported engine: {mlp_engine}")
+
+ backend = jdata.get(key, default)
+ backend = _BACKEND_ALIASES.get(backend, backend)
+ if backend not in _BACKEND_CONFIG:
+ supported = "', '".join(_BACKEND_CONFIG)
+ raise ValueError(
+ f"The backend {backend} is not available. Supported backends are: '{supported}'."
+ )
+ return backend, _BACKEND_CONFIG[backend]
+
+
+def _get_train_backend_config(jdata) -> tuple[str, dict]:
+ """Return the training backend configuration."""
+ return _get_backend(jdata, "train_backend", "tensorflow")
+
+
+def _get_model_backend_config(jdata) -> tuple[str, dict, str]:
+ """Return the training backend and validate its frozen model format."""
+ backend, config = _get_train_backend_config(jdata)
+ default_model_format = config["default_model_format"]
+ if (
+ backend == "pytorch-exportable"
+ and jdata.get("model_devi_engine", "lammps") == "lammps"
+ ):
+ default_model_format = "pt2"
+ model_format = jdata.get("model_format", default_model_format)
+ if model_format not in config["model_formats"]:
+ supported = "', '".join(sorted(config["model_formats"]))
+ raise ValueError(
+ f"The model format {model_format} is not available for backend {backend}. "
+ f"Supported formats are: '{supported}'."
+ )
+ if (
+ backend == "pytorch-exportable"
+ and model_format == "pte"
+ and jdata.get("model_devi_engine", "lammps") == "lammps"
+ ):
+ raise ValueError(
+ "The pte model format is not supported by LAMMPS; use model_format=pt2."
+ )
+ return backend, config, model_format
+
+
+def _get_model_suffix(jdata) -> str:
+ """Return the frozen model suffix."""
+ _, _, model_format = _get_model_backend_config(jdata)
+ return f".{model_format}"
+
+
+def _get_checkpoint_suffix(jdata) -> str:
+ """Return the training checkpoint suffix."""
+ _, config = _get_train_backend_config(jdata)
+ return config["checkpoint_suffix"]
+
+
+def _get_train_backend_flag(jdata) -> str:
+ """Return the DeePMD CLI backend flag."""
+ _, config = _get_train_backend_config(jdata)
+ return config["flag"]
+
+
+def _get_input_model_suffix(models) -> str:
+ """Return the common suffix of input models."""
+ suffixes = {Path(model).suffix.lower() for model in models}
+ if "" in suffixes or len(suffixes) != 1:
+ raise ValueError("Input models must have the same non-empty file suffix.")
+ return suffixes.pop()
+
+
+def _iter_model_sections(training_param):
+ """Yield the top-level model and each model-dictionary branch.
+
+ Parameters
+ ----------
+ training_param : dict
+ DeePMD training parameters.
+
+ Yields
+ ------
+ tuple[str, dict]
+ The dotted configuration path and corresponding model section.
+ """
+ model = training_param.get("model", {})
+ yield "model", model
+ for name, branch in (model.get("model_dict") or {}).items():
+ yield f"model.model_dict.{name}", branch
+
+
+def _get_dpa_model_family(training_param) -> Optional[str]:
+ """Identify the DPA model family in a training configuration.
+
+ Parameters
+ ----------
+ training_param : dict
+ DeePMD training parameters.
+
+ Returns
+ -------
+ str or None
+ "dpa4", "dpa4c", or None when neither family is present.
+
+ Raises
+ ------
+ ValueError
+ If DPA4 and DPA4C branches are mixed in one configuration.
+ """
+ families = set()
+ for _, model in _iter_model_sections(training_param):
+ 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 in {"dpa4", "sezm"} or descriptor_type in {"dpa4", "sezm"}:
+ families.add("dpa4")
+ if len(families) > 1:
+ raise ValueError(
+ "default_training_param cannot mix DPA4 and DPA4C branches because "
+ "they require different DeePMD backends"
+ )
+ return next(iter(families), None)
+
+
+def _validate_dpa_training_config(jdata) -> None:
+ """Validate DPA backend, format, and acceleration-option placement.
+
+ Parameters
+ ----------
+ jdata : dict
+ DP-GEN parameters containing the training and deployment configuration.
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ ValueError
+ If the model family, backend, deployment format, or acceleration
+ options are incompatible.
+ """
+ training_param = jdata.get("default_training_param", {})
+ family = _get_dpa_model_family(training_param)
+ train_backend, _ = _get_train_backend_config(jdata)
+ _, _, model_format = _get_model_backend_config(jdata)
+ if family is None:
+ if train_backend == "pytorch" and model_format == "pt2":
+ raise ValueError(
+ "The regular PyTorch backend only exports pt2 for DPA4/SeZM models."
+ )
+ return
+ expected_backend = "pytorch" if family == "dpa4" else "pytorch-exportable"
+ if train_backend != expected_backend:
+ raise ValueError(
+ f"{family.upper()} training requires train_backend='{expected_backend}', "
+ f"not '{train_backend}'"
+ )
+
+ if jdata.get("model_devi_engine", "lammps") == "lammps":
+ if model_format != "pt2":
+ raise ValueError(
+ f"{family.upper()} LAMMPS model deviation requires model_format='pt2'"
+ )
+
+ if family == "dpa4c":
+ misplaced = []
+ for scope, model in _iter_model_sections(training_param):
+ for key in ("use_compile", "enable_tf32"):
+ if key in model:
+ misplaced.append(f"{scope}.{key}")
+ if misplaced:
raise ValueError(
- f"The backend {backend} is not available. Supported backends are: 'tensorflow', 'pytorch', 'jax'."
+ "DPA4C uses training.enable_compile and training.enable_tf32; "
+ f"remove misplaced {', '.join(misplaced)}"
)
- return suffix
else:
- raise ValueError(f"Unsupported engine: {mlp_engine}")
+ training = training_param.get("training", {})
+ misplaced = [
+ f"training.{key}"
+ for key in ("enable_compile", "enable_tf32")
+ if key in training
+ ]
+ if misplaced:
+ raise ValueError(
+ "DPA4 uses model.use_compile and model.enable_tf32; "
+ f"remove misplaced {', '.join(misplaced)}"
+ )
def get_job_names(jdata):
@@ -667,9 +881,13 @@ def make_train_dp(iter_index, jdata, mdata):
None,
)
if copied_models is not None:
+ input_model_suffix = _get_input_model_suffix(copied_models)
for ii in range(len(copied_models)):
_link_old_models(
- work_path, [copied_models[ii]], ii, basename=f"init{suffix}"
+ work_path,
+ [copied_models[ii]],
+ ii,
+ basename=f"init{input_model_suffix}",
)
# Copy user defined forward files
symlink_user_forward_files(mdata=mdata, task_type="train", work_path=work_path)
@@ -729,8 +947,12 @@ def run_train(iter_index, jdata, mdata):
def run_train_dp(iter_index, jdata, mdata):
# print("debug:run_train:mdata", mdata)
# load json param
+ _validate_dpa_training_config(jdata)
numb_models = jdata["numb_models"]
+ train_backend, _ = _get_train_backend_config(jdata)
+ _, _, model_format = _get_model_backend_config(jdata)
suffix = _get_model_suffix(jdata)
+ checkpoint_suffix = _get_checkpoint_suffix(jdata)
# train_param = jdata['train_param']
train_input_file = default_train_input_file
training_reuse_iter = jdata.get("training_reuse_iter")
@@ -752,6 +974,27 @@ def run_train_dp(iter_index, jdata, mdata):
except KeyError:
mdata = set_version(mdata)
+ if (train_backend == "pytorch-exportable" or model_format == "pt2") and Version(
+ mdata["deepmd_version"]
+ ) < Version("3.2"):
+ raise RuntimeError(
+ "The pytorch-exportable backend and pt2 models require DeePMD-kit 3.2 or later."
+ )
+ if train_backend == "pytorch-exportable" and training_init_frozen_model is not None:
+ raise RuntimeError(
+ "The pytorch-exportable backend does not support training_init_frozen_model; "
+ "use training_finetune_model or a checkpoint instead."
+ )
+ if (
+ train_backend == "pytorch"
+ and model_format == "pt2"
+ and jdata.get("dp_compress", False)
+ ):
+ raise RuntimeError(
+ "The pytorch backend cannot compress pt2 models; use "
+ "pytorch-exportable for a compressible pt2 model."
+ )
+
if (
training_init_model
+ (training_init_frozen_model is not None)
@@ -764,10 +1007,9 @@ def run_train_dp(iter_index, jdata, mdata):
train_command = mdata.get("train_command", "dp").strip()
# assert train_command == "dp", "The 'train_command' should be 'dp'" # the tests should be updated to run this command
- if suffix == ".pth":
- train_command += " --pt"
- elif suffix == ".savedmodel":
- train_command += " --jax"
+ backend_flag = _get_train_backend_flag(jdata)
+ if backend_flag:
+ train_command += f" {backend_flag}"
# paths
iter_name = make_iter_name(iter_index)
@@ -783,6 +1025,7 @@ def run_train_dp(iter_index, jdata, mdata):
task_path = os.path.join(work_path, train_task_fmt % ii)
all_task.append(task_path)
commands = []
+ export_commands = []
if Version(mdata["deepmd_version"]) >= Version("1") and Version(
mdata["deepmd_version"]
) < Version("4"):
@@ -797,25 +1040,43 @@ def run_train_dp(iter_index, jdata, mdata):
if training_init_model:
init_flag = " --init-model old/model.ckpt"
elif training_init_frozen_model is not None:
- init_flag = f" --init-frz-model old/init{suffix}"
+ input_model_suffix = _get_input_model_suffix(training_init_frozen_model)
+ init_flag = f" --init-frz-model old/init{input_model_suffix}"
elif training_finetune_model is not None:
- init_flag = f" --finetune old/init{suffix}"
+ input_model_suffix = _get_input_model_suffix(training_finetune_model)
+ init_flag = f" --finetune old/init{input_model_suffix}"
command = f"{train_command} train {train_input_file}{extra_flags}"
- if suffix == ".pb":
- ckpt_suffix = ".index"
- elif suffix == ".pth":
- ckpt_suffix = ".pt"
- elif suffix == ".savedmodel":
- ckpt_suffix = ".jax"
- else:
- raise RuntimeError(f"Unknown suffix {suffix}")
- command = f"{{ if [ ! -f model.ckpt{ckpt_suffix} ]; then {command}{init_flag}; else {command} --restart model.ckpt; fi }}"
+ command = f"{{ if [ ! -f model.ckpt{checkpoint_suffix} ]; then {command}{init_flag}; else {command} --restart model.ckpt; fi }}"
command = f"/bin/sh -c {shlex.quote(command)}"
commands.append(command)
- command = f"{train_command} freeze"
- commands.append(command)
- if jdata.get("dp_compress", False):
- commands.append(f"{train_command} compress")
+ if model_format == "pt2":
+ if train_backend == "pytorch-exportable":
+ command = (
+ f"{train_command} freeze -c model.ckpt.pt -o frozen_model "
+ "--lower-kind graph"
+ )
+ else:
+ command = f"{train_command} freeze -c model.ckpt.pt -o frozen_model"
+ export_commands.append(command)
+ if jdata.get("dp_compress", False):
+ export_commands.append(
+ f"{train_command} compress -i frozen_model{suffix} "
+ f"-o frozen_model_compressed{suffix}"
+ )
+ else:
+ if train_backend == "pytorch-exportable":
+ command = f"{train_command} freeze -o frozen_model{suffix}"
+ else:
+ command = f"{train_command} freeze"
+ commands.append(command)
+ if jdata.get("dp_compress", False):
+ if train_backend == "pytorch-exportable":
+ commands.append(
+ f"{train_command} compress -i frozen_model{suffix} "
+ f"-o frozen_model_compressed{suffix}"
+ )
+ else:
+ commands.append(f"{train_command} compress")
else:
raise RuntimeError(
"DP-GEN currently only supports for DeePMD-kit 1.x to 3.x version!"
@@ -836,42 +1097,49 @@ def run_train_dp(iter_index, jdata, mdata):
if "srtab_file_path" in jdata.keys():
forward_files.append(zbl_file)
if training_init_model:
- if suffix == ".pb":
+ if checkpoint_suffix == ".index":
forward_files += [
os.path.join("old", "model.ckpt.meta"),
os.path.join("old", "model.ckpt.index"),
os.path.join("old", "model.ckpt.data-00000-of-00001"),
]
- elif suffix == ".pth":
+ elif checkpoint_suffix == ".pt":
forward_files += [os.path.join("old", "model.ckpt.pt")]
- elif suffix == ".savedmodel":
+ elif checkpoint_suffix == ".jax":
forward_files += [os.path.join("old", "model.ckpt.jax")]
else:
- raise RuntimeError(f"Unknown suffix {suffix}")
+ raise RuntimeError(f"Unknown checkpoint suffix {checkpoint_suffix}")
elif training_init_frozen_model is not None or training_finetune_model is not None:
- forward_files.append(os.path.join("old", f"init{suffix}"))
+ input_models = (
+ training_init_frozen_model
+ if training_init_frozen_model is not None
+ else training_finetune_model
+ )
+ input_model_suffix = _get_input_model_suffix(input_models)
+ forward_files.append(os.path.join("old", f"init{input_model_suffix}"))
backward_files = [
- f"frozen_model{suffix}",
"lcurve.out",
"train.log",
"checkpoint",
]
- if jdata.get("dp_compress", False):
- backward_files.append(f"frozen_model_compressed{suffix}")
+ if not export_commands:
+ backward_files.append(f"frozen_model{suffix}")
+ if jdata.get("dp_compress", False):
+ backward_files.append(f"frozen_model_compressed{suffix}")
- if suffix == ".pb":
+ if checkpoint_suffix == ".index":
backward_files += [
"model.ckpt.meta",
"model.ckpt.index",
"model.ckpt.data-00000-of-00001",
]
- elif suffix == ".pth":
+ elif checkpoint_suffix == ".pt":
backward_files += ["model.ckpt.pt"]
- elif suffix == ".savedmodel":
+ elif checkpoint_suffix == ".jax":
backward_files += ["model.ckpt.jax"]
else:
- raise RuntimeError(f"Unknown suffix {suffix}")
+ raise RuntimeError(f"Unknown checkpoint suffix {checkpoint_suffix}")
if not jdata.get("one_h5", False):
init_data_sys_ = jdata["init_data_sys"]
@@ -925,6 +1193,24 @@ def run_train_dp(iter_index, jdata, mdata):
errlog="train.log",
)
submission.run_submission()
+ if export_commands:
+ export_backward_files = [f"frozen_model{suffix}"]
+ if jdata.get("dp_compress", False):
+ export_backward_files.append(f"frozen_model_compressed{suffix}")
+ export_submission = make_submission(
+ mdata["model_devi_machine"],
+ mdata["model_devi_resources"],
+ commands=export_commands,
+ work_path=work_path,
+ run_tasks=run_tasks,
+ group_size=1,
+ forward_common_files=[],
+ forward_files=[f"model.ckpt{checkpoint_suffix}"],
+ backward_files=export_backward_files,
+ outlog="model_export.log",
+ errlog="model_export.log",
+ )
+ export_submission.run_submission()
def post_train(iter_index, jdata, mdata):
@@ -1132,40 +1418,101 @@ def revise_lmp_input_model(
def revise_lmp_input_pair_coeff(lmp_lines, jdata=None):
- """Update pair_coeff lines for D3 support."""
+ """Add explicit DeepMD element mapping and D3 pair coefficients.
+
+ Parameters
+ ----------
+ lmp_lines : list[str]
+ Lines from a LAMMPS input template.
+ jdata : dict, optional
+ DP-GEN parameters, including type_map and optional D3 settings.
+
+ Returns
+ -------
+ list[str]
+ The updated LAMMPS input lines.
+
+ Raises
+ ------
+ RuntimeError
+ If a coefficient must be inserted but the template does not contain
+ exactly one pair_style line.
+ """
if jdata is None:
return lmp_lines
lmp_d3 = jdata.get("lmp_d3", {})
d3_enabled = lmp_d3.get("enable", False) if lmp_d3 else False
+ type_map = jdata.get("type_map", [])
+ type_map_str = " ".join(type_map)
+ type_map_args = f" {type_map_str}" if type_map_str else ""
- if not d3_enabled:
+ if not d3_enabled and not type_map:
return lmp_lines
- # D3 requires type maps (element symbols)
- type_map = jdata.get("type_map", [])
- type_map_str = " ".join(type_map)
+ pair_style_idx = find_only_one_key(lmp_lines, ["pair_style"])
+ pair_style_tokens = lmp_lines[pair_style_idx].partition("#")[0].split()
+ hybrid_pair_style = pair_style_tokens[1].startswith("hybrid")
- # Find pair_coeff line
- pair_coeff_idx = None
+ deepmd_coeff_idx = None
+ fallback_coeff_idx = None
+ d3_coeff_idx = None
for idx, line in enumerate(lmp_lines):
- if line.strip().startswith("pair_coeff") and "* *" in line:
- pair_coeff_idx = idx
- break
-
- if pair_coeff_idx is None:
- # If no pair_coeff found, add them after pair_style
- pair_style_idx = find_only_one_key(lmp_lines, ["pair_style"])
- lmp_lines.insert(pair_style_idx + 1, "pair_coeff * * deepmd\n")
- lmp_lines.insert(
- pair_style_idx + 2, f"pair_coeff * * dispersion/d3 {type_map_str}\n"
+ tokens = line.partition("#")[0].split()
+ if not tokens or tokens[0] != "pair_coeff":
+ continue
+ if "dispersion/d3" in tokens:
+ d3_coeff_idx = idx
+ elif tokens[3:4] == ["deepmd"] and deepmd_coeff_idx is None:
+ deepmd_coeff_idx = idx
+ elif fallback_coeff_idx is None:
+ fallback_coeff_idx = idx
+
+ if deepmd_coeff_idx is None and fallback_coeff_idx is not None:
+ fallback_tokens = lmp_lines[fallback_coeff_idx].partition("#")[0].split()
+ fallback_is_bare = fallback_tokens in (
+ ["pair_coeff"],
+ ["pair_coeff", "*", "*"],
)
+ if not hybrid_pair_style or (d3_enabled and fallback_is_bare):
+ deepmd_coeff_idx = fallback_coeff_idx
+
+ if deepmd_coeff_idx is None:
+ deepmd_coeff_idx = pair_style_idx + 1
+ style = " deepmd" if d3_enabled or hybrid_pair_style else ""
+ line = f"pair_coeff * *{style}{type_map_args}\n"
+ lmp_lines.insert(deepmd_coeff_idx, line)
+ if d3_coeff_idx is not None and d3_coeff_idx >= deepmd_coeff_idx:
+ d3_coeff_idx += 1
else:
- # Replace existing pair_coeff with D3 version
- lmp_lines[pair_coeff_idx] = "pair_coeff * * deepmd\n"
- lmp_lines.insert(
- pair_coeff_idx + 1, f"pair_coeff * * dispersion/d3 {type_map_str}\n"
+ tokens = lmp_lines[deepmd_coeff_idx].partition("#")[0].split()
+ bare_coeff = tokens in (
+ ["pair_coeff"],
+ ["pair_coeff", "*", "*"],
+ ["pair_coeff", "*", "*", "deepmd"],
)
+ if d3_enabled:
+ if tokens[:3] == ["pair_coeff", "*", "*"]:
+ if tokens[3:4] == ["deepmd"]:
+ elements = tokens[4:]
+ else:
+ elements = tokens[3:]
+ else:
+ elements = []
+ element_args = f" {' '.join(elements)}" if elements else type_map_args
+ lmp_lines[deepmd_coeff_idx] = f"pair_coeff * * deepmd{element_args}\n"
+ elif bare_coeff:
+ style = " deepmd" if tokens[3:4] == ["deepmd"] else ""
+ lmp_lines[deepmd_coeff_idx] = f"pair_coeff * *{style}{type_map_args}\n"
+
+ if d3_enabled:
+ d3_line = f"pair_coeff * * dispersion/d3 {type_map_str}\n"
+ if d3_coeff_idx is None:
+ lmp_lines.insert(deepmd_coeff_idx + 1, d3_line)
+ else:
+ d3_tokens = lmp_lines[d3_coeff_idx].partition("#")[0].split()
+ if d3_tokens == ["pair_coeff", "*", "*", "dispersion/d3"]:
+ lmp_lines[d3_coeff_idx] = d3_line
return lmp_lines
@@ -1463,6 +1810,42 @@ def make_model_devi(iter_index, jdata, mdata):
return True
+def _validate_pt2_template_atom_map(lmp_lines):
+ """Validate the atom map required by pt2 LAMMPS templates.
+
+ Parameters
+ ----------
+ lmp_lines : list[str]
+ Lines of the LAMMPS input template.
+
+ Raises
+ ------
+ ValueError
+ If ``atom_modify map yes`` is missing or follows ``read_data`` or
+ ``read_restart``.
+ """
+ atom_map_index = None
+ read_index = None
+ for line_index, line in enumerate(lmp_lines):
+ command = line.partition("#")[0]
+ tokens = command.split()
+ if not tokens:
+ continue
+ if tokens[0] == "atom_modify" and any(
+ tokens[index : index + 2] == ["map", "yes"]
+ for index in range(1, len(tokens) - 1)
+ ):
+ atom_map_index = line_index
+ if read_index is None and re.search(r"\bread_(?:data|restart)\b", command):
+ read_index = line_index
+ if read_index is None:
+ raise ValueError("pt2 LAMMPS templates require read_data or read_restart.")
+ if atom_map_index is None or atom_map_index > read_index:
+ raise ValueError(
+ "pt2 LAMMPS templates require 'atom_modify map yes' before read_data or read_restart."
+ )
+
+
def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems):
model_devi_jobs = jdata["model_devi_jobs"]
if iter_index >= len(model_devi_jobs):
@@ -1558,6 +1941,8 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems):
# revise input of lammps
with open("input.lammps") as fp:
lmp_lines = fp.readlines()
+ if suffix == ".pt2":
+ _validate_pt2_template_atom_map(lmp_lines)
# only revise the line "pair_style deepmd" if the user has not written the full line (checked by then length of the line)
template_has_pair_deepmd = 1
for line_idx, line_context in enumerate(lmp_lines):
@@ -1776,7 +2161,7 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
trj_freq,
mass_map,
tt,
- jdata=jdata,
+ jdata={**jdata, "model_format": suffix[1:]},
tau_t=model_devi_taut,
pres=pp,
tau_p=model_devi_taup,
@@ -2281,7 +2666,12 @@ def run_model_devi(iter_index, jdata, mdata):
if model_devi_engine != "calypso":
run_md_model_devi(iter_index, jdata, mdata)
else:
- run_calypso_model_devi(iter_index, jdata, mdata)
+ run_calypso_model_devi(
+ iter_index,
+ jdata,
+ mdata,
+ model_suffix=_get_model_suffix(jdata),
+ )
def post_model_devi(iter_index, jdata, mdata):
diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py
new file mode 100644
index 000000000..47daac497
--- /dev/null
+++ b/tests/generator/test_deepmd_backend.py
@@ -0,0 +1,366 @@
+import json
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from dpgen.generator.lib.run_calypso import (
+ _find_models,
+ _make_calypso_opt_command,
+)
+from dpgen.generator.run import (
+ _get_checkpoint_suffix,
+ _get_input_model_suffix,
+ _get_model_suffix,
+ _get_train_backend_flag,
+ _validate_dpa_training_config,
+ _validate_pt2_template_atom_map,
+ post_train_dp,
+ run_md_model_devi,
+ run_train_dp,
+)
+
+
+class TestDeepmdBackendConfig(unittest.TestCase):
+ def test_legacy_defaults(self):
+ cases = [
+ ({}, ".pb", ".index", ""),
+ ({"train_backend": "pytorch"}, ".pth", ".pt", "--pt"),
+ ({"train_backend": "jax"}, ".savedmodel", ".jax", "--jax"),
+ ]
+ for jdata, model_suffix, checkpoint_suffix, backend_flag in cases:
+ with self.subTest(jdata=jdata):
+ self.assertEqual(_get_model_suffix(jdata), model_suffix)
+ self.assertEqual(_get_checkpoint_suffix(jdata), checkpoint_suffix)
+ self.assertEqual(_get_train_backend_flag(jdata), backend_flag)
+
+ def test_pytorch_exportable_aliases(self):
+ for backend in ("pytorch-exportable", "pt-expt"):
+ with self.subTest(backend=backend):
+ jdata = {"train_backend": backend}
+ self.assertEqual(_get_model_suffix(jdata), ".pt2")
+ self.assertEqual(_get_checkpoint_suffix(jdata), ".pt")
+ self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt")
+
+ def test_explicit_pt2_formats(self):
+ cases = [
+ {"train_backend": "pytorch", "model_format": "pt2"},
+ {"train_backend": "pytorch-exportable", "model_format": "pt2"},
+ {"train_backend": "pt-expt", "model_format": "pt2"},
+ ]
+ for jdata in cases:
+ with self.subTest(jdata=jdata):
+ self.assertEqual(_get_model_suffix(jdata), ".pt2")
+ self.assertEqual(_get_checkpoint_suffix(jdata), ".pt")
+
+ def test_pte_is_rejected_for_lammps(self):
+ with self.assertRaisesRegex(ValueError, "not supported by LAMMPS"):
+ _get_model_suffix(
+ {"train_backend": "pytorch-exportable", "model_format": "pte"}
+ )
+ self.assertEqual(
+ _get_model_suffix(
+ {"train_backend": "pt-expt", "model_devi_engine": "calypso"}
+ ),
+ ".pte",
+ )
+
+ def test_rejects_incompatible_model_format(self):
+ with self.assertRaisesRegex(ValueError, "not available for backend"):
+ _get_model_suffix({"train_backend": "tensorflow", "model_format": "pt2"})
+
+ def test_input_models_must_share_suffix(self):
+ self.assertEqual(_get_input_model_suffix(["a.pt", "b.pt"]), ".pt")
+ with self.assertRaisesRegex(ValueError, "same non-empty file suffix"):
+ _get_input_model_suffix(["a.pte", "b.pt2"])
+
+ def test_pt2_template_requires_atom_map_before_read(self):
+ _validate_pt2_template_atom_map(
+ ["atom_modify map yes\n", "read_data conf.lmp\n"]
+ )
+ for lines in (
+ ["read_data conf.lmp\n"],
+ ["read_restart restart.100\n", "atom_modify map yes\n"],
+ [
+ 'if "${restart} > 0" then "read_restart restart.100" '
+ 'else "read_data conf.lmp"\n',
+ "atom_modify map yes\n",
+ ],
+ ):
+ with self.subTest(lines=lines):
+ with self.assertRaisesRegex(ValueError, "atom_modify map yes"):
+ _validate_pt2_template_atom_map(lines)
+ with self.assertRaisesRegex(ValueError, "read_data or read_restart"):
+ _validate_pt2_template_atom_map(["atom_modify map yes\n"])
+
+ def test_dpa_backend_and_compile_option_validation(self):
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pytorch",
+ "model_format": "pt2",
+ "default_training_param": {
+ "model": {
+ "type": "DPA4",
+ "use_compile": True,
+ "enable_tf32": True,
+ },
+ "training": {},
+ },
+ }
+ )
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pytorch",
+ "model_format": "pt2",
+ "default_training_param": {"model": {"type": "SeZM"}},
+ }
+ )
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pt-expt",
+ "model_format": "pt2",
+ "default_training_param": {
+ "model": {"descriptor": {"type": "DPA4C"}},
+ "training": {
+ "enable_compile": True,
+ "enable_tf32": True,
+ },
+ },
+ }
+ )
+ with self.assertRaisesRegex(ValueError, "requires train_backend='pytorch'"):
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pytorch-exportable",
+ "model_format": "pt2",
+ "default_training_param": {
+ "model": {"descriptor": {"type": "dpa4"}},
+ "training": {},
+ },
+ }
+ )
+ with self.assertRaisesRegex(ValueError, "training.enable_compile"):
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pytorch-exportable",
+ "model_format": "pt2",
+ "default_training_param": {
+ "model": {
+ "descriptor": {"type": "dpa4c"},
+ "use_compile": True,
+ },
+ "training": {},
+ },
+ }
+ )
+ with self.assertRaisesRegex(ValueError, "cannot mix DPA4 and DPA4C"):
+ _validate_dpa_training_config(
+ {
+ "train_backend": "pytorch-exportable",
+ "model_format": "pt2",
+ "default_training_param": {
+ "model": {
+ "model_dict": {
+ "dpa4": {"descriptor": {"type": "dpa4"}},
+ "dpa4c": {"descriptor": {"type": "dpa4c"}},
+ }
+ },
+ "training": {},
+ },
+ }
+ )
+ with self.assertRaisesRegex(ValueError, "only exports pt2 for DPA4/SeZM"):
+ _validate_dpa_training_config(
+ {"train_backend": "pytorch", "model_format": "pt2"}
+ )
+
+ def test_calypso_discovers_resolved_model_suffix(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ path = Path(tmpdir)
+ for name in ("graph.000.pb", "graph.000.pte", "graph.001.pte"):
+ (path / name).touch()
+ self.assertEqual(len(_find_models(path, ".pb")), 1)
+ self.assertEqual(len(_find_models(path, ".pte")), 2)
+
+ def test_calypso_optimizer_uses_resolved_model(self):
+ command = _make_calypso_opt_command("python", "graph.000.pt2")
+ self.assertIn("calypso_run_opt.py --model ../graph.000.pt2", command)
+
+
+class TestRunTrainDeepmdBackend(unittest.TestCase):
+ def setUp(self):
+ self.tmp = tempfile.TemporaryDirectory()
+ self.addCleanup(self.tmp.cleanup)
+ self.old_cwd = os.getcwd()
+ os.chdir(self.tmp.name)
+ self.addCleanup(os.chdir, self.old_cwd)
+ self.mdata = {
+ "api_version": "1.0",
+ "deepmd_version": "3.2.0",
+ "train_command": "dp",
+ "train_machine": {"name": "train"},
+ "train_resources": {"queue": "train"},
+ "model_devi_machine": {"name": "model-devi"},
+ "model_devi_resources": {"queue": "model-devi"},
+ }
+
+ def _run(self, **updates):
+ jdata = {"numb_models": 1, "one_h5": True}
+ jdata.update(updates)
+ with patch("dpgen.generator.run.make_submission") as make_submission:
+ run_train_dp(0, jdata, self.mdata)
+ calls = []
+ for call in make_submission.call_args_list:
+ details = dict(call.kwargs)
+ details["machine"], details["resources"] = call.args[:2]
+ calls.append(details)
+ return calls[0] if len(calls) == 1 else calls
+
+ def test_legacy_tensorflow_commands_are_preserved(self):
+ call = self._run()
+ self.assertIn("model.ckpt.index", call["commands"][0])
+ self.assertEqual(call["commands"][1], "dp freeze")
+ self.assertIn("frozen_model.pb", call["backward_files"])
+ self.assertIn("model.ckpt.index", call["backward_files"])
+
+ def test_pytorch_dpa4_exports_pt2_on_model_devi_resources(self):
+ train_call, export_call = self._run(
+ train_backend="pytorch",
+ model_format="pt2",
+ default_training_param={"model": {"type": "dpa4"}},
+ )
+ self.assertEqual(train_call["machine"], self.mdata["train_machine"])
+ self.assertEqual(len(train_call["commands"]), 1)
+ self.assertIn("dp --pt train", train_call["commands"][0])
+ self.assertIn("model.ckpt.pt", train_call["backward_files"])
+ self.assertNotIn("frozen_model.pt2", train_call["backward_files"])
+ self.assertEqual(export_call["machine"], self.mdata["model_devi_machine"])
+ self.assertEqual(export_call["resources"], self.mdata["model_devi_resources"])
+ self.assertEqual(
+ export_call["commands"],
+ ["dp --pt freeze -c model.ckpt.pt -o frozen_model"],
+ )
+ self.assertEqual(export_call["forward_files"], ["model.ckpt.pt"])
+ self.assertIn("frozen_model.pt2", export_call["backward_files"])
+
+ def test_legacy_pytorch_commands_are_preserved(self):
+ call = self._run(train_backend="pytorch")
+ self.assertIn("dp --pt train", call["commands"][0])
+ self.assertIn("model.ckpt.pt", call["commands"][0])
+ self.assertEqual(call["commands"][1], "dp --pt freeze")
+ self.assertIn("frozen_model.pth", call["backward_files"])
+ self.assertIn("model.ckpt.pt", call["backward_files"])
+
+ def test_pytorch_exportable_dense_pte(self):
+ call = self._run(
+ train_backend="pytorch-exportable", model_devi_engine="calypso"
+ )
+ self.assertIn("dp --pt-expt train", call["commands"][0])
+ self.assertEqual(
+ call["commands"][1],
+ "dp --pt-expt freeze -o frozen_model.pte",
+ )
+ self.assertIn("frozen_model.pte", call["backward_files"])
+
+ def test_pytorch_exportable_dpa4c_pt2_compression(self):
+ train_call, export_call = self._run(
+ train_backend="pt-expt",
+ model_format="pt2",
+ dp_compress=True,
+ )
+ self.assertIn("dp --pt-expt train", train_call["commands"][0])
+ self.assertEqual(len(train_call["commands"]), 1)
+ self.assertEqual(
+ export_call["commands"],
+ [
+ "dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph",
+ "dp --pt-expt compress -i frozen_model.pt2 -o frozen_model_compressed.pt2",
+ ],
+ )
+ self.assertEqual(export_call["forward_files"], ["model.ckpt.pt"])
+ self.assertIn("frozen_model_compressed.pt2", export_call["backward_files"])
+
+ def test_multiple_pt2_models_are_exported(self):
+ train_call, export_call = self._run(
+ numb_models=4,
+ train_backend="pt-expt",
+ model_format="pt2",
+ )
+ expected_tasks = [f"{index:03d}" for index in range(4)]
+ self.assertEqual(train_call["run_tasks"], expected_tasks)
+ self.assertEqual(export_call["run_tasks"], expected_tasks)
+
+ def test_regular_pytorch_pt2_compression_is_rejected(self):
+ with self.assertRaisesRegex(RuntimeError, "cannot compress pt2"):
+ self._run(
+ train_backend="pytorch",
+ model_format="pt2",
+ dp_compress=True,
+ default_training_param={"model": {"type": "dpa4"}},
+ )
+
+ def test_exportable_init_frozen_model_is_rejected(self):
+ with self.assertRaisesRegex(RuntimeError, "does not support"):
+ self._run(
+ train_backend="pt-expt",
+ training_init_frozen_model=["model.pt2"],
+ )
+
+ def test_deepmd_31_is_rejected_for_pt2(self):
+ self.mdata["deepmd_version"] = "3.1.0"
+ with self.assertRaisesRegex(RuntimeError, "3.2 or later"):
+ self._run(
+ train_backend="pytorch",
+ model_format="pt2",
+ default_training_param={"model": {"type": "dpa4"}},
+ )
+
+ def test_finetune_keeps_source_model_suffix(self):
+ train_call, _ = self._run(
+ train_backend="pt-expt",
+ model_format="pt2",
+ training_finetune_model=["source.pt"],
+ )
+ self.assertIn("--finetune old/init.pt", train_call["commands"][0])
+ self.assertIn(str(Path("old") / "init.pt"), train_call["forward_files"])
+
+ def test_post_train_links_pt2_model(self):
+ jdata = {
+ "numb_models": 1,
+ "train_backend": "pt-expt",
+ "model_format": "pt2",
+ }
+ with patch("dpgen.generator.run.os.symlink") as symlink:
+ post_train_dp(0, jdata, self.mdata)
+ symlink.assert_called_once_with(
+ str(Path("000") / "frozen_model.pt2"),
+ str(Path("iter.000000") / "00.train" / "graph.000.pt2"),
+ )
+
+ def test_model_deviation_forwards_pt2_models(self):
+ work_path = Path("iter.000000") / "01.model_devi"
+ (work_path / "task.000.000000").mkdir(parents=True)
+ (work_path / "graph.000.pt2").touch()
+ (work_path / "cur_job.json").write_text(json.dumps({}), encoding="utf-8")
+ jdata = {
+ "train_backend": "pytorch",
+ "model_format": "pt2",
+ "model_devi_jobs": [{}],
+ }
+ mdata = {
+ "api_version": "1.0",
+ "model_devi_command": "lmp -k on g 1 -sf kk",
+ "model_devi_group_size": 1,
+ "model_devi_machine": {},
+ "model_devi_resources": {},
+ }
+ with patch("dpgen.generator.run.make_submission") as make_submission:
+ run_md_model_devi(0, jdata, mdata)
+ call = make_submission.call_args.kwargs
+ self.assertEqual(call["forward_common_files"], ["graph.000.pt2"])
+ self.assertIn("lmp -k on g 1 -sf kk", call["commands"][0])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/generator/test_lammps.py b/tests/generator/test_lammps.py
index c1b915939..6055d5a1c 100644
--- a/tests/generator/test_lammps.py
+++ b/tests/generator/test_lammps.py
@@ -55,6 +55,65 @@ def test_basic_deepmd_only(self):
# Should NOT contain hybrid/overlay or dispersion/d3
self.assertNotIn("hybrid/overlay", result)
self.assertNotIn("dispersion/d3", result)
+ self.assertNotIn("atom_modify map yes", result)
+
+ def test_pt2_enables_atom_map_before_read(self):
+ """Test that pt2 models enable one atom map before defining the box."""
+ result = make_lammps_input(
+ self.ensemble,
+ self.conf_file,
+ self.graphs,
+ self.nsteps,
+ self.dt,
+ self.neidelay,
+ self.trj_freq,
+ self.mass_map,
+ self.temp,
+ {"model_format": "pt2"},
+ pres=1.0,
+ deepmd_version=self.deepmd_version,
+ )
+
+ atom_map = "atom_modify map yes"
+ self.assertEqual(result.count(atom_map), 1)
+ self.assertLess(result.index(atom_map), result.index("read_restart"))
+ self.assertLess(result.index(atom_map), result.index("read_data"))
+
+ pimd_result = make_lammps_input(
+ self.ensemble,
+ self.conf_file,
+ self.graphs,
+ self.nsteps,
+ self.dt,
+ self.neidelay,
+ self.trj_freq,
+ self.mass_map,
+ self.temp,
+ {"model_format": "pt2"},
+ pres=1.0,
+ deepmd_version=self.deepmd_version,
+ nbeads=4,
+ )
+ self.assertEqual(pimd_result.count(atom_map), 1)
+
+ def test_pt2_uses_explicit_type_mapping(self):
+ """Test that DP-GEN's LAMMPS type order is passed to DeepMD."""
+ result = make_lammps_input(
+ self.ensemble,
+ self.conf_file,
+ self.graphs,
+ self.nsteps,
+ self.dt,
+ self.neidelay,
+ self.trj_freq,
+ self.mass_map,
+ self.temp,
+ {"model_format": "pt2", "type_map": ["C", "Cl", "H", "O"]},
+ pres=1.0,
+ deepmd_version=self.deepmd_version,
+ )
+
+ self.assertIn("pair_coeff * * C Cl H O\n", result)
def test_d3_enabled_basic(self):
"""Test LAMMPS input with D3 dispersion enabled."""
@@ -65,7 +124,8 @@ def test_d3_enabled_basic(self):
"functional": "pbe",
"cutoff": 30.0,
"cn_cutoff": 20.0,
- }
+ },
+ "type_map": ["H", "O"],
}
result = make_lammps_input(
@@ -87,6 +147,9 @@ def test_d3_enabled_basic(self):
self.assertIn("pair_style hybrid/overlay deepmd model.pb", result)
self.assertIn("dispersion/d3 original pbe 30.0 20.0", result)
+ self.assertIn("pair_coeff * * deepmd H O\n", result)
+ self.assertIn("pair_coeff * * dispersion/d3 H O\n", result)
+
# Should contain both pair_coeff lines
lines = result.split("\n")
deepmd_coeff_found = False
diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py
index ab36225cc..66dbf03e1 100644
--- a/tests/generator/test_make_md.py
+++ b/tests/generator/test_make_md.py
@@ -6,11 +6,16 @@
import shutil
import sys
import unittest
+from unittest.mock import patch
import dpdata
import numpy as np
-from dpgen.generator.run import _read_model_devi_file, parse_cur_job_sys_revmat
+from dpgen.generator.run import (
+ _read_model_devi_file,
+ parse_cur_job_sys_revmat,
+ revise_lmp_input_pair_coeff,
+)
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
__package__ = "generator"
@@ -162,6 +167,48 @@ def test_make_model_devi(self):
_check_pt(self, 0, jdata)
# shutil.rmtree('iter.000000')
+ def test_make_model_devi_default_pt2_atom_map(self):
+ test_dir = os.path.dirname(__file__)
+ with open(os.path.join(test_dir, param_file)) as fp:
+ jdata = json.load(fp)
+ with open(os.path.join(test_dir, machine_file)) as fp:
+ mdata = json.load(fp)
+ jdata["train_backend"] = "pytorch-exportable"
+ jdata["sys_configs_prefix"] = test_dir
+ jdata["sys_configs"] = [
+ ["data/al.fcc.02x02x02/01.scale_pert/sys-0032/scale*/000001/POSCAR"]
+ ]
+ jdata["model_devi_jobs"][0]["sys_idx"] = [0]
+ jdata["model_devi_jobs"][0]["temps"] = [50]
+ jdata["model_devi_jobs"][0]["press"] = [1.0]
+ jdata.pop("model_format", None)
+
+ train_dir = os.path.join("iter.000000", "00.train")
+ os.makedirs(train_dir, exist_ok=True)
+ for model_index in range(jdata["numb_models"]):
+ with open(
+ os.path.join(train_dir, f"graph.{model_index:03d}.pt2"), "w"
+ ) as fp:
+ fp.write("model")
+
+ def copy_link(source, target):
+ if not os.path.isabs(source):
+ source = os.path.join(os.path.dirname(target), source)
+ shutil.copyfile(os.path.normpath(source), target)
+
+ with patch("dpgen.generator.run.os.symlink", side_effect=copy_link):
+ make_model_devi(0, jdata, mdata)
+
+ task = sorted(glob.glob("iter.000000/01.model_devi/task.*"))[0]
+ with open(os.path.join(task, "input.lammps")) as fp:
+ lammps_input = fp.read()
+ self.assertEqual(lammps_input.count("atom_modify map yes"), 1)
+ self.assertLess(
+ lammps_input.index("atom_modify map yes"),
+ lammps_input.index("read_data"),
+ )
+ self.assertIn("pair_coeff * * Mg Al\n", lammps_input)
+
def test_make_model_devi_pimd(self):
if os.path.isdir("iter.000000"):
shutil.rmtree("iter.000000")
@@ -515,6 +562,67 @@ def test_make_model_devi_null(self):
)
os.chdir(cwd_)
+ def test_default_pt2_template_with_atom_map(self):
+ test_dir = os.path.dirname(__file__)
+ with open(os.path.join(test_dir, "lmp", "input.lammps")) as fp:
+ template = fp.read()
+ template = template.replace(
+ "read_data conf.lmp",
+ "atom_modify map yes\nread_data conf.lmp",
+ )
+ with tempfile.NamedTemporaryFile(
+ mode="w", suffix=".lammps", dir=".", delete=False
+ ) as fp:
+ fp.write(template)
+ template_path = os.path.abspath(fp.name)
+ self.addCleanup(os.remove, template_path)
+
+ jdata = {
+ "type_map": ["Mg", "Al"],
+ "mass_map": [24, 27],
+ "init_data_prefix": "data",
+ "init_data_sys": ["deepmd"],
+ "init_batch_size": [16],
+ "sys_configs_prefix": test_dir,
+ "sys_configs": [
+ ["data/al.fcc.02x02x02/01.scale_pert/sys-0032/scale*/000001/POSCAR"]
+ ],
+ "numb_models": 1,
+ "shuffle_poscar": False,
+ "model_devi_f_trust_lo": 0.050,
+ "model_devi_f_trust_hi": 0.150,
+ "train_backend": "pytorch-exportable",
+ "model_devi_jobs": [
+ {
+ "sys_idx": [0],
+ "traj_freq": 10,
+ "template": {"lmp": template_path},
+ }
+ ],
+ }
+ train_path = os.path.join("iter.000000", "00.train")
+ os.makedirs(train_path, exist_ok=True)
+ with open(os.path.join(train_path, "graph.000.pt2"), "w") as fp:
+ fp.write("model")
+
+ def copy_link(source, target):
+ if not os.path.isabs(source):
+ source = os.path.join(os.path.dirname(target), source)
+ shutil.copyfile(os.path.normpath(source), target)
+
+ with patch("dpgen.generator.run.os.symlink", side_effect=copy_link):
+ make_model_devi(0, jdata, {"deepmd_version": "3.2"})
+ task = sorted(glob.glob("iter.000000/01.model_devi/task.*"))[0]
+ with open(os.path.join(task, "input.lammps")) as fp:
+ lines = fp.readlines()
+ atom_map_index = next(
+ index for index, line in enumerate(lines) if "atom_modify" in line
+ )
+ read_data_index = next(
+ index for index, line in enumerate(lines) if "read_data" in line
+ )
+ self.assertLess(atom_map_index, read_data_index)
+
class TestParseCurJobRevMat(unittest.TestCase):
def setUp(self):
@@ -580,6 +688,64 @@ def test_parse_cur_job(self):
class MakeModelDeviByReviseMatrix(unittest.TestCase):
+ def test_revise_lmp_input_pair_coeff_adds_type_map(self):
+ jdata = {"type_map": ["C", "Cl", "H", "O"]}
+ cases = (
+ ("pair_coeff\n", "pair_coeff * * C Cl H O\n"),
+ ("pair_coeff * *\n", "pair_coeff * * C Cl H O\n"),
+ (
+ "pair_coeff * * deepmd\n",
+ "pair_coeff * * deepmd C Cl H O\n",
+ ),
+ )
+
+ for pair_coeff, expected in cases:
+ with self.subTest(pair_coeff=pair_coeff):
+ lines = ["pair_style deepmd graph.pb\n", pair_coeff]
+ result = revise_lmp_input_pair_coeff(lines, jdata)
+ self.assertEqual(result[1], expected)
+
+ def test_revise_lmp_input_pair_coeff_preserves_explicit_mapping(self):
+ jdata = {"type_map": ["C", "Cl", "H", "O"]}
+ cases = (
+ "pair_coeff * * H O\n",
+ "pair_coeff * * deepmd H O\n",
+ )
+
+ for pair_coeff in cases:
+ with self.subTest(pair_coeff=pair_coeff):
+ lines = ["pair_style deepmd graph.pb\n", pair_coeff]
+ result = revise_lmp_input_pair_coeff(lines, jdata)
+ self.assertEqual(result[1], pair_coeff)
+
+ def test_revise_lmp_input_pair_coeff_selects_deepmd_in_hybrid(self):
+ jdata = {"type_map": ["C", "Cl", "H", "O"]}
+ lines = [
+ "pair_style hybrid/overlay deepmd graph.pb zero 10.0\n",
+ "pair_coeff * * zero 10.0\n",
+ "pair_coeff * * deepmd\n",
+ ]
+
+ result = revise_lmp_input_pair_coeff(lines, jdata)
+
+ self.assertEqual(result[1], "pair_coeff * * zero 10.0\n")
+ self.assertEqual(result[2], "pair_coeff * * deepmd C Cl H O\n")
+
+ def test_revise_lmp_input_pair_coeff_d3_is_idempotent(self):
+ jdata = {
+ "type_map": ["C", "Cl", "H", "O"],
+ "lmp_d3": {"enable": True},
+ }
+ lines = ["pair_style deepmd graph.pb\n", "pair_coeff * *\n"]
+
+ result = revise_lmp_input_pair_coeff(lines, jdata)
+ result = revise_lmp_input_pair_coeff(result, jdata)
+
+ self.assertEqual(result.count("pair_coeff * * deepmd C Cl H O\n"), 1)
+ self.assertEqual(
+ result.count("pair_coeff * * dispersion/d3 C Cl H O\n"), 1
+ )
+
def test_find_only_one_key_1(self):
lines = ["aaa bbb ccc\n", "bbb ccc\n", "ccc bbb ccc\n"]
idx = find_only_one_key(lines, ["bbb", "ccc"])