What does quantization trade away?

"""Keep the numbers and model shape inside ../widget.html equal to quantize.py.

    python3 embed_widget.py          rewrite the data block in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

import quantize as qz

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
DATA_BLOCK = re.compile(r'(<script type="application/json" id="qnt-data">)(.*?)(</script>)', re.S)


def widget_data() -> dict:
    cfg = qz.load_config()
    return {"model": {"name": "Llama 3.1 8B", "layers": cfg["num_hidden_layers"], "hidden": cfg["hidden_size"],
                      "ffn": cfg["intermediate_size"], "heads": cfg["num_attention_heads"], "kv_heads": cfg["num_key_value_heads"],
                      "head_dim": cfg.get("head_dim"), "vocab": cfg["vocab_size"], "tied": bool(cfg.get("tie_word_embeddings"))},
            "bits": list(qz.BITS), "full_bits": qz.FULL_BITS, "scale_bits": qz.SCALE_BITS, "group": qz.GROUP,
            "sample": {"weights": qz.sample_weights(), "outlier_slot": qz.SAMPLE_OUTLIER_SLOT, "outlier": qz.SAMPLE_OUTLIER,
                       "group": qz.SAMPLE_GROUP},
            "gpus": [[name, int(bw)] for name, bw in qz.BANDWIDTHS.items()]}


def embed(html: str) -> str:
    """html with the data block rewritten from quantize.py."""
    block = json.dumps(widget_data(), indent=1)
    return DATA_BLOCK.sub(lambda m: m.group(1) + "\n" + block + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data block now matches quantize.py")

{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 131072,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": {
    "factor": 8.0,
    "low_freq_factor": 1.0,
    "high_freq_factor": 4.0,
    "original_max_position_embeddings": 8192,
    "rope_type": "llama3"
  },
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.43.0.dev0",
  "use_cache": true,
  "vocab_size": 128256
}

"""Round weights to fewer bits and measure what that costs in error, size and decoding speed.

numpy only. The method is round-to-nearest with a symmetric (absmax) scale: each group of weights
shares one 16-bit scale, and every weight becomes a small signed whole number times that scale.
The model shape comes from llama31_8b_config.json, a copy of the config.json of Llama 3.1 8B on
Hugging Face (read 2026-10-19). Sizes are in bytes unless a name says GiB; bandwidths are in
bytes per second (GB and TB mean 10^9 and 10^12 bytes).
"""
import json
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
GIB = 1024**3

BITS = (16, 8, 4, 3)      # bits stored per weight the lesson compares; 16 is the unquantized reference
FULL_BITS = 16            # the token table and output matrix stay at 16 bits in published 8-bit and 4-bit files
SCALE_BITS = 16           # each scale is stored as a 16-bit float
GROUP = 128               # weights that share one scale in the usual 4-bit files

# Graphics-card memory bandwidth, bytes per second. Sources (read 2026-10-19):
#   H100 SXM  3.35 TB/s  NVIDIA H100 page, specification table
#   H200      4.8 TB/s   NVIDIA H200 page
#   A100 80 GB SXM  2,039 GB/s  NVIDIA A100 data sheet (June 2021 revision), same figure as lesson 19
BANDWIDTHS = {"H100 SXM": 3.35e12, "H200": 4.8e12, "A100 80 GB SXM": 2.039e12}

# Sixteen example weights for the widget: seed, size, and where the optional outlier goes.
SAMPLE_SEED = 18
SAMPLE_GROUP = 4          # the picture shows groups of 4; real files use 128
SAMPLE_OUTLIER_SLOT = 5
SAMPLE_OUTLIER = 1.0


def load_config() -> dict:
    return json.loads((HERE / "llama31_8b_config.json").read_text())


# ---------- the quantizer

def qmax(bits: int) -> int:
    """Largest whole number a signed `bits`-bit code uses: 127 at 8 bits, 7 at 4, 3 at 3.
    The most negative code, -2**(bits-1), goes unused so that the range is symmetric around zero."""
    if not 2 <= bits <= 8:
        raise ValueError("bits must be between 2 and 8")
    return 2 ** (bits - 1) - 1


def _blocks(a: np.ndarray, group: int | None) -> np.ndarray:
    """View `a` as (..., groups, weights per group). group=None makes the whole tensor one group."""
    if group is None:
        return a.reshape(1, -1)
    width = a.shape[-1]
    if width % group:
        raise ValueError(f"group {group} does not divide the row length {width}")
    return a.reshape(*a.shape[:-1], width // group, group)


def quantize(w, bits: int, group: int | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Return whole-number codes (same shape as `w`) and one 16-bit scale per group.

    scale = largest |weight| in the group / qmax, stored as float16. code = round(weight / scale), ties to even,
    kept inside [-qmax, qmax] (the 16-bit scale can round down, which would push the largest weight past qmax).
    group=None uses one scale for the whole tensor; otherwise a group is `group` neighbours along the last axis."""
    w = np.asarray(w, dtype=np.float64)
    top = qmax(bits)
    blocks = _blocks(w, group)
    scales = (np.abs(blocks).max(axis=-1, keepdims=True) / top).astype(np.float16)
    step = scales.astype(np.float64)
    step[step == 0] = 1.0   # a group of zeros has codes of 0 whatever the step is
    codes = np.clip(np.rint(blocks / step), -top, top).astype(np.int32)
    return codes.reshape(w.shape), scales


def dequantize(codes: np.ndarray, scales: np.ndarray, group: int | None = None) -> np.ndarray:
    """The weights the model computes with: code times its group's scale."""
    out = _blocks(codes, group) * scales.astype(np.float64)
    return out.reshape(codes.shape)


def roundtrip(w, bits: int, group: int | None = None) -> np.ndarray:
    codes, scales = quantize(w, bits, group)
    return dequantize(codes, scales, group)


def rms(a) -> float:
    return float(np.sqrt(np.mean(np.asarray(a, dtype=np.float64) ** 2)))


def relative_rms_error(w, w_hat, typical: float | None = None) -> float:
    """Root-mean-square gap between the weights and their rounded versions, as a share of a typical weight's size
    (by default the weights' own root-mean-square size)."""
    w = np.asarray(w, dtype=np.float64)
    return rms(w - w_hat) / (rms(w) if typical is None else typical)


def error(w, bits: int, group: int | None = None) -> float:
    return 0.0 if bits == FULL_BITS else relative_rms_error(w, roundtrip(w, bits, group))


def bits_per_weight(bits: int, group: int | None = GROUP) -> float:
    """Stored bits per weight including the scales: 4 bits and a 16-bit scale per 128 weights is 4.125."""
    if bits == FULL_BITS or group is None:
        return float(bits)   # one scale per whole tensor is a few bytes in all
    return bits + SCALE_BITS / group


# ---------- the data

def synthetic_weights(rows: int = 256, cols: int = 1024, seed: int = 0, outlier: float = 0.0) -> np.ndarray:
    """Bell-curve weights with standard deviation 0.02 (the initializer_range of the Llama config; a stand-in,
    not a trained matrix). outlier > 0 sets the top-left weight to that value."""
    w = np.random.default_rng(seed).normal(0.0, 0.02, size=(rows, cols))
    if outlier:
        w[0, 0] = outlier
    return w


def sample_weights() -> list[float]:
    """Sixteen weights of the widget, three decimals, none larger than 0.3 in size."""
    w = np.random.default_rng(SAMPLE_SEED).normal(0.0, 0.1, size=16)
    return [float(x) for x in np.round(np.clip(w, -0.3, 0.3), 3)]


def sample_with_outlier() -> list[float]:
    w = sample_weights()
    w[SAMPLE_OUTLIER_SLOT] = SAMPLE_OUTLIER
    return w


def sample_error(bits: int, group: int | None, outlier: bool) -> float:
    """The widget's error figure. With the outlier in, one weight of 1.0 would dominate the root-mean-square of
    sixteen, so the gap is always compared with the size of the ordinary sample weights."""
    w = sample_with_outlier() if outlier else sample_weights()
    if bits == FULL_BITS:
        return 0.0
    return relative_rms_error(w, roundtrip(w, bits, group), typical=rms(sample_weights()))


# ---------- the model

def head_dim(cfg: dict) -> int:
    return cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]


def param_count(cfg: dict) -> int:
    """Weights of a Llama-style model from its shape (as in lesson 18): attention, feed-forward and two norms per layer,
    the token table, an output matrix the size of the table unless tied, and the final norm."""
    h, ffn = cfg["hidden_size"], cfg["intermediate_size"]
    kv_width = cfg["num_key_value_heads"] * head_dim(cfg)
    per_layer = 2 * h * h + 2 * h * kv_width + 3 * h * ffn + 2 * h
    return cfg["num_hidden_layers"] * per_layer + embedding_params(cfg) + h


def table_params(cfg: dict) -> int:
    return cfg["vocab_size"] * cfg["hidden_size"]


