run_all_benchmarks.py

#!/usr/bin/env python3
"""Run benchmark_transformer for every (size, context_length) combination
and dump results as JSON for the writeup."""

import json
import gc
import sys

# Ensure the benchmark module is importable
from cs336_systems.benchmark_transformer import (
    get_model_configs,
    create_model,
    generate_random_batch,
    benchmark_model,
    count_params,
    VOCAB_SIZE,
    BATCH_SIZE,
)

import jax

SIZES = list(get_model_configs().keys())
CONTEXT_LENGTHS = [128, 256, 512]
WARM_UP_STEPS = 5
N_STEPS = 10

# (size, ctx_len) combos known to OOM on this GPU – skip gracefully
_SKIP = set()

def main():
    backend = jax.default_backend()
    print(f"JAX backend: {backend}")
    print(f"Sizes: {SIZES}")
    print(f"Context lengths: {CONTEXT_LENGTHS}")
    print(f"Warm-up steps: {WARM_UP_STEPS}, Measurement steps: {N_STEPS}")
    print(f"Batch size: {BATCH_SIZE}, Vocab size: {VOCAB_SIZE}")
    print("=" * 70)

    all_results = []

    for size in SIZES:
        for ctx_len in CONTEXT_LENGTHS:
            tag = f"{size} / ctx={ctx_len}"
            print(f"\n>>> {tag}")
            try:
                model, config = create_model(size=size, context_length=ctx_len)
                n_params = count_params(model)
                print(f"    params = {n_params:,}")

                key = jax.random.key(42)
                k1, k2 = jax.random.split(key)
                inp = generate_random_batch(k1, BATCH_SIZE, ctx_len, VOCAB_SIZE)
                tgt = generate_random_batch(k2, BATCH_SIZE, ctx_len, VOCAB_SIZE)

                res = benchmark_model(model, inp, tgt,
                                      n_steps=N_STEPS,
                                      warm_up_steps=WARM_UP_STEPS)

                entry = {
                    "size": size,
                    "context_length": ctx_len,
                    "n_params": n_params,
                    "forward_mean_ms": res["forward"]["mean"] * 1000,
                    "forward_std_ms": res["forward"]["std"] * 1000,
                    "backward_mean_ms": res["backward"]["mean"] * 1000,
                    "backward_std_ms": res["backward"]["std"] * 1000,
                }
                all_results.append(entry)

                print(f"    forward:  {entry['forward_mean_ms']:.2f} ± {entry['forward_std_ms']:.2f} ms")
                print(f"    backward: {entry['backward_mean_ms']:.2f} ± {entry['backward_std_ms']:.2f} ms")

                # Free memory between runs
                del model, inp, tgt
                gc.collect()

            except Exception as e:
                print(f"    SKIPPED ({e})", file=sys.stderr)
                all_results.append({
                    "size": size,
                    "context_length": ctx_len,
                    "error": str(e),
                })

    out_path = "benchmark_results.json"
    with open(out_path, "w") as f:
        json.dump(all_results, f, indent=2)
    print(f"\nResults written to {out_path}")


if __name__ == "__main__":
    main()