03 - Quantizzazione Modelli: GPTQ, AWQ, INT8

# Quantizzazione uniforme (schema base)
# W_quantized = round(W / scale) + zero_point
# W_dequantized = (W_quantized - zero_point) * scale

import torch
import numpy as np

def quantize_tensor_int8(tensor: torch.Tensor, symmetric: bool = True):
    """
    Quantizzazione INT8 uniforme.
    symmetric=True: zero_point=0, range [-127, 127]
    symmetric=False: range asimmetrico con zero_point
    """
    if symmetric:
        # Scale basato sul valore assoluto massimo
        max_val = tensor.abs().max().item()
        scale = max_val / 127.0
        zero_point = 0
    else:
        # Scale basato su min/max
        min_val = tensor.min().item()
        max_val = tensor.max().item()
        scale = (max_val - min_val) / 255.0
        zero_point = round(-min_val / scale)
        zero_point = max(0, min(255, zero_point))

    # Quantizzazione: FP16 -> INT8
    quantized = torch.clamp(
        torch.round(tensor / scale) + zero_point,
        -128, 127
    ).to(torch.int8)

    # De-quantizzazione: INT8 -> FP16 (per verifica errore)
    dequantized = (quantized.float() - zero_point) * scale

    # Errore di quantizzazione
    error = (tensor - dequantized).abs().mean().item()

    return quantized, scale, zero_point, error

# Esempio pratico
W = torch.randn(1024, 1024, dtype=torch.float16)
q, scale, zp, err = quantize_tensor_int8(W, symmetric=True)

print(f"Originale: {W.dtype}, {W.element_size() * W.numel() / 1024:.1f} KB")
print(f"Quantizzato: {q.dtype}, {q.element_size() * q.numel() / 1024:.1f} KB")
print(f"Riduzione memoria: {(1 - q.element_size()/W.element_size()) * 100:.0f}%")
print(f"Errore medio assoluto: {err:.6f}")
# Output tipico:
# Originale: torch.float16, 2048.0 KB
# Quantizzato: torch.int8, 1024.0 KB
# Riduzione memoria: 50%
# Errore medio assoluto: 0.000394

# Installazione
# pip install bitsandbytes transformers accelerate

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# === CONFIGURAZIONE INT8 ===
config_int8 = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,       # Outlier threshold (default ottimale)
    llm_int8_has_fp16_weight=False
)

# === CONFIGURAZIONE INT4 NF4 (QLoRA style) ===
config_int4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",            # NF4: ottimale per pesi normalmente distribuiti
    bnb_4bit_compute_dtype=torch.bfloat16, # Compute in BF16 (non INT4!)
    bnb_4bit_use_double_quant=True,        # Double quantization: -0.4 bit/param extra
    bnb_4bit_quant_storage=torch.uint8     # Storage format
)

model_name = "meta-llama/Llama-3.1-8B-Instruct"

# Caricamento con INT4 NF4
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=config_int4,
    device_map="auto",           # Distribuisce automaticamente su GPU/CPU
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Verifica memoria usata
mem_gb = model.get_memory_footprint() / 1024**3
print(f"Memoria modello (INT4): {mem_gb:.2f} GB")
# Llama-3.1-8B: ~4.5 GB vs 16 GB in BF16

# Inferenza standard (identica al modello full-precision)
inputs = tokenizer("Spiegami la quantizzazione dei modelli:", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        do_sample=True
    )
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# Installazione AutoGPTQ
# pip install auto-gptq optimum

from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

model_name = "meta-llama/Llama-3.1-8B-Instruct"
output_dir = "./llama-3.1-8b-gptq-int4"

# === DATASET DI CALIBRAZIONE ===
# GPTQ richiede un piccolo dataset (128-512 sample) per calibrare
# la quantizzazione. Più rappresentativo = migliore accuratezza.
from datasets import load_dataset

tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

# Dataset Wikitext per calibrazione (standard per LLM)
calibration_data = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
calibration_texts = [
    text for text in calibration_data["text"]
    if len(text.strip()) > 100
][:128]  # 128 sample di calibrazione

# Tokenizza i testi di calibrazione
calibration_tokens = [
    tokenizer(text, return_tensors="pt", truncation=True, max_length=2048)
    for text in calibration_texts
]

# === CONFIGURAZIONE QUANTIZZAZIONE ===
quantize_config = BaseQuantizeConfig(
    bits=4,              # INT4 quantization
    group_size=128,      # Dimensione gruppo (più piccolo = più accurato, più lento)
    damp_percent=0.01,   # Damping factor per stabilità numerica
    desc_act=True,       # Activation ordering (migliora qualità, opzionale)
    sym=True,            # Quantizzazione simmetrica
)