def embedding_params(cfg: dict) -> int:
    """The token table, plus an output matrix of the same size unless the model reuses the table ("tied")."""
    return table_params(cfg) * (1 if cfg.get("tie_word_embeddings") else 2)


def weight_bytes(cfg: dict, bits: int, group: int | None = GROUP) -> int:
    """Embeddings stay at 16 bits; every other weight uses `bits`, plus one 16-bit scale per `group` of them.
    16 bits means no quantization, so no scales. group=None adds none (a scale per matrix is a few hundred bytes)."""
    if bits not in BITS:
        raise ValueError(f"bits must be one of {BITS}")
    embeddings = embedding_params(cfg)
    others = param_count(cfg) - embeddings
    total = embeddings * FULL_BITS // 8 + others * bits // 8
    if bits != FULL_BITS and group is not None:
        if others % group:
            raise ValueError(f"group {group} does not divide {others}")
        total += others // group * SCALE_BITS // 8
    return total


def bytes_per_token(cfg: dict, bits: int, group: int | None = GROUP) -> int:
    """Bytes read from memory to produce one token for one user.
    Every weight is read once, except that an untied model looks up one row of its token table (hidden x 2 bytes)
    instead of reading the whole table. The attention cache and working memory are not counted."""
    total = weight_bytes(cfg, bits, group)
    if cfg.get("tie_word_embeddings"):
        return total
    return total - table_params(cfg) * FULL_BITS // 8 + cfg["hidden_size"] * FULL_BITS // 8


def tokens_per_second_bound(bytes_read: int, bandwidth: float) -> float:
    """Ceiling for one user: the card cannot read the weights faster than its memory bandwidth."""
    return bandwidth / bytes_read

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import numpy as np

import quantize as qz

CFG = qz.load_config()
OTHERS = qz.param_count(CFG) - qz.embedding_params(CFG)

# Sizes of the .safetensors files of published copies of Llama 3.1 8B Instruct (Hugging Face API, blobs=true, read 2026-10-19).
FILES = (
    ("16-bit", "NousResearch/Meta-Llama-3.1-8B", 16_060_556_376),
    ("8-bit, one scale per output row (w8a16)", "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a16", 9_084_040_952),
    ("4-bit AWQ, group 128", "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4", 5_727_938_576),
    ("4-bit GPTQ, group 128", "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4", 5_735_719_888),
    ("4-bit GPTQ, group 128 (w4a16)", "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w4a16", 5_735_720_552),
)

# Perplexity on WikiText-2, lower is better (AWQ paper, arXiv 2306.00978v6, Table 4; the RTN rows are round-to-nearest, group 128).
PERPLEXITY = (
    ("Llama-2 7B", 5.47, 5.73, 6.66),
    ("Llama-2 13B", 4.88, 4.98, 5.52),
    ("Llama-2 70B", 3.32, 3.46, 3.98),
)

# Model-card results, Llama 3.1 8B Instruct, 16-bit baseline against 4-bit GPTQ (w4a16, group 128): each card reports its own baseline run.
CARD_W4 = (("OpenLLM v1 average", 74.3, 73.5), ("MMLU 5-shot", 68.3, 66.9), ("GSM-8K", 82.8, 82.9),
           ("TruthfulQA", 54.5, 52.8), ("MMLU-Pro", 30.8, 28.8), ("IFEval", 77.9, 76.3))
CARD_W8 = (("OpenLLM v1 average", 74.05, 74.12), ("Hellaswag", 80.47, 80.48), ("Winogrande", 78.06, 77.51), ("TruthfulQA", 54.48, 54.41))


def gib(n):
    return f"{n / qz.GIB:.2f}"


