Linux aarch64 CI updates (#1622)

* Add aarch64 cpu tests and CUDA build to nightly workflow

* aarch64: limit CUDA targets to sm75, sm80, sm90, sm100

* aarch64: limit CUDA targets to sm75, sm80, sm90, sm100

* Update build cpu script

* fix

* Update auditwheel for aarch64

.github/scripts/auditwheel_show.py

@@ -22,7 +22,7 @@ def main():
             print(audit_wheel_output.stdout)
 
         if audit_wheel_output.stderr:
-            print(f"**Error:**\n```{audit_wheel_output.stderr}```")
+            print(f"**Error:**\n```\n{audit_wheel_output.stderr}\n```")
 
         print("---")
 

.github/scripts/build-cpu.sh

@@ -6,12 +6,7 @@ set -xeuo pipefail
 
 pip install cmake==3.28.3
 
-if [ "${build_os:0:6}" == ubuntu ] && [ "${build_arch}" == aarch64 ]; then
-	# Allow cross-compile on aarch64
-	sudo apt-get update
-	sudo apt-get install -y gcc-aarch64-linux-gnu binutils-aarch64-linux-gnu g++-aarch64-linux-gnu
-	cmake -DCMAKE_C_COMPILER=aarch64-linux-gnu-gcc -DCMAKE_CXX_COMPILER=aarch64-linux-gnu-g++ -DCOMPUTE_BACKEND=cpu .
-elif [ "${build_os:0:5}" == macos ] && [ "${build_arch}" == aarch64 ]; then
+if [ "${build_os:0:5}" == macos ] && [ "${build_arch}" == aarch64 ]; then
 	cmake -DCMAKE_OSX_ARCHITECTURES=arm64 -DCOMPUTE_BACKEND=cpu .
 else
 	cmake -DCOMPUTE_BACKEND=cpu .

.github/scripts/build-cuda.sh

@@ -8,6 +8,11 @@ set -xeuo pipefail
 
 if [[ -v cuda_targets ]]; then
     build_capability="${cuda_targets}"
+elif [ "${build_arch}" = "aarch64" ]; then
+    build_capability="75;80;90"
+
+    # CUDA 12.8: Add sm100
+    [[ "${cuda_version}" == 12.8.* ]] && build_capability="75;80;90;100"
 else
     # By default, target Maxwell through Hopper.
     build_capability="50;52;60;61;70;75;80;86;89;90"

.github/workflows/python-package.yml

@@ -198,16 +198,23 @@ jobs:
 
   audit-wheels:
     needs: build-wheels
-    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        os: [ubuntu-22.04, ubuntu-22.04-arm]
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+    runs-on: ${{ matrix.os }}
     env:
       PIP_DISABLE_PIP_VERSION_CHECK: 1
     steps:
       - uses: actions/checkout@v4
-      - name: Download all wheels
+      - name: Download wheel
         uses: actions/download-artifact@v4
         with:
-          merge-multiple: true
-          pattern: "bdist_wheel_*"
+          name: bdist_wheel_${{ matrix.os }}_${{ matrix.arch }}
           path: wheels/
       - name: Set up Python
         uses: actions/setup-python@v5

.github/workflows/tests.yml

@@ -17,8 +17,14 @@ jobs:
   build-cpu:
     strategy:
       matrix:
-        os: [ubuntu-22.04, windows-2025]
-        arch: [x86_64]
+        os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025]
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+          - os: windows-2025
+            arch: x86_64
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4
@@ -44,8 +50,14 @@ jobs:
     strategy:
       matrix:
         cuda_version: ["11.8.0", "12.8.1"]
-        os: [ubuntu-22.04, windows-2025]
-        arch: [x86_64]
+        os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025]
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+          - os: windows-2025
+            arch: x86_64
     runs-on: ${{ matrix.os }}
 
     steps:
@@ -70,7 +82,7 @@ jobs:
         run: bash .github/scripts/build-cuda.sh
         env:
           build_os: ${{ matrix.os }}
-          build_arch: x86_64
+          build_arch: ${{ matrix.arch }}
           cuda_version: ${{ matrix.cuda_version }}
           cuda_targets: "75"
 
@@ -86,9 +98,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-22.04, windows-2025]
-        arch: [x86_64]
+        os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025]
         torch_version: ["2.7.0"]
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+          - os: windows-2025
+            arch: x86_64
     runs-on: ${{ matrix.os }}
     env:
       BNB_TEST_DEVICE: cpu