# === CARICAMENTO E QUANTIZZAZIONE ===
print("Caricamento modello FP16...")
model = AutoGPTQForCausalLM.from_pretrained(
    model_name,
    quantize_config=quantize_config,
    torch_dtype="auto"
)

print("Avvio quantizzazione GPTQ (richiede ~30-60 min su A100)...")
model.quantize(
    calibration_tokens,
    use_triton=False,    # True per inferenza ottimizzata con triton
    batch_size=1,
    cache_examples_on_gpu=True
)

# === SALVATAGGIO ===
model.save_quantized(output_dir, use_safetensors=True)
tokenizer.save_pretrained(output_dir)
print(f"Modello GPTQ salvato in: {output_dir}")

# === CARICAMENTO MODELLO GPTQ GIA QUANTIZZATO ===
print("Caricamento modello GPTQ pre-quantizzato...")
model_gptq = AutoGPTQForCausalLM.from_quantized(
    output_dir,
    use_triton=False,
    device_map="auto",
    inject_fused_attention=True,   # Ottimizzazione memoria attention
    inject_fused_mlp=True          # Ottimizzazione memoria MLP
)

# Inferenza
inputs = tokenizer("La quantizzazione GPTQ funziona:", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model_gptq.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# pip install autoawq

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_name = "meta-llama/Llama-3.1-8B-Instruct"
output_dir = "./llama-3.1-8b-awq-int4"

# Caricamento tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Caricamento modello da quantizzare
print("Caricamento modello per quantizzazione AWQ...")
model = AutoAWQForCausalLM.from_pretrained(
    model_name,
    safetensors=True,
    device_map="cuda",
    trust_remote_code=True
)

# Configurazione AWQ
quant_config = {
    "zero_point": True,     # Quantizzazione asimmetrica (migliore per LLM)
    "q_group_size": 128,    # Dimensione gruppo
    "w_bit": 4,             # 4-bit quantization
    "version": "GEMM"       # GEMM: bilanciamento velocità/qualità
                            # GEMV: ottimizzato per batch size 1 (chatbot)
}

# Dataset di calibrazione personalizzato
calib_data = [
    "La quantizzazione AWQ permette di eseguire modelli grandi su hardware limitato.",
    "I Transformer hanno rivoluzionato il natural language processing con il meccanismo di attention.",
    "Il fine-tuning con LoRA riduce significativamente il numero di parametri addestrabili.",
    # ... aggiungere 128-256 esempi rappresentativi del task target
]

print("Avvio quantizzazione AWQ...")
model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=calib_data
)

# Salvataggio
model.save_quantized(output_dir)
tokenizer.save_pretrained(output_dir)
print(f"Modello AWQ salvato: {output_dir}")

# === CARICAMENTO E INFERENZA MODELLO AWQ ===
model_awq = AutoAWQForCausalLM.from_quantized(
    output_dir,
    fuse_layers=True,          # Ottimizzazione kernel fused
    trust_remote_code=True,
    safetensors=True
)

from transformers import pipeline
pipe = pipeline(
    "text-generation",
    model=model_awq,
    tokenizer=tokenizer,
    device_map="auto"
)

result = pipe("Spiegami AWQ in un paragrafo:", max_new_tokens=150)
print(result[0]["generated_text"])

# Conversione e quantizzazione con llama.cpp
# Prima: installare llama.cpp

# git clone https://github.com/ggerganov/llama.cpp
# cd llama.cpp && make -j4

# 1. Convertire modello HuggingFace -> GGUF FP16
python convert_hf_to_gguf.py \
    meta-llama/Llama-3.1-8B-Instruct \
    --outfile llama-3.1-8b-f16.gguf \
    --outtype f16

# 2. Quantizzare in vari formati
./llama-quantize llama-3.1-8b-f16.gguf llama-3.1-8b-q4_k_m.gguf Q4_K_M
./llama-quantize llama-3.1-8b-f16.gguf llama-3.1-8b-q8_0.gguf Q8_0
./llama-quantize llama-3.1-8b-f16.gguf llama-3.1-8b-q5_k_m.gguf Q5_K_M

# 3. Benchmark performance con llama.cpp
./llama-bench \
    -m llama-3.1-8b-q4_k_m.gguf \
    -p 512 \        # Token prompt
    -n 128 \        # Token da generare
    -t 8            # Thread CPU

# Output tipico su M2 Pro (16 GB):
# Q4_K_M: prompt 45.2 t/s, generate 28.1 t/s
# Q8_0:   prompt 24.1 t/s, generate 15.8 t/s