def main():
    print("== hand-worked examples")
    for bits, w in ((4, [0.32, -0.10, 0.04, -0.70]), (3, [0.32, -0.10, 0.04, -0.70])):
        codes, scales = qz.quantize(w, bits, group=4)
        back = qz.dequantize(codes, scales, 4)
        print(f"{bits} bits: weights {w} scale {float(scales[0, 0]):.7f} codes {codes.tolist()} back {back.round(3).tolist()} "
              f"gaps {(np.array(w) - back).round(3).tolist()}")

    print("== synthetic matrix 256 x 1024, bell curve with std 0.02, seed 0: relative root-mean-square error (%)")
    clean = qz.synthetic_weights()
    print(f"weights {clean.size:,}; groups of 128 per row: {clean.size // 128:,}; largest |weight| {abs(clean).max():.4f} "
          f"= {abs(clean).max() / 0.02:.1f} std")
    for label, w in (("no outlier", clean), ("one weight of 1.0", qz.synthetic_weights(outlier=1.0))):
        for bits in (8, 4, 3):
            tensor, group = qz.error(w, bits), qz.error(w, bits, qz.GROUP)
            zero = np.mean(qz.roundtrip(w, bits) == 0)
            print(f"{label}, {bits} bits: one scale {tensor:.1%} (weights rounded to zero {zero:.1%}); "
                  f"one per 128 {group:.1%} (zero {np.mean(qz.roundtrip(w, bits, qz.GROUP) == 0):.1%})")

    print("== outlier size sweep, 4 bits (std 0.02)")
    for size in (0.0, 0.1, 0.2, 0.5, 1.0):
        w = qz.synthetic_weights(outlier=size)
        print(f"outlier {size:.1f} ({size / 0.02:.0f} std): one scale {qz.error(w, 4):.1%}, one per 128 {qz.error(w, 4, qz.GROUP):.1%}")

    print("== bits per weight including scales (group 128)")
    print("  ".join(f"{b}-bit {qz.bits_per_weight(b):.3f}" for b in qz.BITS))

    print("== Llama 3.1 8B")
    n, e = qz.param_count(CFG), qz.embedding_params(CFG)
    print(f"weights {n:,}; embeddings (table + output) {e:,} = {e * 2:,} B = {gib(e * 2)} GiB kept at 16 bits; others {OTHERS:,}")
    base = qz.weight_bytes(CFG, 16)
    for b in qz.BITS:
        none, grp = qz.weight_bytes(CFG, b, None), qz.weight_bytes(CFG, b)
        print(f"{b:>2} bits: no scales {none:,} B = {gib(none)} GiB; with group-128 scales {grp:,} B = {gib(grp)} GiB "
              f"({base / grp:.2f}x smaller); scales alone {grp - none:,} B = {(grp - none) / 1e6:.1f} MB = {(grp - none) / grp:.1%} of the file")

    print("== decoding ceiling for one user: tokens per second = bandwidth / bytes read per token")
    print(f"token table {qz.table_params(CFG) * 2:,} B; one row {CFG['hidden_size'] * 2:,} B")
    for b in qz.BITS:
        read = qz.bytes_per_token(CFG, b)
        row = "  ".join(f"{name} {qz.tokens_per_second_bound(read, bw):,.0f}" for name, bw in qz.BANDWIDTHS.items())
        print(f"{b:>2} bits: reads {read:,} B ({read / 1e9:.2f} GB) per token; {row}")
    h100 = qz.BANDWIDTHS["H100 SXM"]
    ref = qz.tokens_per_second_bound(qz.bytes_per_token(CFG, 16), h100)
    print("gain over 16 bits on H100 SXM: " + "  ".join(
        f"{b}-bit {qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b), h100) / ref:.2f}x" for b in qz.BITS))
    whole = [qz.weight_bytes(CFG, b) for b in qz.BITS]
    print("if the whole file were read (the simpler whole-file rule): "
          + "  ".join(f"{b}-bit {h100 / w:,.0f}" for b, w in zip(qz.BITS, whole)))

    print("== size formula against published files (bytes)")
    zero_points = OTHERS // qz.GROUP // 2
    for label, repo, size in FILES:
        print(f"{label}: file {size:,} ({repo})")
    f16, w8, awq, gptq, w4 = (f[2] for f in FILES)
    print(f"16-bit: calculator {qz.weight_bytes(CFG, 16):,}; gap {f16 - qz.weight_bytes(CFG, 16):,}")
    h, kv, ffn = CFG["hidden_size"], CFG["num_key_value_heads"] * qz.head_dim(CFG), CFG["intermediate_size"]
    rows = CFG["num_hidden_layers"] * (2 * h + 2 * kv + 2 * ffn + h)   # output rows of q, o, k, v, gate, up, down
    print(f"w8a16: no scales {qz.weight_bytes(CFG, 8, None):,}; group-128 scales {qz.weight_bytes(CFG, 8):,}; file {w8:,}; "
          f"above no-scales {w8 - qz.weight_bytes(CFG, 8, None):,}; one 16-bit scale per output row: {rows:,} scales = {rows * 2:,} B")
    print(f"AWQ: formula + zero points {qz.weight_bytes(CFG, 4) + zero_points:,}; file {awq:,}; gap {awq - qz.weight_bytes(CFG, 4) - zero_points:,} "
          f"({(awq - qz.weight_bytes(CFG, 4) - zero_points) / awq:.4%})")
    columns = CFG["num_hidden_layers"] * (6 * CFG["hidden_size"] + CFG["intermediate_size"])
    rest = gptq - qz.weight_bytes(CFG, 4) - zero_points - columns * 4
    print(f"GPTQ: file above group-128 formula {gptq - qz.weight_bytes(CFG, 4):,}; zero points {zero_points:,}; "
          f"column index 4 B x {columns:,} = {columns * 4:,}; unitemised remainder {rest:,} ({rest / gptq:.3%} of the file)")
    print(f"Red Hat w4a16 vs hugging-quants GPTQ: {w4 - gptq:+,} B")

    print("== perplexity on WikiText-2 (AWQ paper), 16-bit | 4-bit RTN g128 | 3-bit RTN g128; rise over 16-bit")
    for name, f, i4, i3 in PERPLEXITY:
        print(f"{name}: {f:.2f} | {i4:.2f} (+{i4 - f:.2f}, +{(i4 - f) / f:.1%}) | {i3:.2f} (+{i3 - f:.2f}, +{(i3 - f) / f:.1%})")
    print("OPT-6.7B 3-bit g128 (Table 1): 16-bit 10.86, RTN 23.54, keep 1% of channels in 16 bits chosen by activations 11.39, by weight size 22.37, at random 24.23")
    print("Llama-2 7B GSM8K, 4-bit g128 (Table 8): 16-bit 13.87, RTN 11.07, GPTQ 12.13, AWQ 13.57")

    print("== model cards (Llama 3.1 8B Instruct; each card has its own 16-bit run)")
    for title, rows in (("4-bit GPTQ, group 128", CARD_W4), ("8-bit, one scale per output row", CARD_W8)):
        print(title)
        for name, a, b in rows:
            print(f"  {name}: {a} -> {b} ({b - a:+.2f} points, recovery {b / a:.1%})")

    print("== widget sample: 16 weights, groups of 4; error as a share of a typical weight (%)")
    print("weights", qz.sample_weights(), "typical", f"{qz.rms(qz.sample_weights()):.4f}")
    for outlier in (False, True):
        for group, label in ((None, "one scale"), (qz.SAMPLE_GROUP, "one per 4")):
            print(f"outlier {outlier}, {label}: " + "  ".join(f"{b} bits {qz.sample_error(b, group, outlier):.1%}" for b in (8, 4, 3)))


