[r3.0] release v3.0.3 (#4759)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added a new sigmoid activation function compatible with both JAX and
other array types.
- Introduced a JAX configuration to set default matrix multiplication
precision for improved performance.

- **Bug Fixes**
- Enhanced detection of CUDA out-of-memory errors for more robust error
handling.
- Improved validation and error messages for mismatched or incomplete
type maps in data handling.
- Updated logic for stress calculation to better handle cases where cell
information is missing.
- Improved handling of multi-dimensional atom type arrays and model type
checks in TensorFlow evaluation.

- **Chores**
- Upgraded PyTorch and LibTorch versions to 2.7.0 and updated related
download and installation scripts.
- Updated ROCm repository version and CMake minimum version
requirements.
- Improved device placement for tensors in PyTorch to ensure
compatibility.
- Added or updated teardown methods in multiple test classes for better
resource management and cleanup.
- Expanded and refined test coverage for data type mapping and error
scenarios.
  - Updated expected values and test logic for TensorFlow-related tests.
- Updated macOS deployment target to version 11.0 for build environment.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: njzjz

.devcontainer/download_libtorch.sh

@@ -4,5 +4,5 @@ set -ev
 SCRIPT_PATH=$(dirname $(realpath -s $0))
 cd ${SCRIPT_PATH}/..
 
-wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcpu.zip -O ~/libtorch.zip
+wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.7.0%2Bcpu.zip -O ~/libtorch.zip
 unzip ~/libtorch.zip

.github/workflows/build_cc.yml

@@ -52,7 +52,7 @@ jobs:
       env:
         DEBIAN_FRONTEND: noninteractive
     - run: |
-         echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/5.3/ jammy main' | sudo tee /etc/apt/sources.list.d/rocm.list \
+         echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/6.3/ jammy main' | sudo tee /etc/apt/sources.list.d/rocm.list \
          && printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600 \
          && curl -s https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add - \
          && sudo apt-get update \

.github/workflows/test_cc.yml

@@ -36,7 +36,7 @@ jobs:
       run: source/tests/infer/convert-models.sh
     - name: Download libtorch
       run: |
-         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.7.0%2Bcpu.zip -O libtorch.zip
          unzip libtorch.zip
     # https://github.com/actions/runner-images/issues/9491
     - name: Fix kernel mmap rnd bits

.github/workflows/test_cuda.yml

@@ -47,7 +47,7 @@ jobs:
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
     - run: python -m pip install -U uv
-    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.6.0" "jax[cuda12]==0.5.0"
+    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.7.0" "jax[cuda12]==0.5.0"
     - 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)')
@@ -67,7 +67,7 @@ jobs:
       run: source/tests/infer/convert-models.sh
     - name: Download libtorch
       run: |
-         wget https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu124.zip -O libtorch.zip
+         wget https://download.pytorch.org/libtorch/cu126/libtorch-cxx11-abi-shared-with-deps-2.7.0%2Bcu126.zip -O libtorch.zip
          unzip libtorch.zip
     - run: |
         export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch

.github/workflows/test_python.yml

@@ -30,6 +30,7 @@ jobs:
         export TENSORFLOW_ROOT=$(python -c 'import tensorflow;print(tensorflow.__path__[0])')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.5.0;python_version>='3.10'"
+        source/install/uv_with_retry.sh pip install --system -U setuptools
         source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
       env:
         # Please note that uv has some issues with finding
@@ -42,6 +43,8 @@ jobs:
         HOROVOD_WITH_TENSORFLOW: 1
         HOROVOD_WITHOUT_PYTORCH: 1
         HOROVOD_WITH_MPI: 1
+        # https://cmake.org/cmake/help/latest/variable/CMAKE_POLICY_VERSION_MINIMUM.html
+        CMAKE_POLICY_VERSION_MINIMUM: 3.5
     - run: dp --version
     - name: Get durations from cache
       uses: actions/cache@v4

backend/find_pytorch.py

@@ -116,7 +116,7 @@ def get_pt_requirement(pt_version: str = "") -> dict:
         cuda_version = os.environ.get("CUDA_VERSION", "12.2")
         if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
             # CUDA 12.2, cudnn 9
-            pt_version = "2.6.0"
+            pt_version = "2.7.0"
         elif cuda_version in SpecifierSet(">=11,<12"):
             # CUDA 11.8, cudnn 8
             pt_version = "2.3.1"

deepmd/calculator.py

@@ -138,12 +138,13 @@ def calculate(
         self.results["virial"] = v[0].reshape(3, 3)
 
         # convert virial into stress for lattice relaxation
-        if "stress" in properties:
-            if sum(atoms.get_pbc()) > 0:
-                # the usual convention (tensile stress is positive)
-                # stress = -virial / volume
-                stress = -0.5 * (v[0].copy() + v[0].copy().T) / atoms.get_volume()
-                # Voigt notation
-                self.results["stress"] = stress.flat[[0, 4, 8, 5, 2, 1]]
-            else:
-                raise PropertyNotImplementedError
+        if cell is not None:
+            # the usual convention (tensile stress is positive)
+            # stress = -virial / volume
+            stress = -0.5 * (v[0].copy() + v[0].copy().T) / atoms.get_volume()
+            # Voigt notation
+            self.results["stress"] = stress.flat[[0, 4, 8, 5, 2, 1]]
+        elif "stress" in properties:
+            raise PropertyNotImplementedError
+        else:
+            pass

deepmd/dpmodel/utils/network.py

@@ -34,6 +34,19 @@
 )
 
 
+def sigmoid_t(x: np.ndarray) -> np.ndarray:
+    """Sigmoid."""
+    if array_api_compat.is_jax_array(x):
+        from deepmd.jax.env import (
+            jax,
+        )
+
+        # see https://github.com/jax-ml/jax/discussions/15617
+        return jax.nn.sigmoid(x)
+    xp = array_api_compat.array_namespace(x)
+    return 1 / (1 + xp.exp(-x))
+
+
 class Identity(NativeOP):
     def __init__(self) -> None:
         super().__init__()
@@ -313,9 +326,8 @@ def fn(x):
     elif activation_function == "sigmoid":
 
         def fn(x):
-            xp = array_api_compat.array_namespace(x)
             # generated by GitHub Copilot
-            return 1 / (1 + xp.exp(-x))
+            return sigmoid_t(x)
 
         return fn
     elif activation_function.lower() in ("none", "linear"):

deepmd/jax/env.py

@@ -12,6 +12,8 @@
 
 jax.config.update("jax_enable_x64", True)
 # jax.config.update("jax_debug_nans", True)
+# https://github.com/jax-ml/jax/issues/24909
+jax.config.update("jax_default_matmul_precision", "tensorfloat32")
 
 if os.environ.get("DP_DTYPE_PROMOTION_STRICT") == "1":
     jax.config.update("jax_numpy_dtype_promotion", "strict")

deepmd/jax/model/base_model.py

@@ -47,7 +47,7 @@ def forward_common_atomic(
             kk_redu = get_reduce_name(kk)
             model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis)
             kk_derv_r, kk_derv_c = get_deriv_name(kk)
-            if vdef.c_differentiable:
+            if vdef.r_differentiable:
 
                 def eval_output(
                     cc_ext,