12. Quantization-Aware Training — INT8/INT4 Without Quality Loss

graph TB
    subgraph "FP32 → INT8 Quantization (one weight row)"
        FP32["FP32 weights:<br>[-0.82, 0.15, -0.03, 0.67, 1.24, -0.56, 0.91, -1.10]<br>Continuous values, infinite precision"]

        SCALE["Compute scale: s = max(|w|) / 127<br>s = 1.24 / 127 = 0.00976"]

        INT8["INT8 quantized:<br>[-84, 15, -3, 69, 127, -57, 93, -113]<br>Only 256 possible values"]

        DEQ["Dequantized (FP32):<br>[-0.820, 0.146, -0.029, 0.673, 1.240, -0.557, 0.908, -1.103]<br>Errors: [0, 0.004, 0.001, 0.003, 0, 0.003, 0.002, 0.003]"]

        FP32 --> SCALE --> INT8 --> DEQ
    end

    subgraph "FP32 → INT4 Quantization (same weights)"
        FP32_2["Same FP32 weights:<br>[-0.82, 0.15, -0.03, 0.67, 1.24, -0.56, 0.91, -1.10]"]

        SCALE_2["Compute scale: s = max(|w|) / 7<br>s = 1.24 / 7 = 0.177"]

        INT4["INT4 quantized:<br>[-5, 1, 0, 4, 7, -3, 5, -6]<br>Only 16 possible values!"]

        DEQN4["Dequantized (FP32):<br>[-0.886, 0.177, 0.000, 0.709, 1.240, -0.531, 0.886, -1.063]<br>Errors: [0.066, 0.027, 0.030, 0.039, 0, 0.029, 0.024, 0.037]"]

        FP32_2 --> SCALE_2 --> INT4 --> DEQN4
    end

    style INT8 fill:#c8e6c9
    style INT4 fill:#fff9c4
    style DEQN4 fill:#ffcdd2

sequenceDiagram
    participant W as FP32 Weights (master)
    participant FQ as Fake Quantize
    participant FWD as Forward Pass
    participant LOSS as Loss
    participant BWD as Backward (STE)

    Note over W,BWD: QAT Training Loop
    W->>FQ: w_float
    FQ->>FQ: q = round(w/s) · s<br>(simulate INT8 rounding)
    FQ->>FWD: w_quantized (but still FP32 tensor)
    FWD->>FWD: Normal forward with quantized weights<br>Model "sees" what INT8 inference will look like
    FWD->>LOSS: logits → cross-entropy

    LOSS->>BWD: ∂L/∂logits
    BWD->>BWD: Backprop through all layers
    BWD->>FQ: ∂L/∂w_quantized
    FQ->>W: ∂L/∂w ≈ ∂L/∂w_quantized (STE)<br>Pass gradient straight through!
    W->>W: w ← w - η · ∂L/∂w<br>Update FP32 master weights

    Note over W,BWD: Key: FP32 weights adjust to<br>positions where rounding error is minimal

graph LR
    subgraph "GPTQ: Quantize col j, compensate cols j+1...n"
        C1["Col 1<br>✅ Quantized<br>Error: δ₁"]
        C2["Col 2<br>✅ Quantized<br>Error: δ₂<br>+ correction<br>from δ₁"]
        C3["Col 3<br>🔄 Quantizing now<br>Error: δ₃"]
        C4["Col 4<br>⬜ Float<br>Absorbs δ₃ error<br>w₄ -= δ₃·H₃₄/H₃₃"]
        C5["Col 5<br>⬜ Float<br>Absorbs δ₃ error<br>w₅ -= δ₃·H₃₅/H₃₃"]
        CN["Col n<br>⬜ Float<br>Last col absorbs<br>all residual error"]

        C1 --> C2 --> C3 --> C4 --> C5 --> CN
    end

    subgraph "Hessian H = 2XX^T"
        H["H tells us how<br>columns interact.<br>H_jk large → col k can<br>compensate col j's error.<br>H_jj large → col j is<br>sensitive to quantization.<br><br>Calibration: 128 samples<br>from MangaAssist queries"]
    end

    style C1 fill:#c8e6c9
    style C2 fill:#c8e6c9
    style C3 fill:#fff9c4
    style C4 fill:#e3f2fd
    style C5 fill:#e3f2fd

