From 05438d3e4bc613df8ef401c76a247b5fa2b9d9b2 Mon Sep 17 00:00:00 2001
From: Xu Zhao <xzhao9@meta.com>
Date: Thu, 30 May 2024 12:42:26 -0400
Subject: [PATCH] Added torchao nightly workflow

---
 .github/workflows/torchao.yml     | 63 +++++++++++++++++++++++++++++++
 userbenchmark/torchao/__init__.py |  1 +
 userbenchmark/torchao/install.py  | 13 +++++++
 userbenchmark/torchao/run.py      | 30 +++++++++++++++
 4 files changed, 107 insertions(+)
 create mode 100644 .github/workflows/torchao.yml
 create mode 100644 userbenchmark/torchao/__init__.py
 create mode 100644 userbenchmark/torchao/install.py
 create mode 100644 userbenchmark/torchao/run.py

diff --git a/.github/workflows/torchao.yml b/.github/workflows/torchao.yml
new file mode 100644
index 0000000000..222b6c7aae
--- /dev/null
+++ b/.github/workflows/torchao.yml
@@ -0,0 +1,63 @@
+name: Torchao nightly workflow (A100)
+on:
+  workflow_dispatch:
+  # TODO: REMOVE THIS BEFORE THE PR IS MERGED!
+  pull_request:
+
+jobs:
+  run-benchmark:
+    environment: docker-s3-upload
+    env:
+      BASE_CONDA_ENV: "torchbench"
+      CONDA_ENV:  "torchao-nightly"
+      PLATFORM_NAME: "gcp_a100"
+      SETUP_SCRIPT: "/workspace/setup_instance.sh"
+      TORCHBENCH_USERBENCHMARK_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TORCHBENCH_USERBENCHMARK_SCRIBE_GRAPHQL_ACCESS_TOKEN }}
+      AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+      AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+      IS_GHA: 1
+      BUILD_ENVIRONMENT: benchmark-nightly
+    if: ${{ github.repository_owner == 'pytorch' }}
+    runs-on: [a100-runner]
+    steps:
+      - name: Checkout TorchBench
+        uses: actions/checkout@v3
+        with:
+          path: benchmark
+      - name: Tune Nvidia GPU
+        run: |
+          sudo nvidia-smi -pm 1
+          sudo nvidia-smi -ac 1215,1410
+          nvidia-smi
+      - name: Clone and setup conda env
+        run: |
+          CONDA_ENV=${BASE_CONDA_ENV} . "${SETUP_SCRIPT}"
+          conda create --name "${CONDA_ENV}" --clone "${BASE_CONDA_ENV}"
+      - name: Install TorchBench
+        run: |
+          set -x
+          . "${SETUP_SCRIPT}"
+          pushd benchmark
+          python install.py
+      - name: Run the torchao userbenchmark
+        run: |
+          . "${SETUP_SCRIPT}"
+          # remove old results if exists
+          if [ -d benchmark-output ]; then rm -Rf benchmark-output; fi
+          pushd benchmark
+          if [ -d .userbenchmark ]; then rm -Rf .userbenchmark; fi
+          # Install torchao
+          python install.py torchao
+          python run_benchmark.py torchao --ci
+          cp -r ./.userbenchmark/torchao ../benchmark-output
+      - name: Upload result to GH Actions Artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: Torchao nightly result
+          path: benchmark-output/
+      - name: Clean up Conda env
+        if: always()
+        run: |
+          . "${SETUP_SCRIPT}"
+          conda deactivate && conda deactivate
+          conda remove -n "${CONDA_ENV}" --all
diff --git a/userbenchmark/torchao/__init__.py b/userbenchmark/torchao/__init__.py
new file mode 100644
index 0000000000..5f04e41781
--- /dev/null
+++ b/userbenchmark/torchao/__init__.py
@@ -0,0 +1 @@
+BM_NAME = "torchao"
\ No newline at end of file
diff --git a/userbenchmark/torchao/install.py b/userbenchmark/torchao/install.py
new file mode 100644
index 0000000000..62c62a4374
--- /dev/null
+++ b/userbenchmark/torchao/install.py
@@ -0,0 +1,13 @@
+import os
+import subprocess
+
+def install_torchao():
+    # Set ARCH list so that we can build fp16 with SM75+, the logic is copied from
+    # pytorch/builder
+    # https://github.com/pytorch/ao/blob/main/packaging/env_var_script_linux.sh#L16C1-L19
+    torchao_env = os.environ
+    torchao_env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+    subprocess.check_call(["pip", "install", "git+https://github.com/pytorch/ao.git"], env=torchao_env)
+
+if __name__ == "__main__":
+    install_torchao()
\ No newline at end of file
diff --git a/userbenchmark/torchao/run.py b/userbenchmark/torchao/run.py
new file mode 100644
index 0000000000..af8e3d6c3c
--- /dev/null
+++ b/userbenchmark/torchao/run.py
@@ -0,0 +1,30 @@
+import argparse
+
+from typing import List
+
+CI_ARGS = [
+    # TIMM
+    ["--timm", "--inference", "--bfloat16", "--quantization", "noquant", "--output", ".userbenchmark/torchao/torchao_noquant_timm_bfloat16_inference_cuda_performance.csv"],
+    # ["--timm", "--inference", "--bfloat16", "--quantization", "int8dynamic", "--output", ".userbenchmark/torchao/torchao_int8dynamic_timm_bfloat16_inference_cuda_performance.csv"],
+    # ["--timm", "--inference", "--bfloat16", "--quantization", "int8weightonly", "--output", ".userbenchmark/torchao/torchao_int8weightonly_timm_bfloat16_inference_cuda_performance.csv"],
+    # ["--timm", "--inference", "--bfloat16", "--quantization", "autoquant", "--output", ".userbenchmark/torchao/torchao_autoquant_timm_bfloat16_inference_cuda_performance.csv"],
+]
+
+
+def _run_pt2_args(pt2_args: List[str]) -> str:
+    from userbenchmark.dynamo.run import run as run_pt2_benchmark
+    run_pt2_benchmark(pt2_args)
+
+
+def run(args: List[str]):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ci", actions="store_true", help="Run the CI workflow")
+    args, pt2_args = parser.parse_known_args(args)
+
+    if args.ci:
+        group_pt2_args = CI_ARGS
+    else:
+        group_pt2_args = [pt2_args]
+    
+    output_files = [_run_pt2_args(pt2_args) for pt2_args in group_pt2_args]
+    print("\n".join(output_files))