""" 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()