Support torch 2.2.0 (#1268)

k2-fsa · Feb 11, 2024 · aa96d78 · aa96d78
1 parent 415fe1f
commit aa96d78
Show file tree

Hide file tree

Showing 17 changed files with 112 additions and 48 deletions.
diff --git a/.github/workflows/macos-cpu-wheels.yml b/.github/workflows/macos-cpu-wheels.yml
@@ -25,8 +25,11 @@ 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)
           echo "::set-output name=matrix::${MATRIX}"
 
   build_wheels_macos_cpu:
@@ -43,15 +46,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/run-tests-cpu.yml b/.github/workflows/run-tests-cpu.yml
@@ -52,8 +52,8 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest]
-        torch: ["2.1.2"]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        torch: ["2.2.0"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         build_type: ["Release", "Debug"]
 
     steps:

diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
@@ -52,13 +52,22 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         cuda: ["12.1"]
-        torch: ["2.1.2"]
-        python-version: ["3.11"]
+        torch: ["2.2.0"]
+        python-version: ["3.12"]
 
     steps:
       # refer to https://github.com/actions/checkout
       - uses: actions/checkout@v4
 
+      - name: Free space
+        shell: bash
+        run: |
+          df -h
+          rm -rf /opt/hostedtoolcache
+          df -h
+          echo "pwd: $PWD"
+          echo "github.workspace ${{ github.workspace }}"
+
       - name: Run the build process with Docker
         uses: addnab/docker-run-action@v3
         with:

diff --git a/.github/workflows/ubuntu-cpu-wheels.yml b/.github/workflows/ubuntu-cpu-wheels.yml
@@ -27,6 +27,9 @@ jobs:
           # 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 --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
@@ -47,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 }}"
+
       - name: display
         shell: bash
         run: |

diff --git a/.github/workflows/windows-x64-cpu-wheels.yml b/.github/workflows/windows-x64-cpu-wheels.yml
@@ -27,6 +27,9 @@ jobs:
           # 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 --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:
@@ -46,7 +49,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 cmake numpy -f https://download.pytorch.org/whl/torch_stable.html
           CIBW_BUILD: ${{ matrix.python-version }}-win_amd64

diff --git a/cmake/cub.cmake b/cmake/cub.cmake
@@ -21,15 +21,15 @@ function(download_cub)
   include(FetchContent)
 
   set(cub_URL  "https://github.com/NVlabs/cub/archive/1.15.0.tar.gz")
-  set(cub_URL2 "https://huggingface.co/csukuangfj/k2-cmake-deps/resolve/main/cub-1.15.0.tar.gz")
+  set(cub_URL2 "https://hub.nuaa.cf/NVlabs/cub/archive/1.15.0.tar.gz")
   set(cub_HASH "SHA256=1781ee5eb7f00acfee5bff88e3acfc67378f6b3c24281335e18ae19e1f2ff685")
 
   # If you don't have access to the Internet,
   # please pre-download cub
   set(possible_file_locations
     $ENV{HOME}/Downloads/cub-1.15.0.tar.gz
-    ${PROJECT_SOURCE_DIR}/cub-1.15.0.tar.gz
-    ${PROJECT_BINARY_DIR}/cub-1.15.0.tar.gz
+    ${CMAKE_SOURCE_DIR}/cub-1.15.0.tar.gz
+    ${CMAKE_BINARY_DIR}/cub-1.15.0.tar.gz
     /tmp/cub-1.15.0.tar.gz
     /star-fj/fangjun/download/github/cub-1.15.0.tar.gz
   )

diff --git a/cmake/googletest.cmake b/cmake/googletest.cmake
@@ -25,15 +25,15 @@ function(download_googltest)
   include(FetchContent)
 
   set(googletest_URL  "https://github.com/google/googletest/archive/refs/tags/v1.13.0.tar.gz")
