diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index d36c1da026..db0dfb6c61 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -51,8 +51,8 @@ jobs: - run: | export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])') export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') - source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py source/install/uv_with_retry.sh pip install --system --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu123/ + source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py env: DP_VARIANT: cuda DP_ENABLE_NATIVE_OPTIMIZATION: 1 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ae1a0c186f..1c0d197edf 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -65,13 +65,13 @@ repos: - id: clang-format exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$) # markdown, yaml, CSS, javascript - - repo: https://github.com/pre-commit/mirrors-prettier - rev: v4.0.0-alpha.8 - hooks: - - id: prettier - types_or: [markdown, yaml, css] - # workflow files cannot be modified by pre-commit.ci - exclude: ^(source/3rdparty|\.github/workflows|\.clang-format) + # - repo: https://github.com/pre-commit/mirrors-prettier + # rev: v4.0.0-alpha.8 + # hooks: + # - id: prettier + # types_or: [markdown, yaml, css] + # # workflow files cannot be modified by pre-commit.ci + # exclude: ^(source/3rdparty|\.github/workflows|\.clang-format) # Shell - repo: https://github.com/scop/pre-commit-shfmt rev: v3.10.0-1 @@ -83,25 +83,25 @@ repos: hooks: - id: cmake-format #- id: cmake-lint - - repo: https://github.com/njzjz/mirrors-bibtex-tidy - rev: v1.13.0 - hooks: - - id: bibtex-tidy - args: - - --curly - - --numeric - - --align=13 - - --blank-lines - # disable sort: the order of keys and fields has explict meanings - #- --sort=key - - --duplicates=key,doi,citation,abstract - - --merge=combine - #- --sort-fields - #- --strip-comments - - --trailing-commas - - --encode-urls - - --remove-empty-fields - - --wrap=80 + # - repo: https://github.com/njzjz/mirrors-bibtex-tidy + # rev: v1.13.0 + # hooks: + # - id: bibtex-tidy + # args: + # - --curly + # - --numeric + # - --align=13 + # - --blank-lines + # # disable sort: the order of keys and fields has explict meanings + # #- --sort=key + # - --duplicates=key,doi,citation,abstract + # - --merge=combine + # #- --sort-fields + # #- --strip-comments + # - --trailing-commas + # - --encode-urls + # - --remove-empty-fields + # - --wrap=80 # license header - repo: https://github.com/Lucas-C/pre-commit-hooks rev: v1.5.5 diff --git a/deepmd/pd/__init__.py b/deepmd/pd/__init__.py new file mode 100644 index 0000000000..2f4814885c --- /dev/null +++ b/deepmd/pd/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +# import customized OPs globally +from deepmd.pd.cxx_op import ( + ENABLE_CUSTOMIZED_OP, +) +from deepmd.utils.entry_point import ( + load_entry_point, +) + +load_entry_point("deepmd.pd") + +__all__ = [ + "ENABLE_CUSTOMIZED_OP", +] diff --git a/deepmd/pd/model/network/layernorm.py b/deepmd/pd/model/network/layernorm.py deleted file mode 100644 index fdf2433ed2..0000000000 --- a/deepmd/pd/model/network/layernorm.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) - -import numpy as np -import paddle -import paddle.nn as nn - -from deepmd.dpmodel.utils.network import LayerNorm as DPLayerNorm -from deepmd.pd.model.network.init import ( - normal_, - ones_, - zeros_, -) -from deepmd.pd.utils import ( - env, -) -from deepmd.pd.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, -) -from deepmd.pd.utils.utils import ( - get_generator, - to_numpy_array, - to_paddle_tensor, -) - -device = env.DEVICE - - -def empty_t(shape, precision): - return paddle.empty(shape, dtype=precision).to(device=device) - - -class LayerNorm(nn.Layer): - def __init__( - self, - num_in, - eps: float = 1e-5, - uni_init: bool = True, - bavg: float = 0.0, - stddev: float = 1.0, - precision: str = DEFAULT_PRECISION, - trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - ): - super().__init__() - self.eps = eps - self.uni_init = uni_init - self.num_in = num_in - self.precision = precision - self.prec = PRECISION_DICT[self.precision] - self.matrix = self.create_parameter( - shape=[num_in], - dtype=self.prec, - default_initializer=nn.initializer.Assign( - empty_t([num_in], self.prec), - ), - ) - self.bias = self.create_parameter( - shape=[num_in], - dtype=self.prec, - default_initializer=nn.initializer.Assign(empty_t([num_in], self.prec)), - ) - random_generator = get_generator(seed) - if self.uni_init: - ones_(self.matrix.data) - zeros_(self.bias.data) - else: - normal_(self.bias.data, mean=bavg, std=stddev, generator=random_generator) - normal_( - self.matrix.data, - std=stddev / np.sqrt(self.num_in), - generator=random_generator, - ) - self.trainable = trainable - if not self.trainable: - self.matrix.stop_gradient = True - self.bias.stop_gradient = True - - def dim_out(self) -> int: - return self.matrix.shape[0] - - def forward( - self, - xx: paddle.Tensor, - ) -> paddle.Tensor: - """One Layer Norm used by DP model. - - Parameters - ---------- - xx : paddle.Tensor - The input of index. - - Returns - ------- - yy: paddle.Tensor - The output. - """ - # mean = xx.mean(dim=-1, keepdim=True) - # variance = xx.var(dim=-1, unbiased=False, keepdim=True) - # The following operation is the same as above, but will not raise error when using jit model to inference. - # See https://github.com/pytorch/pytorch/issues/85792 - if xx.numel() > 0: - variance, mean = ( - paddle.var(xx, axis=-1, unbiased=False, keepdim=True), - paddle.mean(xx, axis=-1, keepdim=True), - ) - yy = (xx - mean) / paddle.sqrt(variance + self.eps) - else: - yy = xx - if self.matrix is not None and self.bias is not None: - yy = yy * self.matrix + self.bias - return yy - - def serialize(self) -> dict: - """Serialize the layer to a dict. - - Returns - ------- - dict - The serialized layer. - """ - nl = DPLayerNorm( - self.matrix.shape[0], - eps=self.eps, - trainable=self.trainable, - precision=self.precision, - ) - nl.w = to_numpy_array(self.matrix) - nl.b = to_numpy_array(self.bias) - data = nl.serialize() - return data - - @classmethod - def deserialize(cls, data: dict) -> "LayerNorm": - """Deserialize the layer from a dict. - - Parameters - ---------- - data : dict - The dict to deserialize from. - """ - nl = DPLayerNorm.deserialize(data) - obj = cls( - nl["matrix"].shape[0], - eps=nl["eps"], - trainable=nl["trainable"], - precision=nl["precision"], - ) - prec = PRECISION_DICT[obj.precision] - - def check_load_param(ss): - if nl[ss] is not None: - tensor = to_paddle_tensor(nl[ss]) - return paddle.create_parameter( - tensor.shape, - dtype=tensor.dtype, - default_initializer=nn.initializer.Assign(tensor), - ) - return None - - obj.matrix = check_load_param("matrix") - obj.bias = check_load_param("bias") - return obj diff --git a/deepmd/pd/utils/env.py b/deepmd/pd/utils/env.py index 3faf6f22e0..4c104db374 100644 --- a/deepmd/pd/utils/env.py +++ b/deepmd/pd/utils/env.py @@ -78,15 +78,14 @@ def enable_prim(enable: bool = True): """Enable running program in primitive C++ API in eager/static mode.""" - if enable: - from paddle.framework import ( - core, - ) + from paddle.framework import ( + core, + ) - core.set_prim_eager_enabled(True) - core._set_prim_all_enabled(True) - log = logging.getLogger(__name__) - log.info("Enable prim in eager and static mode.") + core.set_prim_eager_enabled(enable) + core._set_prim_all_enabled(enable) + log = logging.getLogger(__name__) + log.info(f"{'Enable' if enable else 'Disable'} prim in eager and static mode.") __all__ = [