if __name__ == "__main__":
    main()

"""Tests for quantize.py: the hand-worked examples of the lesson, the claims it makes, and the sizes it quotes."""
import numpy as np
import pytest

import quantize as qz

CFG = qz.load_config()


def test_hand_worked_four_bit_example():
    """Guards the scale or rounding rule drifting from the table in the lesson: absmax 0.70 / 7 = 0.1, then round."""
    w = [0.32, -0.10, 0.04, -0.70]
    codes, scales = qz.quantize(w, 4, group=4)
    assert codes.tolist() == [3, -1, 0, -7]
    assert float(scales[0, 0]) == pytest.approx(0.1, rel=1e-3)   # stored as a 16-bit float, so 0.0999756
    assert qz.dequantize(codes, scales, group=4).round(2).tolist() == [0.3, -0.1, 0.0, -0.7]


def test_hand_worked_three_bit_example():
    """Guards the 3-bit row of the lesson: only seven codes (-3 to 3), so small weights fall to zero."""
    codes, _ = qz.quantize([0.32, -0.10, 0.04, -0.70], 3, group=4)
    assert codes.tolist() == [1, 0, 0, -3]


def test_ties_round_to_even():
    """Guards round-half-up creeping in (Python's round and numpy's rint agree; JavaScript's Math.round does not)."""
    codes, scales = qz.quantize([3.0, 2.5, 0.5, -0.5, -1.5, -2.5], 3)
    assert float(scales[0, 0]) == 1.0
    assert codes.tolist() == [3, 2, 0, 0, -2, -2]


@pytest.mark.parametrize("bits", [8, 4, 3])
@pytest.mark.parametrize("group", [None, 128])
def test_codes_stay_in_range_and_the_largest_weight_survives(bits, group):
    """Guards a scale cut by qmax = 2**bits (too wide) or a 16-bit scale that rounds down and overflows the range."""
    w = qz.synthetic_weights(8, 256, seed=3, outlier=0.5)
    codes, scales = qz.quantize(w, bits, group)
    assert np.abs(codes).max() == qz.qmax(bits)
    back = qz.dequantize(codes, scales, group)
    assert abs(back).max() == pytest.approx(0.5, rel=2e-3)


