🎲 LAB-04 Playground: Inference Math, Sampling & Quantization¶

import os
import sys
import sys
from pathlib import Path

try:
    import google.colab
    IN_COLAB = True
except ImportError:
    IN_COLAB = False

if IN_COLAB:
    print("🤖 Running in Google Colab. Setting up environment...")
    !pip install -q numpy scikit-learn matplotlib loguru
    
    # Clone repository to access custom local source modules
    if not os.path.exists("production-ai-systems"):
        !git clone https://github.com/RogerioLS/production-ai-systems.git
    
    sys.path.append("production-ai-systems")
    print("🎉 Colab environment successfully initialized!")
# Ensure project root is in the path when running locally
if not IN_COLAB:
    # Localiza dinamicamente a raiz do repositório procurando por 'pyproject.toml'
    current_dir = Path.cwd()
    root_dir = current_dir
    for parent in [current_dir] + list(current_dir.parents):
        if (parent / "pyproject.toml").exists() or (parent / "projects").exists():
            root_dir = parent
            break

    if str(root_dir) not in sys.path:
        sys.path.insert(0, str(root_dir))

    print(f"📁 Project root added to sys.path: {root_dir}")


import math
import torch
import torch.nn.functional as F

from projects.a_llm_basics.src.lab_04_inference import InferenceCostCalculator, SamplingEngine, UniformQuantizer, compute_cross_entropy_loss, compute_perplexity, compute_quantization_error, compute_shannon_entropy

print("✅ LAB-04 Modules imported successfully!")

# Simulated unnormalized logits vector for a 5-token vocabulary
logits = torch.tensor([[12.0, 8.0, 4.0, 1.0, 0.2]])
vocab = ["llama", "transformer", "attention", "token", "cat"]

print("Raw Logits:", logits.numpy()[0])

# Experiment 1: Temperature Scaling
for temp in [0.1, 0.7, 1.0, 2.0]:
    scaled = logits / temp
    probs = F.softmax(scaled, dim=-1)[0]
    entropy = compute_shannon_entropy(probs).item()
    print(f"T={temp:<4.1f} | Probs: {probs.numpy().round(3)} | Entropy: {entropy:.2f} bits")

torch.manual_seed(42)
weights_fp32 = torch.randn(1000)

# Symmetric INT8 Quantization
quantizer_sym = UniformQuantizer(num_bits=8, symmetric=True)
q_sym, scale_sym, zero_sym = quantizer_sym.quantize(weights_fp32)
deq_sym = quantizer_sym.dequantize(q_sym, scale_sym, zero_sym)
metrics_sym = compute_quantization_error(weights_fp32, deq_sym)

print("=== Symmetric INT8 Quantization Metrics ===")
print(f"Scale Factor (S): {scale_sym:.6f}")
print(f"MSE: {metrics_sym['mse']:.6f}")
print(f"Cosine Similarity: {metrics_sym['cosine_similarity']:.6f}")
print(f"SQNR: {metrics_sym['sqnr_db']:.2f} dB")

# Llama-3-8B architecture specifications
calc = InferenceCostCalculator(
    num_layers=32,
    num_heads=32,
    head_dim=128,
    hidden_size=4096,
    num_kv_heads=8  # Grouped-Query Attention (GQA 4:1 ratio)
)

num_params = 8_000_000_000
weight_vram_gb = calc.compute_model_param_bytes(num_params, precision_bytes=2.0) / (1024**3)
kv_vram_mb = calc.compute_kv_cache_bytes(batch_size=1, seq_len=8192, precision_bytes=2.0) / (1024**2)
max_batch = calc.compute_max_batch_size(vram_capacity_bytes=24 * (1024**3), num_params=num_params, seq_len=8192)

print("=== Llama-3-8B Inference Profile (FP16) ===")
print(f"Weights VRAM: {weight_vram_gb:.2f} GB")
print(f"KV Cache (B=1, S=8192): {kv_vram_mb:.2f} MB")
print(f"Max Batch Size on RTX 4090 (24GB VRAM): {max_batch} sequences")