-  set(googletest_URL2 "https://huggingface.co/csukuangfj/k2-cmake-deps/resolve/main/googletest-1.13.0.tar.gz")
+  set(googletest_URL2 "https://hub.nuaa.cf/google/googletest/archive/refs/tags/v1.13.0.tar.gz")
   set(googletest_HASH "SHA256=ad7fdba11ea011c1d925b3289cf4af2c66a352e18d4c7264392fead75e919363")
 
   # If you don't have access to the Internet,
   # please pre-download googletest
   set(possible_file_locations
     $ENV{HOME}/Downloads/googletest-1.13.0.tar.gz
-    ${PROJECT_SOURCE_DIR}/googletest-1.13.0.tar.gz
-    ${PROJECT_BINARY_DIR}/googletest-1.13.0.tar.gz
+    ${CMAKE_SOURCE_DIR}/googletest-1.13.0.tar.gz
+    ${CMAKE_BINARY_DIR}/googletest-1.13.0.tar.gz
     /tmp/googletest-1.13.0.tar.gz
     /star-fj/fangjun/download/github/googletest-1.13.0.tar.gz
   )

diff --git a/cmake/kaldifeat.cmake b/cmake/kaldifeat.cmake
@@ -24,18 +24,18 @@ function(download_kaldifeat)
 
   include(FetchContent)
 
-  set(kaldifeat_URL "https://github.com/csukuangfj/kaldifeat/archive/refs/tags/v1.25.1.tar.gz")
-  set(kaldifeat_URL2 "https://huggingface.co/csukuangfj/k2-cmake-deps/resolve/main/kaldifeat-1.25.1.tar.gz")
-  set(kaldifeat_HASH "SHA256=28a6f77963f390220675614932e1d3c4a49d394f7d192d6af46b383765673a2e")
+  set(kaldifeat_URL  "https://github.com/csukuangfj/kaldifeat/archive/refs/tags/v1.25.4.tar.gz")
+  set(kaldifeat_URL2 "https://hub.nuaa.cf/csukuangfj/kaldifeat/archive/refs/tags/v1.25.4.tar.gz")
+  set(kaldifeat_HASH "SHA256=8a2cb45873a39d92b5e4e1f150269e7ef061c3537ad5acc3cff5ced4997bca96")
 
   # If you don't have access to the Internet,
   # please pre-download kaldifeat
   set(possible_file_locations
-    $ENV{HOME}/Downloads/kaldifeat-1.25.1.tar.gz
-    ${PROJECT_SOURCE_DIR}/kaldifeat-1.25.1.tar.gz
-    ${PROJECT_BINARY_DIR}/kaldifeat-1.25.1.tar.gz
-    /tmp/kaldifeat-1.25.1.tar.gz
-    /star-fj/fangjun/download/github/kaldifeat-1.25.1.tar.gz
+    $ENV{HOME}/Downloads/kaldifeat-1.25.4.tar.gz
+    ${CMAKE_SOURCE_DIR}/kaldifeat-1.25.4.tar.gz
+    ${CMAKE_BINARY_DIR}/kaldifeat-1.25.4.tar.gz
+    /tmp/kaldifeat-1.25.4.tar.gz
+    /star-fj/fangjun/download/github/kaldifeat-1.25.4.tar.gz
   )
 
   foreach(f IN LISTS possible_file_locations)

diff --git a/cmake/moderngpu.cmake b/cmake/moderngpu.cmake
@@ -22,15 +22,15 @@ function(download_moderngpu)
 
   # this is the latest commit of modern gpu as of 2022-04-03
   set(moderngpu_URL  "https://github.com/moderngpu/moderngpu/archive/8ec9ac0de8672de7217d014917eedec5317f75f3.zip")
-  set(moderngpu_URL2 "https://huggingface.co/csukuangfj/k2-cmake-deps/resolve/main/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip")
+  set(moderngpu_URL2 "https://hub.nuaa.cf/moderngpu/moderngpu/archive/8ec9ac0de8672de7217d014917eedec5317f75f3.zip")
   set(moderngpu_HASH "SHA256=1c20ffbb81d6f7bbe6107aaa5ee6d37392677c8a5fc7894935149c3ef0a3c2fb")
 
   # If you don't have access to the Internet,
   # please pre-download moderngpu
   set(possible_file_locations
     $ENV{HOME}/Downloads/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
-    ${PROJECT_SOURCE_DIR}/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
-    ${PROJECT_BINARY_DIR}/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
+    ${CMAKE_SOURCE_DIR}/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
+    ${CMAKE_BINARY_DIR}/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
     /tmp/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
     /star-fj/fangjun/download/github/moderngpu-8ec9ac0de8672de7217d014917eedec5317f75f3.zip
   )

