pd: support dpa2

deepmd/pd/entrypoints/main.py

@@ -92,8 +92,8 @@
     # Initialize DDP
     world_size = dist.get_world_size()
     if world_size > 1:
         assert paddle.version.nccl() != "0"
         fleet.init(is_collective=True)
 
     def prepare_trainer_input_single(
         model_params_single, data_dict_single, rank=0, seed=None
@@ -111,15 +111,15 @@
         # stat files
         stat_file_path_single = data_dict_single.get("stat_file", None)
         if rank != 0:
             stat_file_path_single = None
         elif stat_file_path_single is not None:
             if not Path(stat_file_path_single).exists():
                 if stat_file_path_single.endswith((".h5", ".hdf5")):
                     with h5py.File(stat_file_path_single, "w") as f:
                         pass
                 else:
                     Path(stat_file_path_single).mkdir()
             stat_file_path_single = DPPath(stat_file_path_single, "a")
 
         # validation and training data
         # avoid the same batch sequence among devices
@@ -160,9 +160,9 @@
             seed=data_seed,
         )
     else:
         train_data, validation_data, stat_file_path = {}, {}, {}
         for model_key in config["model"]["model_dict"]:
             (
                 train_data[model_key],
                 validation_data[model_key],
                 stat_file_path[model_key],
@@ -194,24 +194,24 @@
 
     def is_built_with_cuda(self) -> bool:
         """Check if the backend is built with CUDA."""
         return paddle.device.is_compiled_with_cuda()
 
     def is_built_with_rocm(self) -> bool:
         """Check if the backend is built with ROCm."""
         return paddle.device.is_compiled_with_rocm()
 
     def get_compute_device(self) -> str:
         """Get Compute device."""
         return str(DEVICE)
 
     def get_ngpus(self) -> int:
         """Get the number of GPUs."""
         return paddle.device.cuda.device_count()
 
     def get_backend_info(self) -> dict:
         """Get backend information."""
         op_info = {}
         return {
             "Backend": "Paddle",
             "PD ver": f"v{paddle.__version__}-g{paddle.version.commit[:11]}",
             "Enable custom OP": False,
@@ -230,85 +230,85 @@
     use_pretrain_script: bool = False,
     force_load: bool = False,
     output: str = "out.json",
-):
+) -> None:
     log.info("Configuration path: %s", input_file)
     SummaryPrinter()()
     with open(input_file) as fin:
         config = json.load(fin)
     # ensure suffix, as in the command line help, we say "path prefix of checkpoint files"
     if init_model is not None and not init_model.endswith(".pd"):
         init_model += ".pd"
     if restart is not None and not restart.endswith(".pd"):
         restart += ".pd"
 
     # update multitask config
     multi_task = "model_dict" in config["model"]
     shared_links = None
     if multi_task:
         config["model"], shared_links = preprocess_shared_params(config["model"])
         # handle the special key
         assert (
             "RANDOM" not in config["model"]["model_dict"]
         ), "Model name can not be 'RANDOM' in multi-task mode!"
 
     # update fine-tuning config
     finetune_links = None
     if finetune is not None:
         config["model"], finetune_links = get_finetune_rules(
             finetune,
             config["model"],
             model_branch=model_branch,
             change_model_params=use_pretrain_script,
         )
     # update init_model or init_frz_model config if necessary
     if (init_model is not None or init_frz_model is not None) and use_pretrain_script:
         if init_model is not None:
             init_state_dict = paddle.load(init_model)
             if "model" in init_state_dict:
                 init_state_dict = init_state_dict["model"]
             config["model"] = init_state_dict["_extra_state"]["model_params"]
         else:
             raise NotImplementedError("init_frz_model is not supported yet")
 
     # argcheck
     config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
     config = normalize(config, multi_task=multi_task)
 
     # do neighbor stat
     min_nbor_dist = None
     if not skip_neighbor_stat:
         log.info(
             "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)"
         )
 
         if not multi_task:
             type_map = config["model"].get("type_map")
             train_data = get_data(
                 config["training"]["training_data"], 0, type_map, None
             )
             config["model"], min_nbor_dist = BaseModel.update_sel(
                 train_data, type_map, config["model"]
             )
         else:
             min_nbor_dist = {}
             for model_item in config["model"]["model_dict"]:
                 type_map = config["model"]["model_dict"][model_item].get("type_map")
                 train_data = get_data(
                     config["training"]["data_dict"][model_item]["training_data"],
                     0,
                     type_map,
                     None,
                 )
                 config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = (
                     BaseModel.update_sel(
                         train_data, type_map, config["model"]["model_dict"][model_item]
                     )
                 )
 
     with open(output, "w") as fp:
         json.dump(config, fp, indent=4)
 
     trainer = get_trainer(
         config,
         init_model,
         restart,
@@ -319,41 +319,49 @@
         finetune_links=finetune_links,
     )
     # save min_nbor_dist
     if min_nbor_dist is not None:
         if not multi_task:
-            trainer.model.min_nbor_dist = min_nbor_dist
+            trainer.model.min_nbor_dist = paddle.to_tensor(
+                min_nbor_dist,
+                dtype=paddle.float64,
+                place=DEVICE,
+            )
         else:
             for model_item in min_nbor_dist:
-                trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
+                trainer.model[model_item].min_nbor_dist = paddle.to_tensor(
+                    min_nbor_dist[model_item],
+                    dtype=paddle.float64,
+                    place=DEVICE,
+                )
     trainer.run()
 
 
 def freeze(
     model: str,
     output: str = "frozen_model.json",
     head: Optional[str] = None,
-):
+) -> None:
     paddle.set_flags(
         {
             "FLAGS_save_cf_stack_op": 1,
             "FLAGS_prim_enable_dynamic": 1,
             "FLAGS_enable_pir_api": 1,
         }
     )
     model = inference.Tester(model, head=head).model
     model.eval()
     from paddle.static import (
         InputSpec,
     )
 
     """
     ** coord [None, natoms, 3] paddle.float64
     ** atype [None, natoms] paddle.int64
     ** nlist [None, natoms, nnei] paddle.int32
     """
     # NOTE: 'FLAGS_save_cf_stack_op', 'FLAGS_prim_enable_dynamic' and
     # 'FLAGS_enable_pir_api' shoule be enabled when freezing model.
     jit_model = paddle.jit.to_static(
         model.forward_lower,
         full_graph=True,
         input_spec=[
@@ -362,14 +370,14 @@
             InputSpec([-1, -1, -1], dtype="int32", name="nlist"),
         ],
     )
     if output.endswith(".json"):
         output = output[:-5]
     paddle.jit.save(
         jit_model,
         path=output,
         skip_prune_program=True,
     )
     log.info(
         f"Paddle inference model has been exported to: {output}.json and {output}.pdiparams"
     )
 
@@ -383,27 +391,27 @@
     numb_batch: int = 0,
     model_branch: Optional[str] = None,
     output: Optional[str] = None,
