GPTQ and AWQ weight quantisation

Open in Colab ▶️ Run this notebook in Colab

GPTQ and AWQ weight quantisation#

Track 05 - Serving · Notebook 05 · Runtime: ≈15 min on CPU

Prerequisites: 05_serving/01 (roofline), 03_training/01 (mixed precision).

Papers:

  • Frantar et al. 2022, GPTQ.

  • Lin et al. 2023, AWQ.

What#

Quantise a linear layer’s weights from FP16 to 4-bit integer and measure how much output quality is lost. Two methods:

  • Per-group symmetric int4 (round-to-nearest, RTN). Baseline. Split each row of W into groups of 128 columns; compute one scale per group; round. Simple and fast; loses noticeable accuracy on outliers.

  • Activation-aware scaling (AWQ, Lin 2023). Before quantising, rescale each input channel by a factor s_i chosen so the channels with the largest activation magnitudes are preserved at higher precision. (W * s) @ (x / s) = W @ x so the layer output is unchanged by the rescaling; but after quantisation the error on high-activation channels is lower.

Target: AWQ produces lower output MSE than plain RTN on a calibration batch at 4× compression.

Full GPTQ (Frantar 2022) — second- order error correction via Cholesky — is out of scope for a from-scratch notebook; we compare against it qualitatively at the end.

from llm_systems_cookbook.nb import bootstrap

import torch
import torch.nn.functional as F

s = bootstrap("05_serving_05_gptq_awq_weight_quant")

Setup: a linear layer and some calibration data#

FP16 weight matrix of shape (out=4096, in=4096) plus a batch of 256 random activations that stand in for a real calibration set (in practice you’d pass a few hundred tokens through the model and record each layer’s input).

torch.manual_seed(0)
OUT_FEATURES = 4096
IN_FEATURES = 4096
GROUP_SIZE = 128

W = torch.randn(OUT_FEATURES, IN_FEATURES, dtype=torch.float16) * 0.02
# Heavy-tail activations: most channels are well-behaved, a few have
# magnitudes 5-10x larger - mirrors what happens inside a real LLM after
# a few layers of growth.
x = torch.randn(256, IN_FEATURES, dtype=torch.float16)
outlier_channels = torch.randperm(IN_FEATURES)[:64]
x[:, outlier_channels] *= 6.0

y_ref = F.linear(x.float(), W.float()).half()
print(f"W shape = {tuple(W.shape)}  |W|_max = {W.abs().max().item():.3f}")
print(f"x shape = {tuple(x.shape)}  |x|_max = {x.abs().max().item():.3f}")
print(f"y ref shape = {tuple(y_ref.shape)}")

Per-group symmetric int4 (RTN)#

For each row of W, split into groups of group_size columns. For each group compute scale = max(|w|) / 7 (7 is the largest positive int in signed int4). Quantise to q = round(w / scale), clamp to [-8, 7], dequantise by w_hat = q * scale.

