From 1591ecf1997a7a6a7e101e21d9ff91395c6e1064 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Thu, 20 Aug 2026 23:42:12 +0800 Subject: [PATCH 01/13] feat: support DPA4 and DPA4C model formats --- doc/run/param.rst | 28 ++++ dpgen/generator/arginfo.py | 24 +++- dpgen/generator/run.py | 182 +++++++++++++++++++------ tests/generator/test_deepmd_backend.py | 161 ++++++++++++++++++++++ 4 files changed, 353 insertions(+), 42 deletions(-) create mode 100644 tests/generator/test_deepmd_backend.py diff --git a/doc/run/param.rst b/doc/run/param.rst index 74772253e..8aa06a7e0 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -8,3 +8,31 @@ 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. A +regular PyTorch DPA4 model can be exported as an AOTInductor archive with: + +.. code-block:: json + + { + "train_backend": "pytorch", + "model_format": "pt2" + } + +For DPA4C, and for DPA4 using the PyTorch-exportable backend, select the graph +export explicitly: + +.. code-block:: json + + { + "train_backend": "pytorch-exportable", + "model_format": "pt2", + "dp_compress": true + } + +``pt-expt`` is accepted as an alias of ``pytorch-exportable``. The ``pte`` +format remains available for dense PyTorch-exportable models. Training +checkpoints keep the ``.pt`` suffix independently of the frozen model format. diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 0f8abb865..58f8b6ddf 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -94,8 +94,22 @@ 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 depend on ``train_backend``: + ``pb`` for TensorFlow, ``pth`` for PyTorch, ``pte`` for + PyTorch-exportable, and ``savedmodel`` for JAX. PyTorch also supports + ``pt2`` for DPA4. PyTorch-exportable supports ``pte`` and ``pt2``; + use ``pt2`` for the graph export used by DPA4 and DPA4C. + """ ) 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." @@ -139,6 +153,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/run.py b/dpgen/generator/run.py index 0ce513a96..d41be09d0 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -125,22 +125,85 @@ 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_config(jdata) -> tuple[str, dict, str]: + """Return and validate the DeePMD backend and deployment format.""" 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: - raise ValueError( - f"The backend {backend} is not available. Supported backends are: 'tensorflow', 'pytorch', 'jax'." - ) - return suffix - else: + if mlp_engine != "dp": raise ValueError(f"Unsupported engine: {mlp_engine}") + backend = jdata.get("train_backend", "tensorflow") + 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}'." + ) + + config = _BACKEND_CONFIG[backend] + model_format = jdata.get("model_format", config["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}'." + ) + return backend, config, model_format + + +def _get_model_suffix(jdata) -> str: + """Return the frozen model suffix.""" + _, _, model_format = _get_backend_config(jdata) + return f".{model_format}" + + +def _get_checkpoint_suffix(jdata) -> str: + """Return the training checkpoint suffix.""" + _, config, _ = _get_backend_config(jdata) + return config["checkpoint_suffix"] + + +def _get_train_backend_flag(jdata) -> str: + """Return the DeePMD CLI backend flag.""" + _, config, _ = _get_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 get_job_names(jdata): jobkeys = [] @@ -667,9 +730,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) @@ -730,7 +797,9 @@ def run_train_dp(iter_index, jdata, mdata): # print("debug:run_train:mdata", mdata) # load json param numb_models = jdata["numb_models"] + backend, _, model_format = _get_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 +821,27 @@ def run_train_dp(iter_index, jdata, mdata): except KeyError: mdata = set_version(mdata) + if (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 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 ( + 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 +854,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) @@ -797,25 +886,32 @@ 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" + if backend == "pytorch-exportable": + command = f"{train_command} freeze -o frozen_model{suffix}" + if model_format == "pt2": + command += " --lower-kind graph" + elif backend == "pytorch" and model_format == "pt2": + command = f"{train_command} freeze -o frozen_model{suffix}" + else: + command = f"{train_command} freeze" commands.append(command) if jdata.get("dp_compress", False): - commands.append(f"{train_command} compress") + if 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,20 +932,26 @@ 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}", @@ -860,18 +962,18 @@ def run_train_dp(iter_index, jdata, mdata): 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"] diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py new file mode 100644 index 000000000..eb31e4a40 --- /dev/null +++ b/tests/generator/test_deepmd_backend.py @@ -0,0 +1,161 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from dpgen.generator.run import ( + _get_checkpoint_suffix, + _get_input_model_suffix, + _get_model_suffix, + _get_train_backend_flag, + post_train_dp, + 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), ".pte") + self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") + self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt") + + def test_explicit_pt2_formats(self): + for backend in ("pytorch", "pytorch-exportable", "pt-expt"): + with self.subTest(backend=backend): + jdata = {"train_backend": backend, "model_format": "pt2"} + self.assertEqual(_get_model_suffix(jdata), ".pt2") + self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") + + 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"]) + + +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": {}, + "train_resources": {}, + } + + 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) + return make_submission.call_args.kwargs + + 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_regular_pytorch_dpa4_pt2(self): + call = self._run(train_backend="pytorch", model_format="pt2") + self.assertIn("dp --pt train", call["commands"][0]) + self.assertIn("model.ckpt.pt", call["commands"][0]) + self.assertEqual(call["commands"][1], "dp --pt freeze -o frozen_model.pt2") + self.assertIn("frozen_model.pt2", call["backward_files"]) + self.assertIn("model.ckpt.pt", 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") + 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): + call = self._run(train_backend="pt-expt", model_format="pt2", dp_compress=True) + self.assertEqual( + call["commands"][1], + "dp --pt-expt freeze -o frozen_model.pt2 --lower-kind graph", + ) + self.assertEqual( + call["commands"][2], + "dp --pt-expt compress -i frozen_model.pt2 -o frozen_model_compressed.pt2", + ) + self.assertIn("frozen_model_compressed.pt2", call["backward_files"]) + + 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) + + 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") + + def test_finetune_keeps_source_model_suffix(self): + call = self._run( + train_backend="pt-expt", + model_format="pt2", + training_finetune_model=["source.pt"], + ) + self.assertIn("--finetune old/init.pt", call["commands"][0]) + self.assertIn(str(Path("old") / "init.pt"), 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"), + ) + + +if __name__ == "__main__": + unittest.main() From ff667d3f971bac84801cd303ea06a52a75a2cc4f Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 21 Aug 2026 00:01:50 +0800 Subject: [PATCH 02/13] feat: separate training and deployment backends --- doc/run/param.rst | 7 +++ dpgen/generator/arginfo.py | 18 ++++++- dpgen/generator/run.py | 72 +++++++++++++++++-------- tests/generator/test_deepmd_backend.py | 73 +++++++++++++++++++++++--- 4 files changed, 142 insertions(+), 28 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index 8aa06a7e0..486329138 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -19,6 +19,7 @@ regular PyTorch DPA4 model can be exported as an AOTInductor archive with: { "train_backend": "pytorch", + "model_devi_backend": "pytorch-exportable", "model_format": "pt2" } @@ -36,3 +37,9 @@ export explicitly: ``pt-expt`` is accepted as an alias of ``pytorch-exportable``. The ``pte`` format remains available for dense PyTorch-exportable models. Training checkpoints keep the ``.pt`` suffix independently of the frozen model format. +The ``model_devi_backend`` setting makes ``dpgen`` train with ``dp --pt`` but +freeze (and optionally compress) with ``dp --pt-expt``. The resulting ``.pt2`` +models are automatically linked into the model-deviation stage and listed in +the generated LAMMPS input. For Kokkos execution, configure the model-deviation +command to invoke a Kokkos-enabled LAMMPS build, for example +``lmp -k on g 1 -sf kk``. diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 58f8b6ddf..851ad8624 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -104,13 +104,23 @@ def training_args_dp() -> list[Argument]: ) doc_model_format = textwrap.dedent( """\ - The frozen model format. Defaults depend on ``train_backend``: + The frozen model format. Defaults depend on ``model_devi_backend`` (or + ``train_backend`` when no deployment backend is set): ``pb`` for TensorFlow, ``pth`` for PyTorch, ``pte`` for PyTorch-exportable, and ``savedmodel`` for JAX. PyTorch also supports ``pt2`` for DPA4. PyTorch-exportable supports ``pte`` and ``pt2``; use ``pt2`` for the graph export used by DPA4 and DPA4C. """ ) + doc_model_devi_backend = textwrap.dedent( + """\ + The DeePMD backend used to freeze and compress models for model + deviation. It defaults to ``train_backend``. A model trained with + ``pytorch`` can be deployed with ``pytorch-exportable`` to preserve its + ``.pt`` training checkpoint while producing a graph-lowered ``.pt2`` + model for LAMMPS/Kokkos. Other cross-backend exports are 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 `_." @@ -153,6 +163,12 @@ def training_args_dp() -> list[Argument]: default="tensorflow", doc=doc_train_backend, ), + Argument( + "model_devi_backend", + str, + optional=True, + doc=doc_model_devi_backend, + ), Argument( "model_format", str, diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index d41be09d0..1044d0b22 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -154,21 +154,38 @@ } -def _get_backend_config(jdata) -> tuple[str, dict, str]: - """Return and validate the DeePMD backend and deployment format.""" +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": raise ValueError(f"Unsupported engine: {mlp_engine}") - backend = jdata.get("train_backend", "tensorflow") + 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] - config = _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 and validate the deployment backend and model format.""" + train_backend, _ = _get_train_backend_config(jdata) + backend, config = _get_backend(jdata, "model_devi_backend", train_backend) + if backend != train_backend and not ( + train_backend == "pytorch" and backend == "pytorch-exportable" + ): + raise ValueError( + f"Cannot export models trained with backend {train_backend} using " + f"backend {backend}." + ) model_format = jdata.get("model_format", config["default_model_format"]) if model_format not in config["model_formats"]: supported = "', '".join(sorted(config["model_formats"])) @@ -181,19 +198,25 @@ def _get_backend_config(jdata) -> tuple[str, dict, str]: def _get_model_suffix(jdata) -> str: """Return the frozen model suffix.""" - _, _, model_format = _get_backend_config(jdata) + _, _, model_format = _get_model_backend_config(jdata) return f".{model_format}" def _get_checkpoint_suffix(jdata) -> str: """Return the training checkpoint suffix.""" - _, config, _ = _get_backend_config(jdata) + _, 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_backend_config(jdata) + _, config = _get_train_backend_config(jdata) + return config["flag"] + + +def _get_model_backend_flag(jdata) -> str: + """Return the DeePMD CLI deployment backend flag.""" + _, config, _ = _get_model_backend_config(jdata) return config["flag"] @@ -797,7 +820,8 @@ def run_train_dp(iter_index, jdata, mdata): # print("debug:run_train:mdata", mdata) # load json param numb_models = jdata["numb_models"] - backend, _, model_format = _get_backend_config(jdata) + train_backend, _ = _get_train_backend_config(jdata) + model_backend, _, model_format = _get_model_backend_config(jdata) suffix = _get_model_suffix(jdata) checkpoint_suffix = _get_checkpoint_suffix(jdata) # train_param = jdata['train_param'] @@ -821,19 +845,21 @@ def run_train_dp(iter_index, jdata, mdata): except KeyError: mdata = set_version(mdata) - if (backend == "pytorch-exportable" or model_format == "pt2") and Version( - mdata["deepmd_version"] - ) < Version("3.2"): + if ( + train_backend == "pytorch-exportable" + or model_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 backend == "pytorch-exportable" and training_init_frozen_model is not None: + 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 ( - backend == "pytorch" + model_backend == "pytorch" and model_format == "pt2" and jdata.get("dp_compress", False) ): @@ -857,6 +883,10 @@ def run_train_dp(iter_index, jdata, mdata): backend_flag = _get_train_backend_flag(jdata) if backend_flag: train_command += f" {backend_flag}" + model_command = mdata.get("train_command", "dp").strip() + model_backend_flag = _get_model_backend_flag(jdata) + if model_backend_flag: + model_command += f" {model_backend_flag}" # paths iter_name = make_iter_name(iter_index) @@ -895,23 +925,23 @@ def run_train_dp(iter_index, jdata, mdata): 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) - if backend == "pytorch-exportable": - command = f"{train_command} freeze -o frozen_model{suffix}" + if model_backend == "pytorch-exportable": + command = f"{model_command} freeze -o frozen_model{suffix}" if model_format == "pt2": command += " --lower-kind graph" - elif backend == "pytorch" and model_format == "pt2": - command = f"{train_command} freeze -o frozen_model{suffix}" + elif model_backend == "pytorch" and model_format == "pt2": + command = f"{model_command} freeze -o frozen_model{suffix}" else: - command = f"{train_command} freeze" + command = f"{model_command} freeze" commands.append(command) if jdata.get("dp_compress", False): - if backend == "pytorch-exportable": + if model_backend == "pytorch-exportable": commands.append( - f"{train_command} compress -i frozen_model{suffix} " + f"{model_command} compress -i frozen_model{suffix} " f"-o frozen_model_compressed{suffix}" ) else: - commands.append(f"{train_command} compress") + commands.append(f"{model_command} compress") else: raise RuntimeError( "DP-GEN currently only supports for DeePMD-kit 1.x to 3.x version!" diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index eb31e4a40..98176b400 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -1,3 +1,4 @@ +import json import os import tempfile import unittest @@ -7,9 +8,11 @@ from dpgen.generator.run import ( _get_checkpoint_suffix, _get_input_model_suffix, + _get_model_backend_flag, _get_model_suffix, _get_train_backend_flag, post_train_dp, + run_md_model_devi, run_train_dp, ) @@ -36,12 +39,39 @@ def test_pytorch_exportable_aliases(self): self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt") def test_explicit_pt2_formats(self): - for backend in ("pytorch", "pytorch-exportable", "pt-expt"): - with self.subTest(backend=backend): - jdata = {"train_backend": backend, "model_format": "pt2"} + cases = [ + {"train_backend": "pytorch", "model_format": "pt2"}, + {"train_backend": "pytorch-exportable", "model_format": "pt2"}, + {"train_backend": "pt-expt", "model_format": "pt2"}, + { + "train_backend": "pytorch", + "model_devi_backend": "pytorch-exportable", + "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_pytorch_checkpoint_can_use_exportable_deployment(self): + jdata = { + "train_backend": "pytorch", + "model_devi_backend": "pt-expt", + "model_format": "pt2", + } + self.assertEqual(_get_train_backend_flag(jdata), "--pt") + self.assertEqual(_get_model_backend_flag(jdata), "--pt-expt") + + def test_rejects_other_cross_backend_exports(self): + with self.assertRaisesRegex(ValueError, "Cannot export models"): + _get_model_suffix( + { + "train_backend": "tensorflow", + "model_devi_backend": "pytorch-exportable", + } + ) + def test_rejects_incompatible_model_format(self): with self.assertRaisesRegex(ValueError, "not available for backend"): _get_model_suffix({"train_backend": "tensorflow", "model_format": "pt2"}) @@ -81,11 +111,18 @@ def test_legacy_tensorflow_commands_are_preserved(self): self.assertIn("frozen_model.pb", call["backward_files"]) self.assertIn("model.ckpt.index", call["backward_files"]) - def test_regular_pytorch_dpa4_pt2(self): - call = self._run(train_backend="pytorch", model_format="pt2") + def test_pytorch_dpa4_uses_exportable_pt2_deployment(self): + call = self._run( + train_backend="pytorch", + model_devi_backend="pytorch-exportable", + model_format="pt2", + ) self.assertIn("dp --pt train", call["commands"][0]) self.assertIn("model.ckpt.pt", call["commands"][0]) - self.assertEqual(call["commands"][1], "dp --pt freeze -o frozen_model.pt2") + self.assertEqual( + call["commands"][1], + "dp --pt-expt freeze -o frozen_model.pt2 --lower-kind graph", + ) self.assertIn("frozen_model.pt2", call["backward_files"]) self.assertIn("model.ckpt.pt", call["backward_files"]) @@ -156,6 +193,30 @@ def test_post_train_links_pt2_model(self): 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_devi_backend": "pytorch-exportable", + "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() From 4374ce7a90ab55dcc19e1a8332f88cd3b880dd15 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 21 Aug 2026 00:09:23 +0800 Subject: [PATCH 03/13] docs: focus deployment guidance on DP-GEN --- doc/run/param.rst | 4 +--- dpgen/generator/arginfo.py | 2 +- tests/generator/test_deepmd_backend.py | 3 +-- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index 486329138..de3166dd5 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -40,6 +40,4 @@ checkpoints keep the ``.pt`` suffix independently of the frozen model format. The ``model_devi_backend`` setting makes ``dpgen`` train with ``dp --pt`` but freeze (and optionally compress) with ``dp --pt-expt``. The resulting ``.pt2`` models are automatically linked into the model-deviation stage and listed in -the generated LAMMPS input. For Kokkos execution, configure the model-deviation -command to invoke a Kokkos-enabled LAMMPS build, for example -``lmp -k on g 1 -sf kk``. +the generated LAMMPS input. diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 851ad8624..b3cf86578 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -118,7 +118,7 @@ def training_args_dp() -> list[Argument]: deviation. It defaults to ``train_backend``. A model trained with ``pytorch`` can be deployed with ``pytorch-exportable`` to preserve its ``.pt`` training checkpoint while producing a graph-lowered ``.pt2`` - model for LAMMPS/Kokkos. Other cross-backend exports are not supported. + model for model deviation. Other cross-backend exports are 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." diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index 98176b400..545bbcf60 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -206,7 +206,7 @@ def test_model_deviation_forwards_pt2_models(self): } mdata = { "api_version": "1.0", - "model_devi_command": "lmp -k on g 1 -sf kk", + "model_devi_command": "lmp", "model_devi_group_size": 1, "model_devi_machine": {}, "model_devi_resources": {}, @@ -215,7 +215,6 @@ def test_model_deviation_forwards_pt2_models(self): 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__": From 1ad5296e3db8b5041deb5636e24de3026ad26a0f Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Sun, 23 Aug 2026 22:06:29 +0800 Subject: [PATCH 04/13] fix: enable atom map for pt2 model deviation --- dpgen/generator/lib/lammps.py | 2 +- tests/generator/test_lammps.py | 40 ++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/dpgen/generator/lib/lammps.py b/dpgen/generator/lib/lammps.py index b4bf603d7..47dd32d80 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 diff --git a/tests/generator/test_lammps.py b/tests/generator/test_lammps.py index c1b915939..c9afbc77b 100644 --- a/tests/generator/test_lammps.py +++ b/tests/generator/test_lammps.py @@ -55,6 +55,46 @@ 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_d3_enabled_basic(self): """Test LAMMPS input with D3 dispersion enabled.""" From 40adff3065d56234be48b99037434f8e7df5d68c Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Sun, 23 Aug 2026 22:19:52 +0800 Subject: [PATCH 05/13] docs: clarify pt2 deployment format --- dpgen/generator/arginfo.py | 8 +++++--- tests/generator/test_deepmd_backend.py | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index b3cf86578..95cdca4c4 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -116,9 +116,11 @@ def training_args_dp() -> list[Argument]: """\ The DeePMD backend used to freeze and compress models for model deviation. It defaults to ``train_backend``. A model trained with - ``pytorch`` can be deployed with ``pytorch-exportable`` to preserve its - ``.pt`` training checkpoint while producing a graph-lowered ``.pt2`` - model for model deviation. Other cross-backend exports are not supported. + ``pytorch`` can be deployed with ``pytorch-exportable`` and + ``model_format=pt2`` to preserve its ``.pt`` training checkpoint while + producing a graph-lowered ``.pt2`` model for model deviation. If + ``model_format`` is omitted, the deployment backend's default format is + used. Other cross-backend exports are 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." diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index 545bbcf60..98176b400 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -206,7 +206,7 @@ def test_model_deviation_forwards_pt2_models(self): } mdata = { "api_version": "1.0", - "model_devi_command": "lmp", + "model_devi_command": "lmp -k on g 1 -sf kk", "model_devi_group_size": 1, "model_devi_machine": {}, "model_devi_resources": {}, @@ -215,6 +215,7 @@ def test_model_deviation_forwards_pt2_models(self): 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__": From bbdb223ec1fbe0f5e5e3d243b6c6c38083c652b7 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 08:55:10 +0800 Subject: [PATCH 06/13] fix: export pt2 models on deployment resources --- doc/run/param.rst | 30 +++--- dpgen/generator/arginfo.py | 28 +---- dpgen/generator/run.py | 136 +++++++++++++++++-------- tests/generator/test_deepmd_backend.py | 111 +++++++++++--------- tests/generator/test_make_md.py | 63 ++++++++++++ 5 files changed, 244 insertions(+), 124 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index de3166dd5..cd7ef7c11 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -12,19 +12,17 @@ dpgen run param parameters DPA4 and DPA4C --------------- -DPA4 and the PyTorch-exportable backend require DeePMD-kit 3.2 or later. A -regular PyTorch DPA4 model can be exported as an AOTInductor archive with: +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_devi_backend": "pytorch-exportable", "model_format": "pt2" } -For DPA4C, and for DPA4 using the PyTorch-exportable backend, select the graph -export explicitly: +DPA4C uses the PyTorch-exportable backend for both training and graph export: .. code-block:: json @@ -34,10 +32,18 @@ export explicitly: "dp_compress": true } -``pt-expt`` is accepted as an alias of ``pytorch-exportable``. The ``pte`` -format remains available for dense PyTorch-exportable models. Training -checkpoints keep the ``.pt`` suffix independently of the frozen model format. -The ``model_devi_backend`` setting makes ``dpgen`` train with ``dp --pt`` but -freeze (and optionally compress) with ``dp --pt-expt``. The resulting ``.pt2`` -models are automatically linked into the model-deviation stage and listed in -the generated LAMMPS input. +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. + +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 95cdca4c4..23fd68599 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -104,23 +104,11 @@ def training_args_dp() -> list[Argument]: ) doc_model_format = textwrap.dedent( """\ - The frozen model format. Defaults depend on ``model_devi_backend`` (or - ``train_backend`` when no deployment backend is set): - ``pb`` for TensorFlow, ``pth`` for PyTorch, ``pte`` for - PyTorch-exportable, and ``savedmodel`` for JAX. PyTorch also supports - ``pt2`` for DPA4. PyTorch-exportable supports ``pte`` and ``pt2``; - use ``pt2`` for the graph export used by DPA4 and DPA4C. - """ - ) - doc_model_devi_backend = textwrap.dedent( - """\ - The DeePMD backend used to freeze and compress models for model - deviation. It defaults to ``train_backend``. A model trained with - ``pytorch`` can be deployed with ``pytorch-exportable`` and - ``model_format=pt2`` to preserve its ``.pt`` training checkpoint while - producing a graph-lowered ``.pt2`` model for model deviation. If - ``model_format`` is omitted, the deployment backend's default format is - used. Other cross-backend exports are not supported. + 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. """ ) doc_training_iter0_model_path = "The model used to init the first iter training. Number of element should be equal to numb_models." @@ -165,12 +153,6 @@ def training_args_dp() -> list[Argument]: default="tensorflow", doc=doc_train_backend, ), - Argument( - "model_devi_backend", - str, - optional=True, - doc=doc_model_devi_backend, - ), Argument( "model_format", str, diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 1044d0b22..bb2e2f499 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -177,22 +177,28 @@ def _get_train_backend_config(jdata) -> tuple[str, dict]: def _get_model_backend_config(jdata) -> tuple[str, dict, str]: """Return and validate the deployment backend and model format.""" - train_backend, _ = _get_train_backend_config(jdata) - backend, config = _get_backend(jdata, "model_devi_backend", train_backend) - if backend != train_backend and not ( - train_backend == "pytorch" and backend == "pytorch-exportable" + 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" ): - raise ValueError( - f"Cannot export models trained with backend {train_backend} using " - f"backend {backend}." - ) - model_format = jdata.get("model_format", config["default_model_format"]) + 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 @@ -214,12 +220,6 @@ def _get_train_backend_flag(jdata) -> str: return config["flag"] -def _get_model_backend_flag(jdata) -> str: - """Return the DeePMD CLI deployment backend flag.""" - _, config, _ = _get_model_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} @@ -821,7 +821,7 @@ def run_train_dp(iter_index, jdata, mdata): # load json param numb_models = jdata["numb_models"] train_backend, _ = _get_train_backend_config(jdata) - model_backend, _, model_format = _get_model_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'] @@ -845,11 +845,9 @@ def run_train_dp(iter_index, jdata, mdata): except KeyError: mdata = set_version(mdata) - if ( - train_backend == "pytorch-exportable" - or model_backend == "pytorch-exportable" - or model_format == "pt2" - ) and Version(mdata["deepmd_version"]) < Version("3.2"): + 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." ) @@ -859,7 +857,7 @@ def run_train_dp(iter_index, jdata, mdata): "use training_finetune_model or a checkpoint instead." ) if ( - model_backend == "pytorch" + train_backend == "pytorch" and model_format == "pt2" and jdata.get("dp_compress", False) ): @@ -883,10 +881,6 @@ def run_train_dp(iter_index, jdata, mdata): backend_flag = _get_train_backend_flag(jdata) if backend_flag: train_command += f" {backend_flag}" - model_command = mdata.get("train_command", "dp").strip() - model_backend_flag = _get_model_backend_flag(jdata) - if model_backend_flag: - model_command += f" {model_backend_flag}" # paths iter_name = make_iter_name(iter_index) @@ -902,6 +896,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"): @@ -925,23 +920,34 @@ def run_train_dp(iter_index, jdata, mdata): 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) - if model_backend == "pytorch-exportable": - command = f"{model_command} freeze -o frozen_model{suffix}" - if model_format == "pt2": - command += " --lower-kind graph" - elif model_backend == "pytorch" and model_format == "pt2": - command = f"{model_command} freeze -o frozen_model{suffix}" - else: - command = f"{model_command} freeze" - commands.append(command) - if jdata.get("dp_compress", False): - if model_backend == "pytorch-exportable": - commands.append( - f"{model_command} compress -i frozen_model{suffix} " + 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: - commands.append(f"{model_command} compress") + 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!" @@ -984,13 +990,14 @@ def run_train_dp(iter_index, jdata, mdata): 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 checkpoint_suffix == ".index": backward_files += [ @@ -1057,6 +1064,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): @@ -1595,6 +1620,29 @@ 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.""" + atom_map_index = None + read_index = None + for line_index, line in enumerate(lmp_lines): + tokens = line.partition("#")[0].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 + elif tokens[0] in {"read_data", "read_restart"} and read_index is None: + read_index = line_index + if atom_map_index is None or ( + read_index is not None and 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): @@ -1690,6 +1738,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 jdata.get("model_format") == "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): diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index 98176b400..502e3958a 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -8,9 +8,9 @@ from dpgen.generator.run import ( _get_checkpoint_suffix, _get_input_model_suffix, - _get_model_backend_flag, _get_model_suffix, _get_train_backend_flag, + _validate_pt2_template_atom_map, post_train_dp, run_md_model_devi, run_train_dp, @@ -34,7 +34,7 @@ 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), ".pte") + self.assertEqual(_get_model_suffix(jdata), ".pt2") self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt") @@ -43,34 +43,23 @@ def test_explicit_pt2_formats(self): {"train_backend": "pytorch", "model_format": "pt2"}, {"train_backend": "pytorch-exportable", "model_format": "pt2"}, {"train_backend": "pt-expt", "model_format": "pt2"}, - { - "train_backend": "pytorch", - "model_devi_backend": "pytorch-exportable", - "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_pytorch_checkpoint_can_use_exportable_deployment(self): - jdata = { - "train_backend": "pytorch", - "model_devi_backend": "pt-expt", - "model_format": "pt2", - } - self.assertEqual(_get_train_backend_flag(jdata), "--pt") - self.assertEqual(_get_model_backend_flag(jdata), "--pt-expt") - - def test_rejects_other_cross_backend_exports(self): - with self.assertRaisesRegex(ValueError, "Cannot export models"): + def test_pte_is_rejected_for_lammps(self): + with self.assertRaisesRegex(ValueError, "not supported by LAMMPS"): _get_model_suffix( - { - "train_backend": "tensorflow", - "model_devi_backend": "pytorch-exportable", - } + {"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"): @@ -81,6 +70,18 @@ def test_input_models_must_share_suffix(self): 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"], + ): + with self.subTest(lines=lines): + with self.assertRaisesRegex(ValueError, "atom_modify map yes"): + _validate_pt2_template_atom_map(lines) + class TestRunTrainDeepmdBackend(unittest.TestCase): def setUp(self): @@ -93,8 +94,10 @@ def setUp(self): "api_version": "1.0", "deepmd_version": "3.2.0", "train_command": "dp", - "train_machine": {}, - "train_resources": {}, + "train_machine": {"name": "train"}, + "train_resources": {"queue": "train"}, + "model_devi_machine": {"name": "model-devi"}, + "model_devi_resources": {"queue": "model-devi"}, } def _run(self, **updates): @@ -102,7 +105,12 @@ def _run(self, **updates): jdata.update(updates) with patch("dpgen.generator.run.make_submission") as make_submission: run_train_dp(0, jdata, self.mdata) - return make_submission.call_args.kwargs + 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() @@ -111,20 +119,24 @@ def test_legacy_tensorflow_commands_are_preserved(self): self.assertIn("frozen_model.pb", call["backward_files"]) self.assertIn("model.ckpt.index", call["backward_files"]) - def test_pytorch_dpa4_uses_exportable_pt2_deployment(self): - call = self._run( + def test_pytorch_dpa4_exports_pt2_on_model_devi_resources(self): + train_call, export_call = self._run( train_backend="pytorch", - model_devi_backend="pytorch-exportable", model_format="pt2", ) - self.assertIn("dp --pt train", call["commands"][0]) - self.assertIn("model.ckpt.pt", call["commands"][0]) + 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( - call["commands"][1], - "dp --pt-expt freeze -o frozen_model.pt2 --lower-kind graph", + export_call["commands"], + ["dp --pt freeze -c model.ckpt.pt -o frozen_model"], ) - self.assertIn("frozen_model.pt2", call["backward_files"]) - self.assertIn("model.ckpt.pt", call["backward_files"]) + 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") @@ -135,7 +147,9 @@ def test_legacy_pytorch_commands_are_preserved(self): self.assertIn("model.ckpt.pt", call["backward_files"]) def test_pytorch_exportable_dense_pte(self): - call = self._run(train_backend="pytorch-exportable") + 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], @@ -144,16 +158,22 @@ def test_pytorch_exportable_dense_pte(self): self.assertIn("frozen_model.pte", call["backward_files"]) def test_pytorch_exportable_dpa4c_pt2_compression(self): - call = self._run(train_backend="pt-expt", model_format="pt2", dp_compress=True) - self.assertEqual( - call["commands"][1], - "dp --pt-expt freeze -o frozen_model.pt2 --lower-kind graph", + 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( - call["commands"][2], - "dp --pt-expt compress -i frozen_model.pt2 -o frozen_model_compressed.pt2", + 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.assertIn("frozen_model_compressed.pt2", call["backward_files"]) + self.assertEqual(export_call["forward_files"], ["model.ckpt.pt"]) + self.assertIn("frozen_model_compressed.pt2", export_call["backward_files"]) def test_regular_pytorch_pt2_compression_is_rejected(self): with self.assertRaisesRegex(RuntimeError, "cannot compress pt2"): @@ -172,13 +192,13 @@ def test_deepmd_31_is_rejected_for_pt2(self): self._run(train_backend="pytorch", model_format="pt2") def test_finetune_keeps_source_model_suffix(self): - call = self._run( + train_call, _ = self._run( train_backend="pt-expt", model_format="pt2", training_finetune_model=["source.pt"], ) - self.assertIn("--finetune old/init.pt", call["commands"][0]) - self.assertIn(str(Path("old") / "init.pt"), call["forward_files"]) + 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 = { @@ -200,7 +220,6 @@ def test_model_deviation_forwards_pt2_models(self): (work_path / "cur_job.json").write_text(json.dumps({}), encoding="utf-8") jdata = { "train_backend": "pytorch", - "model_devi_backend": "pytorch-exportable", "model_format": "pt2", "model_devi_jobs": [{}], } diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py index ab36225cc..3fb70c890 100644 --- a/tests/generator/test_make_md.py +++ b/tests/generator/test_make_md.py @@ -6,6 +6,7 @@ import shutil import sys import unittest +from unittest.mock import patch import dpdata import numpy as np @@ -515,6 +516,68 @@ def test_make_model_devi_null(self): ) os.chdir(cwd_) + def test_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", + "model_format": "pt2", + "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): From 95c0339182de678e6f60dcf645cbf5c9a9f82e1b Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 09:08:50 +0800 Subject: [PATCH 07/13] fix: validate default pt2 templates --- dpgen/generator/run.py | 16 ++++++++++++++-- tests/generator/test_make_md.py | 5 ++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index bb2e2f499..e854ed2ae 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -1621,7 +1621,19 @@ def make_model_devi(iter_index, jdata, mdata): def _validate_pt2_template_atom_map(lmp_lines): - """Validate the atom map required by pt2 LAMMPS templates.""" + """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): @@ -1738,7 +1750,7 @@ 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 jdata.get("model_format") == "pt2": + 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 diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py index 3fb70c890..5cf799c17 100644 --- a/tests/generator/test_make_md.py +++ b/tests/generator/test_make_md.py @@ -516,7 +516,7 @@ def test_make_model_devi_null(self): ) os.chdir(cwd_) - def test_pt2_template_with_atom_map(self): + 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() @@ -545,8 +545,7 @@ def test_pt2_template_with_atom_map(self): "shuffle_poscar": False, "model_devi_f_trust_lo": 0.050, "model_devi_f_trust_hi": 0.150, - "train_backend": "pytorch", - "model_format": "pt2", + "train_backend": "pytorch-exportable", "model_devi_jobs": [ { "sys_idx": [0], From c4de737a7cf02f866e06002bc6bde0f76e118c12 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 13:21:53 +0800 Subject: [PATCH 08/13] fix: resolve pt2 format for generated inputs --- dpgen/generator/run.py | 2 +- tests/generator/test_make_md.py | 41 +++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index e854ed2ae..bd9f1d960 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -1970,7 +1970,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, diff --git a/tests/generator/test_make_md.py b/tests/generator/test_make_md.py index 5cf799c17..72e3c465e 100644 --- a/tests/generator/test_make_md.py +++ b/tests/generator/test_make_md.py @@ -163,6 +163,47 @@ 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"), + ) + def test_make_model_devi_pimd(self): if os.path.isdir("iter.000000"): shutil.rmtree("iter.000000") From 8c52d30fedad07f7a013f7677001bed0adf05f9d Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 14:39:41 +0000 Subject: [PATCH 09/13] fix: validate DPA deployment workflows --- doc/run/param.rst | 26 ++++++- dpgen/generator/arginfo.py | 11 ++- dpgen/generator/lib/run_calypso.py | 38 ++++++++-- dpgen/generator/run.py | 99 ++++++++++++++++++++++++-- tests/generator/test_deepmd_backend.py | 97 +++++++++++++++++++++++++ 5 files changed, 255 insertions(+), 16 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index cd7ef7c11..a015733e4 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -19,7 +19,14 @@ uses the regular PyTorch backend for both training and export: { "train_backend": "pytorch", - "model_format": "pt2" + "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: @@ -29,7 +36,16 @@ DPA4C uses the PyTorch-exportable backend for both training and graph export: { "train_backend": "pytorch-exportable", "model_format": "pt2", - "dp_compress": true + "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 @@ -37,6 +53,12 @@ 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. +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 diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 23fd68599..57848738e 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -113,7 +113,16 @@ def training_args_dp() -> list[Argument]: ) 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." diff --git a/dpgen/generator/lib/run_calypso.py b/dpgen/generator/lib/run_calypso.py index dbff1f0cb..8c592beeb 100644 --- a/dpgen/generator/lib/run_calypso.py +++ b/dpgen/generator/lib/run_calypso.py @@ -29,8 +29,19 @@ 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 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,7 +61,7 @@ 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") @@ -335,7 +346,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 +364,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 +465,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 +500,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 +516,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 bd9f1d960..018561f75 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -228,6 +228,86 @@ def _get_input_model_suffix(models) -> str: return suffixes.pop() +def _iter_model_sections(training_param): + 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]: + 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 == "dpa4" 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 DPA4/DPA4C backend and acceleration-option placement.""" + training_param = jdata.get("default_training_param", {}) + family = _get_dpa_model_family(training_param) + if family is None: + return + + train_backend, _ = _get_train_backend_config(jdata) + 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": + _, _, model_format = _get_model_backend_config(jdata) + 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( + "DPA4C uses training.enable_compile and training.enable_tf32; " + f"remove misplaced {', '.join(misplaced)}" + ) + else: + 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): jobkeys = [] for ii in jdata.keys(): @@ -819,6 +899,7 @@ 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) @@ -1637,7 +1718,8 @@ def _validate_pt2_template_atom_map(lmp_lines): atom_map_index = None read_index = None for line_index, line in enumerate(lmp_lines): - tokens = line.partition("#")[0].split() + command = line.partition("#")[0] + tokens = command.split() if not tokens: continue if tokens[0] == "atom_modify" and any( @@ -1645,11 +1727,11 @@ def _validate_pt2_template_atom_map(lmp_lines): for index in range(1, len(tokens) - 1) ): atom_map_index = line_index - elif tokens[0] in {"read_data", "read_restart"} and read_index is None: + if read_index is None and re.search(r"\bread_(?:data|restart)\b", command): read_index = line_index - if atom_map_index is None or ( - read_index is not None and atom_map_index > read_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." ) @@ -2475,7 +2557,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 index 502e3958a..d3d0f528c 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -5,11 +5,13 @@ from pathlib import Path from unittest.mock import patch +from dpgen.generator.lib.run_calypso import _find_models 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, @@ -77,10 +79,95 @@ def test_pt2_template_requires_atom_map_before_read(self): 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": "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": {}, + }, + } + ) + + 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) class TestRunTrainDeepmdBackend(unittest.TestCase): @@ -175,6 +262,16 @@ def test_pytorch_exportable_dpa4c_pt2_compression(self): 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) From eb9aeddaf7a876decc92778ed07a56e83829a32f Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 30 Aug 2026 00:21:51 +0800 Subject: [PATCH 10/13] fix: separate DeepMD deployment backend Coding-Agent: Codex Codex-Version: codex-cli 0.151.0 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- doc/run/param.rst | 47 ++++++++++-------- dpgen/generator/arginfo.py | 14 ++++++ dpgen/generator/run.py | 69 ++++++++++++++++++-------- tests/generator/test_deepmd_backend.py | 51 ++++++++++++++++++- 4 files changed, 136 insertions(+), 45 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index a015733e4..43c4ba557 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -19,14 +19,15 @@ uses the regular PyTorch backend for both training and export: { "train_backend": "pytorch", - "model_format": "pt2", - "default_training_param": { - "model": { - "type": "dpa4", - "use_compile": true, - "enable_tf32": true - } - } + "model_devi_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: @@ -35,23 +36,27 @@ DPA4C uses the PyTorch-exportable backend for both training and graph export: { "train_backend": "pytorch-exportable", + "model_devi_backend": "pytorch-exportable", "model_format": "pt2", - "dp_compress": true, - "default_training_param": { - "model": { - "descriptor": {"type": "dpa4c"} - }, - "training": { - "enable_compile": true, - "enable_tf32": true - } - } + "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 +The default ``train_backend`` remains ``tensorflow``. +``model_devi_backend`` defaults to ``train_backend`` but independently selects +the backend flag used by ``freeze`` and ``compress``. ``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. +independently of the frozen model format. Regular PyTorch ``pt2`` export is +accepted only for DPA4/SeZM training templates. The acceleration controls belong to different sections of the DeePMD training template: DPA4 uses ``model.use_compile`` and ``model.enable_tf32``; DPA4C uses diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 57848738e..92a6c4940 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -102,6 +102,14 @@ def training_args_dp() -> list[Argument]: 3.2 or later. """ ) + doc_model_devi_backend = textwrap.dedent( + """\ + The DeePMD-kit backend used to freeze and deploy trained checkpoints for + model deviation. It defaults to ``train_backend`` for backward + compatibility, but may be selected independently when the deployment + toolchain requires a different backend flag. + """ + ) doc_model_format = textwrap.dedent( """\ The frozen model format. Defaults are ``pb`` for TensorFlow, ``pth`` for @@ -162,6 +170,12 @@ def training_args_dp() -> list[Argument]: default="tensorflow", doc=doc_train_backend, ), + Argument( + "model_devi_backend", + str, + optional=True, + doc=doc_model_devi_backend, + ), Argument( "model_format", str, diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 018561f75..6eeaf4508 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -177,7 +177,8 @@ def _get_train_backend_config(jdata) -> tuple[str, dict]: def _get_model_backend_config(jdata) -> tuple[str, dict, str]: """Return and validate the deployment backend and model format.""" - backend, config = _get_train_backend_config(jdata) + train_backend, _ = _get_train_backend_config(jdata) + backend, config = _get_backend(jdata, "model_devi_backend", train_backend) default_model_format = config["default_model_format"] if ( backend == "pytorch-exportable" @@ -220,6 +221,12 @@ def _get_train_backend_flag(jdata) -> str: return config["flag"] +def _get_model_backend_flag(jdata) -> str: + """Return the DeePMD CLI flag used for model export and deployment.""" + _, config, _ = _get_model_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} @@ -265,19 +272,32 @@ def _validate_dpa_training_config(jdata) -> None: """Validate DPA4/DPA4C backend and acceleration-option placement.""" training_param = jdata.get("default_training_param", {}) family = _get_dpa_model_family(training_param) + train_backend, _ = _get_train_backend_config(jdata) + model_backend, _, model_format = _get_model_backend_config(jdata) if family is None: + if ( + train_backend == "pytorch" + and model_backend == "pytorch" + and model_format == "pt2" + ): + raise ValueError( + "The regular PyTorch backend only exports pt2 for DPA4/SeZM models." + ) return - train_backend, _ = _get_train_backend_config(jdata) 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 model_backend != expected_backend: + raise ValueError( + f"{family.upper()} deployment requires model_devi_backend=" + f"'{expected_backend}', not '{model_backend}'" + ) if jdata.get("model_devi_engine", "lammps") == "lammps": - _, _, model_format = _get_model_backend_config(jdata) if model_format != "pt2": raise ValueError( f"{family.upper()} LAMMPS model deviation requires model_format='pt2'" @@ -902,7 +922,7 @@ def run_train_dp(iter_index, jdata, mdata): _validate_dpa_training_config(jdata) numb_models = jdata["numb_models"] train_backend, _ = _get_train_backend_config(jdata) - _, _, model_format = _get_model_backend_config(jdata) + model_backend, _, model_format = _get_model_backend_config(jdata) suffix = _get_model_suffix(jdata) checkpoint_suffix = _get_checkpoint_suffix(jdata) # train_param = jdata['train_param'] @@ -926,9 +946,9 @@ 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"): + if ( + "pytorch-exportable" in {train_backend, model_backend} 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." ) @@ -938,7 +958,7 @@ def run_train_dp(iter_index, jdata, mdata): "use training_finetune_model or a checkpoint instead." ) if ( - train_backend == "pytorch" + model_backend == "pytorch" and model_format == "pt2" and jdata.get("dp_compress", False) ): @@ -957,11 +977,16 @@ def run_train_dp(iter_index, jdata, mdata): "training_init_model, training_init_frozen_model, and training_finetune_model are mutually exclusive." ) - train_command = mdata.get("train_command", "dp").strip() + deepmd_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 - backend_flag = _get_train_backend_flag(jdata) - if backend_flag: - train_command += f" {backend_flag}" + train_command = deepmd_command + train_backend_flag = _get_train_backend_flag(jdata) + if train_backend_flag: + train_command += f" {train_backend_flag}" + model_command = deepmd_command + model_backend_flag = _get_model_backend_flag(jdata) + if model_backend_flag: + model_command += f" {model_backend_flag}" # paths iter_name = make_iter_name(iter_index) @@ -1002,33 +1027,33 @@ def run_train_dp(iter_index, jdata, mdata): command = f"/bin/sh -c {shlex.quote(command)}" commands.append(command) if model_format == "pt2": - if train_backend == "pytorch-exportable": + if model_backend == "pytorch-exportable": command = ( - f"{train_command} freeze -c model.ckpt.pt -o frozen_model " + f"{model_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" + command = f"{model_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"{model_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}" + if model_backend == "pytorch-exportable": + command = f"{model_command} freeze -o frozen_model{suffix}" else: - command = f"{train_command} freeze" + command = f"{model_command} freeze" commands.append(command) if jdata.get("dp_compress", False): - if train_backend == "pytorch-exportable": + if model_backend == "pytorch-exportable": commands.append( - f"{train_command} compress -i frozen_model{suffix} " + f"{model_command} compress -i frozen_model{suffix} " f"-o frozen_model_compressed{suffix}" ) else: - commands.append(f"{train_command} compress") + commands.append(f"{model_command} compress") else: raise RuntimeError( "DP-GEN currently only supports for DeePMD-kit 1.x to 3.x version!" diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index d3d0f528c..b4cb67e43 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -9,6 +9,7 @@ from dpgen.generator.run import ( _get_checkpoint_suffix, _get_input_model_suffix, + _get_model_backend_flag, _get_model_suffix, _get_train_backend_flag, _validate_dpa_training_config, @@ -40,6 +41,17 @@ def test_pytorch_exportable_aliases(self): self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt") + def test_model_deviation_backend_is_independent(self): + jdata = { + "train_backend": "pytorch", + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + } + self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") + self.assertEqual(_get_train_backend_flag(jdata), "--pt") + self.assertEqual(_get_model_backend_flag(jdata), "--pt-expt") + self.assertEqual(_get_model_suffix(jdata), ".pt2") + def test_explicit_pt2_formats(self): cases = [ {"train_backend": "pytorch", "model_format": "pt2"}, @@ -210,6 +222,10 @@ 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": {"descriptor": {"type": "sezm"}}, + "training": {}, + }, ) self.assertEqual(train_call["machine"], self.mdata["train_machine"]) self.assertEqual(len(train_call["commands"]), 1) @@ -225,6 +241,18 @@ def test_pytorch_dpa4_exports_pt2_on_model_devi_resources(self): self.assertEqual(export_call["forward_files"], ["model.ckpt.pt"]) self.assertIn("frozen_model.pt2", export_call["backward_files"]) + def test_separate_model_backend_controls_export_command(self): + train_call, export_call = self._run( + train_backend="pytorch", + model_devi_backend="pytorch-exportable", + model_format="pt2", + ) + self.assertIn("dp --pt train", train_call["commands"][0]) + self.assertEqual( + export_call["commands"], + ["dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph"], + ) + def test_legacy_pytorch_commands_are_preserved(self): call = self._run(train_backend="pytorch") self.assertIn("dp --pt train", call["commands"][0]) @@ -274,7 +302,19 @@ def test_multiple_pt2_models_are_exported(self): 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) + self._run( + train_backend="pytorch", + model_format="pt2", + dp_compress=True, + default_training_param={ + "model": {"descriptor": {"type": "sezm"}}, + "training": {}, + }, + ) + + def test_regular_pytorch_pt2_requires_dpa4_family(self): + with self.assertRaisesRegex(ValueError, "only exports pt2 for DPA4/SeZM"): + self._run(train_backend="pytorch", model_format="pt2") def test_exportable_init_frozen_model_is_rejected(self): with self.assertRaisesRegex(RuntimeError, "does not support"): @@ -286,7 +326,14 @@ def test_exportable_init_frozen_model_is_rejected(self): 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") + self._run( + train_backend="pytorch", + model_format="pt2", + default_training_param={ + "model": {"descriptor": {"type": "sezm"}}, + "training": {}, + }, + ) def test_finetune_keeps_source_model_suffix(self): train_call, _ = self._run( From 5368e5590eabbab5167b987acde46ab3afd48f4b Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Sun, 30 Aug 2026 00:41:28 +0800 Subject: [PATCH 11/13] Revert "fix: separate DeepMD deployment backend" This reverts commit eb9aeddaf7a876decc92778ed07a56e83829a32f. --- doc/run/param.rst | 47 ++++++++---------- dpgen/generator/arginfo.py | 14 ------ dpgen/generator/run.py | 69 ++++++++------------------ tests/generator/test_deepmd_backend.py | 51 +------------------ 4 files changed, 45 insertions(+), 136 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index 43c4ba557..a015733e4 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -19,15 +19,14 @@ uses the regular PyTorch backend for both training and export: { "train_backend": "pytorch", - "model_devi_backend": "pytorch", - "model_format": "pt2", - "default_training_param": { - "model": { - "type": "dpa4", - "use_compile": true, - "enable_tf32": true - } - } + "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: @@ -36,27 +35,23 @@ DPA4C uses the PyTorch-exportable backend for both training and graph export: { "train_backend": "pytorch-exportable", - "model_devi_backend": "pytorch-exportable", "model_format": "pt2", - "dp_compress": true, - "default_training_param": { - "model": { - "descriptor": {"type": "dpa4c"} - }, - "training": { - "enable_compile": true, - "enable_tf32": true - } - } + "dp_compress": true, + "default_training_param": { + "model": { + "descriptor": {"type": "dpa4c"} + }, + "training": { + "enable_compile": true, + "enable_tf32": true + } + } } -The default ``train_backend`` remains ``tensorflow``. -``model_devi_backend`` defaults to ``train_backend`` but independently selects -the backend flag used by ``freeze`` and ``compress``. ``pt-expt`` is accepted as -an alias of ``pytorch-exportable``. PyTorch-exportable model deviation with +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. Regular PyTorch ``pt2`` export is -accepted only for DPA4/SeZM training templates. +independently of the frozen model format. The acceleration controls belong to different sections of the DeePMD training template: DPA4 uses ``model.use_compile`` and ``model.enable_tf32``; DPA4C uses diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 92a6c4940..57848738e 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -102,14 +102,6 @@ def training_args_dp() -> list[Argument]: 3.2 or later. """ ) - doc_model_devi_backend = textwrap.dedent( - """\ - The DeePMD-kit backend used to freeze and deploy trained checkpoints for - model deviation. It defaults to ``train_backend`` for backward - compatibility, but may be selected independently when the deployment - toolchain requires a different backend flag. - """ - ) doc_model_format = textwrap.dedent( """\ The frozen model format. Defaults are ``pb`` for TensorFlow, ``pth`` for @@ -170,12 +162,6 @@ def training_args_dp() -> list[Argument]: default="tensorflow", doc=doc_train_backend, ), - Argument( - "model_devi_backend", - str, - optional=True, - doc=doc_model_devi_backend, - ), Argument( "model_format", str, diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 6eeaf4508..018561f75 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -177,8 +177,7 @@ def _get_train_backend_config(jdata) -> tuple[str, dict]: def _get_model_backend_config(jdata) -> tuple[str, dict, str]: """Return and validate the deployment backend and model format.""" - train_backend, _ = _get_train_backend_config(jdata) - backend, config = _get_backend(jdata, "model_devi_backend", train_backend) + backend, config = _get_train_backend_config(jdata) default_model_format = config["default_model_format"] if ( backend == "pytorch-exportable" @@ -221,12 +220,6 @@ def _get_train_backend_flag(jdata) -> str: return config["flag"] -def _get_model_backend_flag(jdata) -> str: - """Return the DeePMD CLI flag used for model export and deployment.""" - _, config, _ = _get_model_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} @@ -272,32 +265,19 @@ def _validate_dpa_training_config(jdata) -> None: """Validate DPA4/DPA4C backend and acceleration-option placement.""" training_param = jdata.get("default_training_param", {}) family = _get_dpa_model_family(training_param) - train_backend, _ = _get_train_backend_config(jdata) - model_backend, _, model_format = _get_model_backend_config(jdata) if family is None: - if ( - train_backend == "pytorch" - and model_backend == "pytorch" - and model_format == "pt2" - ): - raise ValueError( - "The regular PyTorch backend only exports pt2 for DPA4/SeZM models." - ) return + train_backend, _ = _get_train_backend_config(jdata) 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 model_backend != expected_backend: - raise ValueError( - f"{family.upper()} deployment requires model_devi_backend=" - f"'{expected_backend}', not '{model_backend}'" - ) if jdata.get("model_devi_engine", "lammps") == "lammps": + _, _, model_format = _get_model_backend_config(jdata) if model_format != "pt2": raise ValueError( f"{family.upper()} LAMMPS model deviation requires model_format='pt2'" @@ -922,7 +902,7 @@ def run_train_dp(iter_index, jdata, mdata): _validate_dpa_training_config(jdata) numb_models = jdata["numb_models"] train_backend, _ = _get_train_backend_config(jdata) - model_backend, _, model_format = _get_model_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'] @@ -946,9 +926,9 @@ def run_train_dp(iter_index, jdata, mdata): except KeyError: mdata = set_version(mdata) - if ( - "pytorch-exportable" in {train_backend, model_backend} or model_format == "pt2" - ) and Version(mdata["deepmd_version"]) < Version("3.2"): + 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." ) @@ -958,7 +938,7 @@ def run_train_dp(iter_index, jdata, mdata): "use training_finetune_model or a checkpoint instead." ) if ( - model_backend == "pytorch" + train_backend == "pytorch" and model_format == "pt2" and jdata.get("dp_compress", False) ): @@ -977,16 +957,11 @@ def run_train_dp(iter_index, jdata, mdata): "training_init_model, training_init_frozen_model, and training_finetune_model are mutually exclusive." ) - deepmd_command = mdata.get("train_command", "dp").strip() + 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 - train_command = deepmd_command - train_backend_flag = _get_train_backend_flag(jdata) - if train_backend_flag: - train_command += f" {train_backend_flag}" - model_command = deepmd_command - model_backend_flag = _get_model_backend_flag(jdata) - if model_backend_flag: - model_command += f" {model_backend_flag}" + backend_flag = _get_train_backend_flag(jdata) + if backend_flag: + train_command += f" {backend_flag}" # paths iter_name = make_iter_name(iter_index) @@ -1027,33 +1002,33 @@ def run_train_dp(iter_index, jdata, mdata): command = f"/bin/sh -c {shlex.quote(command)}" commands.append(command) if model_format == "pt2": - if model_backend == "pytorch-exportable": + if train_backend == "pytorch-exportable": command = ( - f"{model_command} freeze -c model.ckpt.pt -o frozen_model " + f"{train_command} freeze -c model.ckpt.pt -o frozen_model " "--lower-kind graph" ) else: - command = f"{model_command} freeze -c model.ckpt.pt -o frozen_model" + 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"{model_command} compress -i frozen_model{suffix} " + f"{train_command} compress -i frozen_model{suffix} " f"-o frozen_model_compressed{suffix}" ) else: - if model_backend == "pytorch-exportable": - command = f"{model_command} freeze -o frozen_model{suffix}" + if train_backend == "pytorch-exportable": + command = f"{train_command} freeze -o frozen_model{suffix}" else: - command = f"{model_command} freeze" + command = f"{train_command} freeze" commands.append(command) if jdata.get("dp_compress", False): - if model_backend == "pytorch-exportable": + if train_backend == "pytorch-exportable": commands.append( - f"{model_command} compress -i frozen_model{suffix} " + f"{train_command} compress -i frozen_model{suffix} " f"-o frozen_model_compressed{suffix}" ) else: - commands.append(f"{model_command} compress") + commands.append(f"{train_command} compress") else: raise RuntimeError( "DP-GEN currently only supports for DeePMD-kit 1.x to 3.x version!" diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index b4cb67e43..d3d0f528c 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -9,7 +9,6 @@ from dpgen.generator.run import ( _get_checkpoint_suffix, _get_input_model_suffix, - _get_model_backend_flag, _get_model_suffix, _get_train_backend_flag, _validate_dpa_training_config, @@ -41,17 +40,6 @@ def test_pytorch_exportable_aliases(self): self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") self.assertEqual(_get_train_backend_flag(jdata), "--pt-expt") - def test_model_deviation_backend_is_independent(self): - jdata = { - "train_backend": "pytorch", - "model_devi_backend": "pytorch-exportable", - "model_format": "pt2", - } - self.assertEqual(_get_checkpoint_suffix(jdata), ".pt") - self.assertEqual(_get_train_backend_flag(jdata), "--pt") - self.assertEqual(_get_model_backend_flag(jdata), "--pt-expt") - self.assertEqual(_get_model_suffix(jdata), ".pt2") - def test_explicit_pt2_formats(self): cases = [ {"train_backend": "pytorch", "model_format": "pt2"}, @@ -222,10 +210,6 @@ 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": {"descriptor": {"type": "sezm"}}, - "training": {}, - }, ) self.assertEqual(train_call["machine"], self.mdata["train_machine"]) self.assertEqual(len(train_call["commands"]), 1) @@ -241,18 +225,6 @@ def test_pytorch_dpa4_exports_pt2_on_model_devi_resources(self): self.assertEqual(export_call["forward_files"], ["model.ckpt.pt"]) self.assertIn("frozen_model.pt2", export_call["backward_files"]) - def test_separate_model_backend_controls_export_command(self): - train_call, export_call = self._run( - train_backend="pytorch", - model_devi_backend="pytorch-exportable", - model_format="pt2", - ) - self.assertIn("dp --pt train", train_call["commands"][0]) - self.assertEqual( - export_call["commands"], - ["dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph"], - ) - def test_legacy_pytorch_commands_are_preserved(self): call = self._run(train_backend="pytorch") self.assertIn("dp --pt train", call["commands"][0]) @@ -302,19 +274,7 @@ def test_multiple_pt2_models_are_exported(self): 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": {"descriptor": {"type": "sezm"}}, - "training": {}, - }, - ) - - def test_regular_pytorch_pt2_requires_dpa4_family(self): - with self.assertRaisesRegex(ValueError, "only exports pt2 for DPA4/SeZM"): - self._run(train_backend="pytorch", model_format="pt2") + self._run(train_backend="pytorch", model_format="pt2", dp_compress=True) def test_exportable_init_frozen_model_is_rejected(self): with self.assertRaisesRegex(RuntimeError, "does not support"): @@ -326,14 +286,7 @@ def test_exportable_init_frozen_model_is_rejected(self): 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": {"descriptor": {"type": "sezm"}}, - "training": {}, - }, - ) + self._run(train_backend="pytorch", model_format="pt2") def test_finetune_keeps_source_model_suffix(self): train_call, _ = self._run( From 2c633ccad435e21fab854e87d454a1f304a9b4dd Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Sun, 30 Aug 2026 00:31:38 +0800 Subject: [PATCH 12/13] fix: enforce same-backend pt2 workflows --- doc/run/param.rst | 36 ++++++++++++++------------ dpgen/generator/arginfo.py | 4 ++- dpgen/generator/run.py | 13 ++++++---- tests/generator/test_deepmd_backend.py | 25 ++++++++++++++++-- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/doc/run/param.rst b/doc/run/param.rst index a015733e4..251557340 100644 --- a/doc/run/param.rst +++ b/doc/run/param.rst @@ -19,14 +19,14 @@ uses the regular PyTorch backend for both training and export: { "train_backend": "pytorch", - "model_format": "pt2", - "default_training_param": { - "model": { - "type": "dpa4", - "use_compile": true, - "enable_tf32": true - } + "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: @@ -36,16 +36,16 @@ DPA4C uses the PyTorch-exportable backend for both training and graph export: { "train_backend": "pytorch-exportable", "model_format": "pt2", - "dp_compress": true, - "default_training_param": { - "model": { - "descriptor": {"type": "dpa4c"} - }, - "training": { - "enable_compile": true, - "enable_tf32": true - } + "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 @@ -53,6 +53,10 @@ 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 diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 57848738e..24cb61362 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -108,7 +108,9 @@ def training_args_dp() -> list[Argument]: 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. + 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." diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py index 018561f75..dbe47f756 100644 --- a/dpgen/generator/run.py +++ b/dpgen/generator/run.py @@ -176,7 +176,7 @@ def _get_train_backend_config(jdata) -> tuple[str, dict]: def _get_model_backend_config(jdata) -> tuple[str, dict, str]: - """Return and validate the deployment backend and model format.""" + """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 ( @@ -251,7 +251,7 @@ def _get_dpa_model_family(training_param) -> Optional[str]: ) if descriptor_type == "dpa4c": families.add("dpa4c") - if model_type == "dpa4" or descriptor_type in {"dpa4", "sezm"}: + if model_type in {"dpa4", "sezm"} or descriptor_type in {"dpa4", "sezm"}: families.add("dpa4") if len(families) > 1: raise ValueError( @@ -265,10 +265,14 @@ def _validate_dpa_training_config(jdata) -> None: """Validate DPA4/DPA4C backend and acceleration-option placement.""" 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 - - train_backend, _ = _get_train_backend_config(jdata) expected_backend = "pytorch" if family == "dpa4" else "pytorch-exportable" if train_backend != expected_backend: raise ValueError( @@ -277,7 +281,6 @@ def _validate_dpa_training_config(jdata) -> None: ) if jdata.get("model_devi_engine", "lammps") == "lammps": - _, _, model_format = _get_model_backend_config(jdata) if model_format != "pt2": raise ValueError( f"{family.upper()} LAMMPS model deviation requires model_format='pt2'" diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index d3d0f528c..22aec88fa 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -106,6 +106,13 @@ def test_dpa_backend_and_compile_option_validation(self): }, } ) + _validate_dpa_training_config( + { + "train_backend": "pytorch", + "model_format": "pt2", + "default_training_param": {"model": {"type": "SeZM"}}, + } + ) _validate_dpa_training_config( { "train_backend": "pt-expt", @@ -160,6 +167,10 @@ def test_dpa_backend_and_compile_option_validation(self): }, } ) + 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: @@ -210,6 +221,7 @@ 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) @@ -274,7 +286,12 @@ def test_multiple_pt2_models_are_exported(self): 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) + 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"): @@ -286,7 +303,11 @@ def test_exportable_init_frozen_model_is_rejected(self): 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") + 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( From d481bfb8b4a9ec6f8381d2517915c843b97be643 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Sun, 30 Aug 2026 07:52:29 +0800 Subject: [PATCH 13/13] fix: pass resolved model to CALYPSO optimization --- dpgen/generator/lib/calypso_run_opt.py | 10 +++++++--- dpgen/generator/lib/run_calypso.py | 12 +++++++++--- tests/generator/test_deepmd_backend.py | 9 ++++++++- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/dpgen/generator/lib/calypso_run_opt.py b/dpgen/generator/lib/calypso_run_opt.py index 1bc765af1..5c190c735 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,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__": diff --git a/dpgen/generator/lib/run_calypso.py b/dpgen/generator/lib/run_calypso.py index 8c592beeb..62b16f7aa 100644 --- a/dpgen/generator/lib/run_calypso.py +++ b/dpgen/generator/lib/run_calypso.py @@ -34,6 +34,14 @@ def _find_models(path, model_suffix=".pb"): 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, @@ -65,9 +73,7 @@ def gen_structures( 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 " diff --git a/tests/generator/test_deepmd_backend.py b/tests/generator/test_deepmd_backend.py index 22aec88fa..47daac497 100644 --- a/tests/generator/test_deepmd_backend.py +++ b/tests/generator/test_deepmd_backend.py @@ -5,7 +5,10 @@ from pathlib import Path from unittest.mock import patch -from dpgen.generator.lib.run_calypso import _find_models +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, @@ -180,6 +183,10 @@ def test_calypso_discovers_resolved_model_suffix(self): 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):