diff --git a/.github/workflows/torchao_experimental_test.yml b/.github/workflows/torchao_experimental_test.yml index a712e7a624..4c56ec0c0e 100644 --- a/.github/workflows/torchao_experimental_test.yml +++ b/.github/workflows/torchao_experimental_test.yml @@ -14,7 +14,7 @@ jobs: test-cpu-ops: strategy: matrix: - runner: [macos-14] + runner: [macos-14, linux.arm64.2xlarge] runs-on: ${{matrix.runner}} defaults: run: @@ -30,7 +30,8 @@ jobs: python-version: "3.10" miniconda-version: "latest" activate-environment: venv - - name: Install requirements + - name: Install requirements mac + if: runner.os == 'macOS' run: | conda activate venv # Install executorch first because it installs its own version @@ -39,6 +40,13 @@ jobs: pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu --force-reinstall pip install -r dev-requirements.txt USE_CPP=1 TORCHAO_BUILD_KLEIDIAI=1 pip install . + - name: Install requirements linux + if: runner.os == 'Linux' + run: | + conda activate venv + pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu --force-reinstall + pip install -r dev-requirements.txt + BUILD_TORCHAO_EXPERIMENTAL=1 TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_BUILD_KLEIDIAI=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 TORCHAO_PARALLEL_BACKEND=OPENMP pip install . - name: Run python tests run: | conda activate venv @@ -46,6 +54,7 @@ jobs: python torchao/experimental/tests/test_embedding_xbit_quantizer.py python torchao/experimental/tests/test_quant_passes.py - name: Run kernels/cpu/aarch64/tests + if: runner.os == 'macOS' run: | conda activate venv pushd torchao/experimental/kernels/cpu/aarch64/tests @@ -53,6 +62,7 @@ jobs: rm -rf /tmp/cmake-out popd - name: Run torchao/experimental/ops/tests + if: runner.os == 'macOS' run: | conda activate venv pushd torchao/experimental/ops/tests @@ -60,6 +70,7 @@ jobs: rm -rf /tmp/cmake-out popd - name: ET ops build + if: runner.os == 'macOS' run: | conda activate venv pushd torchao/experimental