Merge branch 'devel' into econf-embd

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.3.3
+    rev: v0.3.4
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -52,7 +52,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.1
+    rev: v18.1.2
     hooks:
       - id: clang-format
         exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc

backend/dynamic_metadata.py

@@ -77,7 +77,7 @@ def dynamic_metadata(
                 "nvidia-curand-cu11",
                 "nvidia-cusolver-cu11",
                 "nvidia-cusparse-cu11",
-                "nvidia-cudnn-cu11",
+                "nvidia-cudnn-cu11<9",
                 "nvidia-cuda-nvcc-cu11",
             ],
             "cu12": [
@@ -87,7 +87,7 @@ def dynamic_metadata(
                 "nvidia-curand-cu12",
                 "nvidia-cusolver-cu12",
                 "nvidia-cusparse-cu12",
-                "nvidia-cudnn-cu12",
+                "nvidia-cudnn-cu12<9",
                 "nvidia-cuda-nvcc-cu12",
             ],
             "torch": [

deepmd/common.py

@@ -17,8 +17,10 @@
     Dict,
     List,
     Optional,
+    Set,
     TypeVar,
     Union,
+    get_args,
 )
 
 try:
@@ -45,23 +47,28 @@
     "j_loader",
     "expand_sys_str",
     "get_np_precision",
+    "VALID_PRECISION",
+    "VALID_ACTIVATION",
 ]
 
+_PRECISION = Literal["default", "float16", "float32", "float64"]
+_ACTIVATION = Literal[
+    "relu",
+    "relu6",
+    "softplus",
+    "sigmoid",
+    "tanh",
+    "gelu",
+    "gelu_tf",
+    "none",
+    "linear",
+]
+# get_args is new in py38
+VALID_PRECISION: Set[_PRECISION] = set(get_args(_PRECISION))
+VALID_ACTIVATION: Set[_ACTIVATION] = set(get_args(_ACTIVATION))
 
 if TYPE_CHECKING:
     _DICT_VAL = TypeVar("_DICT_VAL")
-    _PRECISION = Literal["default", "float16", "float32", "float64"]
-    _ACTIVATION = Literal[
-        "relu",
-        "relu6",
-        "softplus",
-        "sigmoid",
-        "tanh",
-        "gelu",
-        "gelu_tf",
-        "none",
-        "linear",
-    ]
     __all__.extend(
         [
             "_DICT_VAL",

deepmd/dpmodel/common.py

@@ -4,8 +4,12 @@
     abstractmethod,
 )
 
+import ml_dtypes
 import numpy as np
 
+from deepmd.common import (
+    VALID_PRECISION,
+)
 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
@@ -21,14 +25,22 @@
     "int32": np.int32,
     "int64": np.int64,
     "default": GLOBAL_NP_FLOAT_PRECISION,
+    # NumPy doesn't have bfloat16 (and does't plan to add)
+    # ml_dtypes is a solution, but it seems not supporting np.save/np.load
+    # hdf5 hasn't supported bfloat16 as well (see https://forum.hdfgroup.org/t/11975)
+    "bfloat16": ml_dtypes.bfloat16,
 }
+assert VALID_PRECISION.issubset(PRECISION_DICT.keys())
+
 RESERVED_PRECISON_DICT = {
     np.float16: "float16",
     np.float32: "float32",
     np.float64: "float64",
     np.int32: "int32",
     np.int64: "int64",
+    ml_dtypes.bfloat16: "bfloat16",
 }
+assert set(RESERVED_PRECISON_DICT.keys()) == set(PRECISION_DICT.values())
 DEFAULT_PRECISION = "float64"
 
 

deepmd/infer/model_devi.py