diff --git a/cmake/pybind11.cmake b/cmake/pybind11.cmake
@@ -22,15 +22,15 @@ function(download_pybind11)
 
   # latest commit as of 2022.10.31 that supports python 3.11
   set(pybind11_URL  "https://github.com/pybind/pybind11/archive/5bc0943ed96836f46489f53961f6c438d2935357.zip")
-  set(pybind11_URL2 "https://huggingface.co/csukuangfj/k2-cmake-deps/resolve/main/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip")
+  set(pybind11_URL2 "https://hub.nuaa.cf/pybind/pybind11/archive/5bc0943ed96836f46489f53961f6c438d2935357.zip")
   set(pybind11_HASH "SHA256=ff65a1a8c9e6ceec11e7ed9d296f2e22a63e9ff0c4264b3af29c72b4f18f25a0")
 
   # If you don't have access to the Internet,
   # please pre-download pybind11
   set(possible_file_locations
     $ENV{HOME}/Downloads/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
-    ${PROJECT_SOURCE_DIR}/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
-    ${PROJECT_BINARY_DIR}/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
+    ${CMAKE_SOURCE_DIR}/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
+    ${CMAKE_BINARY_DIR}/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
     /tmp/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
     /star-fj/fangjun/download/github/pybind11-5bc0943ed96836f46489f53961f6c438d2935357.zip
   )

diff --git a/cmake/torch.cmake b/cmake/torch.cmake
@@ -8,7 +8,10 @@ execute_process(
 )
 
 list(APPEND CMAKE_PREFIX_PATH "${TORCH_DIR}")
-find_package(Torch REQUIRED)
+
+if(NOT DEFINED TORCH_LIBRARY)
+  find_package(Torch REQUIRED)
+endif()
 
 # set the global CMAKE_CXX_FLAGS so that
 # k2 uses the same abi flag as PyTorch

diff --git a/k2/torch/csrc/deserialization.cu b/k2/torch/csrc/deserialization.cu
@@ -352,7 +352,7 @@ torch::IValue Load(
                      "when invoking torch.save()";
   }
 
-  auto reader = torch::make_unique<caffe2::serialize::PyTorchStreamReader>(
+  auto reader = std::make_unique<caffe2::serialize::PyTorchStreamReader>(
       std::move(rai));
 
   auto cu = std::make_shared<torch::jit::CompilationUnit>();

diff --git a/scripts/github_actions/build-ubuntu-cpu.sh b/scripts/github_actions/build-ubuntu-cpu.sh
@@ -20,10 +20,10 @@ echo "Installing ${PYTHON_VERSION}.3"
 yum -y install openssl-devel bzip2-devel libffi-devel xz-devel wget redhat-lsb-core
 
 if true; then
-  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}
 

diff --git a/scripts/github_actions/build-ubuntu-cuda.sh b/scripts/github_actions/build-ubuntu-cuda.sh
@@ -22,14 +22,18 @@ 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}
 

diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
@@ -178,10 +178,16 @@ def generate_build_matrix(
             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:
@@ -210,6 +216,8 @@ def generate_build_matrix(
         matrix["1.13.1"] = {"python-version": ["3.7", "3.8", "3.9", "3.10"]}
         matrix["2.0.0"] = {"python-version": ["3.8", "3.9", "3.10", "3.11"]}
         matrix["2.0.1"] = {"python-version": ["3.8", "3.9", "3.10", "3.11"]}
+        # TODO(fangjun): we currently don't support macOS M1 build
+        # since github actions does not support it.
 
     ans = []
     for torch, python_cuda in matrix.items():
@@ -240,9 +248,11 @@ def generate_build_matrix(
                 if p in excluded_python_versions:
                     continue
 
-                if for_windows or for_macos or for_macos_m1:
+                if for_windows:
                     p = "cp" + "".join(p.split("."))
                     ans.append({"torch": torch, "python-version": p})
+                elif for_macos or for_macos_m1:
+                    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