Leave ROCm alone for now

Signed-off-by: Huy Do <huydhn@gmail.com>

.ci/pytorch/macos-test.sh

@@ -177,9 +177,6 @@ checkout_install_torchbench() {
   popd
 
   pip install -r .ci/docker/ci_commit_pins/huggingface-requirements.txt
-  # https://github.com/pytorch/pytorch/issues/160689 to remove torchao because
-  # its current version 0.12.0 doesn't work with transformers 4.54.0
-  pip uninstall -y torchao
 
   echo "Print all dependencies after TorchBench is installed"
   python -mpip freeze

.ci/pytorch/test.sh

@@ -778,11 +778,6 @@ test_single_dynamo_benchmark() {
 }
 
 test_inductor_micro_benchmark() {
-  # torchao requires cuda 8.0 or above for bfloat16 support
-  if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
-    export TORCH_CUDA_ARCH_LIST="8.0;8.6"
-  fi
-  install_torchao
   TEST_REPORTS_DIR=$(pwd)/test/test-reports
   if [[ "${TEST_CONFIG}" == *cpu* ]]; then
     test_inductor_set_cpu_affinity
@@ -1664,37 +1659,50 @@ elif [[ "${TEST_CONFIG}" == *operator_benchmark* ]]; then
     elif [[ "${TEST_CONFIG}" == *all* ]]; then
       TEST_MODE="all"
     fi
-
     test_operator_benchmark cpu ${TEST_MODE}
-
   fi
 elif [[ "${TEST_CONFIG}" == *inductor_distributed* ]]; then
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   test_inductor_distributed
 elif [[ "${TEST_CONFIG}" == *inductor-halide* ]]; then
   test_inductor_halide
 elif [[ "${TEST_CONFIG}" == *inductor-triton-cpu* ]]; then
   test_inductor_triton_cpu
 elif [[ "${TEST_CONFIG}" == *inductor-micro-benchmark* ]]; then
+  install_torchao
   test_inductor_micro_benchmark
 elif [[ "${TEST_CONFIG}" == *huggingface* ]]; then
   install_torchvision
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   id=$((SHARD_NUMBER-1))
   test_dynamo_benchmark huggingface "$id"
 elif [[ "${TEST_CONFIG}" == *timm* ]]; then
   install_torchvision
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   id=$((SHARD_NUMBER-1))
   test_dynamo_benchmark timm_models "$id"
 elif [[ "${TEST_CONFIG}" == cachebench ]]; then
   install_torchaudio
   install_torchvision
+  install_torchao
   PYTHONPATH=/torchbench test_cachebench
 elif [[ "${TEST_CONFIG}" == verify_cachebench ]]; then
   install_torchaudio
   install_torchvision
+  install_torchao
   PYTHONPATH=/torchbench test_verify_cachebench
 elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
   install_torchaudio
   install_torchvision
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   id=$((SHARD_NUMBER-1))
   # https://github.com/opencv/opencv-python/issues/885
   pip_install opencv-python==4.8.0.74
@@ -1714,12 +1722,18 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
   fi
 elif [[ "${TEST_CONFIG}" == *inductor_cpp_wrapper* ]]; then
   install_torchvision
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   PYTHONPATH=/torchbench test_inductor_cpp_wrapper_shard "$SHARD_NUMBER"
   if [[ "$SHARD_NUMBER" -eq "1" ]]; then
     test_inductor_aoti
   fi
 elif [[ "${TEST_CONFIG}" == *inductor* ]]; then
   install_torchvision
+  if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+    install_torchao
+  fi
   test_inductor_shard "${SHARD_NUMBER}"
   if [[ "${SHARD_NUMBER}" == 1 ]]; then
     if [[ "${BUILD_ENVIRONMENT}" != linux-jammy-py3.9-gcc11-build ]]; then

.github/ci_commit_pins/torchao.txt

@@ -1 +1 @@
-51c87b6ead6b7e098ada95d6a7609ee873b854cf
+f32431e593d0e9db86c502d3872dd67ee40a005f

.github/workflows/inductor-micro-benchmark.yml

@@ -35,6 +35,8 @@ jobs:
     needs:
       - get-default-label-prefix
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       build-environment: linux-jammy-cuda12.8-py3.10-gcc9-sm80
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks
@@ -43,6 +45,7 @@ jobs:
         { include: [
           { config: "inductor-micro-benchmark", shard: 1, num_shards: 1, runner: "linux.aws.a100", owners: ["oncall:pt2"] },
         ]}
+      build-additional-packages: "vision audio fbgemm torchao"
     secrets: inherit
 
   test:

.github/workflows/inductor-perf-test-nightly-h100.yml

@@ -137,7 +137,6 @@ jobs:
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 720
-      # disable monitor in perf tests, next step is to enable it
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
@@ -154,7 +153,6 @@ jobs:
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 1440
-      # disable monitor in perf tests, next step is to enable it
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4
@@ -173,7 +171,6 @@ jobs:
       docker-image: ${{ needs.build.outputs.docker-image }}
       test-matrix: ${{ needs.build.outputs.test-matrix }}
       timeout-minutes: 720
-      # disable monitor in perf tests for more investigation
       disable-monitor: false
       monitor-log-interval: 15
       monitor-data-collect-interval: 4

.github/workflows/inductor-periodic.yml

@@ -36,6 +36,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-default-label-prefix
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       build-environment: linux-jammy-cuda12.8-py3.10-gcc9-sm86
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks
@@ -128,6 +130,8 @@ jobs:
     needs:
       - get-default-label-prefix
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       runner_prefix: "${{ needs.get-default-label-prefix.outputs.label-type }}"
       build-environment: linux-jammy-cuda12.8-py3.10-gcc9-sm80
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks

.github/workflows/inductor-unittest.yml

@@ -33,6 +33,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       build-environment: linux-jammy-cuda12.8-py3.10-gcc9-sm86
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks
       cuda-arch-list: '8.6'
@@ -45,6 +47,7 @@ jobs:
           { config: "inductor_cpp_wrapper", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
           { config: "inductor_cpp_wrapper", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g5.4xlarge.nvidia.gpu" },
         ]}
+      build-additional-packages: "vision audio torchao"
     secrets: inherit
 
   inductor-test:

.github/workflows/inductor.yml

@@ -49,6 +49,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       build-environment: linux-jammy-cuda12.8-py3.10-gcc9-sm86
       docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks
       cuda-arch-list: '8.6'

.github/workflows/trunk.yml

@@ -239,6 +239,8 @@ jobs:
     uses: ./.github/workflows/_linux-build.yml
     needs: get-label-type
     with:
+      # More memory is needed to build torchao
+      runner: linux.2xlarge.memory
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       build-environment: linux-jammy-py3.9-gcc11
       docker-image-name: ci-image:pytorch-linux-jammy-py3-gcc11-inductor-benchmarks
@@ -246,6 +248,7 @@ jobs:
         { include: [
           { config: "verify_cachebench", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" },
         ]}
+      build-additional-packages: "vision audio torchao"
     secrets: inherit
 
   verify-cachebench-cpu-test: