compare_models.py

#!/usr/bin/env python3
"""
Compare Baseline vs Tiered HRM — Multi-GPU Benchmark + Comparison Plots.

Usage:
    source venv/bin/activate
    python compare_models.py                              # quick compare
    python compare_models.py --sweep                      # batch-size sweep + plots
    python compare_models.py --iterations 50 --sweep      # thorough
"""

import argparse
import json
import os
import sys
from dataclasses import dataclass, asdict

import torch
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from models.memory_tier import MemoryTierManager
from models.hrm.hrm_tiered import HRM_Tiered
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1


# ═══════════════════════════════════════════════════════════
#  Helpers
# ═══════════════════════════════════════════════════════════

def make_config(batch_size, seq_len, hidden_size, num_heads):
    return {
        "batch_size": batch_size, "seq_len": seq_len,
        "puzzle_emb_ndim": 0, "num_puzzle_identifiers": batch_size,
        "vocab_size": 32,
        "H_cycles": 2, "L_cycles": 2, "H_layers": 4, "L_layers": 4,
        "hidden_size": hidden_size, "expansion": 4.0,
        "num_heads": num_heads, "pos_encodings": "rope",
        "halt_max_steps": 1, "halt_exploration_prob": 0.0,
    }


def make_batch(batch_size, seq_len, device):
    return {
        "inputs": torch.randint(0, 31, (batch_size, seq_len), device=device),
        "labels": torch.randint(0, 31, (batch_size, seq_len), device=device),
        "puzzle_identifiers": torch.arange(batch_size, device=device),
    }


def build_model(arch, config_dict, device):
    """Build using the wrapper classes (same pattern as eval_dummy.py)."""
    if arch == "tiered":
        mm = MemoryTierManager(device=device, enable_tracking=True)
        model = HRM_Tiered(config_dict, memory_manager=mm).to(device)
    else:
        model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
    model.eval()
    return model


@torch.no_grad()
def benchmark(model, batch, device, warmup=5, iterations=20):
    """Time forward pass, return dict of metrics."""
    bs = batch["inputs"].shape[0]

    # Warmup
    for _ in range(warmup):
        carry = model.initial_carry(batch)
        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
        carry.steps = carry.steps.to(device)
        carry.halted = carry.halted.to(device)
        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
        model(carry, batch)

    torch.cuda.reset_peak_memory_stats(device)
    torch.cuda.synchronize()

    latencies = []
    for _ in range(iterations):
        carry = model.initial_carry(batch)
        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
        carry.steps = carry.steps.to(device)
        carry.halted = carry.halted.to(device)
        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        model(carry, batch)
        end.record()
        torch.cuda.synchronize()
        latencies.append(start.elapsed_time(end))

    lat = np.array(latencies)
    return {
        "latency_ms": float(np.mean(lat)),
        "latency_std": float(np.std(lat)),
        "throughput": float(bs / (np.mean(lat) / 1000)),
        "peak_gpu_mb": float(torch.cuda.max_memory_allocated(device) / 1e6),
        "params_m": sum(p.numel() for p in model.parameters()) / 1e6,
    }


# ═══════════════════════════════════════════════════════════
#  Plotting
# ═══════════════════════════════════════════════════════════