def quant_rtn_int4(w: torch.Tensor, group_size: int) -> torch.Tensor:
    '''Row-wise per-group symmetric int4 round-to-nearest.'''
    out_f, in_f = w.shape
    assert in_f % group_size == 0
    w_g = w.view(out_f, in_f // group_size, group_size).float()
    scale = w_g.abs().amax(dim=-1, keepdim=True) / 7.0
    scale = scale.clamp(min=1e-8)
    q = torch.round(w_g / scale).clamp(-8, 7)
    w_hat = (q * scale).view(out_f, in_f).to(w.dtype)
    return w_hat


W_rtn = quant_rtn_int4(W, GROUP_SIZE)
y_rtn = F.linear(x.float(), W_rtn.float()).half()
mse_rtn = (y_ref.float() - y_rtn.float()).pow(2).mean().item()
print(f"RTN:  weight max abs err = {(W - W_rtn).abs().max().item():.5f}")
print(f"      output MSE          = {mse_rtn:.4e}")

Activation-aware scaling (AWQ)#

Key idea: channels with larger activation magnitudes deserve lower quantisation error. Define s_i (one per input channel) from the per-channel magnitude |x|_i, then transform:

W'   = W * s   (scale columns)
x'   = x / s   (scale rows)
y    = (W * s) @ (x / s) = W @ x      (unchanged mathematically)

After this transform, the column magnitudes of W' are larger precisely where it matters; RTN on W' then preserves those columns better, and the overall output error drops.

We use s_i = |x|_i^alpha with alpha ∈ [0, 1]. AWQ’s paper reports alpha ≈ 0.5 as the sweet spot; we sweep and pick the one that minimises output MSE on the calibration set.

def awq_quant_int4(
    w: torch.Tensor, x_calib: torch.Tensor, group_size: int, alpha: float
) -> torch.Tensor:
    '''AWQ-style quantisation: rescale by per-channel activation magnitudes.'''
    # Per-input-channel magnitude from the calibration batch.
    x_mag = x_calib.abs().mean(dim=0).float().clamp(min=1e-8)
    s = x_mag ** alpha
    s = s / s.mean()   # keep scale centred so no overflow.
    w_scaled = w.float() * s[None, :]           # broadcast over rows
    w_scaled_q = quant_rtn_int4(w_scaled.to(w.dtype), group_size).float()
    # Fold s back into weights so inference math matches: w_hat = (w_scaled_q / s)
    w_hat = (w_scaled_q / s[None, :]).to(w.dtype)
    return w_hat


best_mse = float("inf")
best_alpha = None
for alpha in (0.0, 0.25, 0.5, 0.75, 1.0):
    W_awq = awq_quant_int4(W, x, GROUP_SIZE, alpha)
    y_awq = F.linear(x.float(), W_awq.float()).half()
    mse = (y_ref.float() - y_awq.float()).pow(2).mean().item()
    print(f"  alpha={alpha:.2f}   output MSE = {mse:.4e}")
    if mse < best_mse:
        best_mse = mse
        best_alpha = alpha

print(f"AWQ best alpha = {best_alpha}   output MSE = {best_mse:.4e}")
mse_awq = best_mse

Storage math#

Original: 4096 * 4096 * 2 bytes = 32 MiB for this one layer. Int4 with a per-group fp16 scale (one scale per 128 columns):

weight bytes  = 4096 * 4096 * 0.5 = 8 MiB     (half a byte per element)
scale  bytes  = 4096 * 32   * 2   = 0.25 MiB
total         = 8.25 MiB

4× compression in practice (the scales add a tiny constant overhead).

bytes_fp16 = W.numel() * 2
bytes_int4 = W.numel() // 2 + (W.shape[0] * (W.shape[1] // GROUP_SIZE)) * 2
ratio = bytes_fp16 / bytes_int4
print(f"fp16 bytes  = {bytes_fp16 / 1024**2:.2f} MiB")
print(f"int4 bytes  = {bytes_int4 / 1024**2:.2f} MiB  (incl. per-group fp16 scales)")
print(f"compression = {ratio:.2f}x")

s.assert_close("int4_gives_close_to_4x_compression", actual=ratio, expected=4.0, rtol=0.05)
s.check(
    "awq_beats_rtn_on_output_mse",
    lambda: mse_awq < mse_rtn,
    msg=f"RTN MSE={mse_rtn:.3e}  AWQ MSE={mse_awq:.3e}",
)
s.check(
    "awq_outputs_finite",
    lambda: (not torch.isnan(torch.tensor(mse_awq))) and mse_awq > 0,
)

Output MSE vs bit width, RTN vs AWQ#

Sweep the weight bitwidth from 8 down to 2 and compare plain RTN to AWQ with the best calibration alpha. AWQ’s activation-aware rescaling mostly matters at 4 bits and below - above that, RTN already has enough dynamic range.

import matplotlib.pyplot as plt

def rtn_nbit(w, bits, group=GROUP_SIZE):
    out_f, in_f = w.shape
    levels = (1 << (bits - 1)) - 1
    wg = w.view(out_f, in_f // group, group).float()
    scale = (wg.abs().amax(dim=-1, keepdim=True) / levels).clamp(min=1e-8)
    q = torch.round(wg / scale).clamp(-levels - 1, levels)
    return (q * scale).view(out_f, in_f).to(w.dtype)

def awq_nbit(w, x_calib, bits, alpha, group=GROUP_SIZE):
    mag = x_calib.abs().mean(dim=0).float().clamp(min=1e-8) ** alpha
    mag = mag / mag.mean()
    wq = rtn_nbit((w.float() * mag[None, :]).to(w.dtype), bits, group).float()
    return (wq / mag[None, :]).to(w.dtype)

bit_widths = [8, 4, 3, 2]
rtn_mse, awq_mse = [], []
for bits in bit_widths:
    y_r = F.linear(x.float(), rtn_nbit(W, bits).float()).half()
    rtn_mse.append((y_ref.float() - y_r.float()).pow(2).mean().item())
    y_a = F.linear(x.float(), awq_nbit(W, x, bits, best_alpha).float()).half()
    awq_mse.append((y_ref.float() - y_a.float()).pow(2).mean().item())

idx = list(range(len(bit_widths)))
fig, ax = plt.subplots(figsize=(7, 4))
ax.bar([i - 0.2 for i in idx], rtn_mse, width=0.4, color="tab:gray",   label="RTN")
ax.bar([i + 0.2 for i in idx], awq_mse, width=0.4, color="tab:orange", label=f"AWQ (alpha={best_alpha})")
ax.set_yscale("log")
ax.set_xticks(idx)
ax.set_xticklabels([f"{b}-bit" for b in bit_widths])
ax.set_ylabel("output MSE vs fp16 (log)")
ax.set_title("weight quantisation quality: RTN vs AWQ across bit widths")
ax.legend()
ax.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

About GPTQ (context only)#

GPTQ takes a step further: instead of picking the quantisation greedily, it solves the quadratic error-minimisation problem layer by layer. For each column, the quantisation error is added to the remaining unquantised columns (Hessian-weighted), so later columns can compensate for earlier ones. This gives sub-4-bit quality on some layers but needs a Hessian inverse per layer.

Production inference uses either GPTQ- or AWQ-produced weights via gptqmodel, autoawq, or llmcompressor. The notebook’s RTN and AWQ paths represent the “easy” end of the quality axis; GPTQ is the “hard but best” end.

Exercises#

  1. Sweep group_size ∈ {32, 64, 128, 256} and plot MSE vs size. Smaller groups = more scales = more overhead = lower error.

  2. Add per-token clipping: clip W to its 99th-percentile magnitude before quantising. Does it help at alpha=0? At alpha=0.5?

  3. Hook this into transformers: replace a nn.Linear layer in SmolLM2-135M with an AWQ-quantised version and measure perplexity on wikitext-2 (uses the Phase 2 perplexity notebook as the harness).

References#

s.summary()
s.save()