Quantization, Pruning, Distillation & Fast Attention

🎯 Learning Objectives

• Explain model quantization (INT8, FP16, INT4) and its trade-offs
• Implement weight pruning and understand sparsity patterns
• Understand knowledge distillation for model compression
• Describe Flash Attention and KV-cache optimization

📋 Prerequisites

• GPU memory hierarchy
• Floating-point number representation
• Transformer architecture

1. 📖 Core Content

1.1 The Deployment Problem

1.2 Model Quantization

Quantization Formula

• $x$: Original floating-point value
• $\Delta$: Scale factor (step size)
• $z$: Zero point (bias)
• $x_q$: Quantized integer value Dequantization:

Post-Training Quantization (PTQ)

1. Weight-only: Quantize weight matrices (most common)
2. Weight + activation: Quantize both (harder, needs calibration data)
3. Per-tensor: Same scale for entire tensor
4. Per-channel: Different scale per output channel (better)

Quantization-Aware Training (QAT)

python
Run Inline
# runnable
import numpy as np
def quantize_weights(W, bits=8):
    """
    Quantize weights to INT bits
    Args:
        W: Weight matrix (floating-point)
        bits: Target bit-width (8, 4, 2)
    Returns:
        W_q: Quantized integer weights
        scale: Scale factors
    """
    # Find range
    w_max = np.max(np.abs(W))
    # Number of quantization levels
    n_levels = 2 ** (bits - 1)  # For symmetric quantization
    # Scale factor
    scale = w_max / n_levels
    # Quantize
    W_q = np.round(W / scale)
    W_q = np.clip(W_q, -n_levels, n_levels - 1)
    # Memory
    orig_mem = W.size * 32  # bits
    quant_mem = W.size * bits
    compression = orig_mem / quant_mem
    # Dequantize (for evaluation)
    W_deq = W_q * scale
    # Error
    mse = np.mean((W - W_deq)**2)
    return W_q.astype(np.int8), scale, compression, mse
# Example
np.random.seed(42)
W = np.random.randn(4096, 4096) * 0.01
for bits in [8, 4, 2]:
    W_q, scale, compression, mse = quantize_weights(W, bits)
    print(f"INT{bits}: compression={compression:.0f}×, MSE={mse:.6f}")

1.3 Pruning

1.4 Knowledge Distillation

• $\hat{y}_{teacher}$: Teacher's soft probabilities (temperature-scaled)
• $y_{true}$: Ground truth labels
• $\alpha$: Distillation weight Temperature scaling:

1.5 Flash Attention

1.6 KV-Cache

pseudo
Step 1: Compute K₁, V₁ for token 1 → store
Step 2: Compute K₂, V₂ for token 2 → append to cache
         Attention uses cached [K₁, K₂] and [V₁, V₂]
Step 3: Compute K₃, V₃ → append to cache
         Attention uses cached [K₁, K₂, K₃]

• K: 2 bytes × 32 × 2 × 2048 × 4096 = 1.07 GB (FP16)
• V: same = 1.07 GB
• Total: ~2.1 GB per sequence For batch_size=32: 67 GB — often the bottleneck for long contexts!

1.7 Why This Matters

• Quantization: Run 175B models on 2 GPUs instead of 8
• KV-cache: Real-time generation speed
• Flash Attention: Train with 8× longer sequences
• Distillation: Create task-specific small models from large general models

6. 📝 Practice Questions

Q1: GPT-3 has 175B parameters in FP16. What is its memory footprint?

175B × 2 bytes (FP16) = 350 GB

With INT8: 175 GB With INT4: 87.5 GB

This is why quantization is essential: a 175B model at INT4 fits in a single A100 (80GB... barely) or 2 A100s. Q2

<strong>Q2</strong>: For a sequence of length 4096 with d_model=4096 and 32 layers, compute the KV-cache size for batch_size=4.

Per layer: 2 × seq_len × d_model (one K, one V) = 2 × 4096 × 4096 = 33,554,432 elements

All layers: 32 × 33,554,432 = 1,073,741,824 elements

In FP16 (2 bytes): 2,147,483,648 bytes = 2 GB

Batch of 4: 8 GB

This is why generating long sequences requires significant GPU memory! Q3

<strong>Q3</strong>: In Flash Attention, why does tiling the computation reduce memory reads/writes?

Standard attention reads and writes the N×N attention matrix from HBM (slow, high bandwidth but high latency). The attention matrix for N=4096 is 4096² = 16M elements = 128 MB (FP16).

Flash attention processes in tiles (e.g., 128×128 blocks). The tile fits in SRAM (fast, on-chip). Each tile is:

1. Read from HBM to SRAM (only Q, K, V tiles, not the full matrix)
2. Compute attention scores in SRAM
3. Write only the output tile back to HBM

This avoids the expensive HBM round-trip for the full attention matrix. Total I/O is reduced from O(N² + N²) to O(N²) (dominated by reading Q, K, V repeatedly). Q4

<strong>Q4

