benchmark_runpod.py

"""
LUNA 100M - Local Benchmark + RunPod Cost Calculator
=====================================================
Uses PyTorch SDPA (Flash Attention) for realistic training throughput.
Matches the exact LUNA model architecture and training config.
"""

import os
import sys
import time
import math
import json
import gc
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import autocast, GradScaler

# ─── Model Architecture (matches your config exactly) ─────────────────────────

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=1024):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_seq_len).float()
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos())
        self.register_buffer("sin_cached", emb.sin())

    def forward(self, seq_len):
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def apply_rotary(x, cos, sin):
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    return x * cos + rotate_half(x) * sin

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.rotary_dim = int(self.head_dim * rotary_pct)
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
        self.rotary = RotaryEmbedding(self.rotary_dim, block_size)

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)

        cos, sin = self.rotary(T)
        q_rot = apply_rotary(q[..., :self.rotary_dim], cos, sin)
        k_rot = apply_rotary(k[..., :self.rotary_dim], cos, sin)
        q = torch.cat([q_rot, q[..., self.rotary_dim:]], dim=-1)
        k = torch.cat([k_rot, k[..., self.rotary_dim:]], dim=-1)

        # SDPA = Flash Attention / Memory Efficient Attention
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.c_fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(4 * n_embd, n_embd, bias=True)
    def forward(self, x):
        return self.c_proj(self.gelu(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd)
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

class LUNAModel(nn.Module):
    def __init__(self, vocab_size=50254, block_size=1024, n_layer=10, n_embd=768, n_head=12):
        super().__init__()
        self.block_size = block_size
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()

    def forward(self, idx, targets=None):
        B, T = idx.size()
        x = self.wte(idx)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

# ─── Config ────────────────────────────────────────────────────────────────────

DATASET_TOTAL_TOKENS = 4_515_286_950  # Verified from index.json: 270 chunks, sum of all dims
BLOCK_SIZE = 1024
VOCAB_SIZE = 50254
N_LAYER = 10
N_EMBD = 768
N_HEAD = 12
GLOBAL_BATCH_SIZE = 120
MAX_SEQ_LENGTH = 1024

WARMUP_STEPS = 3
BENCHMARK_STEPS = 4

# RunPod GPUs (Community Cloud pricing April 2026)
RUNPOD_GPUS = [
    # (name, $/hr, VRAM_GB, bf16_TF_nonsparse, mem_bw_GBs, arch)
    ("RTX A5000",        0.16, 24,   65,   768, "Ampere"),
    ("RTX 3090",         0.22, 24,   71,   936, "Ampere"),
    ("RTX A6000",        0.33, 48,   77,   768, "Ampere"),
    ("RTX 4090",         0.34, 24,  165,  1008, "Ada"),
    ("A40",              0.35, 48,   75,   696, "Ampere"),
    ("L4",               0.44, 24,  121,   300, "Ada"),
    ("RTX 5090",         0.69, 32,  210,  1792, "Blackwell"),
    ("L40",              0.69, 48,  181,   864, "Ada"),
    ("RTX 6000 Ada",     0.74, 48,  181,   960, "Ada"),
    ("L40S",             0.79, 48,  183,   864, "Ada"),
    ("A100 PCIe 80GB",   1.19, 80,  312, 2039, "Ampere"),
    ("A100 SXM 80GB",    1.39, 80,  312, 2039, "Ampere"),
    ("RTX Pro 6000",     1.69, 96,  260, 1280, "Blackwell"),
    ("H100 PCIe",        1.99, 80,  756, 2039, "Hopper"),
    ("H100 NVL",         2.59, 94,  835, 3938, "Hopper"),
    ("H100 SXM",         2.69, 80,  990, 3352, "Hopper"),
    ("H200",             3.59,141,  990, 4800, "Hopper"),
]

USD_TO_INR = 86.0


def find_max_micro_batch(model, device, seq_len=1024, start=32):
    """Binary search for max micro_batch_size, with 0.65 safety factor."""
    model.train()
    lo, hi, best = 1, start, 1
    opt_tmp = torch.optim.AdamW(model.parameters(), lr=1e-4)

    while hi >= lo:
        mid = (lo + hi) // 2
        try:
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
            opt_tmp.zero_grad(set_to_none=True)
            x = torch.randint(0, VOCAB_SIZE, (mid, seq_len), device=device)
            t = torch.randint(0, VOCAB_SIZE, (mid, seq_len), device=device)
            with autocast(device_type='cuda', dtype=torch.bfloat16):
                _, loss = model(x, t)
            loss.backward()
            opt_tmp.step()
            opt_tmp.zero_grad(set_to_none=True)
            best = mid
            lo = mid + 1
            del x, t, loss
            torch.cuda.empty_cache()
        except (torch.cuda.OutOfMemoryError, RuntimeError):
            try:
                del x, t, loss
            except:
                pass
            torch.cuda.empty_cache()
            opt_tmp.zero_grad(set_to_none=True)
            hi = mid - 1

    safe = max(1, int(best * 0.65))
    del opt_tmp
    torch.cuda.empty_cache()
    gc.collect()
    return safe


def run_benchmark():
    device = torch.device("cuda")
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    print("=" * 72)
    print("  LUNA 100M - TRAINING BENCHMARK & RUNPOD COST CALCULATOR")
    print("=" * 72)

    gpu_name = torch.cuda.get_device_name(0)
    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
    print(f"\n  Local GPU: {gpu_name}")
    print(f"  VRAM: {gpu_mem:.1f} GB")
    print(f"  PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}")

    print(f"\n  Creating LUNA-100M (SDPA/Flash Attention)...")
    model = LUNAModel(VOCAB_SIZE, BLOCK_SIZE, N_LAYER, N_EMBD, N_HEAD).to(device)

    total_params = sum(p.numel() for p in model.parameters())
    # Tied embeddings: wte(50254*768) = 38,595,072 counted once
    unique_params = total_params - model.wte.weight.numel()
    print(f"  Parameters: {total_params:,} total, {unique_params:,} unique")

    print(f"\n  Probing max micro_batch_size...")
    max_mbs = find_max_micro_batch(model, device, MAX_SEQ_LENGTH, start=40)
    print(f"  Safe micro_batch_size: {max_mbs}")

    grad_accum = max(1, GLOBAL_BATCH_SIZE // max_mbs)
    effective_batch = max_mbs * grad_accum
    tokens_per_step = effective_batch * MAX_SEQ_LENGTH
    print(f"  grad_accum={grad_accum}, effective_batch={effective_batch}")
    print(f"  Tokens/step: {tokens_per_step:,}")

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=6e-4, weight_decay=0.1,
        betas=(0.9, 0.95), eps=1e-8
    )
    scaler = GradScaler()

    total_steps = WARMUP_STEPS + BENCHMARK_STEPS
    print(f"\n  Running {WARMUP_STEPS} warmup + {BENCHMARK_STEPS} benchmark steps...")

    model.train()
    step_times = []
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()

    for step in range(total_steps):
        t0 = time.perf_counter()
        optimizer.zero_grad(set_to_none=True)
        step_loss = 0.0

        for _ in range(grad_accum):
            x = torch.randint(0, VOCAB_SIZE, (max_mbs, MAX_SEQ_LENGTH), device=device)
            tgt = torch.randint(0, VOCAB_SIZE, (max_mbs, MAX_SEQ_LENGTH), device=device)
            with autocast(device_type='cuda', dtype=torch.bfloat16):
                _, loss = model(x, tgt)
                loss = loss / grad_accum
            scaler.scale(loss).backward()
            step_loss += loss.item()
            del x, tgt, loss

        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        torch.cuda.synchronize()
        dt = time.perf_counter() - t0

        if step >= WARMUP_STEPS:
            step_times.append(dt)

        tps = tokens_per_step / dt
        phase = "WARM" if step < WARMUP_STEPS else "BENCH"
        print(f"    [{phase}] Step {step:3d} | Loss {step_loss:.4f} | {dt:.2f}s | {tps:,.0f} tok/s")

    # ─── Results ──────────────────────────────────────────────────────────────
    peak_vram_gb = torch.cuda.max_memory_allocated() / 1024**3

    avg_time = sum(step_times) / len(step_times)
    med_time = sorted(step_times)[len(step_times) // 2]
    avg_tps = tokens_per_step / avg_time
    med_tps = tokens_per_step / med_time
    peak_tps = tokens_per_step / min(step_times)

    flops_per_token = 6 * unique_params
    achieved_tf = (avg_tps * flops_per_token) / 1e12
    LOCAL_TF = 88.0   # RTX 4060 Ti BF16 non-sparse
    LOCAL_BW = 288.0   # GB/s
    mfu = achieved_tf / LOCAL_TF

    print("\n" + "=" * 72)
    print("  LOCAL BENCHMARK RESULTS")
    print("=" * 72)
    print(f"  GPU: {gpu_name} ({gpu_mem:.1f} GB)")
    print(f"  Peak VRAM: {peak_vram_gb:.2f} GB ({peak_vram_gb/gpu_mem*100:.0f}%)")
    print(f"  Batch: micro={max_mbs}, accum={grad_accum}, global={effective_batch}")
    print(f"  Step time: avg={avg_time:.3f}s, median={med_time:.3f}s")
    print(f"  Tokens/sec: avg={avg_tps:,.0f}, median={med_tps:,.0f}, peak={peak_tps:,.0f}")
    print(f"  TFLOPS: {achieved_tf:.2f}, MFU: {mfu*100:.1f}%")

    n_steps = math.ceil(DATASET_TOTAL_TOKENS / tokens_per_step)
    local_hrs = (n_steps * avg_time) / 3600

    print(f"  Dataset: 4,515,286,950 tokens | Steps needed: {n_steps:,}")
    print(f"  Local training: {local_hrs:.1f} hrs ({local_hrs/24:.1f} days)")

    # ─── RunPod Estimates ─────────────────────────────────────────────────────
    print("\n" + "=" * 72)
    print("  RUNPOD GPU COMPARISON (Community Cloud, INR/86/USD)")
    print("=" * 72)

    results = []
    for name, price, vram, bf16, bw, arch in RUNPOD_GPUS:
        # Memory estimation for each GPU
        fixed_gb = (unique_params * (2 + 8 + 2)) / 1024**3 + 0.5
        avail_gb = vram - fixed_gb
        act_per_sample = max(0.05, (peak_vram_gb - fixed_gb) / max(max_mbs, 1))
        est_mbs = max(1, min(128, int(avail_gb / act_per_sample)))
        est_ga = max(1, GLOBAL_BATCH_SIZE // est_mbs)
        est_tps_step = est_mbs * est_ga * MAX_SEQ_LENGTH

        # Scaling: 50% compute + 50% bandwidth (validated for small transformers)
        speedup = 0.50 * (bf16 / LOCAL_TF) + 0.50 * (bw / LOCAL_BW)
        est_tps = avg_tps * speedup * 0.90  # 0.90 cloud overhead

        est_steps = math.ceil(DATASET_TOTAL_TOKENS / est_tps_step)
        est_sec = (est_tps_step / est_tps) * est_steps
        est_hrs = est_sec / 3600
        cost_usd = est_hrs * price
        cost_inr = cost_usd * USD_TO_INR

        results.append({
            "gpu": name, "price": price, "vram": vram, "bf16": bf16,
            "bw": bw, "arch": arch, "mbs": est_mbs, "ga": est_ga,
            "tps": round(est_tps), "hours": round(est_hrs, 1),
            "usd": round(cost_usd, 2), "inr": round(cost_inr),
        })

    results.sort(key=lambda r: r["inr"])

    print(f"\n  {'#':<3} {'GPU':<18} {'$/hr':>5} {'VRAM':>5} {'tok/s':>10} "
          f"{'Hours':>7} {'$ USD':>8} {'INR':>10}")
    print("  " + "─" * 72)
    for i, r in enumerate(results):
        s = " *" if i < 3 else ""
        print(f"  {i+1:<3} {r['gpu']:<18} {r['price']:>5.2f} {r['vram']:>4}G "
              f"{r['tps']:>10,} {r['hours']:>7.1f} {r['usd']:>8.2f} {r['inr']:>10,}{s}")

    # Top 5 detailed
    print("\n" + "=" * 72)
    print("  TOP 5 CHEAPEST - DETAILS")
    print("=" * 72)
    for i, r in enumerate(results[:5]):
        sx = r["tps"] / avg_tps if avg_tps > 0 else 0
        print(f"\n  #{i+1}: {r['gpu']} ({r['arch']})")
        print(f"  ├── ${r['price']:.2f}/hr | {r['vram']}GB VRAM | {r['bf16']} TF | {r['bw']} GB/s")
        print(f"  ├── micro_batch: {r['mbs']}, grad_accum: {r['ga']}")
        print(f"  ├── {r['tps']:,} tok/s ({sx:.2f}× local)")
        print(f"  ├── {r['hours']:.1f} hrs ({r['hours']/24:.1f} days)")
        print(f"  +-- ${r['usd']:.2f} = INR {r['inr']:,}")

    # Local reference
    print("\n" + "=" * 72)
    print("  YOUR LOCAL GPU")
    print("=" * 72)
    print(f"  RTX 4060 Ti 16GB: {avg_tps:,.0f} tok/s")
    print(f"  Training: {local_hrs:.1f} hrs ({local_hrs/24:.1f} days)")
    print(f"  Electricity: ~INR {local_hrs * 0.16 * 8:,.0f} (160W x Rs8/kWh)")

    # Recommendation
    best = results[0]
    print("\n" + "=" * 72)
    print("  * RECOMMENDATION")
    print("=" * 72)
    print(f"  Most affordable: {best['gpu']} @ ${best['price']:.2f}/hr")
    print(f"  Time: {best['hours']:.1f} hrs ({best['hours']/24:.1f} days)")
    print(f"  Cost: INR {best['inr']:,} (${best['usd']:.2f})")
    print(f"  Speed: {best['tps']/avg_tps:.1f}× local" if avg_tps > 0 else "")

    fast = [r for r in results if r["hours"] < max(8, local_hrs * 0.15)]
    if fast:
        fb = min(fast, key=lambda r: r["inr"])
        if fb["gpu"] != best["gpu"]:
            print(f"\n  Fastest affordable: {fb['gpu']} @ ${fb['price']:.2f}/hr")
            print(f"  Time: {fb['hours']:.1f} hrs | Cost: INR {fb['inr']:,}")

    print("\n" + "=" * 72)

    # Save JSON
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runpod_cost_estimate.json")
    with open(out, "w") as f:
        json.dump({
            "benchmark": {
                "gpu": gpu_name, "vram_gb": round(gpu_mem, 1),
                "peak_vram_gb": round(peak_vram_gb, 2),
                "micro_batch": max_mbs, "grad_accum": grad_accum,
                "tokens_per_step": tokens_per_step,
                "avg_tok_per_sec": round(avg_tps),
                "median_tok_per_sec": round(med_tps),
                "achieved_tflops": round(achieved_tf, 2),
                "mfu_pct": round(mfu*100, 1),
            },
            "dataset": {"tokens": DATASET_TOTAL_TOKENS, "chunks": 270},
            "model": {"total_params": total_params, "unique_params": unique_params},
            "local_hours": round(local_hrs, 1),
            "runpod": results,
            "usd_to_inr": USD_TO_INR,
        }, f, indent=2)
    print(f"  Saved: {out}")
    print("=" * 72)


if __name__ == "__main__":
    run_benchmark()