def test_a_group_of_zeros_does_not_divide_by_zero():
    """Guards the zero-step guard: 0 / 0 is NaN, which numpy casts to 0 on arm64 but to INT_MIN on x86, so only a raised error shows it."""
    w = np.zeros((2, 8))
    w[1, :4] = [0.1, -0.2, 0.3, 0.05]
    with np.errstate(invalid="raise", divide="raise"):
        codes, scales = qz.quantize(w, 4, group=4)
        back = qz.dequantize(codes, scales, 4)
    assert not np.isnan(back).any()
    assert codes[0].tolist() == [0] * 8


def test_group_must_divide_the_row():
    with pytest.raises(ValueError):
        qz.quantize(np.ones((2, 10)), 4, group=4)


def test_error_grows_as_bits_fall_and_groups_help():
    w = qz.synthetic_weights()
    per_tensor = [qz.error(w, b) for b in (8, 4, 3)]
    per_group = [qz.error(w, b, qz.GROUP) for b in (8, 4, 3)]
    assert per_tensor == sorted(per_tensor) and per_group == sorted(per_group)
    assert all(g < t for g, t in zip(per_group, per_tensor))
    assert qz.error(w, 16) == 0.0


def test_one_outlier_ruins_a_shared_scale_but_not_groups():
    """Guards the lesson's central picture: 1 weight in 262,144 stretches the shared scale until most weights round to zero.
    Also pins the four errors the lesson quotes (19.5, 11.8, 99.5, 11.9), which move if the seed of the made-up matrix changes."""
    clean, spiked = qz.synthetic_weights(), qz.synthetic_weights(outlier=1.0)
    assert [round(100 * qz.error(w, 4, g), 1) for w in (clean, spiked) for g in (None, qz.GROUP)] == [19.5, 11.8, 99.5, 11.9]
    assert qz.error(spiked, 4) > 0.9
    assert qz.error(spiked, 4, qz.GROUP) - qz.error(clean, 4, qz.GROUP) < 0.005
    assert np.count_nonzero(qz.roundtrip(spiked, 4)) < 0.01 * spiked.size


def test_widget_sample_opens_on_the_numbers_the_lesson_quotes():
    """Guards the sample weights or outlier slot drifting from the lesson's 4.7%, 19.1% (big weight) and 34.7% (big weight, one scale)."""
    assert max(map(abs, qz.sample_weights())) == 0.283
    assert [round(100 * qz.sample_error(4, 4, False), 1), round(100 * qz.sample_error(4, 4, True), 1),
            round(100 * qz.sample_error(4, None, True), 1)] == [4.7, 19.1, 34.7]


def test_bits_per_weight_counts_the_scale():
    assert qz.bits_per_weight(4, 128) == 4.125 and qz.bits_per_weight(8, 128) == 8.125 and qz.bits_per_weight(3, 128) == 3.125
    assert qz.bits_per_weight(16, 128) == 16 and qz.bits_per_weight(4, None) == 4


def test_llama_sizes_match_lesson_18_plus_scales():
    """Guards the size formula against lesson 18's weights (8,030,261,248; 16 bits is 14.96 GiB) and the added scales."""
    assert qz.param_count(CFG) == 8_030_261_248
    assert qz.weight_bytes(CFG, 16) == 16_060_522_496
    others = qz.param_count(CFG) - qz.embedding_params(CFG)
    assert qz.weight_bytes(CFG, 4, None) == qz.embedding_params(CFG) * 2 + others // 2
    assert qz.weight_bytes(CFG, 4) - qz.weight_bytes(CFG, 4, None) == others // 128 * 2
    assert [round(qz.weight_bytes(CFG, b) / qz.GIB, 2) for b in qz.BITS] == [14.96, 8.56, 5.31, 4.50]


# Sizes of the `.safetensors` files on Hugging Face (Hub API, read 2026-10-19), Llama 3.1 8B Instruct.
AWQ_INT4_FILE = 5_727_938_576          # AutoAWQ, 4 bits, group 128, scale and 4-bit zero point per group
W8A16_FILE = 9_084_040_952            # llm-compressor, 8 bits, one scale per output row


