Publish pre-built wheels for torch 2.4.0 (#1298)

.github/workflows/run-tests-cpu.yml

@@ -52,7 +52,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest]
-        torch: ["2.2.1"]
+        torch: ["2.4.0"]
         python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         build_type: ["Release", "Debug"]
 

.github/workflows/run-tests.yml

@@ -51,8 +51,8 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
-        cuda: ["12.1"]
-        torch: ["2.2.1"]
+        cuda: ["12.4"]
+        torch: ["2.4.0"]
         python-version: ["3.12"]
 
     steps:
@@ -71,7 +71,7 @@ jobs:
       - name: Run the build process with Docker
         uses: addnab/docker-run-action@v3
         with:
-            image: "pytorch/manylinux-builder:cuda12.1"
+            image: "pytorch/manylinux-builder:cuda12.4"
             options: -v ${{ github.workspace }}:/var/www -e PYTHON_VERSION=${{ matrix.python-version }} -e TORCH_VERSION=${{ matrix.torch }} -e CUDA_VERSION=${{ matrix.cuda }}
             run: |
               echo "pwd: $PWD"

.github/workflows/windows-x64-cpu-wheels.yml

@@ -51,7 +51,7 @@ jobs:
       - name: Build wheels
         uses: pypa/[email protected]
         env:
-          CIBW_BEFORE_BUILD: pip install torch==${{ matrix.torch}}+cpu cmake numpy -f https://download.pytorch.org/whl/torch_stable.html
+          CIBW_BEFORE_BUILD: pip install torch==${{ matrix.torch}}+cpu cmake numpy -f https://download.pytorch.org/whl/torch_stable.html || pip install torch==${{ matrix.torch}}+cpu cmake numpy -f https://download.pytorch.org/whl/torch/
           CIBW_BUILD: ${{ matrix.python-version }}-win_amd64
           CIBW_REPAIR_WHEEL_COMMAND_WINDOWS: ""
           CIBW_BUILD_VERBOSITY: 3

docs/source/installation/from_wheels.rst

@@ -37,11 +37,11 @@ we can use one of the following methods:
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch/
    pip install k2==1.24.3.dev20230719+cpu.torch2.0.1 -f https://k2-fsa.github.io/k2/cpu.html
 
    # method 2
-   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch/
 
    wget https://huggingface.co/csukuangfj/k2/resolve/main/cpu/k2-1.24.3.dev20230719+cpu.torch2.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
@@ -105,11 +105,11 @@ we can use the one of the following methods:
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch/
    pip install k2==1.24.3.dev20230726+cpu.torch2.0.1 -f https://k2-fsa.github.io/k2/cpu.html
 
    # method 2
-   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch/
    wget https://huggingface.co/csukuangfj/k2/resolve/main/windows-cpu/k2-1.24.3.dev20230726+cpu.torch2.0.1-cp38-cp38-win_amd64.whl
 
    # For users from China
@@ -143,11 +143,11 @@ we can use the following methods:
 .. code-block:: bash
 
    # method 1
-   pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/torch/
    pip install k2==1.24.3.dev20230718+cuda11.7.torch2.0.1 -f https://k2-fsa.github.io/k2/cuda.html
 
    # method 2
-   pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+   pip install torch==2.0.1+cu117 -f https://download.pytorch.org/whl/torch/
 
    wget https://huggingface.co/csukuangfj/k2/resolve/main/cuda/k2-1.24.3.dev20230718+cuda11.7.torch2.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 

scripts/github_actions/build-ubuntu-cpu.sh

@@ -82,7 +82,8 @@ python3 -m pip install wheel twine typing_extensions
 python3 -m pip install bs4 requests tqdm auditwheel
 
 echo "Installing torch $TORCH_VERSION"
-python3 -m pip install -qq torch==$TORCH_VERSION+cpu -f https://download.pytorch.org/whl/torch_stable.html
+python3 -m pip install -qq torch==$TORCH_VERSION+cpu -f https://download.pytorch.org/whl/torch_stable.html || \
+python3 -m pip install -qq torch==$TORCH_VERSION+cpu -f https://download.pytorch.org/whl/torch/
 
 rm -rf ~/.cache/pip
 yum clean all

scripts/github_actions/generate_build_matrix.py

@@ -220,10 +220,16 @@ def generate_build_matrix(
             if not for_windows
             else ["11.8.0", "12.1.0"],
         },
+        "2.4.0": {
+            "python-version": ["3.8", "3.9", "3.10", "3.11", "3.12"],
+            "cuda": ["11.8", "12.1", "12.4"]  # default 12.1
+            if not for_windows
+            else ["11.8.0", "12.1.0", "12.4.0"],
+        },
         # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.3.1"
+        latest = "2.4.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
@@ -283,15 +289,20 @@ 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("."))
                     ans.append({"torch": torch, "python-version": p})
                 elif for_macos or for_macos_m1:
                     ans.append({"torch": torch, "python-version": p})
+                elif version_ge(torch, "2.4.0"):
+                    ans.append(
+                        {
+                            "torch": torch,
+                            "python-version": p,
+                            "image": "pytorch/manylinux-builder:cpu-2.4",
+                        }
+                    )
                 elif version_ge(torch, "2.2.0"):
                     ans.append(
                         {

scripts/github_actions/install_cuda.sh

@@ -55,6 +55,9 @@ case "$cuda" in
   12.1)
     url=https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
     ;;
+  12.4)
+    url=https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
+    ;;
   *)
     echo "Unknown cuda version: $cuda"
     exit 1

scripts/github_actions/install_cudnn.sh

@@ -48,6 +48,9 @@ case $cuda in
   12.1)
     filename=cudnn-linux-x86_64-8.9.5.29_cuda12-archive.tar.xz
     ;;
+  12.4)
+    filename=cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz
+    ;;
   *)
     echo "Unsupported cuda version: $cuda"
     exit 1

scripts/github_actions/install_torch.sh

@@ -213,6 +213,23 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.4.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        # https://download.pytorch.org/whl/nightly/torch/
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.1)
+        package="torch==${torch}+cu121"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.4)
+        package="torch==${torch}+cu124"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1