Update the ci file

pytorch · Jun 26, 2024 · eb922e7 · eb922e7
1 parent b14de2b
commit eb922e7
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 35 deletions.
diff --git a/.github/workflows/torchao.yml b/.github/workflows/torchao.yml
@@ -5,15 +5,6 @@ on:
     - cron: '00 18 * * *' # run at 6:00 PM UTC, K8s containers will roll out at 12PM EST
 
 jobs:
-  torchbench:
-    uses: ./.github/workflows/_linux-benchmark-cuda.yml
-    with:
-      userbenchmark: "torchao"
-      userbenchmark-run-args: "--ci --dashboard --torchbench"
-    secrets:
-      HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
-      AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
   timm:
     uses: ./.github/workflows/_linux-benchmark-cuda.yml
     with:
@@ -23,15 +14,24 @@ jobs:
       HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
       AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-  huggingface:
-    uses: ./.github/workflows/_linux-benchmark-cuda.yml
-    with:
-      userbenchmark: "torchao"
-      userbenchmark-run-args: "--ci --dashboard --huggingface"
-    secrets:
-      HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
-      AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+  # torchbench:
+  #   uses: ./.github/workflows/_linux-benchmark-cuda.yml
+  #   with:
+  #     userbenchmark: "torchao"
+  #     userbenchmark-run-args: "--ci --dashboard --torchbench"
+  #   secrets:
+  #     HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+  #     AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+  #     AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+  # huggingface:
+  #   uses: ./.github/workflows/_linux-benchmark-cuda.yml
+  #   with:
+  #     userbenchmark: "torchao"
+  #     userbenchmark-run-args: "--ci --dashboard --huggingface"
+  #   secrets:
+  #     HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+  #     AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+  #     AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}

diff --git a/userbenchmark/torchao/run.py b/userbenchmark/torchao/run.py
@@ -1,4 +1,5 @@
 import argparse
+import itertools
 
 from userbenchmark.utils import get_output_dir
 from typing import List
@@ -8,16 +9,35 @@
 OUTPUT_DIR = get_output_dir(BM_NAME)
 OUTPUT_DIR.mkdir(exist_ok=True, parents=True)
 
-CI_ARGS = [
-    ["--progress", "--performance", "--inference", "--bfloat16", "--quantization", "noquant", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_noquant_timm_models_bfloat16_inference_cuda_performance.csv').resolve())}"],
-    ["--progress", "--accuracy", "--inference", "--bfloat16", "--quantization", "noquant", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_noquant_timm_models_bfloat16_inference_cuda_accuracy.csv').resolve())}"],
-    ["--progress", "--performance", "--inference", "--bfloat16", "--quantization", "int8dynamic", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_int8dynamic_timm_models_bfloat16_inference_cuda_performance.csv').resolve())}"],
-    ["--progress", "--accuracy", "--inference", "--bfloat16", "--quantization", "int8dynamic", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_int8dynamic_timm_models_bfloat16_inference_cuda_accuracy.csv').resolve())}"],
-    ["--progress", "--performance", "--inference", "--bfloat16", "--quantization", "int8weightonly", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_int8weightonly_timm_models_bfloat16_inference_cuda_performance.csv').resolve())}"],
-    ["--progress", "--accuracy", "--inference", "--bfloat16", "--quantization", "int8weightonly", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_int8weightonly_timm_models_bfloat16_inference_cuda_accuracy.csv').resolve())}"],
-    ["--progress", "--performance", "--inference", "--bfloat16", "--quantization", "autoquant", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_autoquant_timm_models_bfloat16_inference_cuda_performance.csv').resolve())}"],
-    ["--progress", "--accuracy", "--inference", "--bfloat16", "--quantization", "autoquant", "--output", f"{str(OUTPUT_DIR.joinpath('torchao_autoquant_timm_models_bfloat16_inference_cuda_accuracy.csv').resolve())}"],
-]
+
+def _get_ci_args(backend: str, modelset: str, dtype, mode: str, device: str, experiment: str) -> List[List[str]]:
+    if modelset == "timm":
+        modelset_full_name = "timm_models"
+    else:
+        modelset_full_name = modelset
+    output_file_name = f"torchao_{backend}_{modelset_full_name}_{dtype}_{mode}_{device}_{experiment}.csv"
+    ci_args = [
+        "--progress",
+        f"--{modelset}",
+        "--quantization",
+        f"{backend}",
+        f"--{mode}",
+        f"--{dtype}",
+        f"--{experiment}",
+        "--output",
+        f"{str(OUTPUT_DIR.joinpath(output_file_name).resolve())}"
+    ]
+    return ci_args
+
+def _get_full_ci_args(modelset: str) -> List[List[str]]:
+    backends = ["autoquant", "int8dynamic", "int8weightonly", "noquant"]
+    modelset = [modelset]
+    dtype = ["bfloat16"]
+    mode = ["inference"]
+    device = ["cuda"]
+    experiment = ["performance", "accuracy"]
+    cfgs = itertools.product(*[backends, modelset, dtype, mode, device, experiment])
+    return [ _get_ci_args(*cfg) for cfg in cfgs]
 
 def _get_output(pt2_args):
     if "--output" in pt2_args:
@@ -42,20 +62,18 @@ def run(args: List[str]):
     args, pt2_args = parser.parse_known_args(args)
 
     if args.ci:
-        group_pt2_args = CI_ARGS
         if args.timm:
-            group_pt2_args.append("--timm")
+            benchmark_args = _get_full_ci_args(modelset="timm")
         elif args.huggingface:
-            group_pt2_args.append("--huggingface")
+            benchmark_args = _get_full_ci_args(modelset="huggingface")
         elif args.torchbench:
-            group_pt2_args.append("--torchbench")
+            benchmark_args = _get_full_ci_args(modelset="torchbench")
         else:
             raise RuntimeError("CI mode must run with --timm, --huggingface, or --torchbench")
     else:
-        group_pt2_args = [pt2_args]
-
+        benchmark_args = [pt2_args]
 
-    output_files = [_run_pt2_args(pt2_args) for pt2_args in group_pt2_args]
+    output_files = [_run_pt2_args(args) for args in benchmark_args]
     # Post-processing
     if args.dashboard:
         post_ci_process(output_files)