14.11 Evaluation & Serving: Honest Benchmarks, int4 Quantization, and Running on a Laptop¶

"""
stacklm/infer.py — the string-in/string-out layer over Ch. 14.4's generate().
Every probe in Section 4, the agent loop in Ch. 14.10, and the CLI in Section 9
call exactly these two functions.
"""
import torch

from stacklm.tokenizer import SPECIAL_TOKENS               # Ch. 14.3
ALL_SPECIALS = frozenset(SPECIAL_TOKENS)


@torch.no_grad()
def generate_text(model, tokenizer, prompt: str, max_new_tokens: int = 64,
                  temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0,
                  stop_id: int | None = None, allowed_special=ALL_SPECIALS,
                  device: str = "cpu", return_n_tokens: bool = False):
    """Encode -> Ch. 14.4's cached generate() -> decode ONLY the new tokens.

    `allowed_special` is not decoration. Ch. 14.3's `StackTokenizer.encode`
    emits the reserved special ids ONLY for strings you explicitly allow;
    forget it and `<|system|>` is tokenized as literal angle-bracket bytes, the
    model never sees the chat frame it was SFT'd on, and your probe measures a
    distribution shift you introduced yourself. Default it ON here so no caller
    can forget. (Ch. 14.3's default is the opposite, deliberately: untrusted
    user text must never be able to forge a role boundary.)

    `return_n_tokens` returns the EXACT count of generated ids. Re-encoding the
    decoded string is not guaranteed to reproduce it, so tok/s computed that way
    is an estimate; the CLI in Section 9 uses this flag instead.
    """
    model.eval()
    ids = torch.tensor([tokenizer.encode(prompt, allowed_special=allowed_special)],
                       dtype=torch.long, device=device)
    assert ids.shape[1] + max_new_tokens <= model.cfg.max_seq_len, (
        "prompt + generation exceeds the context; call model.rebuild_rope() or "
        "truncate — Ch. 14.4's generate() asserts this too, loudly.")
    eos = tokenizer.eos_id if stop_id is None else stop_id
    out = model.generate(ids, max_new_tokens=max_new_tokens, temperature=temperature,
                         top_p=top_p, top_k=top_k, eos_id=eos)
    new = out[0, ids.shape[1]:].tolist()                        # new tokens only
    if eos in new:                                              # 14.4's generate emits the
        new = new[:new.index(eos)]                              # stop token; drop it
    text = tokenizer.decode(new)
    return (text, len(new)) if return_n_tokens else text


def generate_fn(model, tokenizer, prompt, max_new_tokens=64, temperature=0.0, **kw):
    """The exact signature the Section 4 probes expect. Greedy by default:
    a probe score you cannot reproduce bit-for-bit is not a measurement."""
    return generate_text(model, tokenizer, prompt, max_new_tokens=max_new_tokens,
                         temperature=temperature, **kw)

"""
stacklm/eval/ppl.py — held-out perplexity.

Consumes exactly the batch dict Ch. 14.2's `PackedMemmapDataset` yields:
{"input_ids", "position_ids", "targets"} over uint16 memmap shards packed to
seq_len 2048 with document-aware position resets.
"""
import math
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

from stacklm.model.transformer import Stack100M       # Ch. 14.4
from stacklm.data.dataset import PackedMemmapDataset  # Ch. 14.2


def _autocast(device: str):
    """bf16 autocast on CUDA; a no-op context on CPU (where the rest of this
    chapter lives). Hardcoding device_type='cuda' is the classic way to make an
    eval function that only runs on the machine you wrote it on."""
    dev_type = torch.device(device).type
    if dev_type == "cuda":
        return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    return torch.autocast(device_type=dev_type, enabled=False)


@torch.no_grad()
def compute_perplexity(model: Stack100M, val_shard_dir: str, pad_id: int,
                       batch_size: int = 4, max_batches: int | None = None,
                       device: str = "cpu") -> dict:
    """Held-out cross-entropy (nats/token) and perplexity.

    Three details decide whether the number is honest (or fits in RAM):
      1. PAD MASKING. Ch. 14.2 pads the trailing window with `pad_id`, and the
         training loop masks the loss there. If we do not, those positions are
         trivially predictable and drag the mean loss DOWN — a silently
         optimistic perplexity. We remap pad targets to -100 and let
         `ignore_index` do the rest, then count only the surviving tokens.
      2. BATCH SIZE. The all-positions logits tensor is B x T x V x 4 bytes;
         at B=4, T=2047, V=32768 that is 1.07 GB. On CPU the logits, not the
         weights, are the memory ceiling — keep this small.
      3. DOCUMENT MASKING. Ch. 14.2 stores `position_ids` that reset to 0 at
         every document start; `seq_ids = cumsum(position_ids == 0) - 1` turns
         that into the segment id `Stack100M.forward` uses to forbid
         cross-document attention. `position_ids` also drives RoPE, so a
         document starting mid-window is rotated from its OWN position 0,
         exactly as in training. Evaluate without them and you measure the
         model in a geometry it was never trained in.

    Note what is NOT passed: `logits_to_keep`. Its default of 0 is what makes
    `logits` a (B, T, V) tensor instead of the (B, 1, V) a sampling call
    returns — Section 2's pitfall, where it does the most damage.
    """
    model.to(device).eval()
    ds = PackedMemmapDataset(val_shard_dir)          # never seen in ANY stage
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False, drop_last=False)

    total_nll, total_tokens = 0.0, 0
    for i, batch in enumerate(loader):
        if max_batches is not None and i >= max_batches:
            break
        input_ids = batch["input_ids"].to(device)             # (B, T)
        targets = batch["targets"].to(device).clone()         # (B, T), shifted by 1
        position_ids = batch["position_ids"].to(device)
        seq_ids = (position_ids == 0).cumsum(dim=-1) - 1      # (B, T) segment ids

        targets = targets.masked_fill(targets == pad_id, -100)     # rule 1
        # Also drop the boundary target: the last token of doc k "predicts" the
        # BOS of doc k+1, which is a packing artifact, not language modeling.
        boundary = torch.zeros_like(targets, dtype=torch.bool)
        boundary[:, :-1] = position_ids[:, 1:] == 0
        targets = targets.masked_fill(boundary, -100)

        with _autocast(device):
            logits, _ = model(input_ids, position_ids=position_ids,
                              seq_ids=seq_ids)              # (B, T, V)
        assert logits.shape[1] == input_ids.shape[1], "logits_to_keep leaked in"
        # sum (not mean) so we can average correctly over the whole eval set
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)).float(),
                               targets.reshape(-1), reduction="sum", ignore_index=-100)
        total_nll += loss.item()
        total_tokens += int((targets != -100).sum())

    mean_nll = total_nll / max(1, total_tokens)
    return {"loss_nats_per_token": mean_nll,
            "perplexity": math.exp(mean_nll),
            "n_tokens_evaluated": total_tokens}

