GPU & VRAM Optimizacija za LLM Deployment

"""
vram_calculator.py - Kalkulator VRAM potreba za AI Deployment
Modul 2.3: GPU & VRAM Optimizacija
"""

from dataclasses import dataclass
from typing import Optional

# ================================================================
# DEFINICIJE MODELA
# ================================================================

@dataclass
class LLMModel:
    """Opisuje jedan LLM model i njegove specifikacije."""
    naziv: str
    parametri_B: float        # Broj parametara u milijardama
    num_layers: int           # Broj transformer slojeva
    num_attention_heads: int  # Broj attention głava
    hidden_dim: int           # Dimenzija hidden state-a
    max_context_tokens: int   # Maksimalni context window

@dataclass
class QuantizationFormat:
    """Opisuje quantizacijski format modela."""
    naziv: str
    bajta_po_parametru: float
    opis: str
    gubitak_kvaliteta: str

@dataclass
class GPUModel:
    """Opisuje GPU model i njegove specifikacije."""
    naziv: str
    vram_gb: float
    bandwidth_gbs: float
    cijena_eur: int
    tip: str  # "consumer", "professional", "datacenter"

    def fituje_model(self, potreban_vram: float) -> bool:
        return self.vram_gb >= potreban_vram * 1.1  # 10% sigurnosni bufer

# ================================================================
# BAZA PODATAKA MODELA I GPU-OVA
# ================================================================

MODELI = {
    "llama_3_8b": LLMModel("Llama 3.1 8B", 8, 32, 32, 4096, 131072),
    "llama_3_70b": LLMModel("Llama 3.1 70B", 70, 80, 64, 8192, 131072),
    "mistral_7b": LLMModel("Mistral 7B v0.3", 7, 32, 32, 4096, 32768),
    "phi3_mini": LLMModel("Phi-3 Mini (3.8B)", 3.8, 32, 32, 3072, 128000),
}

QUANTIZATION_FORMATS = {
    "fp16": QuantizationFormat("FP16", 2.0, "Puna 16-bit preciznost", "Minimalan"),
    "q8": QuantizationFormat("Q8 (8-bit INT)", 1.0, "8-bit quantizacija", "Zanemariv"),
    "q4": QuantizationFormat("Q4 (4-bit INT)", 0.5, "4-bit quantizacija (Ollama default)", "Prihvatljiv"),
    "q3": QuantizationFormat("Q3 (3-bit INT)", 0.375, "3-bit quantizacija", "Primjetan"),
    "q2": QuantizationFormat("Q2 (2-bit INT)", 0.25, "Agresivna quantizacija", "Značajan"),
}

GPU_MODELI = [
    GPUModel("RTX 3080 (10 GB)", 10, 760, 650, "consumer"),
    GPUModel("RTX 4060 Ti (16 GB)", 16, 288, 450, "consumer"),
    GPUModel("RTX 4090 (24 GB)", 24, 1008, 1800, "consumer"),
    GPUModel("NVIDIA RTX A6000 (48 GB)", 48, 768, 5000, "professional"),
    GPUModel("NVIDIA A100 (80 GB)", 80, 2000, 15000, "datacenter"),
    GPUModel("NVIDIA H100 (80 GB)", 80, 3350, 25000, "datacenter"),
]

# ================================================================
# KALKULACIJE
# ================================================================

def izracunaj_vram_parametri(model: LLMModel, q_format: QuantizationFormat) -> float:
    """Izračunava VRAM za model parametre (bez KV Cache)."""
    # Bazni VRAM za parametre
    VRAM_bazni = model.parametri_B * 1e9 * q_format.bajta_po_parametru
    # Pretvaramo u GB + 10% overhead za runtime
    VRAM_gb = (VRAM_bazni / (1024**3)) * 1.10
    return VRAM_gb

def izracunaj_kv_cache(model: LLMModel, q_format: QuantizationFormat,
                        seq_len: int) -> float:
    """Izračunava VRAM za KV Cache za jednog korisnika sa seq_len tokena."""
    # KV Cache formula:
    # za svaki  sloj (layer), svaka głava (head) drži K i V vektore
    head_dim = model.hidden_dim // model.num_attention_heads
    # 2 (K + V) × num_layers × num_heads × head_dim × seq_len × sizeof(dtype)
    kv_bytes = 2 * model.num_layers * model.num_attention_heads * head_dim * seq_len * q_format.bajta_po_parametru
    return kv_bytes / (1024**3)  # Prebacujemo u GB

