Add benchmark suite

benchmark.py

@@ -0,0 +1,174 @@
+"""
+SmolOmni-MLA Benchmark Suite
+Measures: VRAM usage, KV cache size, throughput, generation quality
+Compares against SmolVLM baseline.
+"""
+import time
+import json
+import argparse
+import os
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForImageTextToText, AutoTokenizer
+
+from smolomni.config import SmolOmniConfig
+from smolomni.model import SmolOmniModel
+from smolomni.svd_init import initialize_mla_from_pretrained
+
+
+def benchmark_kv_cache(config: SmolOmniConfig):
+    """Report KV cache sizes for different configurations."""
+    info = config.kv_cache_size_per_token()
+    return info
+
+
+def benchmark_model_loading(model_variant: str, device: str = "cuda"):
+    """Time model initialization and SVD init."""
+    config = SmolOmniConfig.from_pretrained(f"mla-hybrid-ar-flow-{model_variant}")
+    
+    # Baseline: SmolVLM
+    start = time.time()
+    baseline = AutoModelForImageTextToText.from_pretrained(
+        config.base_model, torch_dtype=torch.bfloat16
+    ).to(device)
+    baseline_time = time.time() - start
+    baseline_params = sum(p.numel() for p in baseline.parameters())
+    
+    # SmolOmni
+    start = time.time()
+    model = SmolOmniModel(config)
+    model = initialize_mla_from_pretrained(model, config.base_model, config)
+    model = model.to(device, dtype=torch.bfloat16)
+    smol_time = time.time() - start
+    smol_params = model.num_parameters()
+    
+    del baseline
+    torch.cuda.empty_cache()
+    
+    return {
+        "baseline": {
+            "load_time_s": baseline_time,
+            "params_M": baseline_params / 1e6,
+        },
+        "smolomni": {
+            "load_time_s": smol_time,
+            "params_M": smol_params / 1e6,
+        },
+    }
+
+
+def benchmark_throughput(model, tokenizer, config, batch_size: int = 1, seq_len: int = 512):
+    """Measure tokens/second for understanding and generation."""
+    device = next(model.parameters()).device
+    
+    # Understanding (AR)
+    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len), device=device)
+    
+    # Warm-up
+    for _ in range(3):
+        with torch.no_grad():
+            _ = model.forward_understanding(input_ids)
+    
+    torch.cuda.synchronize()
+    start = time.time()
+    n_iters = 10
+    for _ in range(n_iters):
+        with torch.no_grad():
+            _ = model.forward_understanding(input_ids)
+    torch.cuda.synchronize()
+    ar_time = (time.time() - start) / n_iters
+    ar_tps = (batch_size * seq_len) / ar_time
+    
+    # Generation (flow-matching, 50 steps)
+    latent_shape = (batch_size, 4, 32, 32)
+    torch.cuda.synchronize()
+    start = time.time()
+    with torch.no_grad():
+        _ = model.generate_image(input_ids[:batch_size], num_steps=50, latent_shape=latent_shape)
+    torch.cuda.synchronize()
+    gen_time = time.time() - start
+    
+    return {
+        "ar_time_ms": ar_time * 1000,
+        "ar_tokens_per_sec": ar_tps,
+        "gen_time_s": gen_time,
+        "gen_steps_per_sec": 50 / gen_time,
+    }
+
+
+def benchmark_vram(model, batch_size: int = 1, seq_len: int = 512):
+    """Measure peak VRAM usage."""
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    
+    device = next(model.parameters()).device
+    input_ids = torch.randint(0, model.config.vocab_size, (batch_size, seq_len), device=device)
+    
+    with torch.no_grad():
+        _ = model.forward_understanding(input_ids)
+    
+    peak_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
+    return {"peak_vram_mb": peak_mb}
+
+
+def run_all_benchmarks(model_variant: str = "256M"):
+    """Run full benchmark suite."""
+    print(f"="*70)
+    print(f"SmolOmni-MLA Benchmark: {model_variant}")
+    print(f"="*70)
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # KV Cache comparison
+    print("\n--- KV Cache Analysis ---")
+    config = SmolOmniConfig.from_pretrained(f"mla-hybrid-ar-flow-{model_variant}")
+    cache_info = benchmark_kv_cache(config)
+    for k, v in cache_info.items():
+        print(f"  {k}: {v}")
+    
+    # Model loading
+    print("\n--- Model Loading ---")
+    load_info = benchmark_model_loading(model_variant, device)
+    print(f"  Baseline (SmolVLM): {load_info['baseline']['load_time_s']:.1f}s, {load_info['baseline']['params_M']:.1f}M params")
+    print(f"  SmolOmni-MLA: {load_info['smolomni']['load_time_s']:.1f}s, {load_info['smolomni']['params_M']:.1f}M params")
+    
+    # Load model for throughput tests
+    model = SmolOmniModel(config)
+    model = initialize_mla_from_pretrained(model, config.base_model, config)
+    model = model.to(device, dtype=torch.bfloat16)
+    model.eval()
+    
+    # VRAM
+    print("\n--- VRAM Usage ---")
+    vram = benchmark_vram(model)
+    print(f"  Peak VRAM: {vram['peak_vram_mb']:.0f} MB")
+    
+    # Throughput
+    print("\n--- Throughput ---")
+    throughput = benchmark_throughput(model, None, config)
+    print(f"  AR forward: {throughput['ar_time_ms']:.1f}ms ({throughput['ar_tokens_per_sec']:.0f} tok/s)")
+    print(f"  Image gen (50 steps): {throughput['gen_time_s']:.1f}s ({throughput['gen_steps_per_sec']:.1f} step/s)")
+    
+    results = {
+        "model_variant": model_variant,
+        "kv_cache": cache_info,
+        "loading": load_info,
+        "vram": vram,
+        "throughput": throughput,
+    }
+    
+    # Save results
+    out_path = f"/app/benchmark_{model_variant}.json"
+    with open(out_path, 'w') as f:
+        json.dump(results, f, indent=2)
+    print(f"\nResults saved to {out_path}")
+    
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_variant", default="256M", choices=["256M", "500M"])
+    args = parser.parse_args()
+    run_all_benchmarks(args.model_variant)