# stacklm/eval/probes.py
import re
import torch

from stacklm.post.chat import SPECIAL            # Ch. 14.9's role-token strings


def chat_prompt(user: str, system: str | None = None) -> str:
    """Byte-identical to `render_conversation(..., add_generation_prompt=True)`:
    <|bos|>[<|system|>SYS<|end|>]<|user|>USER<|end|><|assistant|>

    No newlines between turns — Ch. 14.9's template does not emit any, and a
    stray "\\n" is a token the model never saw in that slot. Ch. 14.9's RLVR
    rollouts render the arithmetic task with NO system turn, so §4.1 passes
    `system=None` to match the RL distribution exactly.
    """
    s = SPECIAL["bos"]
    if system is not None:
        s += f"{SPECIAL['system']}{system}{SPECIAL['end']}"
    return s + f"{SPECIAL['user']}{user}{SPECIAL['end']}{SPECIAL['assistant']}"

import random
from stacklm.post.grpo import make_arithmetic_prompt, exact_match_reward  # Ch. 14.9


def make_arithmetic_probe(n: int = 200, seed: int = 0, max_val: int = 99) -> list[dict]:
    """Freeze n problems from Ch. 14.9's RLVR distribution. Seeded, so the probe
    set is reproducible from the seed alone — a hashed, versioned eval set in
    the cheapest possible form."""
    rng = random.Random(seed)
    return [dict(zip(("question", "answer"), make_arithmetic_prompt(rng, max_val=max_val)))
            for _ in range(n)]


@torch.no_grad()
def eval_arithmetic(model, tokenizer, generate_fn, problems: list[dict],
                    max_new_tokens: int = 64, label: str = "in-distribution") -> dict:
    """Greedy-generate, grade with Ch. 14.9's OWN verifier.

    Two disciplines Ch. 14.9 insists on and an eval must not drop:
      * `max_new_tokens=64` matches the RLVR rollout budget. Truncate to 8 and
        you cut off the '####' line and measure your own budget, not the model.
      * PARSE-FAILURE RATE IS REPORTED SEPARATELY from accuracy. "Wrong answer"
        and "unparseable output" demand opposite fixes; conflating them is how
        teams spend a week tuning RL to fix a formatting bug.
    """
    n_correct = n_parsed = 0
    for p in problems:
        completion = generate_fn(model, tokenizer, prompt=chat_prompt(p["question"]),
                                 max_new_tokens=max_new_tokens, temperature=0.0)
        reward, pred = exact_match_reward(completion, p["answer"])
        n_parsed += int(pred is not None)
        n_correct += int(reward == 1.0)
    n = len(problems)
    return {"label": label, "accuracy": n_correct / n,
            "parse_rate": n_parsed / n, "n": n}


# Two rows, always reported together. The second is the generalization probe
# Ch. 14.9 asks for: same format, operands an order of magnitude larger than
# anything RLVR saw. Expect a large gap — RLVR at 100M generalizes NARROWLY,
# and hiding that behind a single averaged number is the dishonest option.
ARITH_PROBES = [
    (make_arithmetic_probe(n=200, seed=0, max_val=99),  "2-digit (in-distribution)"),
    (make_arithmetic_probe(n=100, seed=1, max_val=999), "3-digit (out-of-distribution)"),
]

import torch.nn.functional as F


@torch.no_grad()
def sequence_logprob(model, tokenizer, prompt: str, continuation: str,
                     device="cpu") -> tuple[float, int]:
    """Return (sum of log p(token | prefix) over `continuation`, n_cont_tokens).

    The subtle part is the boundary. BPE is not guaranteed to satisfy
        encode(prompt + continuation)[:len(encode(prompt))] == encode(prompt)
    — a merge can straddle the join and re-tokenize the last prompt token. When
    that happens, slicing at len(prompt_ids) silently scores the WRONG tokens
    and the whole probe becomes noise. `lm-evaluation-harness` has the identical
    construction (`whole_enc[len(context_enc):]`) and the identical exposure; we
    make the assumption explicit and loud rather than implicit and silent.

    The practical mitigation is baked into TINY_MC_SET below: put the separating
    space at the START of every continuation (" Paris", not "Paris" after a
    trailing space), so the join lands on a whitespace boundary that byte-level
    BPE pre-tokenization does not merge across.
    """
    prompt_ids = tokenizer.encode(prompt)
    full_ids = tokenizer.encode(prompt + continuation)
    n_ctx = len(prompt_ids)
    if full_ids[:n_ctx] != prompt_ids:
        raise ValueError(
            "BPE boundary violation: encode(prompt) is not a prefix of "
            "encode(prompt + continuation). Move the separator into the "
            "continuation, or score with a byte-aligned re-tokenization.")
    if len(full_ids) == n_ctx:
        raise ValueError("empty continuation after tokenization")

    ids = torch.tensor([full_ids], dtype=torch.long, device=device)
    logits, _ = model(ids)                                   # (1, T, V): logits_to_keep=0
    logp = F.log_softmax(logits[0].float(), dim=-1)          # (T, V)

    # logits[i] predicts token i+1. Continuation tokens live at n_ctx .. T-1,
    # so the predicting rows are n_ctx-1 .. T-2.
    targets = ids[0, n_ctx:]                                 # (n_cont,)
    pred_rows = logp[n_ctx - 1: -1, :]                       # (n_cont, V)
    tok_logp = pred_rows.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    return float(tok_logp.sum()), int(tok_logp.numel())


