diff --git a/docs/input.md b/docs/input.md index 6fc03bd1..badafb73 100644 --- a/docs/input.md +++ b/docs/input.md @@ -59,6 +59,76 @@ The `"type" : "dp"` tell the traning method is {dargs:argument}`"dp" `, i The `"config"` key defines the training configs, see {ref}`the full documentation`. The {dargs:argument}`"template_script" ` provides the template training script in `json` format. +For DPA4, use the regular PyTorch training backend and deploy `.pt2` models in +LAMMPS exploration: + +```json +"train": {"type": "dp", "config": {"impl": "pytorch"}}, +"explore": { + "type": "lmp", + "config": {"model_devi_backend": "pytorch", "model_format": "pt2"} +} +``` + +Put DPA4 training acceleration controls in the DeePMD training template under +`model`, not in the DPGEN2 workflow file: + +```json +"model": { + "type": "dpa4", + "use_compile": true, + "enable_tf32": true +} +``` + +For DPA4C, both training and deployment use the PyTorch Exportable backend. +Compression is optional: + +```json +"train": {"type": "dp", "config": {"impl": "pytorch-exportable"}}, +"explore": { + "type": "lmp", + "config": { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + "dp_compress": true + } +} +``` + +Put DPA4C training acceleration controls in the DeePMD training template file +referenced by `train.template_script`, under `training`: + +```json +"model": {"descriptor": {"type": "dpa4c"}}, +"training": { + "training_data": { + "systems": [], + "batch_size": "auto:512" + }, + "numb_steps": 1000000, + "enable_compile": true, + "enable_tf32": true +} +``` + +Do not copy the DPA4 paths `model.use_compile` or `model.enable_tf32` into a +DPA4C template. Conversely, DPA4 does not use the DPA4C paths +`training.enable_compile` or `training.enable_tf32`. DPGEN2 validates these +backend-specific placements before creating the workflow, but it does not inject +or change performance and numerical-policy settings. Therefore setting only +`train.config.impl` is not sufficient to enable compilation: a DPA4 template +must contain `model.use_compile: true`, while a DPA4C template must contain +`training.enable_compile: true`. Run a bounded smoke test and inspect the +generated `task.*/input.json` before launching a long campaign. + +TensorFlow remains the default when `impl` is omitted. +For PyTorch and PyTorch Exportable checkpoints, `model_devi_backend` must match +the training `impl`; checkpoints cannot be frozen across these backends. PT2 +export runs with `run_explore_config`, which must select hardware, the libtorch +version, and other runtime libraries compatible with the LAMMPS exploration +environment. + ### Exploration diff --git a/dpgen2/constants.py b/dpgen2/constants.py index 6d5d0197..7647ef16 100644 --- a/dpgen2/constants.py +++ b/dpgen2/constants.py @@ -4,7 +4,8 @@ train_log_name = "train.log" model_name_pattern = "model.%03d.pb" pytorch_model_name_pattern = "model.%03d.pth" -model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth)" +pt2_model_name_pattern = "model.%03d.pt2" +model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth|\.pt2)" lmp_index_pattern = "%06d" lmp_task_pattern = "task." + lmp_index_pattern lmp_conf_name = "conf.lmp" diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..62f543d8 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -45,6 +45,7 @@ def make_link(content, ref_key): def dp_dist_train_args(): + doc_numb_models = "Number of student models trained for model deviation" doc_config = "Configuration of training" doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. " dock_student_model_path = "The path of student model" @@ -63,6 +64,7 @@ def dp_dist_train_args(): Argument( "template_script", [List[str], str], optional=False, doc=doc_template_script ), + Argument("numb_models", int, optional=True, default=1, doc=doc_numb_models), Argument("student_model_path", str, optional=True, doc=dock_student_model_path), Argument( "student_model_uri", diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 9e07374f..6a97da41 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -107,6 +107,7 @@ RunRelax, RunRelaxHDF5, SelectConfs, + validate_model_backend, ) from dpgen2.op.caly_evo_step_merge import ( CalyEvoStepMerge, @@ -462,6 +463,92 @@ def get_systems_from_data(data, data_prefix=None): return data +def _normalize_model_backend(backend: str) -> str: + return {"pt-expt": "pytorch-exportable"}.get(backend, backend) + + +def _iter_model_sections(template_script: dict): + model = template_script.get("model", {}) + yield "model", model + for name, branch in (model.get("model_dict") or {}).items(): + yield f"model.model_dict.{name}", branch + + +def _model_family(template_script: dict) -> Optional[str]: + families = set() + for _, model in _iter_model_sections(template_script): + model_type = model.get("type") + descriptor = model.get("descriptor", {}) + descriptor_type = ( + descriptor.get("type") if isinstance(descriptor, dict) else None + ) + model_type = model_type.lower() if isinstance(model_type, str) else model_type + descriptor_type = ( + descriptor_type.lower() + if isinstance(descriptor_type, str) + else descriptor_type + ) + if descriptor_type == "dpa4c": + families.add("dpa4c") + if model_type == "dpa4" or descriptor_type in {"dpa4", "sezm"}: + families.add("dpa4") + if len(families) > 1: + raise RuntimeError( + "A training template cannot mix DPA4 and DPA4C branches because " + "they require different DeePMD backends" + ) + return next(iter(families), None) + + +def validate_dpa_training_template( + train_backend: str, + explore_config: dict, + template_script: dict, +) -> None: + """Validate DPA4/DPA4C backend and compile-option placement.""" + family = _model_family(template_script) + if family is None: + return + + train_backend = _normalize_model_backend(train_backend) + expected_backend = "pytorch" if family == "dpa4" else "pytorch-exportable" + if train_backend != expected_backend: + raise RuntimeError( + f"{family.upper()} training requires impl='{expected_backend}', " + f"not '{train_backend}'" + ) + + normalized_explore = RunLmp.normalize_config(explore_config) + if normalized_explore["model_format"] != "pt2": + raise RuntimeError( + f"{family.upper()} LAMMPS exploration requires model_format='pt2'" + ) + + misplaced = [] + if family == "dpa4c": + for scope, model in _iter_model_sections(template_script): + for key in ("use_compile", "enable_tf32"): + if key in model: + misplaced.append(f"{scope}.{key}") + if misplaced: + raise RuntimeError( + "DPA4C uses training.enable_compile and training.enable_tf32; " + f"remove misplaced {', '.join(misplaced)}" + ) + else: + training = template_script.get("training", {}) + misplaced = [ + f"training.{key}" + for key in ("enable_compile", "enable_tf32") + if key in training + ] + if misplaced: + raise RuntimeError( + "DPA4 uses model.use_compile and model.enable_tf32; " + f"remove misplaced {', '.join(misplaced)}" + ) + + def workflow_concurrent_learning( config: Dict, ) -> Step: @@ -472,6 +559,19 @@ def workflow_concurrent_learning( train_style = config["train"]["type"] explore_style = config["explore"]["type"] fp_style = config["fp"]["type"] + template_script_ = config["train"]["template_script"] + if isinstance(template_script_, list): + template_script = [json.loads(Path(ii).read_text()) for ii in template_script_] + else: + template_script = json.loads(Path(template_script_).read_text()) + if train_style in ["dp", "dp-dist"] and explore_style == "lmp": + train_backend = train_config.get("impl", "tensorflow") + validate_model_backend(train_backend, explore_config) + templates = ( + template_script if isinstance(template_script, list) else [template_script] + ) + for template in templates: + validate_dpa_training_template(train_backend, explore_config, template) prep_train_config = config["step_configs"]["prep_train_config"] run_train_config = config["step_configs"]["run_train_config"] prep_explore_config = config["step_configs"]["prep_explore_config"] @@ -493,12 +593,17 @@ def workflow_concurrent_learning( "not match numb_models={numb_models}" ) elif train_style == "dp-dist": + numb_models = config["train"]["numb_models"] + if "student_model_path" in config["train"] and numb_models != 1: + raise RuntimeError( + "student_model_path initializes one model; omit it for multiple " + "from-scratch students or set numb_models=1" + ) init_models_paths = ( [config["train"]["student_model_path"]] if "student_model_path" in config["train"] else None ) - config["train"]["numb_models"] = 1 else: raise RuntimeError(f"unknown params, train_style: {train_style}") @@ -554,11 +659,6 @@ def workflow_concurrent_learning( type_map = config["inputs"]["type_map"] numb_models = config["train"]["numb_models"] - template_script_ = config["train"]["template_script"] - if isinstance(template_script_, list): - template_script = [json.loads(Path(ii).read_text()) for ii in template_script_] - else: - template_script = json.loads(Path(template_script_).read_text()) if ( "teacher_model_path" in explore_config diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index f4fec3a2..74cbd5bc 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -35,8 +35,10 @@ RunDPTrain, ) from .run_lmp import ( + PrepareDPModels, RunLmp, RunLmpHDF5, + validate_model_backend, ) from .run_relax import ( RunRelax, diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 5a9782f4..e660603d 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -61,7 +61,10 @@ def _make_train_command( # find checkpoint if impl == "tensorflow" and os.path.isfile("checkpoint"): checkpoint = "model.ckpt" - elif impl == "pytorch" and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0: + elif ( + impl in ["pytorch", "pytorch-exportable"] + and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0 + ): checkpoint = "model.ckpt-%s.pt" % max( [int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")] ) @@ -184,10 +187,14 @@ def execute( finetune_mode = ip["optional_parameter"]["finetune_mode"] config = ip["config"] if ip["config"] is not None else {} impl = ip["config"].get("impl", "tensorflow") + if impl == "pt-expt": + impl = "pytorch-exportable" dp_command = ip["config"].get("command", "dp").split() - assert impl in ["tensorflow", "pytorch"] + assert impl in ["tensorflow", "pytorch", "pytorch-exportable"] if impl == "pytorch": dp_command.append("--pt") + elif impl == "pytorch-exportable": + dp_command.append("--pt-expt") finetune_args = config.get("finetune_args", "") train_args = config.get("train_args", "") config = RunDPTrain.normalize_config(config) @@ -318,7 +325,7 @@ def clean_before_quit(): shutil.copy2("input_v2_compat.json", train_script_name) # freeze model - if impl == "pytorch": + if impl in ["pytorch", "pytorch-exportable"]: model_file = "model.ckpt.pt" else: ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"]) @@ -339,10 +346,10 @@ def clean_before_quit(): ) raise FatalError("dp freeze failed") model_file = "frozen_model.pb" - fplog.write("#=================== freeze std out ===================\n") - fplog.write(out) - fplog.write("#=================== freeze std err ===================\n") - fplog.write(err) + fplog.write("#=================== freeze std out ===================\n") + fplog.write(out) + fplog.write("#=================== freeze std err ===================\n") + fplog.write(err) clean_before_quit() @@ -503,7 +510,7 @@ def decide_init_model( @staticmethod def training_args(): doc_command = "The command for DP, 'dp' for default" - doc_impl = "The implementation/backend of DP. It can be 'tensorflow' or 'pytorch'. 'tensorflow' for default." + doc_impl = "The implementation/backend of DP. It can be 'tensorflow', 'pytorch', or 'pytorch-exportable' (alias 'pt-expt'). 'tensorflow' for default." doc_init_model_policy = "The policy of init-model training. It can be\n\n\ - 'no': No init-model training. Traing from scratch.\n\n\ - 'yes': Do init-model training.\n\n\ diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 60cd9305..5f474f54 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -41,6 +41,7 @@ model_name_match_pattern, model_name_pattern, plm_output_name, + pt2_model_name_pattern, pytorch_model_name_pattern, ) from dpgen2.utils import ( @@ -51,6 +52,34 @@ run_command, ) +_MODEL_BACKEND_ALIASES = {"pt-expt": "pytorch-exportable"} +_MODEL_BACKEND_FLAGS = { + "pytorch": "--pt", + "pytorch-exportable": "--pt-expt", +} + + +class PrepareDPModels(OP): + """Freeze DP checkpoints once before exploration tasks are fanned out.""" + + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "config": BigParameter(dict), + "models": Artifact(List[Path]), + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign({"models": Artifact(List[Path])}) + + @OP.exec_sign_check + def execute(self, ip: OPIO) -> OPIO: + config = RunLmp.normalize_config(ip["config"] or {}) + return OPIO({"models": prepare_dp_models(ip["models"], config)}) + class RunLmp(OP): r"""Execute a LAMMPS task. @@ -152,10 +181,26 @@ def execute( if ext == ".pb": mname = model_name_pattern % (idx) Path(mname).symlink_to(mm) + elif ext == ".pth": + mname = pytorch_model_name_pattern % (idx) + Path(mname).symlink_to(mm) + elif ext == ".pt2": + mname = pt2_model_name_pattern % (idx) + Path(mname).symlink_to(mm) elif ext == ".pt": # freeze model - mname = pytorch_model_name_pattern % (idx) - freeze_model(mm, mname, config.get("model_frozen_head")) + backend = _model_backend(config) + mname = _model_name(idx, config["model_format"]) + freeze_model( + mm, + mname, + config.get("model_frozen_head"), + backend, + ) + if config["dp_compress"]: + compressed = _compressed_model_name(idx, config["model_format"]) + compress_model(mname, compressed, backend) + mname = compressed else: raise RuntimeError( "Model file with extension '%s' is not supported" % ext @@ -166,6 +211,8 @@ def execute( random.shuffle(model_names) set_models(lmp_input_name, model_names) + if any(Path(name).suffix == ".pt2" for name in model_names): + ensure_pt2_atom_map(lmp_input_name) # run lmp command = " ".join([command, "-i", lmp_input_name, "-log", lmp_log_name]) @@ -232,6 +279,11 @@ def lmp_args(): doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature" doc_use_hdf5 = "Use HDF5 to store trajs and model_devis" doc_extra_output_files = "Extra output file names, support wildcards" + doc_model_devi_backend = ( + "The DeePMD backend used to freeze models for exploration" + ) + doc_model_format = "The frozen model format. Use 'pt2' for DPA4 and DPA4C" + doc_dp_compress = "Compress the frozen model before exploration" return [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), Argument( @@ -269,6 +321,27 @@ def lmp_args(): default=[], doc=doc_extra_output_files, ), + Argument( + "model_devi_backend", + str, + optional=True, + default="pytorch", + doc=doc_model_devi_backend, + ), + Argument( + "model_format", + str, + optional=True, + default="pth", + doc=doc_model_format, + ), + Argument( + "dp_compress", + bool, + optional=True, + default=False, + doc=doc_dp_compress, + ), ] @staticmethod @@ -306,7 +379,7 @@ def set_models(lmp_input_name: str, model_names: List[str]): break if match_first == -1: raise RuntimeError( - f"cannot file model pattern {pattern} in line " f" {lmp_input_lines[idx]}" + f"cannot file model pattern {pattern} in line {lmp_input_lines[idx]}" ) if match_last == -1: raise RuntimeError(f"last matching index should not be -1, terribly wrong ") @@ -323,6 +396,56 @@ def set_models(lmp_input_name: str, model_names: List[str]): f.write("".join(lmp_input_lines)) +def ensure_pt2_atom_map(lmp_input_name: str): + """Ensure a PT2 LAMMPS input enables the atom map before reading atoms. + + Parameters + ---------- + lmp_input_name : str + Path to the LAMMPS input file. + + Raises + ------ + RuntimeError + If an existing ``atom_modify map yes`` follows the first + ``read_data`` or ``read_restart`` command, or neither read command is + present. + """ + with open(lmp_input_name, encoding="utf8") as f: + lmp_input_lines = f.readlines() + + read_index = next( + ( + index + for index, line in enumerate(lmp_input_lines) + if re.search(r"\bread_(?:data|restart)\b", line.partition("#")[0]) + ), + None, + ) + if read_index is None: + raise RuntimeError("PT2 LAMMPS inputs require read_data or read_restart") + + atom_map_index = next( + ( + index + for index, line in enumerate(lmp_input_lines) + if re.match(r"^\s*atom_modify\s+.*\bmap\s+yes\b", line.partition("#")[0]) + ), + None, + ) + if atom_map_index is not None: + if atom_map_index > read_index: + raise RuntimeError( + "PT2 LAMMPS inputs require 'atom_modify map yes' before " + "read_data or read_restart" + ) + return + + lmp_input_lines.insert(read_index, "atom_modify map yes\n") + with open(lmp_input_name, "w", encoding="utf8") as f: + f.write("".join(lmp_input_lines)) + + def find_only_one_key(lmp_lines, key, raise_not_found=True): found = [] for idx in range(len(lmp_lines)): @@ -362,31 +485,144 @@ def get_ele_temp(lmp_log_name): return None -def freeze_model(input_model, frozen_model, head=None): - freeze_args = "-o %s" % frozen_model +def _model_backend(config): + backend = _MODEL_BACKEND_ALIASES.get( + config["model_devi_backend"], config["model_devi_backend"] + ) + model_format = config["model_format"] + if backend not in _MODEL_BACKEND_FLAGS: + raise RuntimeError(f"Unsupported model-deviation backend '{backend}'") + if model_format not in ["pth", "pt2"]: + raise RuntimeError(f"Unsupported model format '{model_format}'") + if model_format == "pth" and backend != "pytorch": + raise RuntimeError("The pth model format requires the pytorch backend") + if config["dp_compress"] and not ( + backend == "pytorch-exportable" and model_format == "pt2" + ): + raise RuntimeError( + "Compressed pt2 models require the pytorch-exportable backend" + ) + return backend + + +def validate_model_backend(train_backend, config): + """Validate that a PyTorch checkpoint is frozen by its training backend. + + Parameters + ---------- + train_backend : str + DeePMD training backend. + config : dict + LAMMPS exploration configuration. + + Raises + ------ + RuntimeError + If PyTorch training and deployment backends differ. + """ + train_backend = _MODEL_BACKEND_ALIASES.get(train_backend, train_backend) + if train_backend not in _MODEL_BACKEND_FLAGS: + return + model_backend = _model_backend(RunLmp.normalize_config(config)) + if model_backend != train_backend: + raise RuntimeError( + f"The model-deviation backend '{model_backend}' cannot freeze a " + f"checkpoint trained by '{train_backend}'; use the same backend " + "for training and model deployment" + ) + + +def _model_name(index, model_format): + if model_format == "pt2": + return pt2_model_name_pattern % index + return pytorch_model_name_pattern % index + + +def _compressed_model_name(index, model_format): + return "model.%03d.compressed.%s" % (index, model_format) + + +def prepare_dp_models(models, config): + """Return frozen models, exporting checkpoints once when needed.""" + backend = _model_backend(config) + prepared = [] + output_dir = Path("prepared_models") + for idx, model in enumerate(models): + model = Path(model).resolve() + ext = model.suffix + if ext != ".pt": + if ext not in [".pb", ".pth", ".pt2"]: + raise RuntimeError( + "Model file with extension '%s' is not supported" % ext + ) + prepared.append(model) + continue + output_dir.mkdir(exist_ok=True) + frozen_model = output_dir / _model_name(idx, config["model_format"]) + freeze_model( + model, + frozen_model, + config.get("model_frozen_head"), + backend, + ) + if config["dp_compress"]: + compressed_model = output_dir / _compressed_model_name( + idx, config["model_format"] + ) + compress_model(frozen_model, compressed_model, backend) + frozen_model = compressed_model + prepared.append(frozen_model) + return prepared + + +def freeze_model(input_model, frozen_model, head=None, backend="pytorch"): + backend = _MODEL_BACKEND_ALIASES.get(backend, backend) + freeze_cmd = [ + "dp", + _MODEL_BACKEND_FLAGS[backend], + "freeze", + "-c", + str(input_model), + "-o", + str(frozen_model), + ] if head is not None: - freeze_args += " --head %s" % head - freeze_cmd = "dp --pt freeze -c %s %s" % (input_model, freeze_args) - ret, out, err = run_command(freeze_cmd, shell=True) + freeze_cmd.extend(["--head", str(head)]) + if backend == "pytorch-exportable" and Path(frozen_model).suffix == ".pt2": + freeze_cmd.extend(["--lower-kind", "graph"]) + ret, out, err = run_command(freeze_cmd) if ret != 0: logging.error( - "".join( - ( - "freeze failed\n", - "command was", - freeze_cmd, - "out msg", - out, - "\n", - "err msg", - err, - "\n", - ) - ) + "freeze failed\ncommand was %s\nout msg%s\nerr msg%s\n", + freeze_cmd, + out, + err, ) raise TransientError("freeze failed") +def compress_model(input_model, output_model, backend="pytorch-exportable"): + backend = _MODEL_BACKEND_ALIASES.get(backend, backend) + compress_cmd = [ + "dp", + _MODEL_BACKEND_FLAGS[backend], + "compress", + "-i", + str(input_model), + "-o", + str(output_model), + ] + ret, out, err = run_command(compress_cmd) + if ret != 0: + logging.error( + "compress failed\ncommand was%s\nout msg%s\nerr msg%s\n", + compress_cmd, + out, + err, + ) + raise TransientError("compress failed") + + def merge_pimd_files(): traj_files = glob.glob("traj.*.dump") if len(traj_files) > 0: diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 3e0a0a0f..699716c1 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -40,6 +40,9 @@ from dpgen2.constants import ( lmp_index_pattern, ) +from dpgen2.op.run_lmp import ( + PrepareDPModels, +) from dpgen2.utils.step_config import ( init_executor, ) @@ -91,10 +94,12 @@ def __init__( ), ) - self._keys = ["prep-lmp", "run-lmp"] + self._keys = ["prep-lmp", "prepare-models", "run-lmp"] self.step_keys = {} ii = "prep-lmp" self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + ii = "prepare-models" + self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) ii = "run-lmp" self.step_keys[ii] = "--".join( ["%s" % self.inputs.parameters["block_id"], ii + "-{{item}}"] @@ -148,6 +153,10 @@ def _prep_run_lmp( run_executor = init_executor(run_config.pop("executor")) template_slice_config = run_config.pop("template_slice_config", {}) + prepare_models_config = deepcopy(run_config) + prepare_models_config.pop("continue_on_num_success", None) + prepare_models_config.pop("continue_on_success_ratio", None) + prep_lmp = Step( "prep-lmp", template=PythonOPTemplate( @@ -166,6 +175,26 @@ def _prep_run_lmp( ) prep_run_steps.add(prep_lmp) + prepare_models = Step( + "prepare-models", + template=PythonOPTemplate( + PrepareDPModels, + output_artifact_archive={"models": None}, + python_packages=upload_python_packages, + **run_template_config, + ), + parameters={ + "config": prep_run_steps.inputs.parameters["explore_config"], + }, + artifacts={ + "models": prep_run_steps.inputs.artifacts["models"], + }, + key=step_keys["prepare-models"], + executor=run_executor, + **prepare_models_config, + ) + prep_run_steps.add(prepare_models) + run_lmp = Step( "run-lmp", template=PythonOPTemplate( @@ -193,7 +222,7 @@ def _prep_run_lmp( }, artifacts={ "task_path": prep_lmp.outputs.artifacts["task_paths"], - "models": prep_run_steps.inputs.artifacts["models"], + "models": prepare_models.outputs.artifacts["models"], }, with_sequence=argo_sequence( argo_len(prep_lmp.outputs.parameters["task_names"]), diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py index bc551a2c..88f2d548 100644 --- a/tests/entrypoint/test_submit.py +++ b/tests/entrypoint/test_submit.py @@ -23,6 +23,7 @@ print_list_steps, submit_concurrent_learning, update_reuse_step_scheduler, + validate_dpa_training_template, ) from dpgen2.exploration.render import ( TrajRenderLammps, @@ -106,6 +107,117 @@ def modify_output_parameter(self, key, scheduler): class TestSubmit(unittest.TestCase): + def test_validate_dpa4_training_template(self): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + { + "model": { + "type": "dpa4", + "descriptor": {"type": "dpa4"}, + "use_compile": True, + "enable_tf32": True, + }, + "training": {}, + }, + ) + with self.assertRaisesRegex(RuntimeError, "model.use_compile"): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + { + "model": {"type": "dpa4", "descriptor": {"type": "dpa4"}}, + "training": {"enable_compile": True}, + }, + ) + + def test_validate_dpa4c_training_template(self): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": {"descriptor": {"type": "dpa4c"}}, + "training": {"enable_compile": True, "enable_tf32": True}, + }, + ) + with self.assertRaisesRegex(RuntimeError, "training.enable_compile"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "descriptor": {"type": "dpa4c"}, + "use_compile": True, + }, + "training": {}, + }, + ) + + with self.assertRaisesRegex(RuntimeError, "training.enable_tf32"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "descriptor": {"type": "DPA4C"}, + "enable_tf32": True, + }, + "training": {"enable_compile": True}, + }, + ) + + def test_validate_dpa_training_backend_and_format(self): + dpa4c = { + "model": {"descriptor": {"type": "dpa4c"}}, + "training": {}, + } + with self.assertRaisesRegex(RuntimeError, "requires impl='pytorch-exportable'"): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + dpa4c, + ) + with self.assertRaisesRegex(RuntimeError, "requires model_format='pt2'"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pth", + }, + dpa4c, + ) + + def test_validate_mixed_dpa4_and_dpa4c_branches(self): + with self.assertRaisesRegex(RuntimeError, "cannot mix DPA4 and DPA4C"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "model_dict": { + "teacher": {"descriptor": {"type": "dpa4"}}, + "student": {"descriptor": {"type": "dpa4c"}}, + } + }, + "training": { + "enable_compile": True, + "enable_tf32": True, + }, + }, + ) + def test_expand_idx(self): ilist = ["1", "3-5", "10-20:2"] olist = expand_idx(ilist) diff --git a/tests/entrypoint/test_submit_args.py b/tests/entrypoint/test_submit_args.py index ec7e4582..518a154c 100644 --- a/tests/entrypoint/test_submit_args.py +++ b/tests/entrypoint/test_submit_args.py @@ -126,6 +126,7 @@ def test(self): # self.assertEqual(old_data['default_training_param'], new_data['train']['template_script']) self.assertEqual(new_data["train"]["template_script"], "dp_input_template") self.assertEqual(RunDPTrain.normalize_config({}), new_data["train"]["config"]) + self.assertEqual(new_data["train"]["config"]["impl"], "tensorflow") self.assertEqual( RunLmp.normalize_config(old_data.get("lmp_config", {})), new_data["explore"]["config"], diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 7649b520..d9a03800 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -434,6 +434,39 @@ def test_exec_v1(self, mocked_run): jdata = json.load(fp) self.assertDictEqual(jdata, self.expected_odict_v1) + @patch("dpgen2.op.run_dp_train.run_command") + def test_exec_pytorch_exportable(self, mocked_run): + mocked_run.return_value = (0, "foo\n", "") + config = self.config.copy() + config.update({"impl": "pt-expt", "init_model_policy": "no"}) + Path(self.task_path).mkdir(exist_ok=True) + with open(Path(self.task_path) / train_script_name, "w") as fp: + json.dump(self.idict_v2, fp, indent=4) + + out = RunDPTrain().execute( + OPIO( + { + "config": config, + "task_name": self.task_name, + "task_path": Path(self.task_path), + "init_model": Path(self.init_model), + "init_data": [Path(ii) for ii in self.init_data], + "iter_data": [Path(ii) for ii in self.iter_data], + } + ) + ) + + self.assertEqual(out["model"], Path(self.task_name) / "model.ckpt.pt") + mocked_run.assert_called_once_with( + ["dp", "--pt-expt", "train", train_script_name] + ) + self.assertEqual( + out["log"].read_text(), + "#=================== train std out ===================\n" + "foo\n" + "#=================== train std err ===================\n", + ) + @patch("dpgen2.op.run_dp_train.run_command") def test_exec_v2(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", ""), (0, "bar\n", "")] diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 650fd82e..743fb307 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -32,12 +32,19 @@ lmp_model_devi_name, lmp_traj_name, model_name_pattern, + pt2_model_name_pattern, ) from dpgen2.op.run_lmp import ( + PrepareDPModels, RunLmp, + _model_backend, + compress_model, + ensure_pt2_atom_map, get_ele_temp, merge_pimd_files, + prepare_dp_models, set_models, + validate_model_backend, ) from dpgen2.utils import ( BinaryFileInput, @@ -102,6 +109,40 @@ def test_success(self, mocked_run): (work_dir / (model_name_pattern % ii)).read_text(), f"model{ii}" ) + @patch("dpgen2.op.run_lmp.run_command") + def test_pt2_enables_atom_map_before_read(self, mocked_run): + mocked_run.return_value = (0, "", "") + (self.task_path / lmp_input_name).write_text( + "atom_style atomic\n" + 'if "${restart} > 0" then "read_restart dpgen.restart.*" ' + 'else "read_data conf.lmp"\n' + "pair_style deepmd model.000.pb model.001.pb out_freq 10\n" + ) + models = [self.model_path / f"model_{index}.pt2" for index in range(2)] + for model in models: + model.write_text("model") + + def copy_link(source, target, target_is_directory=False): + shutil.copyfile(source, target) + + with patch("os.symlink", side_effect=copy_link): + RunLmp().execute( + OPIO( + { + "config": {"command": "mylmp"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": models, + } + ) + ) + + lmp_input = (Path(self.task_name) / lmp_input_name).read_text() + atom_map = "atom_modify map yes" + self.assertEqual(lmp_input.count(atom_map), 1) + self.assertLess(lmp_input.index(atom_map), lmp_input.index("read_restart")) + self.assertLess(lmp_input.index(atom_map), lmp_input.index("read_data")) + @patch("dpgen2.op.run_lmp.run_command") def test_error(self, mocked_run): mocked_run.side_effect = [(1, "foo\n", "")] @@ -258,6 +299,121 @@ def swap_element(arg): arg[0] = bk[1] +class TestPrepareDPModels(unittest.TestCase): + def setUp(self): + self.model_dir = Path("checkpoint_models") + self.model_dir.mkdir() + self.models = [] + for idx in range(2): + model = self.model_dir / f"model.{idx}.pt" + model.write_text("checkpoint") + self.models.append(model) + + def tearDown(self): + shutil.rmtree(self.model_dir, ignore_errors=True) + shutil.rmtree("prepared_models", ignore_errors=True) + + @patch("dpgen2.op.run_lmp.run_command") + def test_dpa4_pt2(self, mocked_run): + mocked_run.return_value = (0, "", "") + models = PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pytorch", + "model_format": "pt2", + }, + "models": self.models, + } + ) + )["models"] + self.assertEqual( + models, + [ + Path("prepared_models/model.000.pt2"), + Path("prepared_models/model.001.pt2"), + ], + ) + mocked_run.assert_has_calls( + [ + call( + [ + "dp", + "--pt", + "freeze", + "-c", + str(model.resolve()), + "-o", + str(Path("prepared_models") / f"model.{idx:03d}.pt2"), + ] + ) + for idx, model in enumerate(self.models) + ] + ) + + @patch("dpgen2.op.run_lmp.run_command") + def test_dpa4c_compressed_pt2(self, mocked_run): + mocked_run.return_value = (0, "", "") + models = PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pt-expt", + "model_format": "pt2", + "dp_compress": True, + }, + "models": self.models[:1], + } + ) + )["models"] + self.assertEqual(models, [Path("prepared_models/model.000.compressed.pt2")]) + mocked_run.assert_has_calls( + [ + call( + [ + "dp", + "--pt-expt", + "freeze", + "-c", + str(self.models[0].resolve()), + "-o", + str(Path("prepared_models/model.000.pt2")), + "--lower-kind", + "graph", + ] + ), + call( + [ + "dp", + "--pt-expt", + "compress", + "-i", + str(Path("prepared_models/model.000.pt2")), + "-o", + str(Path("prepared_models/model.000.compressed.pt2")), + ] + ), + ] + ) + + def test_training_and_deployment_backends_must_match(self): + with self.assertRaisesRegex(RuntimeError, "cannot freeze a checkpoint"): + validate_model_backend( + "pytorch", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + ) + validate_model_backend( + "pt-expt", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + ) + + class TestSetModels(unittest.TestCase): def setUp(self): self.input_name = Path("lmp.input") @@ -274,6 +430,16 @@ def test(self): set_models(input_name, self.model_names) self.assertEqual(input_name.read_text(), expected_output) + def test_pt2(self): + lmp_config = "pair_style deepmd model.000.pb model.001.pb out_freq 10\n" + expected_output = "pair_style deepmd model.000.pt2 model.001.pt2 out_freq 10\n" + self.input_name.write_text(lmp_config) + set_models( + self.input_name, + [pt2_model_name_pattern % 0, pt2_model_name_pattern % 1], + ) + self.assertEqual(self.input_name.read_text(), expected_output) + def test_failed(self): lmp_config = "pair_style deepmd model.000.pb model.001.pb out_freq 10 out_file model_devi.out model.002.pb\n" input_name = self.input_name @@ -372,3 +538,116 @@ def tearDown(self): ]: if os.path.exists(f): os.remove(f) + + +class TestPrepareDPModelsPassthrough(unittest.TestCase): + def setUp(self): + self.model_dir = Path("_test_models") + self.model_dir.mkdir(exist_ok=True) + + def tearDown(self): + shutil.rmtree(self.model_dir, ignore_errors=True) + shutil.rmtree("prepared_models", ignore_errors=True) + + def test_pth_passthrough(self): + model = self.model_dir / "model.000.pth" + model.write_text("frozen") + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pth"} + ) + result = prepare_dp_models([model], config) + self.assertEqual(result, [model.resolve()]) + + def test_pt2_passthrough(self): + model = self.model_dir / "model.000.pt2" + model.write_text("frozen") + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) + result = prepare_dp_models([model], config) + self.assertEqual(result, [model.resolve()]) + + def test_unsupported_extension_raises(self): + model = self.model_dir / "model.onnx" + model.write_text("bad") + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) + with self.assertRaisesRegex(RuntimeError, "not supported"): + prepare_dp_models([model], config) + + +class TestModelBackendValidation(unittest.TestCase): + def test_unsupported_backend(self): + with self.assertRaisesRegex( + RuntimeError, "Unsupported model-deviation backend" + ): + _model_backend( + { + "model_devi_backend": "bogus", + "model_format": "pt2", + "dp_compress": False, + } + ) + + def test_unsupported_format(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "xyz", + "dp_compress": False, + } + ) + + def test_pth_requires_pytorch(self): + with self.assertRaisesRegex( + RuntimeError, "pth model format requires the pytorch" + ): + _model_backend( + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pth", + "dp_compress": False, + } + ) + + def test_compress_requires_exportable_pt2(self): + with self.assertRaisesRegex(RuntimeError, "Compressed pt2"): + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "pt2", + "dp_compress": True, + } + ) + + def test_validate_non_pytorch_backend_skips(self): + validate_model_backend( + "tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) + + +class TestEnsurePt2AtomMap(unittest.TestCase): + def test_map_already_present_before_read(self): + lines = "atom_modify map yes\nread_data conf.lmp\n" + Path("test_input.lammps").write_text(lines) + ensure_pt2_atom_map("test_input.lammps") + result = Path("test_input.lammps").read_text() + self.assertIn("atom_modify map yes", result) + self.assertEqual(result.count("atom_modify"), 1) + os.remove("test_input.lammps") + + def test_no_read_command_raises(self): + Path("test_input.lammps").write_text("atom_modify map yes\npair_style deepmd\n") + with self.assertRaisesRegex(RuntimeError, "read_data or read_restart"): + ensure_pt2_atom_map("test_input.lammps") + os.remove("test_input.lammps") + + +class TestCompressModelFailure(unittest.TestCase): + @patch("dpgen2.op.run_lmp.run_command") + def test_compress_failure_raises(self, mocked_run): + mocked_run.return_value = (1, "", "compress error") + with self.assertRaisesRegex(TransientError, "compress failed"): + compress_model("input.pt2", "output.pt2", "pytorch-exportable") diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py index fa0571a7..f1392984 100644 --- a/tests/test_block_cl.py +++ b/tests/test_block_cl.py @@ -225,6 +225,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", @@ -405,6 +406,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py index 49b4c873..bf49cd9d 100644 --- a/tests/test_dpgen_loop.py +++ b/tests/test_dpgen_loop.py @@ -270,6 +270,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", @@ -502,6 +503,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", diff --git a/tests/test_prep_run_lmp_config.py b/tests/test_prep_run_lmp_config.py new file mode 100644 index 00000000..7da46cb8 --- /dev/null +++ b/tests/test_prep_run_lmp_config.py @@ -0,0 +1,46 @@ +import unittest + +from dflow import ( + Step, +) + +from dpgen2.op.prep_lmp import ( + PrepLmp, +) +from dpgen2.op.run_lmp import ( + RunLmp, +) +from dpgen2.superop.prep_run_lmp import ( + PrepRunLmp, +) +from dpgen2.utils.step_config import normalize as normalize_step_dict + + +class TestPrepRunLmpConfig(unittest.TestCase): + def test_non_sliced_model_preparation_ignores_slice_success_controls(self): + config = normalize_step_dict( + { + "continue_on_num_success": 1, + "continue_on_success_ratio": 0.5, + } + ) + steps = PrepRunLmp( + "prep-run-lmp", + PrepLmp, + RunLmp, + prep_config=normalize_step_dict({}), + run_config=config, + ) + prepare_models = next( + step for step in steps.steps if step.name == "prepare-models" + ) + run_lmp = next(step for step in steps.steps if step.name == "run-lmp") + + self.assertIsNone(prepare_models.continue_on_num_success) + self.assertIsNone(prepare_models.continue_on_success_ratio) + self.assertEqual(run_lmp.continue_on_num_success, 1) + self.assertEqual(run_lmp.continue_on_success_ratio, 0.5) + + +if __name__ == "__main__": + unittest.main()