<strong>Q4</strong>: Compare weight-only INT4 quantization with weight+activation INT8 quantization for inference speed.

Weight-only INT4:

• Weights stored as INT4, dequantized to FP16 on-the-fly
• Computations still in FP16
• Memory: 4× reduction
• Speed: Memory-bandwidth-bound → faster (less data to load)
• Quality: Minimal degradation (modern LLMs are robust to weight quantization)

Weight+Activation INT8:

• Both weights and activations in INT8
• Matrix multiplication in INT8 (2-4× faster than FP16 on modern GPUs)
• Memory: 4× reduction for weights, 2× for activations
• Quality: More degradation (activations vary more than weights)

For most use cases, weight-only INT4 is preferred for quality preservation. INT8 compute is only beneficial for compute-bound operations (very large batches). Q5

<strong>Q5

<strong>Q5</strong>: With 50% unstructured sparsity (half the weights pruned), what speedup can you expect?

Unstructured sparsity gives little to no speedup on standard GPUs because:

• Dense matrix multiplication (cuBLAS) is highly optimized
• Sparse operations require irregular memory access
• Only ~15-20% pruned weights can be effectively skipped

Speedup requires structured sparsity (NVIDIA's 2:4 sparsity gives 2× speedup on A100/H100 with n:m support).

This is why unstructured pruning for LLMs is less popular than quantization — you get memory reduction but not speedup. Q6

<strong>Q6

<strong>Q6</strong>: In knowledge distillation with temperature T, what happens as T → ∞?

As T → ∞:

• All logits are divided by infinity → all near 0
• Softmax output becomes uniform distribution: p_i = 1/\text{vocab_size}
• The teacher's "soft targets" convey no information about relative probabilities
• The student only learns from the hard labels

This is why intermediate temperatures (T=2-10) are used: they preserve the teacher's knowledge about relative similarities between classes while softening the distribution enough for efficient learning. Q7

<strong>Q7

<strong>Q7</strong>: A quantized model shows accuracy drops mostly on factual recall tasks but not on reasoning tasks. Why might this be?

Factual recall tasks (e.g., "What is the capital of France?") depend on exact weight values — the model must output exactly "Paris." Quantization introduces small errors that can shift the output distribution just enough to change the argmax token.

Reasoning tasks (e.g., math, logic) involve computation patterns where small errors in intermediate steps may cancel out. The final answer is more robust to quantization noise.

This is why factuality benchmarks (MMLU, TriviaQA) are more sensitive to quantization than reasoning benchmarks (GSM8K). Q8

<strong>Q8

<strong>Q8</strong>: A 7B model has 32 layers. If we quantize only every other layer to INT4 (keeping others at FP16), what is the effective precision and memory?

Half the layers (16) in INT4, half (16) in FP16:

Average precision: (16 × 4 + 16 × 16) / 32 = 10 bits average

Memory:

• FP16 layers: 16 layers × d_model² × 4 projections × 2 bytes = 16 × 16M × 8 × 2 = 16 × 256M = 4GB
• INT4 layers: 16 × 128M = 2GB
• Total: ~6GB (vs 8GB for full FP16)

This mixed precision approach can preserve quality better than full INT4 while still saving memory. It's based on the observation that different layers have different sensitivity to quantization. Q9

<strong>Q9

<strong>Q9

<strong>Q9</strong>: Grouped-query attention (GQA) reduces KV-cache by sharing keys/values across query heads. How does this work?

In standard multi-head attention: each of h heads has its own K and V. GQA divides heads into g groups (g < h). Heads in the same group share K and V:

• Standard MHA: h independent Q, K, V heads
• GQA: h Q heads, g K/V heads (g < h)

For g=8, h=32: 4× KV-cache reduction (32 KV heads → 8 KV heads)

GQA is used in LLaMA 2 70B, Mistral, and other modern LLMs to enable longer context windows without proportionally increasing KV-cache memory. Q10

<strong>Q10

<strong>Q10

<strong>Q10</strong>: Speculative decoding generates tokens 2-3× faster without quality loss. How does it work?

Speculative decoding uses a draft model (small, fast) and a target model (large, accurate):

1. Draft model generates k tokens quickly (e.g., 5 tokens)
2. Target model validates all k tokens in one forward pass
3. If all k consistent → accept all (huge speedup)
4. If discrepancy → reject from divergence point, resample

Speedup factor: \text{speedup} = \frac{k}{\text{acceptance_rate} + (1 - \text{acceptance_rate}) \cdot k}

With k=5, acceptance_rate=0.8: speedup = 5 / (0.8 + 0.2×5) = 5/1.8 ≈ 2.8×

The key insight: the target model's single forward pass is only slightly slower than the draft model's sequential passes, but verifies multiple tokens at once. This gives ~2-3× wall-clock speedup without any loss in output quality.

7. 🔗 Cross-References

• Next: Advanced Positional Encodings (Week 11)
• Previous: RLHF & Alignment
• Video: BSDA5004 Week 10 transcripts Join Discord PreviousRLHF & AlignmentNextAdvanced PE: RoPE & ALiBi