[proof-of-concept] benchmarks dashboard

.github/workflows/perf-dashboard.yml

@@ -0,0 +1,113 @@
+name: Perf dashboard
+
+# Runs the microbenchmarks on a self-hosted AMD GPU runner, appends the results
+# to the per-family CSV shards, and publishes the static dashboard to gh-pages.
+# Prerequisite: the `gh-pages` branch must already exist (seed it once with the
+# dashboard files + an empty data/ dir).
+#
+# Opt-in only -- the expensive GPU benchmarks do NOT fire on every push/PR. The
+# `bench` job runs when explicitly requested:
+#   * a PR carries the `ci-perf-test` label, or
+#   * a push's head commit title contains `[ci-perf-test]`, or
+#   * a manual `workflow_dispatch`.
+#
+# TODO(perf-dashboard): prune PR shards on close. PR runs write
+# perf-<family>-pr<N>.csv and currently accumulate indefinitely. Add a companion
+# job triggered on `pull_request: closed` that deletes perf-*-pr<N>.csv and
+# removes its rows from index.csv (e.g. a `dashboard_prune.py --pr <N>`), then
+# commits gh-pages -- so merged/closed PRs don't leave dead data behind.
+
+on:
+  # Triggers are intentionally broad; the `bench` job's `if:` is the opt-in gate.
+  push:
+    branches: [dev]
+    paths:
+      - 'transformer_engine/**'
+      - 'benchmarks/microbenchmarks/**'
+  pull_request:
+    # `labeled` so that adding the `ci-perf-test` label triggers a run; no path
+    # filter so the label works on any PR.
+    types: [opened, synchronize, reopened, labeled]
+  workflow_dispatch:
+
+permissions:
+  contents: write            # push commits to gh-pages
+
+# Serialize every run globally: they all push to the single mutable gh-pages
+# branch, so we must not push concurrently.
+concurrency:
+  group: perf-dashboard
+  cancel-in-progress: false
+
+jobs:
+  bench:
+    # Opt-in gate: only spend a GPU runner when a perf run was explicitly asked
+    # for (PR label, [ci-perf-test] commit tag, or manual dispatch). A skipped
+    # job consumes no runner.
+    if: >-
+      github.event_name == 'workflow_dispatch'
+      || (github.event_name == 'pull_request'
+          && contains(github.event.pull_request.labels.*.name, 'ci-perf-test'))
+      || (github.event_name == 'push'
+          && contains(github.event.head_commit.message, '[ci-perf-test]'))
+    runs-on: [self-hosted, linux, gfx950]      # your AMD GPU runner label(s)
+    # GitHub's hosted runners have no AMD GPU; a self-hosted runner is required.
+    # If TE isn't in the runner's base env, run the job inside your TE container:
+    # container:
+    #   image: <your-rocm-te-image>
+    #   options: --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 16g
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Checkout published site (gh-pages)
+        uses: actions/checkout@v4
+        with:
+          ref: gh-pages
+          path: site
+
+      - name: Run microbenchmarks
+        working-directory: benchmarks/microbenchmarks
+        run: |
+          # Keep going if one benchmark fails; publish whatever produced a CSV.
+          for b in gemm gemm_fp8 casting normalization grouped_gemm; do
+            python "benchmark_${b}.py" --csv || echo "::warning::benchmark_${b} failed"
+          done
+
+      - name: Ingest into per-family CSV shards
+        working-directory: benchmarks/microbenchmarks
+        run: |
+          set -euo pipefail
+          if [ "${{ github.event_name }}" = "pull_request" ]; then
+            python dashboard_ingest.py benchmark_*.csv \
+              --out-dir "$GITHUB_WORKSPACE/site/data" \
+              --pr "${{ github.event.number }}" \
+              --commit "${{ github.event.pull_request.head.sha }}" \
+              --arch gfx950 --runner "${{ runner.name }}"
+          else
+            python dashboard_ingest.py benchmark_*.csv \
+              --out-dir "$GITHUB_WORKSPACE/site/data" \
+              --ref dev \
+              --commit "$GITHUB_SHA" \
+              --arch gfx950 --runner "${{ runner.name }}"
+          fi
+
+      - name: Refresh the static front-end (keep data/ as-is)
+        run: |
+          rsync -a --exclude=data \
+            benchmarks/microbenchmarks/dashboard/ "$GITHUB_WORKSPACE/site/"
+
+      - name: Commit & push to gh-pages
+        working-directory: site
+        run: |
+          set -euo pipefail
+          git config user.name  "perf-bot"
+          git config user.email "perf-bot@users.noreply.github.com"
+          git add -A
+          if git diff --cached --quiet; then
+            echo "no changes to publish"; exit 0
+          fi
+          git commit -m "perf: ${GITHUB_SHA::8} (${{ github.event_name }}) $(date -u +%FT%TZ)"
+          # Serialized by `concurrency`, but rebase-then-push is cheap insurance.
+          git pull --rebase origin gh-pages || true
+          git push origin gh-pages