-):
+) -> None:
     if input_file.endswith(".pd"):
         old_state_dict = paddle.load(input_file)
         model_state_dict = copy.deepcopy(old_state_dict.get("model", old_state_dict))
         model_params = model_state_dict["_extra_state"]["model_params"]
     else:
         raise RuntimeError(
             "Paddle now do not support change bias directly from a freezed model file"
             "Please provided a checkpoint file with a .pd extension"
         )
     multi_task = "model_dict" in model_params
     bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic"
     if multi_task:
         assert (
             model_branch is not None
         ), "For multitask model, the model branch must be set!"
         assert model_branch in model_params["model_dict"], (
             f"For multitask model, the model branch must be in the 'model_dict'! "
             f"Available options are : {list(model_params['model_dict'].keys())}."
         )
         log.info(f"Changing out bias for model {model_branch}.")
     model = training.get_model_for_wrapper(model_params)
     type_map = (
         model_params["type_map"]
@@ -415,7 +423,7 @@
         wrapper = ModelWrapper(model)
         wrapper.set_state_dict(old_state_dict["model"])
     else:
         raise NotImplementedError("Only support .pd file")
 
     if bias_value is not None:
         # use user-defined bias
@@ -468,7 +476,7 @@
     if not multi_task:
         model = updated_model
     else:
         model[model_branch] = updated_model
 
     if input_file.endswith(".pd"):
         output_path = (
@@ -479,18 +487,18 @@
             old_state_dict["model"] = wrapper.state_dict()
             old_state_dict["model"]["_extra_state"] = model_state_dict["_extra_state"]
         else:
             old_state_dict = wrapper.state_dict()
             old_state_dict["_extra_state"] = model_state_dict["_extra_state"]
         paddle.save(old_state_dict, output_path)
     else:
         raise NotImplementedError("Only support .pd file now")
 
     log.info(f"Saved model to {output_path}")
 
 
 def main(args: Optional[Union[list[str], argparse.Namespace]] = None):
     if not isinstance(args, argparse.Namespace):
         FLAGS = parse_args(args=args)
     else:
         FLAGS = args
 
@@ -503,7 +511,7 @@
     log.info("DeePMD version: %s", __version__)
 
     if FLAGS.command == "train":
         train(
             input_file=FLAGS.INPUT,
             init_model=FLAGS.init_model,
             restart=FLAGS.restart,
@@ -516,14 +524,14 @@
             output=FLAGS.output,
         )
     elif FLAGS.command == "freeze":
         if Path(FLAGS.checkpoint_folder).is_dir():
             checkpoint_path = Path(FLAGS.checkpoint_folder)
             latest_ckpt_file = (checkpoint_path / "checkpoint").read_text()
             FLAGS.model = str(checkpoint_path.joinpath(latest_ckpt_file))
         else:
             FLAGS.model = FLAGS.checkpoint_folder
         FLAGS.output = str(Path(FLAGS.output).with_suffix(".json"))
         freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
     elif FLAGS.command == "change-bias":
         change_bias(
             input_file=FLAGS.INPUT,
@@ -536,8 +544,8 @@
             output=FLAGS.output,
         )
     else:
         raise RuntimeError(f"Invalid command {FLAGS.command}!")
 
 
 if __name__ == "__main__":
     main()

deepmd/pd/loss/ener.py

@@ -10,7 +10,6 @@
     TaskLoss,
 )
 from deepmd.pd.utils import (
-    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -114,7 +113,7 @@
         self.enable_atom_ener_coeff = enable_atom_ener_coeff
         self.numb_generalized_coord = numb_generalized_coord
         if self.has_gf and self.numb_generalized_coord < 1:
             raise RuntimeError(
                 "When generalized force loss is used, the dimension of generalized coordinates should be larger than 0"
             )
         self.use_l1_all = use_l1_all
@@ -188,13 +187,13 @@
                 )
                 # more_loss['log_keys'].append('rmse_e')
             else:  # use l1 and for all atoms
                 l1_ener_loss = F.l1_loss(
                     energy_pred.reshape([-1]),
                     energy_label.reshape([-1]),
                     reduction="sum",
                 )
                 loss += pref_e * l1_ener_loss
                 more_loss["mae_e"] = self.display_if_exist(
                     F.l1_loss(
                         energy_pred.reshape([-1]),
                         energy_label.reshape([-1]),
@@ -204,10 +203,10 @@
                 )
                 # more_loss['log_keys'].append('rmse_e')
             if mae:
                 mae_e = paddle.mean(paddle.abs(energy_pred - energy_label)) * atom_norm
                 more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
                 mae_e_all = paddle.mean(paddle.abs(energy_pred - energy_label))
                 more_loss["mae_e_all"] = self.display_if_exist(
                     mae_e_all.detach(), find_energy
                 )
 
@@ -224,10 +223,7 @@
 
             if self.relative_f is not None:
                 force_label_3 = force_label.reshape([-1, 3])
-                # norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
-                norm_f = (
-                    decomp.norm(force_label_3, axis=1, keepdim=True) + self.relative_f
-                )
+                norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
                 diff_f_3 = diff_f.reshape([-1, 3])
                 diff_f_3 = diff_f_3 / norm_f
                 diff_f = diff_f_3.reshape([-1])
@@ -245,15 +241,15 @@
                         rmse_f.detach(), find_force
                     )
                 else:
                     l1_force_loss = F.l1_loss(force_label, force_pred, reduction="none")
                     more_loss["mae_f"] = self.display_if_exist(
                         l1_force_loss.mean().detach(), find_force
                     )
                     l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
                     loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
                 if mae:
                     mae_f = paddle.mean(paddle.abs(diff_f))
                     more_loss["mae_f"] = self.display_if_exist(
                         mae_f.detach(), find_force
                     )
 
@@ -322,8 +318,8 @@
             rmse_v = l2_virial_loss.sqrt() * atom_norm
             more_loss["rmse_v"] = self.display_if_exist(rmse_v.detach(), find_virial)
             if mae:
                 mae_v = paddle.mean(paddle.abs(diff_v)) * atom_norm
                 more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial)
 
         if self.has_ae and "atom_energy" in model_pred and "atom_ener" in label:
             atom_ener = model_pred["atom_energy"]