# === UTILIZZO CON PYTHON via llama-cpp-python ===
# pip install llama-cpp-python

from llama_cpp import Llama

llm = Llama(
    model_path="./llama-3.1-8b-q4_k_m.gguf",
    n_ctx=4096,         # Context window
    n_threads=8,        # Thread CPU
    n_gpu_layers=35,    # Offload 35 layer su GPU (0 = solo CPU)
    verbose=False
)

response = llm(
    "Q: Come funziona la quantizzazione dei modelli? A:",
    max_tokens=256,
    stop=["Q:", "\n\n"],
    echo=True
)
print(response["choices"][0]["text"])

# === GGUF CON OLLAMA (più semplice) ===
# Crea Modelfile per Ollama
modelfile_content = """
FROM ./llama-3.1-8b-q4_k_m.gguf
PARAMETER temperature 0.7
PARAMETER num_ctx 4096
SYSTEM "Sei un assistente tecnico esperto di deep learning."
"""

with open("Modelfile", "w") as f:
    f.write(modelfile_content)

# ollama create mio-llama -f Modelfile
# ollama run mio-llama

# Benchmark automatico con lm-evaluation-harness
# pip install lm-eval

# Valutazione modello BF16 (baseline)
lm_eval --model hf \
    --model_args "pretrained=meta-llama/Llama-3.1-8B-Instruct" \
    --tasks hellaswag,mmlu \
    --batch_size 4 \
    --output_path results_bf16/

# Valutazione modello quantizzato GPTQ
lm_eval --model hf \
    --model_args "pretrained=./llama-3.1-8b-gptq-int4,use_auto_gptq=True" \
    --tasks hellaswag,mmlu \
    --batch_size 4 \
    --output_path results_gptq_int4/

# === SCRIPT BENCHMARK MEMORIA E VELOCITA ===
import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

def benchmark_model(model_name_or_path, quant_config=None, n_tokens=100, n_runs=5):
    """Benchmark completo: memoria, latenza, throughput."""
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path,
        quantization_config=quant_config,
        device_map="cuda",
        torch_dtype=torch.bfloat16 if quant_config is None else None
    )

    # Memoria usata
    mem_gb = model.get_memory_footprint() / 1024**3

    # Prompt di test
    prompt = "Explain the transformer architecture in detail:" * 3
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    # Warm-up
    with torch.no_grad():
        model.generate(**inputs, max_new_tokens=10)

    # Benchmark
    torch.cuda.synchronize()
    latencies = []
    for _ in range(n_runs):
        start = time.perf_counter()
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=n_tokens)
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start
        latencies.append(elapsed)

    avg_latency = sum(latencies) / len(latencies)
    throughput = n_tokens / avg_latency

    return {
        "memoria_gb": round(mem_gb, 2),
        "latenza_sec": round(avg_latency, 3),
        "throughput_tps": round(throughput, 1)
    }

# Confronto BF16 vs INT8 vs INT4
results = {}

results["BF16"] = benchmark_model("meta-llama/Llama-3.1-8B-Instruct")

config_8bit = BitsAndBytesConfig(load_in_8bit=True)
results["INT8"] = benchmark_model("meta-llama/Llama-3.1-8B-Instruct", config_8bit)

config_4bit = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True
)
results["INT4-NF4"] = benchmark_model("meta-llama/Llama-3.1-8B-Instruct", config_4bit)

for name, r in results.items():
    print(f"{name:10} | Mem: {r['memoria_gb']:5.2f} GB | "
          f"Latenza: {r['latenza_sec']:.3f}s | "
          f"Throughput: {r['throughput_tps']:.1f} t/s")

# Risultati tipici Llama-3.1-8B su RTX 3090:
# BF16       | Mem: 16.02 GB | Latenza: 3.821s | Throughput: 26.2 t/s
# INT8       | Mem:  8.51 GB | Latenza: 4.103s | Throughput: 24.4 t/s
# INT4-NF4   | Mem:  4.89 GB | Latenza: 3.412s | Throughput: 29.3 t/s

import torch
import torch.nn as nn
from torch.quantization import (
    prepare_qat, convert, get_default_qat_qconfig,
    QConfigMapping
)
from torch.ao.quantization import get_default_qat_qconfig_mapping

# === DEFINIZIONE MODELLO SEMPLICE ===
class SimpleTransformerBlock(nn.Module):
    def __init__(self, d_model=256, nhead=4, ff_dim=1024):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        ff_out = self.ff(x)
        return self.norm2(x + ff_out)