benchmarks/microbenchmarks/.gitignore

@@ -0,0 +1,10 @@
+# Generated microbenchmark result CSVs (wide format; ingested into the dashboard,
+# then discarded). Not source -- never commit these.
+benchmark_*.csv
+
+# Generated single-file bundle (dashboard/dist/dashboard.html embeds the CSV data).
+dashboard/dist/
+
+# Local gh-pages checkout used by dashboard_redeploy.sh / build_bundle.py.
+# It's a separate repo clone full of generated data -- keep it out of TE.
+te-dash/

benchmarks/microbenchmarks/benchmark_gemm.py

@@ -10,7 +10,7 @@
 import transformer_engine.pytorch as te
 from utils import (
     generate_gemm_test_cases,
-    time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    time_forward_backward, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
 )
 
 BENCHMARK_LABEL = "GEMM"
@@ -26,20 +26,23 @@ def bench_gemm(Case, M, N, K, dtype):
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
+    def zero_grads():
+        x.grad = None
+        linear.weight.grad = None
+
     def fwd_bwd_func():
         out = linear(x)
         out.backward(grad_out)
-        x.grad = None
-        linear.weight.grad = None
+        zero_grads()
 
     fwd_bwd_func()
 
     fwd_flops = 2 * M * N * K
     bwd_flops = 2 * fwd_flops  # dX + dW
 
-    fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_bwd_ms, fwd_bwd_measurement = time_func(fwd_bwd_func)
-    bwd_ms = fwd_bwd_ms - fwd_ms
+    fwd_ms, bwd_ms, record_kwargs = time_forward_backward(
+        fwd_func, fwd_bwd_func, grad_out
+    )
 
     fwd_tflops = compute_tflops(fwd_flops, fwd_ms)
     bwd_tflops = compute_tflops(bwd_flops, bwd_ms)
@@ -51,9 +54,7 @@ def fwd_bwd_func():
         fwd_tflops,
         bwd_ms,
         bwd_tflops,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        **record_kwargs,
     )
 
 

benchmarks/microbenchmarks/benchmark_gemm_fp8.py

@@ -16,7 +16,7 @@
 from transformer_engine.common.recipe import DelayedScaling, Format
 from utils import (
     generate_gemm_test_cases,
-    time_func, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
+    time_forward_backward, compute_tflops, make_forward_backward_metric_records, run_benchmarks,
 )
 
 RECIPES = {
@@ -43,19 +43,22 @@ def fwd_func():
         with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE):
             return linear(x)
 
+    def zero_grads():
+        x.grad = None
+        linear.weight.grad = None
+
     def fwd_bwd_func():
         with te.fp8_autocast(enabled=True, fp8_recipe=FP8_RECIPE):
             out = linear(x)
             out.backward(grad_out)
-        x.grad = None
-        linear.weight.grad = None
+        zero_grads()
 
     fwd_flops = 2 * M * N * K
     bwd_flops = 2 * fwd_flops
 
-    fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_bwd_ms, fwd_bwd_measurement = time_func(fwd_bwd_func)
-    bwd_ms = fwd_bwd_ms - fwd_ms
+    fwd_ms, bwd_ms, record_kwargs = time_forward_backward(
+        fwd_func, fwd_bwd_func, grad_out
+    )
 
     fwd_tflops = compute_tflops(fwd_flops, fwd_ms)
     bwd_tflops = compute_tflops(bwd_flops, bwd_ms)
