Quantization

Layer 3 · Code

Quantization

Implement group-wise int4 quantization and dequantization from scratch, and measure the actual reconstruction error on realistic weight data.

11 min read110 XP

Below is a from-scratch, dependency-light implementation of group-wise affine (asymmetric) quantization — the same core arithmetic underlying GPTQ, AWQ, and GGUF's k-quants, minus their calibration-based cleverness. We quantize to int4, measure reconstruction error, and see directly why per-group scales beat a single per-tensor scale on data with outliers.

python
import numpy as np

def quantize_groupwise(weights, bits=4, group_size=128):
    """
    weights: 1D or 2D numpy array of fp32/fp16 weights.
    Returns: (int_codes, scales, zero_points, orig_shape)
    Quantization is applied along the last axis, in chunks of group_size.
    """
    orig_shape = weights.shape
    flat = weights.reshape(-1, orig_shape[-1]) if weights.ndim > 1 else weights.reshape(1, -1)
    n_rows, n_cols = flat.shape
    assert n_cols % group_size == 0, "pad or choose a divisor group size in practice"
    n_groups = n_cols // group_size

    qmax = 2 ** bits - 1  # e.g. 15 for 4-bit unsigned
    codes = np.zeros_like(flat, dtype=np.uint8)
    scales = np.zeros((n_rows, n_groups), dtype=np.float32)
    zero_points = np.zeros((n_rows, n_groups), dtype=np.float32)

    for r in range(n_rows):
        for g in range(n_groups):
            chunk = flat[r, g*group_size:(g+1)*group_size]
            w_min, w_max = chunk.min(), chunk.max()
            # Asymmetric affine mapping: scale maps [w_min, w_max] -> [0, qmax]
            scale = (w_max - w_min) / qmax if w_max > w_min else 1.0
            zero_point = w_min
            q = np.round((chunk - zero_point) / scale).clip(0, qmax)
            codes[r, g*group_size:(g+1)*group_size] = q.astype(np.uint8)
            scales[r, g] = scale
            zero_points[r, g] = zero_point

    return codes, scales, zero_points, orig_shape

def dequantize_groupwise(codes, scales, zero_points, group_size=128):
    n_rows, n_cols = codes.shape
    n_groups = n_cols // group_size
    out = np.zeros_like(codes, dtype=np.float32)
    for r in range(n_rows):
        for g in range(n_groups):
            s, z = scales[r, g], zero_points[r, g]
            out[r, g*group_size:(g+1)*group_size] = codes[r, g*group_size:(g+1)*group_size] * s + z
    return out


# --- Build realistic-ish weight data: mostly small values, a few outliers ---
np.random.seed(0)
n_rows, n_cols = 8, 512
weights = np.random.normal(0, 0.02, size=(n_rows, n_cols)).astype(np.float32)
# Inject outlier "channels": a few columns with much larger magnitude, as
# observed empirically in real transformer weight matrices.
outlier_cols = np.random.choice(n_cols, size=5, replace=False)
weights[:, outlier_cols] *= 25.0

def measure_error(weights, group_size):
    codes, scales, zps, shape = quantize_groupwise(weights, bits=4, group_size=group_size)
    recon = dequantize_groupwise(codes, scales, zps, group_size)
    mse = np.mean((weights - recon) ** 2)
    max_err = np.max(np.abs(weights - recon))
    return mse, max_err

for gsize in [512, 128, 32]:  # 512 == "per-tensor" (whole row is one group)
    mse, max_err = measure_error(weights, gsize)
    bits_overhead = 32 / gsize  # one fp32 scale + zero point per group, rough accounting
    print(f"group_size={gsize:4d}  MSE={mse:.6f}  max_err={max_err:.4f}  scale_overhead≈{bits_overhead:.3f} bits/weight")
Group-wise int4 quantize/dequantize with scale and zero-point.