Provide pre-built wheels for torch 2.3.0 (#1283)

k2-fsa · Apr 25, 2024 · 8905c6b · 8905c6b
1 parent 50cdb51
commit 8905c6b
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 38 deletions.
diff --git a/.github/workflows/macos-cpu-wheels.yml b/.github/workflows/macos-cpu-wheels.yml
@@ -25,17 +25,17 @@ jobs:
         id: set-matrix
         run: |
           # outputting for debugging purposes
-          python ./scripts/github_actions/generate_build_matrix.py --for-macos
-          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-macos)
+          # python ./scripts/github_actions/generate_build_matrix.py --for-macos
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-macos)
 
-          # python ./scripts/github_actions/generate_build_matrix.py --for-macos --test-only-latest-torch
-          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-macos --test-only-latest-torch)
+          python ./scripts/github_actions/generate_build_matrix.py --for-macos --test-only-latest-torch
+          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-macos --test-only-latest-torch)
           echo "::set-output name=matrix::${MATRIX}"
 
   build_wheels_macos_cpu:
     needs: generate_build_matrix
     name: ${{ matrix.torch }} ${{ matrix.python-version }}
-    runs-on: macos-latest
+    runs-on: macos-14
     strategy:
       fail-fast: false
       matrix:
@@ -47,14 +47,14 @@ jobs:
           fetch-depth: 0
 
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
       - name: Install dependencies
         shell: bash
         run: |
-          pip install -q torch==${{ matrix.torch}} cmake numpy wheel twine setuptools
+          pip install -q torch==${{ matrix.torch}} cmake numpy wheel>=0.40.0 twine setuptools
 
       - name: Build wheel
         shell: bash
@@ -63,7 +63,21 @@ jobs:
           mkdir wheelhouse
           cp -v dist/* wheelhouse
 
-      - name: Display wheels
+      - name: Display wheels (before fix)
+        shell: bash
+        run: |
+          ls -lh ./wheelhouse/
+
+      - name: Fix wheel platform tag
+        run: |
+          # See https://github.com/glencoesoftware/zeroc-ice-py-macos-x86_64/pull/3/files
+          # See:
+          #  * https://github.com/pypa/wheel/issues/406
+          python -m wheel tags \
+            --platform-tag=macosx_11_0_arm64 \
+            --remove wheelhouse/*.whl
+
+      - name: Display wheels (after fix)
         shell: bash
         run: |
           ls -lh ./wheelhouse/

diff --git a/.github/workflows/ubuntu-cpu-wheels.yml b/.github/workflows/ubuntu-cpu-wheels.yml
@@ -25,11 +25,11 @@ jobs:
         id: set-matrix
         run: |
           # outputting for debugging purposes
-          python ./scripts/github_actions/generate_build_matrix.py
-          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py)
+          # python ./scripts/github_actions/generate_build_matrix.py
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py)
 
-          # python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch
-          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch)
+          python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch
+          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --test-only-latest-torch)
           echo "::set-output name=matrix::${MATRIX}"
 
   build-manylinux-wheels:

diff --git a/.github/workflows/ubuntu-cuda-wheels.yml b/.github/workflows/ubuntu-cuda-wheels.yml
@@ -25,11 +25,11 @@ jobs:
         id: set-matrix
         run: |
           # outputting for debugging purposes
-          # python ./scripts/github_actions/generate_build_matrix.py --enable-cuda --test-only-latest-torch
-          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --enable-cuda --test-only-latest-torch)
+          python ./scripts/github_actions/generate_build_matrix.py --enable-cuda --test-only-latest-torch
+          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --enable-cuda --test-only-latest-torch)
 
-          python ./scripts/github_actions/generate_build_matrix.py --enable-cuda
-          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --enable-cuda)
+          # python ./scripts/github_actions/generate_build_matrix.py --enable-cuda
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --enable-cuda)
 
           echo "::set-output name=matrix::${MATRIX}"
 

diff --git a/.github/workflows/windows-x64-cpu-wheels.yml b/.github/workflows/windows-x64-cpu-wheels.yml
@@ -25,11 +25,11 @@ jobs:
         id: set-matrix
         run: |
           # outputting for debugging purposes
-          python ./scripts/github_actions/generate_build_matrix.py --for-windows
-          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-windows)
+          # python ./scripts/github_actions/generate_build_matrix.py --for-windows
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-windows)
 
-          # python ./scripts/github_actions/generate_build_matrix.py --for-windows --test-only-latest-torch
-          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-windows --test-only-latest-torch)
+          python ./scripts/github_actions/generate_build_matrix.py --for-windows --test-only-latest-torch
+          MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --for-windows --test-only-latest-torch)
           echo "::set-output name=matrix::${MATRIX}"
 
   build_wheels_win64_cpu:

diff --git a/cmake/torch.cmake b/cmake/torch.cmake
@@ -8,6 +8,8 @@ execute_process(
 )
 
 list(APPEND CMAKE_PREFIX_PATH "${TORCH_DIR}")
+include_directories(${TORCH_DIR}/include/torch/csrc/api/include)
+include_directories(${TORCH_DIR}/include)
 
 if(NOT DEFINED TORCH_LIBRARY)
   find_package(Torch REQUIRED)

diff --git a/scripts/github_actions/build-ubuntu-cuda.sh b/scripts/github_actions/build-ubuntu-cuda.sh
@@ -67,6 +67,8 @@ python3 -m pip install bs4 requests tqdm auditwheel
 echo "Installing torch"
 ./install_torch.sh
 
+sed -i.bak /9.0a/d /Python-*/py-3.*/lib/python3.*/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
+
 rm -rf ~/.cache/pip >/dev/null 2>&1
 yum clean all >/dev/null 2>&1
 