def create_plots(base_res, tier_res, sweep_data, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    c_base, c_tier = "#4A90D9", "#E85D75"
    bg, text, grid = "#1a1a2e", "#e0e0e0", "#333355"

    plt.rcParams.update({
        "figure.facecolor": bg, "axes.facecolor": "#16213e",
        "axes.edgecolor": grid, "axes.labelcolor": text,
        "text.color": text, "xtick.color": text, "ytick.color": text,
        "grid.color": grid, "grid.alpha": 0.3,
        "font.family": "sans-serif", "font.size": 11,
    })

    n_plots = 6 if sweep_data else 5
    fig = plt.figure(figsize=(16, 10))
    fig.suptitle("HRM Baseline vs Tiered (SRAM/DRAM) Comparison",
                 fontsize=18, fontweight="bold", y=0.98)
    gs = GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.35)
    labels = ["Baseline", "Tiered"]

    def bar_plot(ax, title, ylabel, vals, fmt=".2f"):
        bars = ax.bar(labels, vals, color=[c_base, c_tier],
                      edgecolor="white", linewidth=0.5, width=0.5)
        ax.set_title(title, fontweight="bold")
        ax.set_ylabel(ylabel)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width()/2, b.get_height() * 1.02,
                    f"{v:{fmt}}", ha="center", fontsize=10, color=text)
        ax.grid(axis="y")

    # 1. Latency
    bar_plot(fig.add_subplot(gs[0, 0]), "Inference Latency", "ms",
             [base_res["latency_ms"], tier_res["latency_ms"]])

    # 2. Throughput
    bar_plot(fig.add_subplot(gs[0, 1]), "Throughput", "samples/sec",
             [base_res["throughput"], tier_res["throughput"]], fmt=".0f")

    # 3. GPU Memory
    bar_plot(fig.add_subplot(gs[0, 2]), "Peak GPU Memory", "MB",
             [base_res["peak_gpu_mb"], tier_res["peak_gpu_mb"]], fmt=".0f")

    # 4. Parameters
    bar_plot(fig.add_subplot(gs[1, 0]), "Model Parameters", "Millions",
             [base_res["params_m"], tier_res["params_m"]], fmt=".1f")

    # 5. Summary text
    ax5 = fig.add_subplot(gs[1, 1])
    speedup = base_res["latency_ms"] / tier_res["latency_ms"]
    mem_diff = tier_res["peak_gpu_mb"] - base_res["peak_gpu_mb"]
    tp_gain = (tier_res["throughput"] / base_res["throughput"] - 1) * 100
    summary = (
        f"Speedup: {speedup:.2f}x\n"
        f"Throughput: {tp_gain:+.1f}%\n"
        f"Memory Δ: {mem_diff:+.0f} MB\n"
        f"Params: identical"
    )
    ax5.text(0.5, 0.5, summary, transform=ax5.transAxes,
             ha="center", va="center", fontsize=14, fontfamily="monospace",
             bbox=dict(boxstyle="round,pad=0.5", facecolor="#0f3460", alpha=0.8))
    ax5.set_title("Summary", fontweight="bold")
    ax5.axis("off")

    # 6. Sweep plot
    ax6 = fig.add_subplot(gs[1, 2])
    if sweep_data:
        bs_list = [s["batch_size"] for s in sweep_data["baseline"]]
        ax6.plot(bs_list, [s["latency_ms"] for s in sweep_data["baseline"]],
                 "o-", color=c_base, label="Baseline", linewidth=2, markersize=6)
        ax6.plot(bs_list, [s["latency_ms"] for s in sweep_data["tiered"]],
                 "s-", color=c_tier, label="Tiered", linewidth=2, markersize=6)
        ax6.set_xlabel("Batch Size")
        ax6.set_ylabel("Latency (ms)")
        ax6.set_title("Latency vs Batch Size", fontweight="bold")
        ax6.legend(facecolor="#16213e", edgecolor=grid)
        ax6.grid(True)
    else:
        ax6.text(0.5, 0.5, "Run with --sweep\nfor batch size\ncomparison",
                 transform=ax6.transAxes, ha="center", va="center", fontsize=12)
        ax6.set_title("Latency vs Batch Size", fontweight="bold")
        ax6.axis("off")

    path = os.path.join(output_dir, "model_comparison.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"  Plot saved → {path}")
    return path


# ═══════════════════════════════════════════════════════════
#  Main
# ═══════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="Compare Baseline vs Tiered HRM")
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--seq-len", type=int, default=81)
    parser.add_argument("--hidden-size", type=int, default=512)
    parser.add_argument("--num-heads", type=int, default=8)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--sweep", action="store_true", help="Batch-size sweep")
    parser.add_argument("--output-dir", type=str, default="benchmark_results")
    args = parser.parse_args()

    device = torch.device("cuda")
    cfg = make_config(args.batch_size, args.seq_len, args.hidden_size, args.num_heads)

    print("=" * 64)
    print("  HRM Model Comparison: Baseline vs Tiered (SRAM/DRAM)")
    print(f"  Device: {torch.cuda.get_device_name(0)}")
    print(f"  Config: bs={args.batch_size}, seq={args.seq_len}, hidden={args.hidden_size}")
    print("=" * 64)

    # ── Build ──
    print("\n  Building Baseline...")
    base_model = build_model("baseline", cfg, device)
    batch = make_batch(args.batch_size, args.seq_len, device)

    print("  Building Tiered...")
    tier_model = build_model("tiered", cfg, device)

    # ── Benchmark ──
    print(f"\n  Benchmarking Baseline ({args.iterations} iters)...")
    base_res = benchmark(base_model, batch, device, args.warmup, args.iterations)
    print(f"    → {base_res['latency_ms']:.2f} ms | {base_res['throughput']:.0f} samp/s | {base_res['peak_gpu_mb']:.0f} MB")

    batch_t = make_batch(args.batch_size, args.seq_len, device)
    print(f"  Benchmarking Tiered ({args.iterations} iters)...")
    tier_res = benchmark(tier_model, batch_t, device, args.warmup, args.iterations)
    print(f"    → {tier_res['latency_ms']:.2f} ms | {tier_res['throughput']:.0f} samp/s | {tier_res['peak_gpu_mb']:.0f} MB")

    speedup = base_res["latency_ms"] / tier_res["latency_ms"]
    print(f"\n  Speedup: {speedup:.2f}x")

    # ── Sweep ──
    sweep_data = None
    if args.sweep:
        print("\n  Running batch-size sweep...")
        sweep_data = {"baseline": [], "tiered": []}
        del base_model, tier_model
        torch.cuda.empty_cache()

        for bs in [1, 4, 8, 16, 32, 64]:
            print(f"    bs={bs}...", end=" ", flush=True)
            c = make_config(bs, args.seq_len, args.hidden_size, args.num_heads)
            b = make_batch(bs, args.seq_len, device)

            bm = build_model("baseline", c, device)
            br = benchmark(bm, b, device, warmup=3, iterations=10)
            br["batch_size"] = bs
            sweep_data["baseline"].append(br)
            del bm

            tm = build_model("tiered", c, device)
            tr = benchmark(tm, b, device, warmup=3, iterations=10)
            tr["batch_size"] = bs
            sweep_data["tiered"].append(tr)
            del tm
            torch.cuda.empty_cache()
            print(f"base={br['latency_ms']:.2f}ms, tier={tr['latency_ms']:.2f}ms")

    # ── Plots ──
    print("\n  Generating plots...")
    create_plots(base_res, tier_res, sweep_data, args.output_dir)

    # ── Save JSON ──
    results = {"baseline": base_res, "tiered": tier_res, "speedup": speedup}
    if sweep_data:
        results["sweep"] = sweep_data
    json_path = os.path.join(args.output_dir, "comparison_results.json")
    os.makedirs(args.output_dir, exist_ok=True)
    with open(json_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"  Results saved → {json_path}")

    print("\n" + "=" * 64)
    print("  Done!")
    print("=" * 64)


if __name__ == "__main__":
    main()