def preporuci_gpu(ukupni_vram_gb: float) -> list:
    """Vraća listu GPU-ova koji mogu podnijeti potrebni VRAM."""
    return [gpu for gpu in GPU_MODELI if gpu.fituje_model(ukupni_vram_gb)]

# ================================================================
# IZVJEŠTAJ
# ================================================================

def generiraj_izvjestaj(
    model_key: str,
    format_key: str,
    concurrent_users: int,
    avg_context_tokens: int
):
    """Generiše potpuni VRAM izvještaj za zadani scenarij."""
    model = MODELI[model_key]
    qformat = QUANTIZATION_FORMATS[format_key]

    SEPARATOR = "═" * 65

    print(f"\n{SEPARATOR}")
    print(f"  🧮 VRAM KALKULATOR — AI DEPLOYMENT PLAN")
    print(SEPARATOR)
    print(f"\n📋 SCENARIJ:")
    print(f"   Model:               {model.naziv}")
    print(f"   Quantizacija:        {qformat.naziv}")
    print(f"   Simultani korisnici: {concurrent_users}")
    print(f"   Prosj. kontekst:     {avg_context_tokens:,} tokena")

    # Komponente VRAM
    vram_model = izracunaj_vram_parametri(model, qformat)
    vram_kv_jednog = izracunaj_kv_cache(model, qformat, avg_context_tokens)
    vram_kv_ukupno = vram_kv_jednog * concurrent_users
    vram_ukupno = vram_model + vram_kv_ukupno

    print(f"\n💾 VRAM BREAKDOWN:")
    print(f"   Model parametri ({qformat.naziv}): {vram_model:.2f} GB")
    print(f"   KV Cache / korisnik:               {vram_kv_jednog:.3f} GB")
    print(f"   KV Cache ukupno ({concurrent_users} kor.):        {vram_kv_ukupno:.2f} GB")
    print(f"   ──────────────────────────────────────────")
    print(f"   🎯 UKUPNO POTREBNO:                {vram_ukupno:.2f} GB VRAM")

    print(f"\n🖥️  PREPORUČENI GPU MODELI (sortiran po cijeni):")
    preporuke = preporuci_gpu(vram_ukupno)
    if preporuke:
        for gpu in sorted(preporuke, key=lambda g: g.cijena_eur):
            marker = "✅" if gpu.cijena_eur < 10000 else "💎"
            print(f"   {marker} {gpu.naziv}")
            print(f"      VRAM: {gpu.vram_gb} GB | BW: {gpu.bandwidth_gbs} GB/s | Cijena: ~€{gpu.cijena_eur:,}")
    else:
        print("   ❌ Ni jedan GPU ne zadovoljava — razmotrite Multi-GPU setup!")

    print(f"\n⚠️  NAPOMENA O KVALITETU ({qformat.naziv}):")
    print(f"   Gubitak vs FP16: {qformat.gubitak_kvaliteta}")

    return vram_ukupno

# ================================================================
# POKRETANJE KALKULATORA ZA VIŠE SCENARIJA
# ================================================================

if __name__ == "__main__":
    print("🧮 AI Infrastruktura VRAM Kalkulator\n")

    # SCENARIJ 1: Mali tim (10 korisnika, 8B model)
    generiraj_izvjestaj("llama_3_8b", "q4", concurrent_users=10, avg_context_tokens=8192)

    # SCENARIJ 2: Srednja firma (50 korisnika, 8B model)
    generiraj_izvjestaj("llama_3_8b", "q4", concurrent_users=50, avg_context_tokens=8192)

    # SCENARIJ 3: Enterprise (30 korisnika, 70B model, dugi kontekst)
    generiraj_izvjestaj("llama_3_70b", "q4", concurrent_users=30, avg_context_tokens=32768)

    print("\n" + "═" * 65)
    print("  💡 SAVJETI ZA PLANIRANJE AI INFRASTRUKTURE")
    print("═" * 65)
    print("""
  1. Uvijek planirajte s 20-30% VRAM buffera za peak opterećenje
  2. Q4 quantizacija daje dobar quality/cost tradeoff za interne alate
  3. Za dulje kontekste (100k+ tokena), KV Cache dominira nad modelom!
  4. Pre-kupovinu GPU-a: testirajte sa Ollama CPU modo za validaciju use-case-a
  5. RTX 4060 Ti 16GB = idealan startni lokalni AI server za manje timove
""")

python vram_calculator.py