@@ -303,7 +303,7 @@ def calc_model_devi(
     energies = []
     forces = []
     virials = []
-    natom = atype.shape[-1]
+    natom = np.array(atype).shape[-1]
     for dp in models:
         ret = dp.eval(
             coord,

deepmd/pt/model/network/mlp.py

@@ -30,6 +30,8 @@
 )
 from deepmd.pt.utils.utils import (
     ActivationFn,
+    to_numpy_array,
+    to_torch_tensor,
 )
 
 try:
@@ -151,9 +153,9 @@ def serialize(self) -> dict:
             precision=self.precision,
         )
         nl.w, nl.b, nl.idt = (
-            self.matrix.detach().cpu().numpy(),
-            self.bias.detach().cpu().numpy() if self.bias is not None else None,
-            self.idt.detach().cpu().numpy() if self.idt is not None else None,
+            to_numpy_array(self.matrix),
+            to_numpy_array(self.bias),
+            to_numpy_array(self.idt),
         )
         return nl.serialize()
 
@@ -180,7 +182,7 @@ def deserialize(cls, data: dict) -> "MLPLayer":
 
         def check_load_param(ss):
             return (
-                nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device))
+                nn.Parameter(data=to_torch_tensor(nl[ss]))
                 if nl[ss] is not None
                 else None
             )

deepmd/pt/model/task/polarizability.py

@@ -230,12 +230,12 @@ def compute_output_stats(
 
                 sys_constant_matrix = []
                 for sys in range(len(sampled)):
-                    nframs = sampled[sys]["type"].shape[0]
+                    nframs = sampled[sys]["atype"].shape[0]
 
                     if sampled[sys]["find_atomic_polarizability"] > 0.0:
                         sys_atom_polar = compute_stats_from_atomic(
                             sampled[sys]["atomic_polarizability"].numpy(force=True),
-                            sampled[sys]["type"].numpy(force=True),
+                            sampled[sys]["atype"].numpy(force=True),
                         )[0]
                     else:
                         if not sampled[sys]["find_polarizability"] > 0.0:
@@ -244,7 +244,7 @@ def compute_output_stats(
                             (nframs, self.ntypes), dtype=env.GLOBAL_NP_FLOAT_PRECISION
                         )
                         for itype in range(self.ntypes):
-                            type_mask = sampled[sys]["type"] == itype
+                            type_mask = sampled[sys]["atype"] == itype
                             sys_type_count[:, itype] = type_mask.sum(dim=1).numpy(
                                 force=True
                             )

deepmd/pt/utils/env.py

@@ -4,6 +4,9 @@
 import numpy as np
 import torch
 
+from deepmd.common import (
+    VALID_PRECISION,
+)
 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
@@ -40,20 +43,24 @@
     "double": torch.float64,
     "int32": torch.int32,
     "int64": torch.int64,
+    "bfloat16": torch.bfloat16,
 }
 GLOBAL_PT_FLOAT_PRECISION = PRECISION_DICT[np.dtype(GLOBAL_NP_FLOAT_PRECISION).name]
 GLOBAL_PT_ENER_FLOAT_PRECISION = PRECISION_DICT[
     np.dtype(GLOBAL_ENER_FLOAT_PRECISION).name
 ]
 PRECISION_DICT["default"] = GLOBAL_PT_FLOAT_PRECISION
+assert VALID_PRECISION.issubset(PRECISION_DICT.keys())
 # cannot automatically generated
 RESERVED_PRECISON_DICT = {
     torch.float16: "float16",
     torch.float32: "float32",
     torch.float64: "float64",
     torch.int32: "int32",
     torch.int64: "int64",
+    torch.bfloat16: "bfloat16",
 }
+assert set(PRECISION_DICT.values()) == set(RESERVED_PRECISON_DICT.keys())
 DEFAULT_PRECISION = "float64"
 
 # throw warnings if threads not set

deepmd/pt/utils/finetune.py