@@ -91,33 +93,33 @@ auditwheel --verbose repair \
   --exclude libtorch_cuda.so \
   --exclude libtorch_python.so \
   \
-  --exclude libcudnn.so.8 \
+  --exclude libcublas.so \
   --exclude libcublas.so.11 \
-  --exclude libcublasLt.so.11 \
   --exclude libcublas.so.12 \
-  --exclude libcublas.so \
-  --exclude libcublasLt.so.12 \
   --exclude libcublasLt.so \
+  --exclude libcublasLt.so.11 \
+  --exclude libcublasLt.so.12 \
   --exclude libcudart.so.11.0 \
   --exclude libcudart.so.12 \
-  --exclude libnvrtc.so.11.2 \
-  --exclude libnvrtc.so.12 \
-  --exclude libnvrtc.so \
-  --exclude libcupti.so.12 \
+  --exclude libcudnn.so.8 \
+  --exclude libcufft.so \
+  --exclude libcufft.so.11 \
   --exclude libcupti.so \
-  --exclude libcusparse.so.12 \
-  --exclude libcusparse.so \
-  --exclude libnvJitLink.so.12 \
-  --exclude libnvJitLink.so \
-  --exclude libcurand.so.10 \
+  --exclude libcupti.so.12 \
   --exclude libcurand.so \
-  --exclude libcufft.so.11 \
-  --exclude libcufft.so \
-  --exclude libnccl.so.2 \
+  --exclude libcurand.so.10 \
+  --exclude libcusparse.so \
+  --exclude libcusparse.so.12 \
   --exclude libnccl.so \
+  --exclude libnccl.so.2 \
+  --exclude libnvJitLink.so \
+  --exclude libnvJitLink.so.12 \
+  --exclude libnvrtc.so \
+  --exclude libnvrtc.so.11.2 \
+  --exclude libnvrtc.so.12 \
   --exclude libshm.so \
-  --exclude libtorch_cuda_cu.so \
   --exclude libtorch_cuda_cpp.so \
+  --exclude libtorch_cuda_cu.so \
   --plat manylinux_2_17_x86_64 \
   -w /var/www/wheelhouse \
   dist/*.whl

diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
@@ -208,10 +208,16 @@ def generate_build_matrix(
             if not for_windows
             else ["11.8.0", "12.1.0"],
         },
+        "2.3.0": {
+            "python-version": ["3.8", "3.9", "3.10", "3.11", "3.12"],
+            "cuda": ["11.8", "12.1"]  # default 12.1
+            if not for_windows
+            else ["11.8.0", "12.1.0"],
+        },
         # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.2.2"
+        latest = "2.3.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
@@ -271,6 +277,9 @@ def generate_build_matrix(
             for p in python_versions:
                 if p in excluded_python_versions:
                     continue
+                if for_macos and p in ["3.8", "3.9"]:
+                    # macOS arm64 in github actions does not support python 3.8 or 3.9
+                    continue
 
                 if for_windows:
                     p = "cp" + "".join(p.split("."))

diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+set -x
+
 if [ x"$TORCH_VERSION" != x"" ] && [ x"$CUDA_VERSION" != x"" ]; then
     torch=$TORCH_VERSION
     cuda=$CUDA_VERSION
@@ -198,6 +200,19 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.3.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        url=https://download.pytorch.org/whl/torch_stable.html
+        ;;
+      12.1)
+        package="torch==${torch}"
+        # Leave it empty to use PyPI.
+        url=
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1