cc: add torch backend support for Multiple backend (#3162)

need to test in union environment (tf and pt) see #3119 --------- Signed-off-by: Lysithea <[email protected]> Signed-off-by: Jinzhe Zeng <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jinzhe Zeng <[email protected]> Co-authored-by: Han Wang <[email protected]>
deepmodeling · Feb 28, 2024 · d377ccb · d377ccb
1 parent fd17e2e
commit d377ccb
Show file tree

Hide file tree

Showing 13 changed files with 2,201 additions and 11 deletions.
diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml
@@ -56,7 +56,7 @@ jobs:
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
-        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib
+        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib
       if: ${{ !matrix.check_memleak }}
     # test ipi
     - run: pytest --cov=deepmd source/ipi/tests
@@ -65,7 +65,7 @@ jobs:
         TF_INTRA_OP_PARALLELISM_THREADS: 1
         TF_INTER_OP_PARALLELISM_THREADS: 1
         PATH: ${{ github.workspace }}/dp_test/bin:$PATH
-        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib
+        LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib:${{ github.workspace }}/libtorch/lib
       if: ${{ !matrix.check_memleak }}
     - uses: codecov/codecov-action@v4
       env:

diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
@@ -38,6 +38,8 @@ jobs:
       with:
         useLocalCache: true
         useCloudCache: false
+    - name: Install wget and unzip
+      run: apt-get update && apt-get install -y wget unzip
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
@@ -53,7 +55,13 @@ jobs:
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
     - run: dp --version
     - run: python -m pytest source/tests --durations=0
-    - run: source/install/test_cc_local.sh
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.1%2Bcu121.zip -O libtorch.zip
+         unzip libtorch.zip
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1
@@ -63,7 +71,7 @@ jobs:
         DP_VARIANT: cuda
         DP_USE_MPICH2: 1
     - run: |
-        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$GITHUB_WORKSPACE/libtorch/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
         export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
         python -m pytest source/lmp/tests
         python -m pytest source/ipi/tests

diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
@@ -17,6 +17,9 @@
     communicate_extended_output,
     fit_output_to_model_output,
 )
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.pt.utils.nlist import (
     extend_input_and_build_neighbor_list,
     nlist_distinguish_types,
@@ -115,6 +118,9 @@ def forward_common(
                 The keys are defined by the `ModelOutputDef`.
 
             """
+            coord = coord.to(env.GLOBAL_PT_FLOAT_PRECISION)
+            if box is not None:
+                box = box.to(env.GLOBAL_PT_FLOAT_PRECISION)
             (
                 extended_coord,
                 extended_atype,
@@ -183,6 +189,7 @@ def forward_common_lower(
                 the result dict, defined by the `FittingOutputDef`.
 
             """
+            extended_coord = extended_coord.to(env.GLOBAL_PT_FLOAT_PRECISION)
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
             nlist = self.format_nlist(extended_coord, extended_atype, nlist)