def test_formula_lands_on_published_files():
    """Guards the size formula against what real files weigh. AWQ stores a 4-bit zero point beside each scale (half a byte)."""
    others = qz.param_count(CFG) - qz.embedding_params(CFG)
    awq = qz.weight_bytes(CFG, 4) + others // 128 // 2
    assert abs(awq - AWQ_INT4_FILE) / AWQ_INT4_FILE < 1e-4
    assert qz.weight_bytes(CFG, 8, None) < W8A16_FILE < qz.weight_bytes(CFG, 8, 128)


def test_decode_reads_one_row_of_the_token_table_not_all_of_it():
    """Guards counting the 1.05 GB token table as read per token when only one 8 KB row is."""
    table = qz.table_params(CFG) * 2
    assert qz.bytes_per_token(CFG, 16) == qz.weight_bytes(CFG, 16) - table + 4096 * 2
    tied = dict(CFG, tie_word_embeddings=True)
    assert qz.bytes_per_token(tied, 16) == qz.weight_bytes(tied, 16)   # a tied table is also the output matrix: read in full


def test_speed_ceiling_rises_as_bits_fall():
    h100 = qz.BANDWIDTHS["H100 SXM"]
    ceilings = [qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b), h100) for b in qz.BITS]
    assert ceilings == sorted(ceilings)
    assert [round(c) for c in ceilings] == [223, 412, 721, 887]

"""Guards for ../widget.html: its data block, its size, and that its JavaScript computes what quantize.py does."""
import json
import re
import subprocess

import numpy as np

import quantize as qz
from embed_widget import WIDGET, embed

HTML = WIDGET.read_text(encoding="utf-8")
CFG = qz.load_config()

# Reads the number code out of the widget between its marker comments, runs it, and compares with the
# values Python wrote into the fixture. Rounding must agree to the last bit; only sums may differ by rounding noise.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const DATA = JSON.parse(html.match(/<script type="application\/json" id="qnt-data">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the numbers'), html.indexOf('/* ---------- the page'));
const api = new Function('DATA', "'use strict';" + js + '; return {rne, f16, quantize, sampleError, weightBytes, bytesPerToken, ceiling};')(DATA);
const bad = [];
const close = (a, b) => Math.abs(a - b) <= 1e-12 * Math.max(1, Math.abs(b));
fx.f16.forEach(([x, want]) => { if (api.f16(x) !== want) bad.push(['f16', x, api.f16(x), want]); });
fx.rne.forEach(([x, want]) => { if (api.rne(x) !== want) bad.push(['rne', x, api.rne(x), want]); });
fx.quant.forEach((c) => {
  const q = api.quantize(c.w, c.bits, c.group);
  if (JSON.stringify(q.codes) !== JSON.stringify(c.codes) || JSON.stringify(q.scales) !== JSON.stringify(c.scales)) bad.push(['quantize', c.bits, c.group, c.w, q.codes, c.codes]);
  if (!q.back.every((v, i) => close(v, c.back[i]))) bad.push(['back', c.bits, c.group]);
});
fx.sample.forEach((c) => { if (!close(api.sampleError(c.bits, c.group, c.outlier), c.error)) bad.push(['sampleError', c]); });
fx.model.forEach((c) => {
  if (api.weightBytes(c.bits, c.group) !== c.weights || api.bytesPerToken(c.bits, c.group) !== c.read) bad.push(['bytes', c]);
  c.gpus.forEach((g) => { if (api.ceiling(c.bits, c.group, g.bandwidth) !== g.tps) bad.push(['tps', c.bits, c.group, g]); });
});
console.log(JSON.stringify({ f16: fx.f16.length, quant: fx.quant.length, sample: fx.sample.length, model: fx.model.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def test_widget_data_block_matches_quantize():
    """Guards an edit to the constants or sample weights that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_every_element_the_script_looks_up_exists():
    """Guards a widget that never starts: a lookup of an id the page does not contain throws in the browser."""
    page = HTML[:HTML.index("<script")]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):]))
    assert lookups and all(f'id="{i}"' in page for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in page)
    assert 'id="qnt-data"' in HTML and 'id="qnt-root"' in page


def test_widget_is_small_and_asks_for_nothing_online():
    assert len(HTML.encode("utf-8")) < 40 * 1024
    for banned in ("http://", "https://", "fetch(", "XMLHttpRequest", "WebSocket", "sendBeacon", "import("):
        assert banned not in HTML, banned
    assert " src=" not in HTML and "@import" not in HTML


def test_widget_numbers_match_the_python_code(tmp_path):
    """Guards the widget's JavaScript drifting from quantize.py: rounding (ties to even), 16-bit scales, errors, sizes and speeds."""
    rng = np.random.default_rng(5)
    f16 = [[float(x), float(np.float16(x))] for x in np.concatenate([10 ** rng.uniform(-9, 3, 400), [0.0, 1.0, 2.0**-14, 2.0**-24, 1.0009765625, 65504.0]])]
    rne = [[x, float(np.rint(x))] for x in [-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, -0.4, 0.51, 7.0] + list(rng.uniform(-130, 130, 200))]
    quant = []
    vectors = [qz.sample_weights(), qz.sample_with_outlier(), [3.0, 2.5, 0.5, -0.5, -1.5, -2.5, 1.0, 0.0]]
    vectors += [list(np.round(rng.normal(0, s, 16), 4)) for s in (0.001, 0.1, 5.0) for _ in range(20)]
    for w in vectors:
        for bits in (8, 4, 3):
            for group in (None, 4, 8):
                codes, scales = qz.quantize(w, bits, group)
                quant.append({"w": [float(x) for x in w], "bits": bits, "group": group, "codes": codes.ravel().tolist(),
                              "scales": scales.astype(np.float64).ravel().tolist(), "back": qz.dequantize(codes, scales, group).ravel().tolist()})
    sample = [{"bits": b, "group": g, "outlier": o, "error": qz.sample_error(b, g, o)}
              for b in qz.BITS for g in (None, qz.SAMPLE_GROUP) for o in (False, True)]
    model = [{"bits": b, "group": g, "weights": qz.weight_bytes(CFG, b, g), "read": qz.bytes_per_token(CFG, b, g),
              "gpus": [{"bandwidth": bw, "tps": qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b, g), bw)} for bw in qz.BANDWIDTHS.values()]}
             for b in qz.BITS for g in (None, qz.GROUP)]
    (tmp_path / "fixture.json").write_text(json.dumps({"f16": f16, "rne": rne, "quant": quant, "sample": sample, "model": model}), encoding="utf-8")
    (tmp_path / "parity.mjs").write_text(PARITY, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "parity.mjs"), str(tmp_path / "fixture.json"), str(WIDGET)],
                         capture_output=True, text=True, timeout=120)
    assert run.returncode == 0, run.stdout + run.stderr
    assert json.loads(run.stdout)["quant"] == len(vectors) * 9 == 567
    assert any(c["codes"].count(0) > 10 for c in quant) and any(max(map(abs, c["codes"])) == 127 for c in quant)