graph TB
    subgraph "Llama 3 8B: Per-Layer Bit Assignment"
        EMB["Embedding layer<br>INT8 (higher precision)<br>Maps tokens → continuous space<br>Error here amplifies 32×"]

        L1["Layers 1-4 (early)<br>INT8 (8-bit)<br>Error amplification: 4.5×<br>Capture low-level patterns"]

        L5["Layers 5-16 (middle)<br>INT4 (4-bit, GPTQ)<br>Error amplification: 2.1×<br>Redundant representations"]

        L17["Layers 17-28 (late-middle)<br>INT4 (4-bit, GPTQ)<br>Error amplification: 1.3×<br>High-level features"]

        L29["Layers 29-32 (final)<br>INT8 (8-bit)<br>Error amplification: 1.0×<br>Direct output influence"]

        HEAD["LM Head<br>FP16 (highest precision)<br>Maps to vocabulary logits<br>Precision critical for next-token"]

        EMB --> L1 --> L5 --> L17 --> L29 --> HEAD
    end

    subgraph "Memory Summary"
        MEM["Embedding: 512MB (INT8)<br>Layers 1-4: 800MB (INT8)<br>Layers 5-28: 2.4GB (INT4)<br>Layers 29-32: 800MB (INT8)<br>LM Head: 512MB (FP16)<br>─────────────<br>Total: ~5.0GB<br>(vs 16GB FP16, 8GB uniform INT8)"]
    end

    style EMB fill:#c8e6c9
    style L1 fill:#c8e6c9
    style L5 fill:#fff9c4
    style L17 fill:#fff9c4
    style L29 fill:#c8e6c9
    style HEAD fill:#bbdefb

graph TB
    subgraph "AWQ: Scale important channels before quantization"
        ACT["Calibration activations X<br>128 MangaAssist queries"]

        CH1["Channel 1: mean(|X₁|) = 0.02<br>Low activation → low importance<br>Scale: s₁ = 0.14^0.5 = 0.37"]
        CH2["Channel 2: mean(|X₂|) = 0.15<br>High activation → HIGH importance<br>Scale: s₂ = 1.0^0.5 = 1.0"]
        CH3["Channel 3: mean(|X₃|) = 0.08<br>Medium activation<br>Scale: s₃ = 0.53^0.5 = 0.73"]

        ACT --> CH1 & CH2 & CH3

        Q1["w₁ scaled by 0.37<br>→ fewer quant bins<br>→ more error (OK, low impact)"]
        Q2["w₂ scaled by 1.0<br>→ full quant bins<br>→ less error (good, high impact)"]
        Q3["w₃ scaled by 0.73<br>→ moderate bins<br>→ moderate error"]

        CH1 --> Q1
        CH2 --> Q2
        CH3 --> Q3
    end

    style Q2 fill:#c8e6c9
    style Q1 fill:#ffcdd2

import torch
import torch.nn as nn
import torch.quantization as quant
from transformers import DistilBertModel, DistilBertTokenizer
from torch.quantization import QuantStub, DeQuantStub, prepare_qat, convert


class QuantizedIntentClassifier(nn.Module):
    """
    DistilBERT intent classifier with quantization-aware training.
    Simulates INT8 during training so the model learns quantization-robust weights.
    """

    def __init__(self, num_intents: int = 10, model_name: str = "distilbert-base-uncased"):
        super().__init__()
        self.quant = QuantStub()     # Quantize input
        self.dequant = DeQuantStub() # Dequantize output

        self.bert = DistilBertModel.from_pretrained(model_name)
        self.classifier = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, num_intents),
        )

    def forward(self, input_ids, attention_mask):
        # Quantize activations entering the classifier head
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]  # [CLS] token
        x = self.quant(cls_output)
        logits = self.classifier(x)
        logits = self.dequant(logits)
        return logits