# A hand-written, deliberately tiny probe (2 items shown; the set used in
# practice is ~50-100 items spanning basic science, geography, and word sense —
# small enough to eyeball every item for leakage).
TINY_MC_SET = [
    {"question": "The chemical symbol for water is",
     "choices": [" H2O", " CO2", " NaCl", " O2"], "answer_idx": 0},
    {"question": "The capital of France is",
     "choices": [" Paris", " Berlin", " Madrid", " Rome"], "answer_idx": 0},
    # ... remaining items omitted; same {question, choices, answer_idx} shape.
]


def eval_mc_probe(model, tokenizer, mc_set: list[dict] = TINY_MC_SET,
                  device="cpu") -> dict:
    """Reports BOTH metrics `lm-evaluation-harness` reports:
      acc      — raw summed log-prob (the harness's `acc`)
      acc_norm — summed log-prob divided by the continuation's UTF-8 BYTE length
                 (the harness's `acc_norm`), removing the length bias that
                 otherwise favours whichever option tokenizes shortest.
    If they disagree, your choices are length-imbalanced and the raw number is
    measuring string length as much as knowledge."""
    n_raw = n_norm = 0
    for item in mc_set:
        raw, norm = [], []
        for choice in item["choices"]:
            s, _ = sequence_logprob(model, tokenizer, item["question"], choice, device)
            raw.append(s)
            norm.append(s / max(1, len(choice.encode("utf-8"))))
        n_raw += int(max(range(len(raw)), key=lambda i: raw[i]) == item["answer_idx"])
        n_norm += int(max(range(len(norm)), key=lambda i: norm[i]) == item["answer_idx"])
    n = len(mc_set)
    return {"acc": n_raw / n, "acc_norm": n_norm / n, "n": n}

def normalize_answer(s: str) -> str:
    """Standard SQuAD/NQ-style EM normalization: lowercase, strip punctuation,
    articles, and extra whitespace."""
    s = s.lower().strip()
    s = re.sub(r"\b(a|an|the)\b", " ", s)
    s = re.sub(r"[^\w\s]", "", s)
    return re.sub(r"\s+", " ", s).strip()


@torch.no_grad()
def eval_retrieval_qa(model, tokenizer, generate_fn, retriever,
                      qa_pairs: list[dict]) -> dict:
    """qa_pairs: [{"question": str, "gold_answer": str}, ...]
    retriever: BM25Retriever or HashEmbedRetriever from Ch. 14.10.

    Ch. 14.10's `.search(query, k)` returns list[tuple[Passage, float]] — a
    (passage, score) PAIR, and `Passage` is a dataclass with `.doc_id`/`.text`.
    Interpolating the tuple straight into the prompt (a genuinely easy mistake)
    feeds the model a Python repr and scores ~0 for reasons that have nothing to
    do with the model.

    We also report RETRIEVER RECALL@1 separately. With k=1 this is a joint test
    of two components; without the recall number a bad EM is unattributable —
    you cannot tell whether the model failed to read or the retriever failed to
    fetch."""
    n_correct, n_retrieved = 0, 0
    for pair in qa_pairs:
        hits = retriever.search(pair["question"], k=1)
        passage_text = hits[0][0].text if hits else ""        # (Passage, score) -> .text
        n_retrieved += int(normalize_answer(pair["gold_answer"])
                           in normalize_answer(passage_text))
        prompt = chat_prompt(f"Passage: {passage_text}\nQuestion: {pair['question']}",
                             system="Answer using only the passage below.")
        completion = generate_fn(model, tokenizer, prompt=prompt, max_new_tokens=16,
                                 temperature=0.0)   # specials enabled by default (§2)
        n_correct += int(normalize_answer(completion) == normalize_answer(pair["gold_answer"]))
    n = len(qa_pairs)
    return {"exact_match": n_correct / n, "retriever_recall@1": n_retrieved / n, "n": n}

"""
stacklm/eval/agent.py — evaluate the Ch. 14.10 ReAct loop end to end.
Runs the REAL loop (model drives, tools execute) and decomposes the outcome.
"""
from dataclasses import dataclass
from collections import Counter
import statistics

from stacklm.agent.react import parse_assistant_step
from stacklm.agent.loop import run_agent          # returns (answer, trace)
from stacklm.agent.distill import ToolEnv, normalize


@dataclass
class AgentTask:
    question: str
    gold: str                    # verifiable final answer (exact match)
    gold_tool: str               # "search" or "calc" — the tool a competent agent picks first
    gold_evidence: str = ""      # substring that MUST appear in an observation if
                                 # retrieval worked; enables the "read it and
                                 # ignored it" diagnosis