@@ -67,9 +70,7 @@ def fwd_bwd_func():
         fwd_tflops,
         bwd_ms,
         bwd_tflops,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        **record_kwargs,
     )
 
 

benchmarks/microbenchmarks/benchmark_grouped_gemm.py

@@ -9,7 +9,7 @@
 import transformer_engine.pytorch as te
 from utils import (
     DTYPE_LIST,
-    time_func,
+    time_forward_backward,
     compute_tflops,
     make_forward_backward_metric_records,
     run_benchmarks,
@@ -122,21 +122,24 @@ def fwd_func_te():
     out_te = fwd_func_te()
     grad_out = torch.randn_like(out_te)
 
-    def fwd_bwd_func_te():
-        out = grouped_linear(x, m_splits, m_splits_tensor=m_splits_tensor)
-        out.backward(grad_out)
+    def zero_grads():
         x.grad = None
         for param in grouped_linear.parameters():
             param.grad = None
 
+    def fwd_bwd_func_te():
+        out = grouped_linear(x, m_splits, m_splits_tensor=m_splits_tensor)
+        out.backward(grad_out)
+        zero_grads()
+
     fwd_bwd_func_te()
 
     fwd_total_flops = 2 * sum_M * N * K
     bwd_total_flops = 2 * fwd_total_flops
 
-    fwd_te_ms, fwd_measurement = time_func(fwd_func_te)
-    fwd_bwd_te_ms, fwd_bwd_measurement = time_func(fwd_bwd_func_te)
-    bwd_te_ms = fwd_bwd_te_ms - fwd_te_ms
+    fwd_te_ms, bwd_te_ms, record_kwargs = time_forward_backward(
+        fwd_func_te, fwd_bwd_func_te, grad_out
+    )
 
     fwd_te_tflops = compute_tflops(fwd_total_flops, fwd_te_ms)
     bwd_te_tflops = compute_tflops(bwd_total_flops, bwd_te_ms)
@@ -148,9 +151,7 @@ def fwd_bwd_func_te():
         fwd_te_tflops,
         bwd_te_ms,
         bwd_te_tflops,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        **record_kwargs,
     )
 
 

benchmarks/microbenchmarks/benchmark_normalization.py

@@ -17,7 +17,7 @@
 import transformer_engine.pytorch as te
 from utils import (
     DTYPE_LIST, MODEL_HIDDEN_SIZES, M_SIZE_LIST,
-    time_func, compute_gbps, make_forward_backward_metric_records, run_benchmarks,
+    time_forward_backward, compute_gbps, make_forward_backward_metric_records, run_benchmarks,
 )
 
 NORM_TYPES = [
@@ -55,22 +55,25 @@ def bench_norm(Case, M, hidden_size, norm_name, norm_cls, dtype):
     out = fwd_func()
     grad_out = torch.randn_like(out)
 
-    def fwd_bwd_func():
-        out = norm(x)
-        out.backward(grad_out)
+    def zero_grads():
         x.grad = None
         for p in norm.parameters():
             p.grad = None
 
+    def fwd_bwd_func():
+        out = norm(x)
+        out.backward(grad_out)
+        zero_grads()
+
     fwd_bwd_func()
 
     elem_bytes = x.element_size()
     fwd_bytes = 2 * M * hidden_size * elem_bytes   # read x, write y
     bwd_bytes = 4 * M * hidden_size * elem_bytes   # read grad+x+y, write grad_x
 
-    fwd_ms, fwd_measurement = time_func(fwd_func)
-    fwd_bwd_ms, fwd_bwd_measurement = time_func(fwd_bwd_func)
-    bwd_ms = fwd_bwd_ms - fwd_ms
+    fwd_ms, bwd_ms, record_kwargs = time_forward_backward(
+        fwd_func, fwd_bwd_func, grad_out
+    )
 
     fwd_gbps = compute_gbps(fwd_bytes, fwd_ms)
     bwd_gbps = compute_gbps(bwd_bytes, bwd_ms)
@@ -82,9 +85,7 @@ def fwd_bwd_func():
         fwd_gbps,
         bwd_ms,
         bwd_gbps,
-        backward_derived=True,
-        fwd_measurement=fwd_measurement,
-        fwd_bwd_measurement=fwd_bwd_measurement,
+        **record_kwargs,
     )