The affine map, the error it leaves, and why one outlier ruins a whole row.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
quantization.ipynb, or run it
locally: python3 quantization.py.
python3 quantization.py
Memory, rounding, the outlier and the repair.
python3 quantization.py --round
The same eight weights at int8 and int4. Halving the bit width multiplies the rms error by 20.6, because the step grows from 0.004016 to 0.072857.
python3 quantization.py --outlier
One weight of +4.20 makes the scale 8.2 times coarser. The outlier is then stored exactly while the small weights carry up to 17.32 per cent error.
python3 quantization.py --channel
Two scales instead of one cut the rms error thirteenfold, at a cost of a quarter of a bit per weight, so int4 with group scales really costs 4.25 bits.
Download quantization.py
· served verbatim at https://nlp.jcrlabz.com/code/quantization.py
"""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()