diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
index 25dded26aa..0ac29a7d9b 100644
--- a/.github/workflows/test_python.yml
+++ b/.github/workflows/test_python.yml
@@ -5,48 +5,43 @@ name: Test Python
 jobs:
   testpython:
     name: Test Python
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         include:
           - python: 3.7
-            gcc: 5
-            tf: 1.14
-          - python: 3.7
-            gcc: 6
-            tf: 1.14
-          - python: 3.7
-            gcc: 7
-            tf: 1.14
-          - python: 3.7
-            gcc: 8
             tf: 1.14
           - python: 3.8
-            gcc: 8
             tf:
           - python: "3.11"
-            gcc: 8
             tf:
 
-    container: ghcr.io/deepmodeling/deepmd-kit-test-environment:py${{ matrix.python }}-gcc${{ matrix.gcc }}-tf${{ matrix.tf }}
     steps:
-    - name: work around permission issue
-      run: git config --global --add safe.directory /__w/deepmd-kit/deepmd-kit
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python }}
+        cache: 'pip'
+    - uses: mpi4py/setup-mpi@v1
+      if: ${{ matrix.tf == '' }}
+      with:
+        mpi: openmpi
     # https://github.com/pypa/pip/issues/11770
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
     - run: pip install -e .[cpu,test]
       env:
-        CC: gcc-${{ matrix.gcc }}
-        CXX: g++-${{ matrix.gcc }}
         TENSORFLOW_VERSION: ${{ matrix.tf }}
         DP_BUILD_TESTING: 1
+    - run: pip install horovod mpi4py
+      if: ${{ matrix.tf == '' }}
+      env:
+        HOROVOD_WITH_TENSORFLOW: 1
+        HOROVOD_WITHOUT_GLOO: 1
     - run: dp --version
     - run: pytest --cov=deepmd --cov=deepmd_cli source/tests --durations=0
     - uses: codecov/codecov-action@v3
       with:
         gcov: true
-        gcov_executable: gcov-${{ matrix.gcc }}
   pass:
     name: Pass testing Python
     needs: [testpython]
diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py
index 8fe3cedb63..aa75d5ecb4 100644
--- a/backend/find_tensorflow.py
+++ b/backend/find_tensorflow.py
@@ -112,16 +112,31 @@ def get_tf_requirement(tf_version: str = "") -> dict:
     if tf_version == "":
         tf_version = os.environ.get("TENSORFLOW_VERSION", "")
 
+    extra_requires = []
+    extra_select = {}
+    if not (tf_version == "" or tf_version in SpecifierSet(">=2.12")):
+        extra_requires.append("protobuf<3.20")
+    if tf_version == "" or tf_version in SpecifierSet(">=1.15"):
+        extra_select["mpi"] = [
+            "horovod",
+            "mpi4py",
+        ]
+    else:
+        extra_select["mpi"] = []
+
     if tf_version == "":
         return {
             "cpu": [
                 "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
+                *extra_requires,
             ],
             "gpu": [
                 "tensorflow",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
+                *extra_requires,
             ],
+            **extra_select,
         }
     elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(
         ">=2.0,<2.1"
@@ -129,22 +144,28 @@ def get_tf_requirement(tf_version: str = "") -> dict:
         return {
             "cpu": [
                 f"tensorflow=={tf_version}",
+                *extra_requires,
             ],
             "gpu": [
                 f"tensorflow-gpu=={tf_version}; platform_machine!='aarch64'",
                 f"tensorflow=={tf_version}; platform_machine=='aarch64'",
+                *extra_requires,
             ],
+            **extra_select,
         }
     else:
         return {
             "cpu": [
                 f"tensorflow-cpu=={tf_version}; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 f"tensorflow=={tf_version}; platform_machine=='aarch64'  or (platform_machine=='arm64' and platform_system == 'Darwin')",
+                *extra_requires,
             ],
             "gpu": [
                 f"tensorflow=={tf_version}",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
+                *extra_requires,
             ],
+            **extra_select,
         }