"""Quantisation, worked as Chapter 23 works it. Four things, all reproducing the chapter's tables: the memory table what each precision costs for a 7B model one tensor the affine map, rounded and dequantised, with its error the outlier why one large weight ruins the whole tensor per-channel the fix, and what it costs to store python3 quantization.py # all four python3 quantization.py --memory # bytes per parameter python3 quantization.py --round # int8 and int4 on eight weights python3 quantization.py --outlier # one big weight, and the damage python3 quantization.py --channel # per-channel scales Install: nothing, the Python standard library is enough """ import argparse # Eight weights from one row of a hidden layer. Small, spread either side of # zero, and deliberately not symmetric: this is what most of a trained model # looks like. WEIGHTS = [0.12, -0.45, 0.33, -0.08, 0.51, -0.22, 0.07, -0.37] # The same row after one outlier arrives. Large activations and weights in a # few coordinates are a measured property of trained transformers, not an # invention for this example. OUTLIER = WEIGHTS + [4.20] PRECISIONS = [("fp32", 4.0), ("fp16", 2.0), ("int8", 1.0), ("int4", 0.5)] PARAMS_7B = 7_000_000_000 def show_memory(): print("\nSTEP 1 what a precision costs\n") print(" precision bytes/param a 7B model versus fp16") print(" " + "-" * 56) fp16 = PARAMS_7B * 2.0 / 1e9 for name, b in PRECISIONS: gb = PARAMS_7B * b / 1e9 print(f" {name:9} {b:11.1f} {gb:9.1f} GB {gb / fp16:9.2f}x") print(" " + "-" * 56) print("\n Read the int4 row. A model that needed a data-centre card at fp16") print(" fits on a consumer GPU at 3.5 GB, and nothing was retrained.") print(" Weights are stored in 16 bits and rarely need 16 bits to be used.") def quantise(ws, bits): """Symmetric affine quantisation. Returns scale, codes, dequantised values.""" qmax = 2 ** (bits - 1) - 1 # 127 for int8, 7 for int4 scale = max(abs(w) for w in ws) / qmax codes = [max(-qmax - 1, min(qmax, round(w / scale))) for w in ws] deq = [c * scale for c in codes] return scale, codes, deq def errors(ws, deq): errs = [abs(w - d) for w, d in zip(ws, deq)] rms = (sum(e * e for e in errs) / len(errs)) ** 0.5 return max(errs), rms def show_round(ws=WEIGHTS, title="STEP 2 one tensor, rounded and put back"): print(f"\n{title}\n") print(" A weight is stored as an integer code and a shared scale:\n") print(" q = round(w / S), S = max|w| / q_max, w_hat = q * S\n") for bits in (8, 4): scale, codes, deq = quantise(ws, bits) qmax = 2 ** (bits - 1) - 1 mx, rms = errors(ws, deq) print(f" int{bits}: q_max = {qmax}, S = max|w| / {qmax} = {scale:.6f}\n") print(" w w/S q w_hat error") print(" " + "-" * 44) for w, c, d in zip(ws, codes, deq): print(f" {w:+7.2f} {w / scale:+9.2f} {c:+5d} {d:+8.4f} {abs(w - d):8.5f}") print(" " + "-" * 44) print(f" max error {mx:.5f} rms error {rms:.5f}\n") s8, _, d8 = quantise(ws, 8) s4, _, d4 = quantise(ws, 4) _, r8 = errors(ws, d8) _, r4 = errors(ws, d4) print(f" Halving the bit width multiplies the rms error by {r4 / r8:.1f}, because") print(f" the step between neighbouring codes grows from {s8:.6f} to {s4:.6f}.") def show_outlier(): print("\nSTEP 3 one large weight ruins the row\n") s_clean, _, _ = quantise(WEIGHTS, 8) s_out, codes, deq = quantise(OUTLIER, 8) print(" Add a single outlier of +4.20 to the same eight weights. The scale") print(" is set by the largest value, so every other weight now shares a") print(f" step {s_out / s_clean:.1f} times coarser.\n") print(f" int8 scale without the outlier S = {s_clean:.6f}") print(f" int8 scale with it S = {s_out:.6f} ({s_out / s_clean:.1f}x coarser)\n") print(" w q w_hat error error %") print(" " + "-" * 46) for w, c, d in zip(OUTLIER, codes, deq): pct = 100 * abs(w - d) / abs(w) if w else 0.0 flag = " <- outlier" if abs(w) > 1 else "" print(f" {w:+7.2f} {c:+5d} {d:+8.4f} {abs(w - d):8.5f} {pct:7.2f}%{flag}") print(" " + "-" * 46) small = [(w, d) for w, d in zip(OUTLIER, deq) if abs(w) <= 1] worst = max(100 * abs(w - d) / abs(w) for w, d in small) print(f"\n The outlier itself is fine, at {100 * abs(4.20 - deq[-1]) / 4.20:.2f} per cent. The small weights") print(f" now carry up to {worst:.2f} per cent error, and there are far more of them.") print(" This is why naive per-tensor quantisation of a transformer fails.") def show_channel(): print("\nSTEP 4 give each group its own scale\n") print(" Split the row into groups and quantise each with its own scale, so") print(" one extreme value cannot set the step for everything else.\n") groups = [("the eight normal weights", WEIGHTS), ("the outlier alone", [4.20])] _, _, deq_all = quantise(OUTLIER, 8) _, rms_all = errors(OUTLIER, deq_all) rms_parts, n = 0.0, 0 for name, g in groups: s, _, d = quantise(g, 8) mx, rms = errors(g, d) rms_parts += sum((a - b) ** 2 for a, b in zip(g, d)) n += len(g) print(f" {name:26} S = {s:.6f} max error {mx:.5f}") rms_grouped = (rms_parts / n) ** 0.5 print(f"\n rms error, one scale for the row {rms_all:.5f}") print(f" rms error, two scales {rms_grouped:.5f}") print(f" improvement {rms_all / rms_grouped:.0f}x\n") print(" The cost is the scales themselves. One fp16 scale per group of 64") print(" weights adds 2 bytes per 64, which is 0.25 bits per weight, so int4") print(" with group scales really costs 4.25 bits. That is the trade every") print(" quantisation scheme is negotiating.") def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--memory", action="store_true") ap.add_argument("--round", action="store_true") ap.add_argument("--outlier", action="store_true") ap.add_argument("--channel", action="store_true") a = ap.parse_args() run_all = not (a.memory or a.round or a.outlier or a.channel) if a.memory or run_all: show_memory() if a.round or run_all: show_round() if a.outlier or run_all: show_outlier() if a.channel or run_all: show_channel() if __name__ == "__main__": main()