@@ -146,7 +146,9 @@ def change_finetune_model_params(finetune_model, model_config, model_branch=""):
             model_branch_from=model_branch,
         )
         finetune_links["Default"] = (
-            model_branch if finetune_from_multi_task else "Default"
+            model_config["model_branch_chosen"]
+            if finetune_from_multi_task
+            else "Default"
         )
     else:
         assert model_branch == "", (

deepmd/pt/utils/stat.py

@@ -61,12 +61,17 @@ def make_stat_input(datasets, dataloaders, nbatches):
                         if dd not in sys_stat:
                             sys_stat[dd] = []
                         sys_stat[dd].append(stat_data[dd])
+                    elif isinstance(stat_data[dd], np.float32):
+                        sys_stat[dd] = stat_data[dd]
                     else:
                         pass
+
         for key in sys_stat:
-            if sys_stat[key] is None or sys_stat[key][0] is None:
+            if isinstance(sys_stat[key], np.float32):
+                pass
+            elif sys_stat[key] is None or sys_stat[key][0] is None:
                 sys_stat[key] = None
-            else:
+            elif isinstance(stat_data[dd], torch.Tensor):
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
         dict_to_device(sys_stat)
         lst.append(sys_stat)

deepmd/pt/utils/utils.py

@@ -5,6 +5,7 @@
     overload,
 )
 
+import ml_dtypes
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -85,6 +86,9 @@ def to_numpy_array(
     prec = NP_PRECISION_DICT.get(prec, None)
     if prec is None:
         raise ValueError(f"unknown precision {xx.dtype}")
+    if xx.dtype == torch.bfloat16:
+        # https://github.com/pytorch/pytorch/issues/109873
+        xx = xx.float()
     return xx.detach().cpu().numpy().astype(prec)
 
 
@@ -109,13 +113,20 @@ def to_torch_tensor(
     prec = PT_PRECISION_DICT.get(prec, None)
     if prec is None:
         raise ValueError(f"unknown precision {xx.dtype}")
+    if xx.dtype == ml_dtypes.bfloat16:
+        # https://github.com/pytorch/pytorch/issues/109873
+        xx = xx.astype(np.float32)
     return torch.tensor(xx, dtype=prec, device=DEVICE)
 
 
 def dict_to_device(sample_dict):
     for key in sample_dict:
         if isinstance(sample_dict[key], list):
             sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]]
+        if isinstance(sample_dict[key], np.float32):
+            sample_dict[key] = (
+                torch.ones(1, dtype=torch.float32, device=DEVICE) * sample_dict[key]
+            )
         else:
             if sample_dict[key] is not None:
                 sample_dict[key] = sample_dict[key].to(DEVICE)

deepmd/tf/common.py

@@ -18,6 +18,8 @@
 )
 
 from deepmd.common import (
+    VALID_ACTIVATION,
+    VALID_PRECISION,
     add_data_requirement,
     data_requirement,
     expand_sys_str,
@@ -69,6 +71,7 @@
     "float64": tf.float64,
     "bfloat16": tf.bfloat16,
 }
+assert VALID_PRECISION.issubset(PRECISION_DICT.keys())
 
 
 def gelu(x: tf.Tensor) -> tf.Tensor:
@@ -138,6 +141,7 @@ def gelu_wrapper(x):
     "linear": lambda x: x,
     "none": lambda x: x,
 }
+assert VALID_ACTIVATION.issubset(ACTIVATION_FN_DICT.keys())
 
 
 def get_activation_func(

deepmd/tf/fit/polar.py

@@ -208,7 +208,9 @@ def compute_output_stats(self, all_stat):
 
                         polar_bias.append(
                             np.sum(
-                                all_stat["atomic_polarizability"][ss][:, index_lis, :]
+                                all_stat["atomic_polarizability"][ss].reshape(
+                                    nframes, len(atom_has_polar), -1
+                                )[:, index_lis, :]
                                 / nframes,
                                 axis=(0, 1),
                             ).reshape((1, 9))

deepmd/tf/model/tensor.py

@@ -86,7 +86,7 @@ def data_stat(self, data):
         all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False)
         m_all_stat = merge_sys_stat(all_stat)
         self._compute_input_stat(m_all_stat, protection=self.data_stat_protect)
-        self._compute_output_stat(all_stat)
+        self._compute_output_stat(m_all_stat)
 
     def _compute_input_stat(self, all_stat, protection=1e-2):
         self.descrpt.compute_input_stats(