def eval_agent(model, tok, env: ToolEnv, tasks: list[AgentTask],
               max_steps: int = 6) -> dict:
    n_calls = n_parseable = n_tool_right = n_solved = n_capped = 0
    turns: list[int] = []
    taxonomy: Counter[str] = Counter()

    for task in tasks:
        answer, trace = run_agent(model, tok, task.question, env, max_steps=max_steps)
        assistant = [t for kind, t in trace if kind == "assistant"]
        observations = [t for kind, t in trace if kind == "observation"]
        acts = [parse_assistant_step(t) for t in assistant]
        calls = [a for a in acts if a.kind == "tool"]

        n_calls += len(calls)
        n_parseable += sum(1 for a in calls if a.tool != "__malformed__")
        first_tool = calls[0].tool if calls else None
        n_tool_right += int(first_tool == task.gold_tool)

        # CAPPING IS A TRACE-SHAPE PROPERTY, NOT A LAST-ACTION PROPERTY.
        # Ch. 14.10's run_agent returns the instant an emission parses as
        # "final", so every terminating path leaves AT MOST max_steps assistant
        # emissions. The only way to get max_steps + 1 is the forced-synthesis
        # fallback appended after the step budget is exhausted — and that
        # fallback ("Thought: I must answer now.\nAnswer: ...") itself parses as
        # kind == "final". Testing `acts[-1].kind != "final"` therefore never
        # fires: it is dead code, and cap_rate would be a structural zero.
        capped = len(assistant) > max_steps
        n_capped += int(capped)
        turns.append(max_steps if capped else len(assistant))   # don't count the fallback

        solved = normalize(answer) == normalize(task.gold)
        n_solved += int(solved)

        # Failure taxonomy, in priority order — each task lands in exactly one bucket.
        evidence_seen = bool(task.gold_evidence) and any(
            task.gold_evidence.lower() in o.lower() for o in observations)
        if solved:
            taxonomy["ok"] += 1
        elif any(a.tool == "__malformed__" for a in calls):
            taxonomy["malformed_call"] += 1
        elif any(o.startswith("RepeatedCall:") for o in observations):
            taxonomy["looped"] += 1
        elif capped:
            taxonomy["non_termination"] += 1
        elif not calls:
            taxonomy["never_called_a_tool"] += 1
        elif first_tool != task.gold_tool:
            taxonomy["wrong_tool"] += 1
        elif evidence_seen:
            taxonomy["retrieved_but_ignored"] += 1
        else:
            taxonomy["other"] += 1

    n = len(tasks)
    return {
        "n_tasks": n,
        "tool_format_validity": (n_parseable / n_calls) if n_calls else float("nan"),
        "tool_choice_acc": n_tool_right / n,
        "turns_mean": statistics.mean(turns),
        "turns_median": statistics.median(turns),
        "cap_rate": n_capped / n,
        "exact_match": n_solved / n,
        "taxonomy": dict(taxonomy),
    }

import torch


def gptq_quantize_column_by_column(W: torch.Tensor, X_calib: torch.Tensor, bits: int = 4,
                                   damp: float = 1e-2) -> torch.Tensor:
    """Simplified reference implementation of GPTQ's per-layer reconstruction.
    W: (d_out, d_in) weight matrix. X_calib: (n_samples, d_in) calibration
    activations captured from a real forward pass on held-out text.

    Pedagogical: real GPTQ batches columns and uses a running Cholesky
    factorization for speed. For production use GPTQModel or llm-compressor
    (Section 9) — do not ship this."""
    d_out, d_in = W.shape
    W = W.clone().float()

    # Hessian of the layer's local least-squares reconstruction objective.
    H = 2 * (X_calib.T @ X_calib) / X_calib.shape[0]
    H += damp * torch.eye(d_in, device=W.device)          # damping for stability
    H_inv = torch.linalg.inv(H)

    qmax = 2 ** (bits - 1) - 1
    for q in range(d_in):
        col = W[:, q]
        scale = col.abs().max() / qmax if col.abs().max() > 0 else 1.0
        q_col = torch.clamp(torch.round(col / scale), -qmax, qmax)
        w_hat = q_col * scale
        error = (col - w_hat) / H_inv[q, q]                # per-output-row error
        W[:, q] = w_hat
        if q + 1 < d_in:
            # Spread the reconstruction error onto not-yet-quantized columns.
            W[:, q + 1:] -= torch.outer(error, H_inv[q, q + 1:])
    return W

