ptq_int8.py

"""
quantise/ptq_int8.py

Post-Training Quantisation (PTQ) pipeline for GPT-2.
Converts BF16 fine-tuned model to INT8 for efficient inference.

Why PTQ over QAT?
- PTQ requires only a small calibration dataset (~512 samples)
- No retraining required — calibration takes minutes vs hours
- INT8 inference: ~2x memory reduction, ~2-4x throughput on supported hardware
- Accuracy drop typically < 1% on language models with good calibration data

Pipeline:
  1. Load BF16 checkpoint
  2. Calibrate with representative data (captures activation ranges)
  3. Quantise linear layers to INT8 (weights + activations)
  4. Evaluate accuracy vs FP32 baseline
  5. Benchmark latency and throughput

Note on Graphcore IPU:
  IPU supports low-precision compute natively — this pipeline demonstrates
  the quantisation reasoning applicable to any accelerator with INT8/FP16 support.
"""

import argparse
import time
from pathlib import Path
from typing import Dict

import torch
import torch.nn as nn
from torch.quantization import (
    quantize_dynamic,
    prepare,
    convert,
    get_default_qconfig,
    QConfig,
)
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from datasets import load_dataset


# ── Calibration ───────────────────────────────────────────────────────────────

def get_calibration_data(tokenizer, n_samples: int = 512, seq_len: int = 128):
    """
    Load calibration dataset.

    Calibration data must be REPRESENTATIVE of real inference inputs.
    Poor calibration data → poor activation range estimation → accuracy drop.
    We use WikiText validation set — same distribution as fine-tuning data.
    """
    dataset = load_dataset("wikitext", "wikitext-103-v1", split="validation")
    texts = [t for t in dataset["text"] if len(t.strip()) > 50][:n_samples]

    encodings = tokenizer(
        texts,
        truncation=True,
        max_length=seq_len,
        padding="max_length",
        return_tensors="pt",
    )
    return encodings


# ── Dynamic quantisation (fastest, no calibration needed) ────────────────────

def apply_dynamic_quantisation(model: nn.Module) -> nn.Module:
    """
    Dynamic INT8 quantisation — quantises weights statically,
    activations dynamically at inference time.

    Pros: No calibration data needed, very fast to apply
    Cons: Less optimal than static quant — activation ranges computed per-batch
    Best for: Transformer models where matmul dominates compute
    """
    quantised = quantize_dynamic(
        model,
        qconfig_spec={nn.Linear},
        dtype=torch.qint8,
        inplace=False,
    )
    return quantised


# ── Static quantisation (better accuracy, requires calibration) ──────────────

def apply_static_quantisation(model: nn.Module, calibration_data: dict,
                               device: str = "cpu") -> nn.Module:
    """
    Static INT8 quantisation with calibration.
    Captures activation ranges during calibration → better INT8 mapping.
    """
    model.eval()
    model.to(device)

    # Set quantisation config — fbgemm for x86, qnnpack for ARM
    backend = "fbgemm"  # swap to "qnnpack" for ARM/mobile
    model.qconfig = get_default_qconfig(backend)
    torch.backends.quantized.engine = backend

    # Prepare — inserts observers to record activation ranges
    prepared = prepare(model, inplace=False)

    # Calibrate — run ~100-512 samples through to capture ranges
    print("Calibrating...")
    with torch.no_grad():
        for i in range(0, min(100, calibration_data["input_ids"].shape[0]), 8):
            batch = {k: v[i:i+8].to(device)
                     for k, v in calibration_data.items()
                     if k in ("input_ids", "attention_mask")}
            prepared(**batch)

    # Convert — replaces float ops with INT8 quantised ops
    quantised = convert(prepared, inplace=False)
    print("Static INT8 quantisation complete.")
    return quantised


# ── Evaluation ────────────────────────────────────────────────────────────────

def compute_perplexity(model: nn.Module, tokenizer, n_samples: int = 100,
                        seq_len: int = 128) -> float:
    """
    Compute perplexity on WikiText validation set.
    Lower perplexity = better language model.
    Acceptable INT8 accuracy degradation: < 2 perplexity points vs FP32 baseline.
    """
    model.eval()
    dataset = load_dataset("wikitext", "wikitext-103-v1", split="validation")
    texts = [t for t in dataset["text"] if len(t.strip()) > 50][:n_samples]

    total_loss = 0.0
    count = 0

    with torch.no_grad():
        for text in texts:
            inputs = tokenizer(text, return_tensors="pt", truncation=True,
                               max_length=seq_len)
            labels = inputs["input_ids"].clone()
            try:
                outputs = model(**inputs, labels=labels)
                total_loss += outputs.loss.item()
                count += 1
            except Exception:
                continue

    avg_loss = total_loss / max(count, 1)
    return torch.exp(torch.tensor(avg_loss)).item()


