Support torch 2.2.0 (#94)

csukuangfj · Feb 10, 2024 · 2e042b3 · 2e042b3
1 parent b75abef
commit 2e042b3
Show file tree

Hide file tree

Showing 9 changed files with 101 additions and 54 deletions.
diff --git a/.github/workflows/macos-cpu-wheels.yml b/.github/workflows/macos-cpu-wheels.yml
@@ -3,7 +3,7 @@ name: build-wheels-cpu-macos
 on:
   push:
     branches:
-      - fix-wheel
+      - fix-wheel-2
     tags:
       - '*'
   workflow_dispatch:
@@ -28,6 +28,9 @@ jobs:
           # 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 --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:
@@ -44,15 +47,22 @@ jobs:
         with:
           fetch-depth: 0
 
-      # see https://cibuildwheel.readthedocs.io/en/stable/changelog/
-      # for a list of versions
-      - name: Build wheels
-        uses: pypa/[email protected]
-        env:
-          CIBW_BEFORE_BUILD: pip install torch==${{ matrix.torch}} cmake numpy
-          CIBW_BUILD: ${{ matrix.python-version }}-*
-          CIBW_REPAIR_WHEEL_COMMAND_MACOS: ""
-          CIBW_BUILD_VERBOSITY: 3
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        shell: bash
+        run: |
+          pip install -q torch==${{ matrix.torch}} cmake numpy wheel twine setuptools
+
+      - name: Build wheel
+        shell: bash
+        run: |
+          python3 setup.py bdist_wheel
+          mkdir wheelhouse
+          cp -v dist/* wheelhouse
 
       - name: Display wheels
         shell: bash

diff --git a/.github/workflows/ubuntu-cpu-wheels.yml b/.github/workflows/ubuntu-cpu-wheels.yml
@@ -3,7 +3,7 @@ name: build-wheels-cpu-ubuntu
 on:
   push:
     branches:
-      - fix-wheel
+      - fix-wheel-2
     tags:
       - '*'
   workflow_dispatch:

diff --git a/.github/workflows/ubuntu-cuda-wheels.yml b/.github/workflows/ubuntu-cuda-wheels.yml
@@ -3,7 +3,7 @@ name: build-wheels-cuda-ubuntu
 on:
   push:
     branches:
-      - fix-wheel
+      - fix-wheel-2
     tags:
       - '*'
   workflow_dispatch:
@@ -28,6 +28,9 @@ jobs:
           # outputting for debugging purposes
           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 --test-only-latest-torch
+          # MATRIX=$(python ./scripts/github_actions/generate_build_matrix.py --enable-cuda --test-only-latest-torch)
           echo "::set-output name=matrix::${MATRIX}"
 
   build-manylinux-wheels:
@@ -44,6 +47,15 @@ jobs:
         with:
           fetch-depth: 0
 
+      - name: Free space
+        shell: bash
+        run: |
+          df -h
+          rm -rf /opt/hostedtoolcache
+          df -h
+          echo "pwd: $PWD"
+          echo "github.workspace ${{ github.workspace }}"
+
       # see https://github.com/pytorch/test-infra/blob/9e3d392690719fac85bad0c9b67f530e48375ca1/tools/scripts/generate_binary_build_matrix.py
       # https://github.com/pytorch/builder/tree/main/manywheel
       # https://github.com/pytorch/builder/pull/476

diff --git a/.github/workflows/windows-x64-cpu-wheels.yml b/.github/workflows/windows-x64-cpu-wheels.yml
@@ -3,7 +3,7 @@ name: build-wheels-cpu-win64
 on:
   push:
     branches:
-      - fix-wheel
+      - fix-wheel-2
     tags:
       - '*'
   workflow_dispatch:
@@ -47,7 +47,7 @@ jobs:
       # see https://cibuildwheel.readthedocs.io/en/stable/changelog/
       # for a list of versions
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.11.4
+        uses: pypa/cibuildwheel@v2.16.5
         env:
           CIBW_BEFORE_BUILD: pip install torch==${{ matrix.torch}}+cpu -f https://download.pytorch.org/whl/torch_stable.html cmake numpy
           CIBW_BUILD: ${{ matrix.python-version }}-win_amd64

diff --git a/doc/source/installation/from_wheels.rst b/doc/source/installation/from_wheels.rst
@@ -16,21 +16,20 @@ Suppose you want to install the following wheel:
 
 .. code-block:: bash
 
-   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cpu/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cpu/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
 you can use one of the following methods:
 
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.1.2+cpu -f https://download.pytorch.org/whl/torch_stable.html
-   pip install kaldifeat==1.25.3.dev20231221+cpu.torch2.1.2 -f https://csukuangfj.github.io/kaldifeat/cpu.html
-
+   pip install torch==2.2.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install kaldifeat==1.25.4.dev20240210+cpu.torch2.2.0 -f https://csukuangfj.github.io/kaldifeat/cpu.html
 
    # method 2
-   pip install torch==2.1.2+cpu -f https://download.pytorch.org/whl/torch_stable.html
-   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cpu/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
-   pip install ./kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   pip install torch==2.2.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cpu/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   pip install ./kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
 
 Windows (CPU)
@@ -40,20 +39,20 @@ Suppose you want to install the following wheel:
 
 .. code-block:: bash
 
-   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/windows-cpu/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-win_amd64.whl
+   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/windows-cpu/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-win_amd64.whl
 
 you can use one of the following methods:
 
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.1.2+cpu -f https://download.pytorch.org/whl/torch_stable.html
-   pip install kaldifeat==1.25.3.dev20231221+cpu.torch2.1.2 -f https://csukuangfj.github.io/kaldifeat/cpu.html
+   pip install torch==2.2.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install kaldifeat==1.25.4.dev20240210+cpu.torch2.2.0 -f https://csukuangfj.github.io/kaldifeat/cpu.html
 
    # method 2
-   pip install torch==2.1.2+cpu -f https://download.pytorch.org/whl/torch_stable.html
-   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/windows-cpu/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-win_amd64.whl
-   pip install ./kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-win_amd64.whl
+   pip install torch==2.2.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/windows-cpu/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-win_amd64.whl
+   pip install ./kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-win_amd64.whl
 
 macOS (CPU)
 -----------
@@ -62,20 +61,20 @@ Suppose you want to install the following wheel:
 
 .. code-block:: bash
 
-   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/macos/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-macosx_10_9_x86_64.whl
+   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/macos/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-macosx_12_0_universal2.whl
 
 you can use one of the following methods:
 
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.1.2
-   pip install kaldifeat==1.25.3.dev20231221+cpu.torch2.1.2 -f https://csukuangfj.github.io/kaldifeat/cpu.html
+   pip install torch==2.2.0
+   pip install kaldifeat==1.25.4.dev20240210+cpu.torch2.2.0 -f https://csukuangfj.github.io/kaldifeat/cpu.html
 
    # method 2
-   pip install torch==2.1.2 -f https://download.pytorch.org/whl/torch_stable.html
-   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/macos/kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-macosx_10_9_x86_64.whl
-   pip install ./kaldifeat-1.25.3.dev20231221+cpu.torch2.1.2-cp311-cp311-macosx_10_9_x86_64.whl
+   pip install torch==2.2.0 -f https://download.pytorch.org/whl/torch_stable.html
+   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/macos/kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-macosx_12_0_universal2.whl
+   pip install ./kaldifeat-1.25.4.dev20240210+cpu.torch2.2.0-cp312-cp312-macosx_12_0_universal2.whl
 
 Linux (CUDA)
 ------------
@@ -84,17 +83,17 @@ Suppose you want to install the following wheel:
 
 .. code-block:: bash
 
-   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cuda/kaldifeat-1.25.3.dev20231221+cuda12.1.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cuda/kaldifeat-1.25.4.dev20240210+cuda12.1.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
 you can use one of the following methods:
 
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.1.2+cu121 -f https://download.pytorch.org/whl/torch_stable.html
-   pip install kaldifeat==1.25.3.dev20231221+cuda12.1.torch2.1.2 -f https://csukuangfj.github.io/kaldifeat/cuda.html
+   pip install torch==2.2.0+cu121 -f https://download.pytorch.org/whl/torch_stable.html
+   pip install kaldifeat==1.25.4.dev20240210+cuda12.1.torch2.2.0 -f https://csukuangfj.github.io/kaldifeat/cuda.html
 
    # method 2
-   pip install torch==2.1.2+cu121 -f https://download.pytorch.org/whl/torch_stable.html
-   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cuda/kaldifeat-1.25.3.dev20231221+cuda12.1.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
-   pip install ./kaldifeat-1.25.3.dev20231221+cuda12.1.torch2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   pip install torch==2.2.0+cu121 -f https://download.pytorch.org/whl/torch_stable.html
+   wget https://huggingface.co/csukuangfj/kaldifeat/resolve/main/ubuntu-cuda/kaldifeat-1.25.4.dev20240210+cuda12.1.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+   pip install ./kaldifeat-1.25.4.dev20240210+cuda12.1.torch2.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
diff --git a/scripts/github_actions/build-ubuntu-cpu.sh b/scripts/github_actions/build-ubuntu-cpu.sh
@@ -17,11 +17,11 @@ fi
 
 yum -y install openssl-devel bzip2-devel libffi-devel xz-devel wget redhat-lsb-core
 
-echo "Installing ${PYTHON_VERSION}.3"
+echo "Installing ${PYTHON_VERSION}.2"
 
-curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.3/Python-${PYTHON_VERSION}.3.tgz
-tar xf Python-${PYTHON_VERSION}.3.tgz
-pushd Python-${PYTHON_VERSION}.3
+curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.2/Python-${PYTHON_VERSION}.2.tgz
+tar xf Python-${PYTHON_VERSION}.2.tgz
+pushd Python-${PYTHON_VERSION}.2
 
 PYTHON_INSTALL_DIR=$PWD/py-${PYTHON_VERSION}
 

diff --git a/scripts/github_actions/build-ubuntu-cuda.sh b/scripts/github_actions/build-ubuntu-cuda.sh
@@ -22,14 +22,19 @@ if [ -z $CUDA_VERSION ]; then
   exit 1
 fi
 
+if [[ $TORCH_VERSION =~ 2.2.* && $CUDA_VERSION =~ 12.* ]]; then
+  # see https://github.com/pytorch/pytorch/issues/113948
+  export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
+fi
+
 
 yum -y install openssl-devel bzip2-devel libffi-devel xz-devel wget redhat-lsb-core
 
 
-echo "Installing ${PYTHON_VERSION}.3"
-curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.3/Python-${PYTHON_VERSION}.3.tgz
-tar xf Python-${PYTHON_VERSION}.3.tgz
-pushd Python-${PYTHON_VERSION}.3
+echo "Installing ${PYTHON_VERSION}.2"
+curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.2/Python-${PYTHON_VERSION}.2.tgz
+tar xf Python-${PYTHON_VERSION}.2.tgz
+pushd Python-${PYTHON_VERSION}.2
 
 PYTHON_INSTALL_DIR=$PWD/py-${PYTHON_VERSION}
 
@@ -44,7 +49,7 @@ make install >/dev/null 2>&1
 popd
 
 echo "pwd: $PWD"
-# rm -rf Python-${PYTHON_VERSION}.3
+# rm -rf Python-${PYTHON_VERSION}.2
 
 export PATH=$PYTHON_INSTALL_DIR/bin:$PATH
 export LD_LIBRARY_PATH=$PYTHON_INSTALL_DIR/lib:$LD_LIBRARY_PATH

diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
@@ -53,10 +53,10 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
         #      "python-version": ["3.6", "3.7", "3.8"],
         #      "cuda": ["10.1", "10.2"],
         #  },
-        "1.6.0": {
-            "python-version": ["3.6", "3.7", "3.8"],
-            "cuda": ["10.1", "10.2"] if not for_windows else ["10.1.243", "10.2.89"],
-        },
+        #  "1.6.0": {
+        #      "python-version": ["3.6", "3.7", "3.8"],
+        #      "cuda": ["10.1", "10.2"] if not for_windows else ["10.1.243", "10.2.89"],
+        #  },
         "1.7.0": {
             "python-version": ["3.6", "3.7", "3.8"],
             "cuda": ["10.1", "10.2", "11.0"]
@@ -165,10 +165,16 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
             if not for_windows
             else ["11.8.0", "12.1.0"],
         },
+        "2.2.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.1.2"
+        latest = "2.2.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
@@ -205,9 +211,11 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
                 if p in excluded_python_versions:
                     continue
 
-                if for_windows or for_macos:
+                if for_windows:
                     p = "cp" + "".join(p.split("."))
                     ans.append({"torch": torch, "python-version": p})
+                elif for_macos:
+                    ans.append({"torch": torch, "python-version": p})
                 else:
                     ans.append(
                         {

diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
@@ -185,6 +185,19 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.2.*)
+    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