Skip to content

Commit

Permalink
Added torchao nightly workflow
Browse files Browse the repository at this point in the history
  • Loading branch information
xuzhao9 committed May 30, 2024
1 parent 68d7736 commit 05438d3
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 0 deletions.
63 changes: 63 additions & 0 deletions .github/workflows/torchao.yml
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions userbenchmark/torchao/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
BM_NAME = "torchao"
13 changes: 13 additions & 0 deletions userbenchmark/torchao/install.py
Original file line number Diff line number Diff line change
@@ -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()
30 changes: 30 additions & 0 deletions userbenchmark/torchao/run.py
Original file line number Diff line number Diff line change
@@ -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))

0 comments on commit 05438d3

Please sign in to comment.