def benchmark_latency(model: nn.Module, tokenizer, n_runs: int = 50,
                       seq_len: int = 128) -> Dict[str, float]:
    """
    Benchmark inference latency and throughput.
    Returns p50, p95, p99 latency and tokens/sec throughput.
    """
    model.eval()
    dummy = tokenizer("Benchmark input text for latency measurement",
                      return_tensors="pt", truncation=True,
                      max_length=seq_len, padding="max_length")

    latencies = []

    # Warmup
    with torch.no_grad():
        for _ in range(5):
            model(**dummy)

    # Measure
    with torch.no_grad():
        for _ in range(n_runs):
            t0 = time.perf_counter()
            model(**dummy)
            latencies.append((time.perf_counter() - t0) * 1000)

    latencies_sorted = sorted(latencies)
    n = len(latencies_sorted)
    tokens_per_sec = (seq_len * 1000) / (sum(latencies) / n)

    return {
        "p50_ms":       round(latencies_sorted[n // 2], 2),
        "p95_ms":       round(latencies_sorted[int(n * 0.95)], 2),
        "p99_ms":       round(latencies_sorted[int(n * 0.99)], 2),
        "mean_ms":      round(sum(latencies) / n, 2),
        "tokens_per_s": round(tokens_per_sec, 1),
    }


def model_size_mb(model: nn.Module) -> float:
    """Estimate model size in MB from parameter count and dtype."""
    total_bytes = 0
    for p in model.parameters():
        total_bytes += p.numel() * p.element_size()
    return round(total_bytes / 1e6, 1)


# ── Main ─────────────────────────────────────────────────────────────────────

def main(args):
    print("=" * 60)
    print("  INT8 Post-Training Quantisation Pipeline")
    print("=" * 60)

    tokenizer = GPT2Tokenizer.from_pretrained(args.model)
    tokenizer.pad_token = tokenizer.eos_token

    # Load BF16 checkpoint or pretrained weights
    print(f"\n[1] Loading model: {args.model}")
    if args.checkpoint:
        state = torch.load(args.checkpoint, map_location="cpu")
        model = GPT2LMHeadModel.from_pretrained(args.model)
        model.load_state_dict(state["model_state_dict"])
    else:
        model = GPT2LMHeadModel.from_pretrained(args.model, torch_dtype=torch.float32)

    model.eval()

    # FP32 baseline
    print("\n[2] FP32 baseline benchmark")
    fp32_ppl    = compute_perplexity(model, tokenizer)
    fp32_lat    = benchmark_latency(model, tokenizer)
    fp32_size   = model_size_mb(model)
    print(f"    Perplexity: {fp32_ppl:.2f} | Size: {fp32_size}MB | "
          f"p50: {fp32_lat['p50_ms']}ms | Throughput: {fp32_lat['tokens_per_s']} tok/s")

    # Dynamic INT8
    print("\n[3] Applying dynamic INT8 quantisation")
    dyn_model   = apply_dynamic_quantisation(model)
    dyn_ppl     = compute_perplexity(dyn_model, tokenizer)
    dyn_lat     = benchmark_latency(dyn_model, tokenizer)
    dyn_size    = model_size_mb(dyn_model)
    print(f"    Perplexity: {dyn_ppl:.2f} | Size: {dyn_size}MB | "
          f"p50: {dyn_lat['p50_ms']}ms | Throughput: {dyn_lat['tokens_per_s']} tok/s")

    # Results table
    print("\n" + "=" * 60)
    print("  Quantisation Results Summary")
    print("=" * 60)
    print(f"  {'Metric':<25} {'FP32':>10} {'INT8 Dynamic':>14}")
    print(f"  {'-'*50}")
    print(f"  {'Perplexity':<25} {fp32_ppl:>10.2f} {dyn_ppl:>14.2f}")
    print(f"  {'Model Size (MB)':<25} {fp32_size:>10.1f} {dyn_size:>14.1f}")
    print(f"  {'P50 Latency (ms)':<25} {fp32_lat['p50_ms']:>10.2f} {dyn_lat['p50_ms']:>14.2f}")
    print(f"  {'P95 Latency (ms)':<25} {fp32_lat['p95_ms']:>10.2f} {dyn_lat['p95_ms']:>14.2f}")
    print(f"  {'Throughput (tok/s)':<25} {fp32_lat['tokens_per_s']:>10.1f} {dyn_lat['tokens_per_s']:>14.1f}")
    print(f"  {'Accuracy Drop (PPL)':<25} {'—':>10} {dyn_ppl - fp32_ppl:>+14.2f}")
    print(f"  {'Size Reduction':<25} {'—':>10} {(1 - dyn_size/fp32_size)*100:>13.1f}%")
    print(f"  {'Speedup':<25} {'—':>10} {fp32_lat['p50_ms']/dyn_lat['p50_ms']:>13.2f}x")
    print("=" * 60)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model",      default="gpt2")
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--n_samples",  type=int, default=100)
    main(parser.parse_args())