@@ -405,7 +401,7 @@
                 )
             )
         if self.has_gf > 0:
             label_requirement.append(
                 DataRequirementItem(
                     "drdq",
                     ndof=self.numb_generalized_coord * 3,
@@ -415,7 +411,7 @@
                 )
             )
         if self.enable_atom_ener_coeff:
             label_requirement.append(
                 DataRequirementItem(
                     "atom_ener_coeff",
                     ndof=1,

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import functools
 import logging
 from typing import (
@@ -52,7 +51,7 @@
         fitting,
         type_map: list[str],
         **kwargs,
-    ):
+    ) -> None:
         super().__init__(type_map, **kwargs)
         ntypes = len(type_map)
         self.type_map = type_map
@@ -116,12 +115,12 @@
 
     def set_eval_descriptor_hook(self, enable: bool) -> None:
         """Set the hook for evaluating descriptor and clear the cache for descriptor list."""
         self.enable_eval_descriptor_hook = enable
         self.eval_descriptor_list = []
 
     def eval_descriptor(self) -> paddle.Tensor:
         """Evaluate the descriptor."""
         return paddle.concat(self.eval_descriptor_list)
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
@@ -144,7 +143,7 @@
         Set the case embedding of this atomic model by the given case_idx,
         typically concatenated with the output of the descriptor and fed into the fitting net.
         """
         self.fitting_net.set_case_embd(case_idx)
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -179,7 +178,7 @@
 
     def has_message_passing(self) -> bool:
         """Returns whether the atomic model has message passing."""
         return self.descriptor.has_message_passing()
 
     def need_sorted_nlist_for_lower(self) -> bool:
         """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
@@ -201,7 +200,7 @@
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)
@@ -212,6 +211,37 @@
         obj = super().deserialize(data)
         return obj
 
+    def enable_compression(
+        self,
+        min_nbor_dist: float,
+        table_extrapolate: float = 5,
+        table_stride_1: float = 0.01,
+        table_stride_2: float = 0.1,
+        check_frequency: int = -1,
+    ) -> None:
+        """Call descriptor enable_compression().
+
+        Parameters
+        ----------
+        min_nbor_dist
+            The nearest distance between atoms
+        table_extrapolate
+            The scale of model extrapolation
+        table_stride_1
+            The uniform stride of the first table
+        table_stride_2
+            The uniform stride of the second table
+        check_frequency
+            The overflow check frequency
+        """
+        self.descriptor.enable_compression(
+            min_nbor_dist,
+            table_extrapolate,
+            table_stride_1,
+            table_stride_2,
+            check_frequency,
+        )
+
     def forward_atomic(
         self,
         extended_coord,
@@ -258,7 +288,7 @@
         )
         assert descriptor is not None
         if self.enable_eval_descriptor_hook:
             self.eval_descriptor_list.append(descriptor)
         # energy, force
         fit_ret = self.fitting_net(
             descriptor,
@@ -278,7 +308,7 @@
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-    ):
+    ) -> None:
         """
         Compute or load the statistics parameters of the model,
         such as mean and standard deviation of descriptors or the energy bias of the fitting net.
@@ -297,19 +327,19 @@
         if stat_file_path is not None and self.type_map is not None:
             # descriptors and fitting net with different type_map
             # should not share the same parameters
             stat_file_path /= " ".join(self.type_map)
 
         @functools.lru_cache
         def wrapped_sampler():
             sampled = sampled_func()
             if self.pair_excl is not None:
                 pair_exclude_types = self.pair_excl.get_exclude_types()
                 for sample in sampled:
                     sample["pair_exclude_types"] = list(pair_exclude_types)
             if self.atom_excl is not None:
                 atom_exclude_types = self.atom_excl.get_exclude_types()
                 for sample in sampled:
                     sample["atom_exclude_types"] = list(atom_exclude_types)
             return sampled
 
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
@@ -330,11 +360,11 @@
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
         return self.fitting_net.get_sel_type()
 
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 
         If False, the shape is (nframes, nloc, ndim).
         """
         return False

deepmd/pd/model/descriptor/__init__.py

@@ -5,18 +5,38 @@
 from .descriptor import (
     DescriptorBlock,
 )
+from .dpa1 import (
+    DescrptBlockSeAtten,
+    DescrptDPA1,
+)
+from .dpa2 import (
+    DescrptDPA2,
+)
 from .env_mat import (
     prod_env_mat,
 )
+from .repformers import (
+    DescrptBlockRepformers,
+)
 from .se_a import (
     DescrptBlockSeA,
     DescrptSeA,
 )
+from .se_t_tebd import (
+    DescrptBlockSeTTebd,
+    DescrptSeTTebd,
+)
 
 __all__ = [
     "BaseDescriptor",
     "DescriptorBlock",
+    "DescrptBlockRepformers",
     "DescrptBlockSeA",
+    "DescrptBlockSeAtten",
+    "DescrptBlockSeTTebd",
+    "DescrptDPA1",
+    "DescrptDPA2",
     "DescrptSeA",
+    "DescrptSeTTebd",
     "prod_env_mat",
 ]