Build wheels for cu124

ghstack-source-id: 8d4638774fb9bc169f5a272a6b6e56ff4996effa
Pull Request resolved: fairinternal/xformers#1195

__original_commit__ = fairinternal/xformers@15ac528

.github/actions/setup-windows-runner/action.yml

@@ -14,7 +14,7 @@ runs:
   using: composite
   steps:
     - name: Install cuda
-      uses: okazunori2013/cuda-toolkit@v0.3.6
+      uses: Jimver/cuda-toolkit@v0.2.16
       with:
         cuda: ${{ inputs.cuda }}
         method: network

.github/workflows/wheels.yml

@@ -25,29 +25,48 @@ jobs:
       run: |
         import os
         import json
+        import itertools
         environ = os.environ
 
         PY_VERSIONS = ['3.8', '3.9', '3.10', '3.11', '3.12']
+        CU_VERSIONS = ['118', '121', '124']
+        ROCM_VERSIONS = ["6.0", "6.1"]
+        PY_CU = list(itertools.product(PY_VERSIONS, CU_VERSIONS))
+        PY_ROCM = list(itertools.product(PY_VERSIONS, ROCM_VERSIONS))
+        print("Full matrix PY_CU", PY_CU)
         if os.environ["GITHUB_EVENT_NAME"] == "pull_request":
-          print("pull-request: only keeping 2 python versions")
-          PY_VERSIONS = [PY_VERSIONS[0], PY_VERSIONS[-1]]
+          print("pull-request: limiting matrix to save resources")
+          PY_CU = [(PY_VERSIONS[0], CU_VERSIONS[0])]
+          for cu in CU_VERSIONS[1:]:
+            PY_CU.append((PY_VERSIONS[-1], cu))
+          print("Limited matrix PY_CU", PY_CU)
+          PY_ROCM = [(PY_VERSIONS[-1], ROCM_VERSIONS[-1])]
 
         include = []
         for os in ['8-core-ubuntu', 'windows-8-core']:
-          for python in PY_VERSIONS:
-            for torch_version in ['2.4.0']:
-              for toolkit_type, toolkit_short_versions in {'cuda': ["118", "121"], 'rocm': ["6.0", "6.1"]}.items():
-                if os == 'windows-8-core' and toolkit_type == 'rocm':
-                  continue
-                for toolkit_short_version in toolkit_short_versions:
-                  include.append(dict(
-                    os=os,
-                    python=python,
-                    torch_version=torch_version,
-                    toolkit_type=toolkit_type,
-                    toolkit_short_version=toolkit_short_version,
-                  ))
-                  print(include[-1])
+          for torch_version in ['2.4.0']:
+            # CUDA builds
+            for python, cuda_short_version in PY_CU:
+              include.append(dict(
+                os=os,
+                python=python,
+                torch_version=torch_version,
+                toolkit_type="cuda",
+                toolkit_short_version=cuda_short_version,
+              ))
+              print(include[-1])
+            # ROCM builds
+            for python, rocm_short_version in PY_ROCM:
+              if os == 'windows-8-core':
+                continue
+              include.append(dict(
+                os=os,
+                python=python,
+                torch_version=torch_version,
+                toolkit_type="rocm",
+                toolkit_short_version=rocm_short_version,
+              ))
+              print(include[-1])
         matrix = {'include': include}
         print(json.dumps(matrix))
         with open(environ["GITHUB_OUTPUT"], "a") as fd:

.github/workflows/wheels_build.yml

@@ -66,31 +66,23 @@ jobs:
           import sys
           print(sys.version)
           cushort = "${{ inputs.toolkit_short_version }}"
-          TORCH_CUDA_DEFAULT = "121"  # pytorch 2.1.0
+          TORCH_CUDA_DEFAULT = "121"  # pytorch 2.4.0
           # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/linux-links.ts
           full_version, install_script = {
+            "124": ("12.4.1", "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run"),
             "121": ("12.1.0", "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"),
             "118": ("11.8.0", "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"),
-            "117": ("11.7.1", "https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run"),
-            "116": ("11.6.2", "https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_510.47.03_linux.run"),
-
             "6.0": ("6.0.2", "https://repo.radeon.com/amdgpu-install/6.0.2/rhel/7.9/amdgpu-install-6.0.60002-1.el7.noarch.rpm"),
             "6.1": ("6.1.2", "https://repo.radeon.com/amdgpu-install/6.1.2/el/7/amdgpu-install-6.1.60102-1.el7.noarch.rpm"),
           }[cushort]
           with open(os.environ['GITHUB_OUTPUT'], "r+") as fp:
             fp.write("CUDA_VERSION=" + full_version + "\n")
             if cushort == TORCH_CUDA_DEFAULT:
               fp.write("CUDA_VERSION_SUFFIX=\n")
-              fp.write("TORCH_ORG_S3_PATH=s3://pytorch/whl\n")
-              fp.write("PUBLISH_PYPI=1\n")
             else:
               fp.write("CUDA_VERSION_SUFFIX=+" + ("cu" if "cuda" == "${{ inputs.toolkit_type }}" else "rocm") + cushort + "\n")
-              fp.write("TORCH_ORG_S3_PATH=s3://pytorch/whl/" + cushort + "\n")
-              fp.write("PUBLISH_PYPI=0\n")
             fp.write("CUDA_INSTALL_SCRIPT=" + install_script + "\n")
       - run: echo "CUDA_VERSION_SUFFIX=${{ steps.cuda_info.outputs.CUDA_VERSION_SUFFIX }}"
-      - run: echo "TORCH_ORG_S3_PATH=${{ steps.cuda_info.outputs.TORCH_ORG_S3_PATH }}"
-      - run: echo "PUBLISH_PYPI=${{ steps.cuda_info.outputs.PUBLISH_PYPI }}"
 
       - name: Recursive checkout
         uses: actions/checkout@v3

CHANGELOG.md

@@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.0.28] - TBD
 ### Added
+- Added wheels for cuda 12.4
 ### Improved
 - Profiler: Fix computation of FLOPS for the attention when using xFormers
 - Profiler: Fix MFU/HFU calculation when multiple dtypes are used