# Runs the widget's page code under Node against a minimal stand-in for the document, reads the three readouts, then clicks
# the outlier checkbox and picks "One for all" as the lesson tells the reader to.
RENDER = r"""
import fs from 'node:fs';
const html = fs.readFileSync(process.argv[2], 'utf8');
const json = html.match(/<script type="application\/json" id="qnt-data">([\s\S]*?)<\/script>/)[1];
const code = html.slice(html.indexOf('<script>', html.indexOf('</script>')) + 8, html.lastIndexOf('</script>'));
const els = {}, handlers = {};
const el = (id) => els[id] ||= { textContent: id === 'qnt-data' ? json : '', innerHTML: '', style: {}, checked: false,
  addEventListener(type, fn) { (handlers[id] ||= {})[type] = fn; }, getBoundingClientRect: () => ({ height: 0 }), querySelectorAll: () => [] };
const win = { addEventListener() {} };
new Function('document', 'window', 'parent', code)({ getElementById: el }, win, win);
const read = () => ({ size: els['qnt-size'].textContent, speed: els['qnt-speed'].textContent, error: els['qnt-err'].textContent });
const out = { open: read() };
handlers['qnt-big'].change({ target: { checked: true } });
out.outlier = read();
handlers['qnt-mode'].change({ target: { value: '0' } });
out.oneForAll = read();
console.log(JSON.stringify(out));
"""


def test_widget_readouts_follow_the_lesson(tmp_path):
    """Guards the page code the parity test cannot reach: size and speed always use groups of 128 (picking "One for all" changed
    them before), GiB is 2**30 bytes, and the opening state, the big-weight box and the scale choice reach the readouts."""
    (tmp_path / "render.mjs").write_text(RENDER, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "render.mjs"), str(WIDGET)], capture_output=True, text=True, timeout=60)
    assert run.returncode == 0, run.stdout + run.stderr
    out = json.loads(run.stdout)
    assert out["open"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "4.7%"}
    assert out["outlier"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "19.1%"}
    assert out["oneForAll"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "34.7%"}