Merge branch 'add_paddle_backend_core_and_water_se_e2_a' of https://g…

…ithub.com/HydrogenSulfate/deepmd-kit into add_paddle_backend_core_and_water_se_e2_a
deepmodeling · Nov 14, 2024 · 8cdef03 · 8cdef03
2 parents 6e6dedc + 1f108bd
commit 8cdef03
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 202 deletions.
diff --git 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
@@ -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
@@ -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
diff --git 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__ = [