class SimpleModel(nn.Module):
    def __init__(self, vocab_size=1000, d_model=256, n_layers=4):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.Sequential(*[
            SimpleTransformerBlock(d_model) for _ in range(n_layers)
        ])
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        x = self.embed(x)
        x = self.blocks(x)
        return self.head(x)

# === TRAINING BASELINE (FP32) ===
model = SimpleModel()
model.train()

# === CONFIGURAZIONE QAT ===
# QConfig specifica come quantizzare activations e weights
qconfig_mapping = get_default_qat_qconfig_mapping("x86")

# Prepara il modello per QAT: inserisce FakeQuantize nodes
# che simulano la quantizzazione durante il forward pass
from torch.ao.quantization import prepare_qat_fx

# Traccia il modello con esempio di input
example_input = torch.randint(0, 1000, (2, 32))
model_prepared = prepare_qat_fx(model, qconfig_mapping, example_input)

# === QAT TRAINING LOOP ===
optimizer = torch.optim.Adam(model_prepared.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

def train_qat(model, n_epochs=10, freeze_quantizer_epoch=8):
    for epoch in range(n_epochs):
        # Dati sintetici (sostituire con dataset reale)
        x = torch.randint(0, 1000, (32, 64))
        y = torch.randint(0, 1000, (32, 64))

        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out.view(-1, 1000), y.view(-1))
        loss.backward()
        optimizer.step()

        # Freeze quantizer dopo N epoche: stabilizza le scale
        if epoch == freeze_quantizer_epoch:
            model.apply(torch.quantization.disable_observer)
            print(f"Epoch {epoch}: FakeQuantize observers disabilitati")

        if epoch % 2 == 0:
            print(f"Epoch {epoch}/{n_epochs}, Loss: {loss.item():.4f}")

train_qat(model_prepared)

# === CONVERSIONE FINALE INT8 ===
model_prepared.eval()
model_int8 = convert(model_prepared)

# Salvataggio
torch.save(model_int8.state_dict(), "model_qat_int8.pt")

# Confronto dimensioni
fp32_size = sum(p.numel() * 4 for p in model.parameters()) / 1024**2
int8_size = sum(p.numel() * 1 for p in model_int8.parameters()) / 1024**2
print(f"Modello FP32: {fp32_size:.1f} MB")
print(f"Modello QAT INT8: {int8_size:.1f} MB")
print(f"Riduzione: {(1 - int8_size/fp32_size)*100:.0f}%")

# Concetto di SmoothQuant (implementazione semplificata)
import torch
import torch.nn as nn

def compute_smooth_scale(activations: torch.Tensor, weights: torch.Tensor,
                          alpha: float = 0.5) -> torch.Tensor:
    """
    Calcola il fattore di smoothing per SmoothQuant.
    alpha: bilanciamento tra difficolta attivazioni e pesi (0.5 = equo)
    """
    # Max assoluto per canale nelle attivazioni
    act_max = activations.abs().max(dim=0).values  # [hidden_dim]

    # Max assoluto per canale nei pesi
    w_max = weights.abs().max(dim=0).values  # [hidden_dim]

    # Fattore di scaling: riduce attivazioni, aumenta pesi proporzionalmente
    scale = (act_max.pow(alpha) / w_max.pow(1 - alpha)).clamp(min=1e-5)

    return scale

def smooth_layer(layer: nn.Linear, calibration_acts: torch.Tensor,
                  alpha: float = 0.5):
    """
    Applica SmoothQuant a un layer Linear.
    Divide le attivazioni in input per scale, moltiplica i pesi per scale.
    """
    scale = compute_smooth_scale(calibration_acts, layer.weight.T, alpha)

    # Modifica in-place (equivalente matematico ma più facile da quantizzare)
    layer.weight.data = layer.weight.data * scale.unsqueeze(0)

    # Ritorna la scala inversa da applicare alle attivazioni
    return scale

# Uso pratico: SmoothQuant e tipicamente applicato con librerie come
# quanto (Hugging Face) o mediante ottimum:
# pip install optimum[quanto]

from transformers import AutoModelForCausalLM, AutoTokenizer
from optimum.quanto import quantize, freeze, qint8

model_name = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Quantizzazione INT8 con Quanto (include SmoothQuant internamente)
quantize(model, weights=qint8, activations=qint8)

# Calibrazione con alcuni esempi
with torch.no_grad():
    for text in ["Calibration text 1", "Calibration text 2"]:
        inputs = tokenizer(text, return_tensors="pt")
        model(**inputs)

# Freeze: consolida i parametri quantizzati
freeze(model)
print("Modello quantizzato INT8 con Quanto/SmoothQuant pronto!")