benchmark.py

"""
Core benchmark harness.

For a given ExperimentConfig, this:
  1. loads the model at the requested dtype (optionally wrapped in torch.compile)
  2. runs a few warm-up passes (and separately times the first / compiling pass)
  3. times `n_iters` steady-state forward passes
  4. records peak GPU memory
  5. checks correctness against a cached fp32-eager reference (same inputs)

Every step is wrapped so a failing config (e.g. OOM, unsupported dtype on
CPU) returns a BenchmarkResult with `error` set instead of crashing the
whole sweep -- the orchestrating agent can just move on.
"""
from __future__ import annotations

import time
from functools import lru_cache
from typing import Dict, Tuple

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from config import BenchmarkResult, CompileMode, ExperimentConfig

DEFAULT_MODEL_NAME = "gpt2" 

_reference_cache: Dict[Tuple[str, int, int, int], torch.Tensor] = {}

@lru_cache(maxsize=4)
def _get_tokenizer(model_name: str):
    tok = AutoTokenizer.from_pretrained(model_name)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    return tok

def _make_inputs(tokenizer, batch_size: int, seq_len: int, device: str, seed: int = 42):
    g = torch.Generator(device="cpu").manual_seed(seed)
    vocab_size = tokenizer.vocab_size
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), generator=g).to(device)
    attention_mask = torch.ones((batch_size, seq_len), dtype=torch.long, device=device)
    return input_ids, attention_mask

def _timed_forward(model, inputs: dict, device: str):
    if device == "cuda":
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        with torch.no_grad():
            out = model(**inputs)
        end.record()
        torch.cuda.synchronize()
        elapsed_ms = start.elapsed_time(end)
    else:
        t0 = time.perf_counter()
        with torch.no_grad():
            out = model(**inputs)
        elapsed_ms = (time.perf_counter() - t0) * 1000
    return out, elapsed_ms

def _get_reference_logits(
    model_name: str, config: ExperimentConfig, input_ids, attention_mask, device: str, seed: int
) -> torch.Tensor:
    # `seed` must be part of the key: _make_inputs is seed-dependent, so a
    # different seed produces different input_ids for the same
    # (model_name, batch_size, seq_len), and a stale cached reference would
    # silently be compared against the wrong inputs.
    key = (model_name, config.batch_size, config.seq_len, seed)
    if key not in _reference_cache:
        ref_model = (
            AutoModelForCausalLM.from_pretrained(model_name)
            .to(device=device, dtype=torch.float32)
            .eval()
        )
        with torch.no_grad():
            ref_out = ref_model(input_ids=input_ids, attention_mask=attention_mask)
        _reference_cache[key] = ref_out.logits.detach().float().cpu()
        del ref_model
        if device == "cuda":
            torch.cuda.empty_cache()
    return _reference_cache[key]

def run_benchmark(
    config: ExperimentConfig,
    model_name: str = DEFAULT_MODEL_NAME,
    device: str = "cuda",
    n_warmup: int = 5,
    n_iters: int = 20,
    seed: int = 42,
    check_correctness: bool = True,
) -> BenchmarkResult:
    try:
        torch.manual_seed(seed)
        torch_dtype = config.dtype.to_torch()

        if device == "cpu" and torch_dtype in (torch.float16, torch.bfloat16):
            torch_dtype = torch.float32

        tokenizer = _get_tokenizer(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        model = model.to(device=device, dtype=torch_dtype).eval()

        if config.compile_mode != CompileMode.EAGER:
            model = torch.compile(model, mode=config.compile_mode.value)

        input_ids, attention_mask = _make_inputs(
            tokenizer, config.batch_size, config.seq_len, device, seed=seed
        )
        inputs = {"input_ids": input_ids, "attention_mask": attention_mask}

        compile_overhead_s = None
        t_first_start = time.perf_counter()
        _timed_forward(model, inputs, device)
        if config.compile_mode != CompileMode.EAGER:
            compile_overhead_s = time.perf_counter() - t_first_start

        for _ in range(max(0, n_warmup - 1)):
            _timed_forward(model, inputs, device)

        if device == "cuda":
            torch.cuda.reset_peak_memory_stats(device)

        latencies_ms = []
        last_logits = None
        for _ in range(n_iters):
            out, elapsed_ms = _timed_forward(model, inputs, device)
            latencies_ms.append(elapsed_ms)
            last_logits = out.logits

        mean_latency_ms = float(np.mean(latencies_ms))
        std_latency_ms = float(np.std(latencies_ms))
        seconds = mean_latency_ms / 1000.0
        throughput = (config.batch_size * config.seq_len) / seconds if seconds > 0 else 0.0

        peak_memory_mb = 0.0
        if device == "cuda":
            peak_memory_mb = torch.cuda.max_memory_allocated(device) / (1024**2)

        max_abs_error = None
        top1_agreement = None
        correctness_passed = None
        if check_correctness:
            reference = _get_reference_logits(model_name, config, input_ids, attention_mask, device, seed)
            current = last_logits.detach().float().cpu()
            max_abs_error = (current - reference).abs().max().item()
            top1_agreement = (current.argmax(dim=-1) == reference.argmax(dim=-1)).float().mean().item()
            correctness_passed = top1_agreement >= 0.95

        return BenchmarkResult(
            config=config,
            mean_latency_ms=mean_latency_ms,
            std_latency_ms=std_latency_ms,
            throughput_tokens_per_s=throughput,
            peak_memory_mb=peak_memory_mb,
            compile_overhead_s=compile_overhead_s,
            max_abs_error=max_abs_error,
            top1_agreement=top1_agreement,
            correctness_passed=correctness_passed,
        )

    except torch.cuda.OutOfMemoryError as e: 
        if device == "cuda":
            torch.cuda.empty_cache()
        return BenchmarkResult(
            config=config,
            mean_latency_ms=0.0,
            std_latency_ms=0.0,
            throughput_tokens_per_s=0.0,
            peak_memory_mb=0.0,
            error=f"OOM: {e}",
        )
    except Exception as e:  
        return BenchmarkResult(
            config=config,
            mean_latency_ms=0.0,
            std_latency_ms=0.0,
            throughput_tokens_per_s=0.0,
            peak_memory_mb=0.0,
            error=f"{type(e).__name__}: {e}",
        )