class QATTrainer:
    """Train with fake-quantization to prepare for INT8 deployment."""

    def __init__(self, num_intents: int = 10):
        self.model = QuantizedIntentClassifier(num_intents)
        self.tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

    def prepare_qat(self):
        """Insert fake-quantization modules for QAT."""
        # Define quantization config
        self.model.qconfig = quant.get_default_qat_qconfig("x86")

        # Fuse modules for efficiency (Conv+BN+ReLU, Linear+ReLU)
        # For our classifier head:
        torch.quantization.fuse_modules(
            self.model.classifier,
            [["0", "1"]],  # Linear + ReLU
            inplace=True,
        )

        # Insert fake-quant observers
        prepare_qat(self.model, inplace=True)
        print("QAT prepared. Fake-quantization modules inserted.")

    def train(self, train_loader, val_loader, epochs: int = 5, lr: float = 2e-5):
        """
        Train with quantization simulation.
        The model "sees" quantized weights during forward pass
        but updates FP32 master weights via STE during backward pass.
        """
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        self.model.train()
        for epoch in range(epochs):
            total_loss = 0
            correct = 0
            total = 0

            # Freeze batch norm stats after epoch 3 for stable quantization
            if epoch >= 3:
                self.model.apply(torch.quantization.disable_observer)

            for batch in train_loader:
                input_ids = batch["input_ids"]
                attention_mask = batch["attention_mask"]
                labels = batch["labels"]

                logits = self.model(input_ids, attention_mask)
                loss = criterion(logits, labels)

                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

                total_loss += loss.item()
                correct += (logits.argmax(dim=-1) == labels).sum().item()
                total += labels.size(0)

            train_acc = correct / total
            val_acc = self._evaluate(val_loader)
            print(f"Epoch {epoch+1}: loss={total_loss/len(train_loader):.4f}, "
                  f"train_acc={train_acc:.4f}, val_acc={val_acc:.4f}")

    def convert_to_int8(self):
        """Convert QAT model to actual INT8 for deployment."""
        self.model.eval()
        quantized_model = convert(self.model, inplace=False)
        return quantized_model

    def _evaluate(self, loader):
        self.model.eval()
        correct = total = 0
        with torch.no_grad():
            for batch in loader:
                logits = self.model(batch["input_ids"], batch["attention_mask"])
                correct += (logits.argmax(dim=-1) == batch["labels"]).sum().item()
                total += batch["labels"].size(0)
        self.model.train()
        return correct / total

    def benchmark_size_and_speed(self, quantized_model, sample_input):
        """Compare FP32 vs INT8 model."""
        import time

        # Size
        torch.save(self.model.state_dict(), "/tmp/fp32_model.pt")
        torch.save(quantized_model.state_dict(), "/tmp/int8_model.pt")

        import os
        fp32_size = os.path.getsize("/tmp/fp32_model.pt") / 1024 / 1024
        int8_size = os.path.getsize("/tmp/int8_model.pt") / 1024 / 1024
        print(f"FP32: {fp32_size:.1f}MB, INT8: {int8_size:.1f}MB, "
              f"Compression: {fp32_size/int8_size:.1f}×")

        # Speed (CPU inference)
        input_ids, attention_mask = sample_input
        self.model.eval()
        quantized_model.eval()

        # Warmup
        for _ in range(10):
            self.model(input_ids, attention_mask)
            quantized_model(input_ids, attention_mask)

        # Benchmark
        start = time.time()
        for _ in range(100):
            self.model(input_ids, attention_mask)
        fp32_time = (time.time() - start) / 100

        start = time.time()
        for _ in range(100):
            quantized_model(input_ids, attention_mask)
        int8_time = (time.time() - start) / 100

        print(f"FP32: {fp32_time*1000:.1f}ms, INT8: {int8_time*1000:.1f}ms, "
              f"Speedup: {fp32_time/int8_time:.1f}×")

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch


class GPTQQuantizer:
    """
    Apply GPTQ quantization to Llama 3 8B.
    Uses calibration data from MangaAssist to optimize quantization.
    """

    def __init__(self, model_name: str = "meta-llama/Llama-3-8b-hf"):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

    def quantize(
        self,
        calibration_texts: list[str],
        bits: int = 4,
        group_size: int = 128,
        desc_act: bool = True,
    ):
        """
        Quantize model using GPTQ with MangaAssist calibration data.

        Args:
            calibration_texts: 128 representative MangaAssist conversations
            bits: 4 for INT4, 8 for INT8
            group_size: Quantize in groups of 128 weights (more groups = more precision)
            desc_act: Sort columns by activation magnitude (descending) before quantizing
        """
        quantize_config = BaseQuantizeConfig(
            bits=bits,
            group_size=group_size,
            desc_act=desc_act,
            sym=True,  # Symmetric quantization
            model_file_base_name="mangaassist-llama3-8b",
        )

        # Load model in FP16
        model = AutoGPTQForCausalLM.from_pretrained(
            self.model_name,
            quantize_config=quantize_config,
            torch_dtype=torch.float16,
            device_map="auto",
        )

        # Prepare calibration data
        calibration_dataset = [
            self.tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
            for text in calibration_texts
        ]

        # Run GPTQ quantization
        model.quantize(
            calibration_dataset,
            batch_size=1,
        )

        return model

    def save_quantized(self, model, output_dir: str):
        model.save_quantized(output_dir)
        self.tokenizer.save_pretrained(output_dir)
        print(f"Saved GPTQ model to {output_dir}")

    def evaluate_quality(self, model, test_prompts: list[str]) -> dict:
        """Compare quantized model quality against FP16 baseline."""
        results = {"perplexity_samples": [], "generation_samples": []}

        for prompt in test_prompts[:20]:
            inputs = self.tokenizer(prompt, return_tensors="pt").to(model.device)

            with torch.no_grad():
                outputs = model(**inputs, labels=inputs["input_ids"])
                perplexity = torch.exp(outputs.loss).item()
                results["perplexity_samples"].append(perplexity)

            generated = model.generate(**inputs, max_new_tokens=100)
            text = self.tokenizer.decode(generated[0], skip_special_tokens=True)
            results["generation_samples"].append({
                "prompt": prompt,
                "response": text[len(prompt):],
            })

        results["mean_perplexity"] = sum(results["perplexity_samples"]) / len(results["perplexity_samples"])
        return results

from awq import AutoAWQForCausalLM


class AWQQuantizer:
    """
    Activation-Aware Weight Quantization.
    Preserves channels with large activations at higher precision.
    """

    def __init__(self, model_name: str = "meta-llama/Llama-3-8b-hf"):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

    def quantize(self, calibration_texts: list[str]):
        """
        Apply AWQ quantization.
        Automatically determines per-channel scaling factors
        from calibration data.
        """
        model = AutoAWQForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16,
            device_map="auto",
        )

        quant_config = {
            "zero_point": True,
            "q_group_size": 128,
            "w_bit": 4,
        }

        # Tokenize calibration data
        calib_data = [
            self.tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
            for text in calibration_texts[:128]
        ]

        model.quantize(
            tokenizer=self.tokenizer,
            quant_config=quant_config,
            calib_data=calib_data,
        )

        return model


class QuantizationBenchmark:
    """Compare GPTQ vs AWQ vs QAT on MangaAssist tasks."""

    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8b-hf")

    def run_benchmark(self, models: dict, test_data: list[dict]) -> dict:
        """
        Benchmark multiple quantization methods head-to-head.
        """
        results = {}
        for name, model in models.items():
            perplexities = []
            latencies = []
            import time

            for item in test_data:
                inputs = self.tokenizer(
                    item["prompt"], return_tensors="pt", max_length=256, truncation=True,
                ).to(model.device)

                # Perplexity
                with torch.no_grad():
                    outputs = model(**inputs, labels=inputs["input_ids"])
                    perplexities.append(torch.exp(outputs.loss).item())

                # Latency
                start = time.time()
                with torch.no_grad():
                    model.generate(**inputs, max_new_tokens=50)
                latencies.append(time.time() - start)

            results[name] = {
                "mean_perplexity": sum(perplexities) / len(perplexities),
                "mean_latency_ms": sum(latencies) / len(latencies) * 1000,
                "model_size_gb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1e9,
            }

        return results