[Zentorch] Benchmarks and CI

.github/workflows/benchmark_zentorch_plugin.yaml

@@ -0,0 +1,56 @@
+name: ZenTorch Plugin - Benchmark
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [main]
+    paths:
+      - "benchmarks/zentorch/**/*.py"
+      - ".github/workflows/benchmark_zentorch_plugin.yaml"
+  pull_request:
+    branches: [main]
+    paths:
+      - "benchmarks/zentorch/**/*.py"
+      - ".github/workflows/benchmark_zentorch_plugin.yaml"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build_and_benchmark:
+    runs-on: [self-hosted, amd-cpu, epyc, genoa]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Copy zentorch binary to current directory
+        run: cp
+          /home/github_actions/actions-runner/zentorch-0.1.0-cp38-cp38-manylinux2014_x86_64.whl
+          docker/transformers-pytorch-amd-cpu-zentorch/zentorch-0.1.0-cp38-cp38-manylinux2014_x86_64.whl
+
+      - name: Build Docker image
+        run: docker build
+          --build-arg USER_ID=$(id -u)
+          --build-arg GROUP_ID=$(id -g)
+          --tag optimum-amd-zentorch:2.2.1
+          docker/transformers-pytorch-amd-cpu-zentorch
+
+      - name: Run benchmarks
+        uses: addnab/docker-run-action@v3
+        env:
+          HF_WRITE_TOKEN: ${{ secrets.HF_WRITE_TOKEN }}
+        with:
+          image: optimum-amd-zentorch:2.2.1
+          options: |
+            --rm
+            --shm-size 64G
+            --env HF_WRITE_TOKEN
+            --volume ${{ github.workspace }}:/workspace
+            --workdir /workspace
+          run: |
+            pip install .[zentorch,benchmarks]
+            huggingface-cli login --token $HF_WRITE_TOKEN
+            amdrun python benchmarks/zentorch/benchmark_encoder_models.py
+            amdrun python benchmarks/zentorch/benchmark_text_generation_models.py
+            amdrun python benchmarks/zentorch/benchmark_image_diffusion_pipelines.py

benchmarks/zentorch/benchmark_common_models.py

