Upload benchmark_runpod.py with huggingface_hub

benchmark_runpod.py

@@ -1,412 +1,412 @@
-"""
-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()
+"""
+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()