def _fake_quant_grouped(W: torch.Tensor, bits: int = 4, group_size: int = 64):
    """Symmetric per-group RTN, dequantized straight back to fp32 ('fake quant').
    Same grid the reference path in Section 7 ships, so the search optimizes the
    quantizer we actually use rather than a different, coarser one."""
    d_out, d_in = W.shape
    g = W.view(d_out, d_in // group_size, group_size)
    qmax = 2 ** (bits - 1) - 1
    s = (g.abs().amax(dim=2, keepdim=True) / qmax).clamp(min=1e-8)
    return (torch.round(g / s).clamp(-qmax, qmax) * s).view(d_out, d_in)


def awq_search_channel_scales(W: torch.Tensor, X_calib: torch.Tensor, bits: int = 4,
                              group_size: int = 64, n_grid: int = 20) -> torch.Tensor:
    """AWQ's salient-channel search. W: (d_out, d_in). X_calib: (n_samples, d_in).

    THE OBJECTIVE MUST BE OUTPUT ERROR MEASURED AFTER DE-SCALING. A tempting
    shortcut — the squared error of the SCALED matrix against its quantization —
    is monotonically increasing in alpha (larger salient channels => larger
    group max => larger step size), so its argmin is always alpha = 0 and the
    whole search silently degenerates into a no-op. Measure what you care about:
    the layer's output.

    Because alpha = 0 (s = 1, plain RTN) is in the grid, the returned scale can
    never be worse than RTN on the calibration set. That is the property that
    makes a grid search over one scalar a legitimate algorithm rather than a
    hopeful heuristic.
    """
    act = X_calib.abs().mean(dim=0).clamp(min=1e-5)   # per-input-channel saliency
    act = act / act.mean()                            # >1 on salient channels
    ref = X_calib @ W.T                               # fp32 reference output

    best_s, best_err = torch.ones_like(act), float("inf")
    for i in range(n_grid + 1):
        alpha = i / n_grid
        s = act.pow(alpha)
        W_hat = _fake_quant_grouped(W * s.unsqueeze(0), bits, group_size) / s.unsqueeze(0)
        err = (X_calib @ W_hat.T - ref).pow(2).mean().item()
        if err < best_err:
            best_err, best_s = err, s
    return best_s     # fold s into W; fold 1/s into the PRECEDING norm/layer

"""
stacklm/serve/quantize.py — round-to-nearest int8/int4 weight-only PTQ, with a
packed on-disk format for CPU inference (Ch. 14.11 reference path).
"""
import json
import torch
import torch.nn as nn


def quantize_int8_per_row(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Symmetric per-output-channel int8. weight: (d_out, d_in).
    Returns (int8 codes, fp32 scales of shape (d_out,))."""
    w = weight.float()
    scale = w.abs().amax(dim=1).clamp(min=1e-8) / 127.0          # (d_out,)
    q = torch.round(w / scale.unsqueeze(1)).clamp(-127, 127)
    return q.to(torch.int8), scale


def dequantize_int8_per_row(q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return q.float() * scale.unsqueeze(1)


def quantize_int4_grouped(weight: torch.Tensor, group_size: int = 64):
    """Asymmetric per-group int4, packed two values per uint8 byte.
    weight: (d_out, d_in), d_in divisible by group_size.
    Returns (packed uint8 (d_out, d_in//2), scale, zero_point),
    scale/zero_point of shape (d_out, d_in // group_size)."""
    d_out, d_in = weight.shape
    assert d_in % group_size == 0, f"d_in={d_in} not divisible by group_size={group_size}"
    n_groups = d_in // group_size

    w = weight.float().view(d_out, n_groups, group_size)
    w_min, w_max = w.amin(dim=2), w.amax(dim=2)                  # (d_out, n_groups)
    scale = ((w_max - w_min) / 15.0).clamp(min=1e-8)             # 4-bit unsigned: 0..15
    zero_point = torch.round(-w_min / scale)                     # integer zero-point

    q = torch.round(w / scale.unsqueeze(2) + zero_point.unsqueeze(2)).clamp(0, 15)
    q = q.view(d_out, d_in).to(torch.uint8)

    # Pack two nibbles per byte: low nibble = even column, high nibble = odd column.
    q_even, q_odd = q[:, 0::2], q[:, 1::2]
    packed = (q_even | (q_odd << 4)).to(torch.uint8)             # (d_out, d_in // 2)
    return packed, scale, zero_point


def dequantize_int4_grouped(packed, scale, zero_point, d_in: int, group_size: int = 64):
    d_out = packed.shape[0]
    n_groups = d_in // group_size
    q_even = (packed & 0x0F).to(torch.float32)
    q_odd = ((packed >> 4) & 0x0F).to(torch.float32)
    q = torch.empty(d_out, d_in, device=packed.device)
    q[:, 0::2], q[:, 1::2] = q_even, q_odd
    q = q.view(d_out, n_groups, group_size)
    w = (q - zero_point.unsqueeze(2)) * scale.unsqueeze(2)
    return w.view(d_out, d_in)

class QuantizedLinear(nn.Module):
    """Memory-efficient stand-in for nn.Linear. Stores weights packed at int8 or
    int4; on each forward, dequantizes to fp32 and does a standard matmul.

    This is the *pedagogically* correct approach — it is exactly what the weight
    compression buys you — but it does NOT give a matmul speedup, since PyTorch
    still computes in fp32 and the unpack is extra work. Section 9.1 measures
    that honestly. Real speedups need integer-native kernels."""

    def __init__(self, bits: int, group_size: int = 64):
        super().__init__()
        assert bits in (4, 8)
        self.bits, self.group_size = bits, group_size
        self.d_in = self.d_out = None

    # ---- construction ----------------------------------------------------
    @classmethod
    def from_float(cls, linear: nn.Linear, bits: int, group_size: int = 64):
        layer = cls(bits, group_size)
        layer.d_in, layer.d_out = linear.in_features, linear.out_features
        # register (not plain-assign) so a bias round-trips through state_dict;
        # Stack-100M is LLaMA-style bias-free, so this is normally None.
        layer.register_buffer(
            "bias", linear.bias.detach().clone() if linear.bias is not None else None)
        if bits == 8:
            q, scale = quantize_int8_per_row(linear.weight.data)
            layer.register_buffer("q_weight", q)
            layer.register_buffer("scale", scale)
        else:
            packed, scale, zp = quantize_int4_grouped(linear.weight.data, group_size)
            layer.register_buffer("q_weight", packed)
            layer.register_buffer("scale", scale)
            layer.register_buffer("zero_point", zp)
        return layer

    @classmethod
    def empty(cls, in_features: int, out_features: int, bits: int,
              group_size: int = 64, device="meta"):
        """Allocate correctly-shaped buffers WITHOUT ever materializing the fp32
        weight. This is what you want at LOAD time: building a real fp32
        Stack100M just to quantize and throw it away costs 405 MB of peak RSS on
        a machine whose whole point is that it only has to hold 63 MB."""
        layer = cls(bits, group_size)
        layer.d_in, layer.d_out = in_features, out_features
        layer.register_buffer("bias", None)
        if bits == 8:
            layer.register_buffer("q_weight", torch.empty(out_features, in_features,
                                                          dtype=torch.int8, device=device))
            layer.register_buffer("scale", torch.empty(out_features, dtype=torch.float32,
                                                       device=device))
        else:
            n_g = in_features // group_size
            layer.register_buffer("q_weight", torch.empty(out_features, in_features // 2,
                                                          dtype=torch.uint8, device=device))
            layer.register_buffer("scale", torch.empty(out_features, n_g,
                                                       dtype=torch.float32, device=device))
            layer.register_buffer("zero_point", torch.empty(out_features, n_g,
                                                            dtype=torch.float32, device=device))
        return layer

    # ---- dequantization --------------------------------------------------
    def dequantize(self) -> torch.Tensor:
        if self.bits == 8:
            return dequantize_int8_per_row(self.q_weight, self.scale)
        return dequantize_int4_grouped(self.q_weight, self.scale, self.zero_point,
                                       self.d_in, self.group_size)

    def dequantize_rows(self, idx: torch.Tensor) -> torch.Tensor:
        """Dequantize ONLY the requested output rows — the embedding lookup path.
        Gathering 2048 rows of a 32768-row table costs 6% of a full dequantize."""
        if self.bits == 8:
            return self.q_weight[idx].float() * self.scale[idx].unsqueeze(1)
        return dequantize_int4_grouped(self.q_weight[idx], self.scale[idx],
                                       self.zero_point[idx], self.d_in, self.group_size)

    @property
    def weight(self) -> torch.Tensor:
        """Compatibility shim for Ch. 14.4, which reaches for `.weight` in two
        places: `fused_ce_z_loss(x, self.lm_head.weight, ...)` when
        `cfg.loss_chunk > 0`, and `estimate_params()`'s `tok_emb.weight.numel()`.
        Without this property both raise AttributeError the moment the model is
        quantized — a sibling chapter's documented API breaking silently.

        It MATERIALIZES the full fp32 tensor, so it is a correctness shim, not a
        fast path. `quantize_stacklm` forces `cfg.loss_chunk = 0` for exactly
        that reason."""
        return self.dequantize()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.dequantize()
        return torch.nn.functional.linear(x, w.to(x.dtype), self.bias)


class QuantizedEmbedding(nn.Module):
    """Token lookup that reads the SAME packed buffers as the quantized, tied
    `lm_head`, dequantizing only the gathered rows.

    Why this class must exist: Ch. 14.4 ties the weights by aliasing
    (`lm_head.weight = tok_emb.weight`). Replacing `lm_head` with a
    QuantizedLinear BREAKS the alias — `tok_emb` keeps its own fp32 Parameter,
    it stays in state_dict(), and the "quantized" export silently carries a
    16.78M x 4 B = 67 MB fp32 table alongside the packed weights, roughly
    doubling every headline in this chapter."""

    def __init__(self, head: QuantizedLinear):
        super().__init__()
        # Deliberately NOT a submodule: nn.Module.__setattr__ would register it,
        # duplicating every key as `tok_emb._head.*` in state_dict(). We want the
        # packed buffers to appear exactly once, under `lm_head.*`.
        object.__setattr__(self, "_head", head)

    @property
    def weight(self) -> torch.Tensor:
        return self._head.dequantize()          # same shim, same reason

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        rows = self._head.dequantize_rows(idx.reshape(-1))     # (n, d_model)
        return rows.view(*idx.shape, -1)

def _swap_linears(module: nn.Module, bits: int, group_size: int) -> None:
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Linear):
            assert bits == 8 or child.in_features % group_size == 0, (
                f"{name}: in_features={child.in_features} does not tile at "
                f"group_size={group_size}")
            setattr(module, name, QuantizedLinear.from_float(child, bits, group_size))
        else:
            _swap_linears(child, bits, group_size)


def quantize_stacklm(model: nn.Module, bits: int, group_size: int = 64,
                     embedding_bits: int | None = None) -> nn.Module:
    """Replace every nn.Linear with a QuantizedLinear, then re-establish the
    embedding tie against the quantized head.

    RMSNorm scales stay fp32 (35,072 params ~= 0.14 MB — negligible and
    error-sensitive; every production format does the same).

    `embedding_bits` keeps the tied embedding/head at a HIGHER precision than
    the body — e.g. `bits=4, embedding_bits=8`. Not a hedge: the 32768 x 512
    table is 17% of the parameters and its rare-token rows are thinly trained,
    so it is disproportionately fragile at 4 bits. llama.cpp's mixed "K-quant"
    recipes make the same choice. Section 8 lets you measure whether it matters
    for YOUR checkpoint instead of guessing."""
    tied = (getattr(model, "cfg", None) is not None
            and model.cfg.tie_embeddings
            and model.lm_head.weight.data_ptr() == model.tok_emb.weight.data_ptr())

    head_linear = model.lm_head if tied else None
    if tied and embedding_bits is not None and embedding_bits != bits:
        # Quantize the head separately, at its own bit-width, BEFORE the walk.
        model.lm_head = QuantizedLinear.from_float(head_linear, embedding_bits, group_size)

    _swap_linears(model, bits, group_size)

    if tied:
        model.tok_emb = QuantizedEmbedding(model.lm_head)   # drops the fp32 Parameter

    if getattr(model, "cfg", None) is not None and model.cfg.loss_chunk:
        # Ch. 14.4's chunked fused CE path calls lm_head.weight per chunk, which
        # on a quantized model means a 67 MB materialization per chunk. We never
        # need a training loss at serve time; turn it off explicitly rather than
        # letting the shim hide a 100x slowdown.
        model.cfg.loss_chunk = 0
    return model


def build_quantized_shell(cfg, bits: int, group_size: int = 64,
                          embedding_bits: int | None = None):
    """Construct the quantized module tree on the `meta` device — right shapes,
    zero bytes — ready for `load_state_dict(..., assign=True)`."""
    from stacklm.model.transformer import Stack100M
    with torch.device("meta"):
        model = Stack100M(cfg)
    tied = cfg.tie_embeddings
    eb = embedding_bits or bits
    model.lm_head = QuantizedLinear.empty(cfg.d_model, cfg.vocab_size, eb, group_size)
    for blk in model.blocks:
        a, m = blk.attn, blk.mlp
        a.wq = QuantizedLinear.empty(cfg.d_model, cfg.n_heads * cfg.head_dim, bits, group_size)
        a.wk = QuantizedLinear.empty(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bits, group_size)
        a.wv = QuantizedLinear.empty(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bits, group_size)
        a.wo = QuantizedLinear.empty(cfg.n_heads * cfg.head_dim, cfg.d_model, bits, group_size)
        m.gate = QuantizedLinear.empty(cfg.d_model, cfg.intermediate, bits, group_size)
        m.up = QuantizedLinear.empty(cfg.d_model, cfg.intermediate, bits, group_size)
        m.down = QuantizedLinear.empty(cfg.intermediate, cfg.d_model, bits, group_size)
    if tied:
        model.tok_emb = QuantizedEmbedding(model.lm_head)
    return model


def state_dict_bytes(model: nn.Module) -> float:
    """Actual bytes, in MB, deduplicated by storage.

    Deduplication is not pedantry: an UNQUANTIZED tied model has BOTH
    `tok_emb.weight` and `lm_head.weight` in its state_dict pointing at one
    16.78M-element tensor, so the naive sum reports ~472 MB for a 405 MB model
    and every compression ratio you compute against it is wrong."""
    seen, total = set(), 0
    for t in model.state_dict().values():
        if t.data_ptr() in seen:
            continue
        seen.add(t.data_ptr())
        total += t.numel() * t.element_size()
    return total / 1e6

>>> m = Stack100M(StackConfig())
>>> [k for k in m.state_dict() if "tok_emb" in k or k == "lm_head.weight"]
['tok_emb.weight', 'lm_head.weight']              # aliased, but BOTH keys exist
>>> round(state_dict_bytes(m), 1)
405.4                                             # MB, counting the tie ONCE
>>> _ = quantize_stacklm(m, bits=4, group_size=64)
>>> [k for k in m.state_dict() if "tok_emb" in k or "lm_head" in k]
['lm_head.q_weight', 'lm_head.scale', 'lm_head.zero_point']   # no fp32 table anywhere
>>> round(state_dict_bytes(m), 1)
63.5                                              # 63.3 packed + 0.14 fp32 norms

from safetensors.torch import save_file, load_file


def export_quantized(model: nn.Module, path_prefix: str, bits: int, group_size: int,
                     config: dict, embedding_bits: int | None = None) -> None:
    """Serialize to `<prefix>.safetensors` + `<prefix>.json`.

    safetensors (huggingface/safetensors) over torch.save: no pickle (so loading
    an untrusted checkpoint cannot execute code), a readable header, and
    zero-copy mmap on load — which is why cold-start load time in Section 9 is
    dominated by page faults rather than deserialization.

    safetensors REFUSES aliased tensors, which is a feature here: it is a
    machine-checked version of the Section 7.3 tie assertion."""
    sd = model.state_dict()
    ptrs = [t.data_ptr() for t in sd.values()]
    assert len(set(ptrs)) == len(ptrs), (
        "aliased tensors in state_dict — the embedding tie was not collapsed")
    save_file({k: v.contiguous() for k, v in sd.items()}, path_prefix + ".safetensors")
    with open(path_prefix + ".json", "w") as f:
        json.dump({"bits": bits, "group_size": group_size,
                   "embedding_bits": embedding_bits or bits,
                   "format": "stacklm-rtn-v1", "architecture": config}, f, indent=2)


def load_quantized(path_prefix: str, device="cpu"):
    """Load without ever allocating the fp32 model."""
    from stacklm.config import StackConfig
    with open(path_prefix + ".json") as f:
        meta = json.load(f)
    cfg = StackConfig(**meta["architecture"])
    model = build_quantized_shell(cfg, meta["bits"], meta["group_size"],
                                  meta.get("embedding_bits"))
    sd = load_file(path_prefix + ".safetensors", device=device)
    # assign=True (PyTorch >= 2.1) REPLACES the meta params/buffers with the
    # loaded tensors rather than copying into them — required here, because you
    # cannot copy into a meta tensor.
    model.load_state_dict(sd, assign=True, strict=True)
    # RoPE tables are persistent=False, so they are NOT in the state_dict and are
    # still on `meta`. Rebuild them or the first forward dies on a meta tensor.
    model.rebuild_rope(cfg.max_seq_len, cfg.rope_theta, device=device)
    model.eval()
    return model, meta

"""
stacklm/serve/export.py — fp32 post-trained checkpoint -> quantized artifact.
Run: `python -m stacklm.serve.export --bits 4 --embedding_bits 8`
"""
import argparse
from dataclasses import asdict

import torch

from stacklm.config import StackConfig
from stacklm.model.transformer import Stack100M
from stacklm.serve.quantize import (quantize_stacklm, export_quantized,
                                    state_dict_bytes)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", default="checkpoints/stack100m/post_final.pt")
    p.add_argument("--out_dir", default="checkpoints/stack100m")
    p.add_argument("--bits", type=int, default=4, choices=[4, 8])
    p.add_argument("--group_size", type=int, default=64)
    p.add_argument("--embedding_bits", type=int, default=8,
                   help="keep the tied table at higher precision than the body")
    args = p.parse_args()

    # Ch. 14.7's save_checkpoint stores {"model", "config", "step", ...} with
    # `config` already asdict()-ed, so weights_only=True is enough to load it.
    ck = torch.load(args.ckpt, map_location="cpu", weights_only=True)
    cfg = StackConfig(**ck["config"])
    model = Stack100M(cfg)                     # the constructor re-establishes the tie
    model.load_state_dict(ck["model"])
    model.eval()

    quantize_stacklm(model, bits=args.bits, group_size=args.group_size,
                     embedding_bits=args.embedding_bits)
    prefix = f"{args.out_dir}/stack100m_int{args.bits}"
    export_quantized(model, prefix, args.bits, args.group_size,
                     asdict(cfg), embedding_bits=args.embedding_bits)
    print(f"wrote {prefix}.safetensors  ({state_dict_bytes(model):.1f} MB)")
    print(f"wrote {prefix}.json")


if __name__ == "__main__":
    main()

"""stacklm/eval/sweep.py — run the full battery across quantization settings."""
import copy

from stacklm.serve.quantize import quantize_stacklm, state_dict_bytes


def evaluate_quantization_sweep(fp32_model, tok, val_shard_dir, probes, *,
                                configs=((None, None, None), (8, None, None),
                                         (4, 128, None), (4, 64, None), (4, 64, 8)),
                                device="cpu", max_batches=20):
    """configs: (bits, group_size, embedding_bits). `probes` bundles the
    Section 4 fixtures:
       {"arith": [(problems, label), ...], "mc": [...], "qa": [...],
        "retriever": ..., "env": ..., "agent_tasks": [...], "pad_id": int}
    Returns one row per config, with deltas against the fp32 baseline."""
    rows = []
    for bits, gs, eb in configs:
        # deepcopy preserves ALIASING, so the tied lm_head/tok_emb stays tied in
        # the copy — which is exactly what `quantize_stacklm`'s data_ptr check
        # relies on. Rebuilding via Stack100M(cfg)+load_state_dict also re-ties,
        # because the constructor performs the tie; a naive tensor-by-tensor
        # "clone" helper would NOT.
        m = copy.deepcopy(fp32_model)
        if bits is not None:
            quantize_stacklm(m, bits=bits, group_size=gs or 64, embedding_bits=eb)
        m.to(device).eval()

        name = "fp32" if bits is None else f"int{bits}/g{gs or 64}" + (f"/e{eb}" if eb else "")
        ppl = compute_perplexity(m, val_shard_dir, probes["pad_id"],
                                 device=device, max_batches=max_batches)
        arith_in = eval_arithmetic(m, tok, generate_fn, probes["arith"][0][0])
        agent = eval_agent(m, tok, probes["env"], probes["agent_tasks"])
        rows.append({
            "config": name,
            "bytes_mb": state_dict_bytes(m),          # dedup'd: fp32 reads 405, not 472
            "ppl": ppl["perplexity"],
            "arith": arith_in["accuracy"],
            "parse": arith_in["parse_rate"],
            "mc": eval_mc_probe(m, tok, probes["mc"], device=device)["acc_norm"],
            "qa": eval_retrieval_qa(m, tok, generate_fn, probes["retriever"],
                                    probes["qa"])["exact_match"],
            "fmt": agent["tool_format_validity"],
            "agent_em": agent["exact_match"],
        })
        del m

    base = rows[0]
    for r in rows:
        r["d_ppl"] = r["ppl"] - base["ppl"]
        for k in ("arith", "mc", "qa", "fmt", "agent_em"):
            r["d_" + k] = r[k] - base[k]
    return rows


def print_sweep(rows) -> None:
    hdr = (f"{'config':<15}{'MB':>7}{'PPL':>8}{'dPPL':>8}{'arith':>8}"
           f"{'MC':>7}{'QA':>7}{'fmt':>7}{'agent':>7}")
    print(hdr, "\n", "-" * len(hdr), sep="")
    for r in rows:
        print(f"{r['config']:<15}{r['bytes_mb']:>7.1f}{r['ppl']:>8.2f}{r['d_ppl']:>+8.2f}"
              f"{r['arith']:>8.1%}{r['mc']:>7.1%}{r['qa']:>7.1%}"
              f"{r['fmt']:>7.1%}{r['agent_em']:>7.1%}")

"""
stacklm/serve/cli.py — CPU text generation from a quantized checkpoint, with
measured latency and memory.
Run (after `python -m stacklm.serve.export --bits 4`):
    python -m stacklm.serve.cli --bits 4
"""
import argparse
import os
import resource
import time

import torch

from stacklm.tokenizer import StackTokenizer                     # Ch. 14.3
from stacklm.serve.quantize import load_quantized                # Section 7.4
from stacklm.infer import generate_text                          # Section 2


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--bits", type=int, default=4, choices=[4, 8])
    p.add_argument("--checkpoint_dir", type=str, default="checkpoints/stack100m")
    p.add_argument("--prompt", type=str, default="The mitochondria is")
    p.add_argument("--max_new_tokens", type=int, default=64)
    p.add_argument("--threads", type=int, default=0, help="0 = all cores")
    args = p.parse_args()

    # CPU decode at batch 1 is memory-bandwidth bound, so more threads help only
    # until they saturate DRAM. Sweep this; on many laptops the optimum is the
    # PHYSICAL core count, not the SMT thread count.
    torch.set_num_threads(args.threads or os.cpu_count())

    t0 = time.perf_counter()
    model, meta = load_quantized(f"{args.checkpoint_dir}/stack100m_int{args.bits}")
    tok = StackTokenizer.load(f"{args.checkpoint_dir}/tokenizer.json")
    load_time = time.perf_counter() - t0

    # Warm up once: the first forward pays lazy-init and page-fault costs that
    # would otherwise be charged to your tok/s number.
    generate_text(model, tok, args.prompt, max_new_tokens=4, temperature=0.0)

    t1 = time.perf_counter()
    text, n_new = generate_text(model, tok, args.prompt,
                                max_new_tokens=args.max_new_tokens,
                                temperature=0.7, top_p=0.95, return_n_tokens=True)
    gen_time = time.perf_counter() - t1

    # Peak RSS counts the mmap'd weights and the KV cache, which a Python-level
    # tracer like tracemalloc cannot see at all. (POSIX only; `ru_maxrss` is KiB
    # on Linux and BYTES on macOS — check yours before reporting a number.)
    peak_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024

    print(f"--- Stack-100M (int{args.bits}, group={meta['group_size']}) on CPU ---")
    print(f"Threads:          {torch.get_num_threads()}")
    print(f"Load time:        {load_time:.2f} s")
    print(f"Generation:       {gen_time:.2f} s  ({n_new / gen_time:.1f} tok/s)")
    print(f"Peak RSS:         {peak_rss_mb:.0f} MB")
    print(f"Output: {args.prompt}{text}")


if __name__ == "__main__":
    main()

# pip install torchao
import torch
from torchao.quantization import quantize_, Int8WeightOnlyConfig, Int4WeightOnlyConfig

model = load_fp32_stack100m().eval()

quantize_(model, Int8WeightOnlyConfig())                    # W8A16, works broadly
# quantize_(model, Int4WeightOnlyConfig(group_size=64))     # W4A16, grouped like ours

# torch.compile is not optional here: torchao's quantized tensor subclasses are
# designed to be traced and fused, and the eager-mode numbers are misleading.
model = torch.compile(model, mode="max-autotune")

# pip install llmcompressor
from llmcompressor import oneshot                          # entrypoint moved across
from llmcompressor.modifiers.quantization import GPTQModifier  # releases; check yours

recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
oneshot(model=model, dataset=calibration_dataset, recipe=recipe,
        max_seq_length=2048, num_calibration_samples=512)
model.save_pretrained("Stack-100M-W4A16", save_compressed=True)  # compressed-tensors

# In a llama.cpp checkout, after exporting Stack-100M in HF layout:
python convert_hf_to_gguf.py ./stack100m-hf --outfile stack100m-f16.gguf --outtype f16
./llama-quantize stack100m-f16.gguf stack100m-q4_k_m.gguf Q4_K_M
./llama-cli -m stack100m-q4_k_m.gguf -p "The mitochondria is" -n 64
./llama-server -m stack100m-q4_k_m.gguf --port 8080   # OpenAI-compatible endpoint

python -m stacklm.serve.export --bits 4 --embedding_bits 8
python -m stacklm.serve.cli    --bits 4 --prompt "The mitochondria is"