@@ -0,0 +1,62 @@
+from optimum_benchmark.backends.pytorch.config import PyTorchConfig
+from optimum_benchmark.benchmarks.inference.config import InferenceConfig
+from optimum_benchmark.experiment import ExperimentConfig, launch
+from optimum_benchmark.launchers.process.config import ProcessConfig
+
+
+REPO_ID = "optimum-amd/zentorch-benchmarks"
+EXPERIMENT_NAME = "common_models"
+
+COMMON_MODELS_LIST = [
+    "google-bert/bert-base-uncased",
+]
+INPUT_SHAPES = {
+    "batch_size": 1,
+    "sequence_length": 2,
+}
+TORCH_COMPILE_CONFIG = {
+    "backend": "zentorch",
+}
+
+
+def benchmark_common_models():
+    for model in COMMON_MODELS_LIST:
+        launcher_config = ProcessConfig(start_method="spawn")  # isolated process
+        benchmark_config = InferenceConfig(
+            memory=True,
+            latency=True,
+            input_shapes=INPUT_SHAPES,
+        )
+        backend_config = PyTorchConfig(
+            model=model,
+            device="cpu",
+            no_weights=True,
+            torch_compile=True,
+            torch_compile_config=TORCH_COMPILE_CONFIG,
+        )
+
+        experiment_config = ExperimentConfig(
+            experiment_name=EXPERIMENT_NAME,
+            benchmark=benchmark_config,
+            launcher=launcher_config,
+            backend=backend_config,
+        )
+
+        benchmark_report = launch(experiment_config)
+
+        experiment_config.push_to_hub(
+            commit_message="Added experiment config",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+        benchmark_report.push_to_hub(
+            commit_message="Added benchmark report",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+
+
+if __name__ == "__main__":
+    benchmark_common_models()

benchmarks/zentorch/benchmark_image_diffusion_pipelines.py

@@ -0,0 +1,65 @@
+from optimum_benchmark.backends.pytorch.config import PyTorchConfig
+from optimum_benchmark.benchmarks.inference.config import InferenceConfig
+from optimum_benchmark.experiment import ExperimentConfig, launch
+from optimum_benchmark.launchers.process.config import ProcessConfig
+
+
+REPO_ID = "optimum-amd/zentorch-benchmarks"
+EXPERIMENT_NAME = "image_diffusion_pipelines"
+
+IMAGE_DIFFUSION_PIPELINES_LIST = [
+    "stabilityai/stable-diffusion-2-1",
+]
+INPUT_SHAPES = {
+    "batch_size": 1,
+}
+CALL_KWARGS = {
+    "num_inference_steps": 2,
+    "num_images_per_prompt": 1,
+}
+TORCH_COMPILE_CONFIG = {
+    "backend": "zentorch",
+}
+
+
+def benchmark_image_diffusion_pipelines():
+    for model in IMAGE_DIFFUSION_PIPELINES_LIST:
+        launcher_config = ProcessConfig(start_method="spawn")  # isolated process
+        benchmark_config = InferenceConfig(
+            memory=True,
+            latency=True,
+            input_shapes=INPUT_SHAPES,
+            call_kwargs=CALL_KWARGS,
+        )
+        backend_config = PyTorchConfig(
+            model=model,
+            device="cpu",
+            torch_compile=True,
+            torch_compile_config=TORCH_COMPILE_CONFIG,
+        )
+
+        experiment_config = ExperimentConfig(
+            experiment_name=EXPERIMENT_NAME,
+            benchmark=benchmark_config,
+            launcher=launcher_config,
+            backend=backend_config,
+        )
+
+        benchmark_report = launch(experiment_config)
+
+        experiment_config.push_to_hub(
+            commit_message="Added experiment config",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+        benchmark_report.push_to_hub(
+            commit_message="Added benchmark report",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+
+
+if __name__ == "__main__":
+    benchmark_image_diffusion_pipelines()

benchmarks/zentorch/benchmark_text_generation_with_static_cache_models.py

@@ -0,0 +1,75 @@
+from optimum_benchmark.backends.pytorch.config import PyTorchConfig
+from optimum_benchmark.benchmarks.inference.config import InferenceConfig
+from optimum_benchmark.experiment import ExperimentConfig, launch
+from optimum_benchmark.launchers.process.config import ProcessConfig
+
+
+REPO_ID = "optimum-amd/zentorch-benchmarks"
+EXPERIMENT_NAME = "text_generation_with_static_cache_models"
+
+# for list with static cache support
+# https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+_setup_cache%28self&type=code
+TEXT_GENERATION_WITH_STATIC_CACHE_MODELS_LIST = [
+    "google/gemma-2b",
+    # "google/gemma-7b",
+    # "huggyllama/llama-7b",
+    # "meta-llama/Llama-2-7b-hf",
+    # "mistralai/Mistral-7B-v0.1",
+]
+INPUT_SHAPES = {
+    "batch_size": 1,
+    "sequence_length": 2,
+}
+GENERATE_KWARGS = {
+    "max_new_tokens": 2,
+    "min_new_tokens": 2,
+}
+TORCH_COMPILE_CONFIG = {
+    "backend": "zentorch",
+}
+CACHE_IMPLEMENTATION = "static"
+
+
+def benchmark_text_generation_with_static_cache_models():
+    for model in TEXT_GENERATION_WITH_STATIC_CACHE_MODELS_LIST:
+        launcher_config = ProcessConfig(start_method="spawn")  # isolated process
+        benchmark_config = InferenceConfig(
+            memory=True,
+            latency=True,
+            input_shapes=INPUT_SHAPES,
+            generate_kwargs=GENERATE_KWARGS,
+        )
+        backend_config = PyTorchConfig(
+            model=model,
+            device="cpu",
+            no_weights=True,
+            torch_compile=True,
+            torch_compile_config=TORCH_COMPILE_CONFIG,
+            cache_implementation=CACHE_IMPLEMENTATION,
+        )
+
+        experiment_config = ExperimentConfig(
+            experiment_name=EXPERIMENT_NAME,
+            benchmark=benchmark_config,
+            launcher=launcher_config,
+            backend=backend_config,
+        )
+
+        benchmark_report = launch(experiment_config)
+
+        experiment_config.push_to_hub(
+            commit_message="Added experiment config",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+        benchmark_report.push_to_hub(
+            commit_message="Added benchmark report",
+            subfolder=f"{EXPERIMENT_NAME}/{model}",
+            repo_id=REPO_ID,
+            private=True,
+        )
+
+
+if __name__ == "__main__":
+    benchmark_text_generation_with_static_cache_models()