Upload folder using huggingface_hub

src/eaimath/__init__.py

@@ -0,0 +1,6 @@
+"""CS6013 Efficient AI — compression recipe for Qwen3.5-4B on the math domain."""
+
+from .model import DEFAULT_MODEL
+
+__all__ = ["DEFAULT_MODEL"]
+__version__ = "0.1.0"

src/eaimath/adaptive_artifact.py

@@ -0,0 +1,423 @@
+"""Physical packing for the Round-14 block-adaptive Bucket-C representation.
+
+The research builder selects one bit width per output-row block, but evaluates
+the resulting values as an ordinary BF16 Hugging Face checkpoint.  This module
+turns that selection into the actual submission object: densely packed integer
+codes, FP16 group scales, one uint8 width selector per adaptive block, and BF16
+passthrough tensors for protected parameters.
+
+Unlike the older nested Bucket-B artifact, ``state_dict`` here is deliberately a
+flat mapping of names to tensors.  The course starter's size audit counts only
+immediate tensor values, so this layout makes its tensor-payload report equal to
+the exact representation accounting instead of accidentally reporting zero.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import re
+import time
+from pathlib import Path
+from typing import Any, Mapping
+
+import torch
+
+from .model import load_model, save_checkpoint
+from .pack import pack_tensor, unpack_tensor
+from .quantize import EXCLUDED, QuantSpec, quantize_dequantize
+
+FORMAT = "eaimath-block-adaptive-c-v1"
+EXCLUDED_RE = re.compile(EXCLUDED)
+PROTECTED_RE = re.compile(r"A_log|dt_bias|conv1d|in_proj_a|in_proj_b|norm|bias")
+ADAPTIVE_RE = re.compile(r"linear_attn|self_attn|\.mlp\.")
+
+
+def tensor_payload_bytes(tensors: Mapping[str, torch.Tensor]) -> int:
+    """Bytes counted by the starter evaluator's immediate state-dict audit."""
+    return sum(int(value.numel()) * int(value.element_size()) for value in tensors.values())
+
+
+def sha256_file(path: str | Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        while chunk := handle.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _is_adaptive(name: str, tensor: torch.Tensor) -> bool:
+    return (
+        tensor.ndim == 2
+        and name.endswith(".weight")
+        and not PROTECTED_RE.search(name)
+        and bool(ADAPTIVE_RE.search(name))
+    )
+
+
+def _store_entry(
+    entry: dict[str, Any],
+    tensors: dict[str, torch.Tensor],
+    prefix: str,
+) -> dict[str, Any]:
+    """Move entry tensors into the flat payload and retain JSON-safe metadata."""
+    metadata: dict[str, Any] = {}
+    tensor_keys: dict[str, str] = {}
+    for field, value in entry.items():
+        if isinstance(value, torch.Tensor):
+            key = f"{prefix}.{field}"
+            if key in tensors:
+                raise KeyError(f"duplicate artifact tensor key: {key}")
+            tensors[key] = value.detach().cpu().contiguous()
+            tensor_keys[field] = key
+        elif isinstance(value, tuple):
+            metadata[field] = list(value)
+        else:
+            metadata[field] = value
+    metadata["tensor_keys"] = tensor_keys
+    return metadata
+
+
+def _load_entry(metadata: Mapping[str, Any], tensors: Mapping[str, torch.Tensor]) -> dict[str, Any]:
+    entry = {key: value for key, value in metadata.items() if key != "tensor_keys"}
+    for field, key in metadata.get("tensor_keys", {}).items():
+        if key not in tensors:
+            raise KeyError(f"artifact tensor is missing: {key}")
+        entry[field] = tensors[key]
+    if "shape" in entry:
+        entry["shape"] = tuple(entry["shape"])
+    return entry
+
+
+def _row_ids(widths: torch.Tensor, bit_width: int, rows: int, block: int) -> torch.Tensor:
+    selected: list[int] = []
+    for block_index, width in enumerate(widths.tolist()):
+        if int(width) != bit_width:
+            continue
+        start = block_index * block
+        selected.extend(range(start, min(rows, start + block)))
+    return torch.tensor(selected, dtype=torch.int64)
+
+
+def pack_block_adaptive_state(
+    state_dict: Mapping[str, torch.Tensor],
+    allocation_report: Mapping[str, Any],
+    *,
+    reference_state_dict: Mapping[str, torch.Tensor] | None = None,
+    scale_normalization: str = "none",
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Pack a block-adaptive BF16 state dict and verify exact byte accounting.
+
+    ``state_dict`` should be the unquantized learned source.  When supplied,
+    ``reference_state_dict`` is the expanded fake-quant checkpoint used during
+    research evaluation; the reported repack error then proves that physical
+    codes reproduce that evaluated model without applying a second quantizer.
+    """
+    if allocation_report.get("format") != FORMAT:
+        raise ValueError(
+            f"expected allocation format {FORMAT!r}, got {allocation_report.get('format')!r}"
+        )
+    selected_bits = allocation_report.get("selected_bits")
+    if not isinstance(selected_bits, dict) or not selected_bits:
+        raise ValueError("allocation report has no selected_bits mapping")
+    block = int(allocation_report.get("row_block", 0))
+    group_size = int(allocation_report.get("group_size", 0))
+    if block <= 0 or group_size <= 0:
+        raise ValueError(f"invalid block/group sizes: {block}/{group_size}")
+
+    tensors: dict[str, torch.Tensor] = {}
+    entries: dict[str, dict[str, Any]] = {}
+    aliases: dict[str, str] = {}
+    seen: dict[int, str] = {}
+    reported_names: set[str] = set()
+    max_abs_error = 0.0
+    max_rmse = 0.0
+    max_abs_scale_error = 0.0
+    max_scale_rmse = 0.0
+    worst_tensor = ""
+    canonical_index = 0
+
+    for name, tensor in state_dict.items():
+        if not isinstance(tensor, torch.Tensor) or EXCLUDED_RE.search(name):
+            continue
+        pointer = tensor.data_ptr()
+        if pointer and pointer in seen:
+            aliases[name] = seen[pointer]
+            continue
+        if pointer:
+            seen[pointer] = name
+        widths_raw = selected_bits.get(name)
+        if widths_raw is None:
+            raise KeyError(
+                f"text tensor {name!r} has no allocation; refusing uncharged passthrough"
+            )
+        widths = torch.tensor([int(value) for value in widths_raw], dtype=torch.uint8)
+        if widths.numel() == 0:
+            raise ValueError(f"tensor {name!r} has an empty width selection")
+        reported_names.add(name)
+        prefix = f"t{canonical_index:04d}"
+        canonical_index += 1
+
+        adaptive = _is_adaptive(name, tensor)
+        expected_blocks = math.ceil(tensor.shape[0] / block) if adaptive else 1
+        if widths.numel() != expected_blocks:
+            raise ValueError(
+                f"width count for {name} is {widths.numel()}, expected {expected_blocks}"
+            )
+
+        if adaptive:
+            selector_key = f"{prefix}.selector"
+            tensors[selector_key] = widths.contiguous()
+            groups: dict[str, dict[str, Any]] = {}
+            reconstructed = torch.empty_like(tensor, device="cpu")
+            for bit_width in sorted(set(int(value) for value in widths.tolist())):
+                ids = _row_ids(widths, bit_width, tensor.shape[0], block)
+                selected = tensor.detach().cpu().index_select(0, ids)
+                packed = pack_tensor(
+                    selected,
+                    bit_width,
+                    group_size=group_size,
+                    scale_normalization=scale_normalization,
+                )
+                groups[str(bit_width)] = {
+                    "rows": int(ids.numel()),
+                    "entry": _store_entry(packed, tensors, f"{prefix}.b{bit_width}"),
+                }
+                unpacked = unpack_tensor(packed).to(reconstructed.dtype)
+                reconstructed.index_copy_(0, ids, unpacked)
+                if bit_width < 16:
+                    ideal = quantize_dequantize(
+                        selected,
+                        QuantSpec(".*", bit_width, group_size),
+                    )
+                    scale_delta = ideal.float() - unpacked.float()
+                    if scale_delta.numel():
+                        max_abs_scale_error = max(
+                            max_abs_scale_error, float(scale_delta.abs().max())
+                        )
+                        max_scale_rmse = max(
+                            max_scale_rmse,
+                            float(scale_delta.square().mean().sqrt()),
+                        )
+            entries[name] = {
+                "kind": "adaptive_rows",
+                "shape": list(tensor.shape),
+                "block": block,
+                "selector_key": selector_key,
+                "groups": groups,
+            }
+        else:
+            bit_width = int(widths[0])
+            packed = pack_tensor(
+                tensor.detach().cpu(),
+                bit_width,
+                group_size=group_size,
+                scale_normalization=scale_normalization,
+            )
+            entries[name] = {
+                "kind": "tensor",
+                "bits": bit_width,
+                "entry": _store_entry(packed, tensors, prefix),
+            }
+            reconstructed = unpack_tensor(packed)
+            if bit_width < 16:
+                ideal = quantize_dequantize(
+                    tensor.detach().cpu(),
+                    QuantSpec(".*", bit_width, group_size),
+                )
+                scale_delta = ideal.float() - reconstructed.float()
+                if scale_delta.numel():
+                    max_abs_scale_error = max(
+                        max_abs_scale_error, float(scale_delta.abs().max())
+                    )
+                    max_scale_rmse = max(
+                        max_scale_rmse,
+                        float(scale_delta.square().mean().sqrt()),
+                    )
+
+        comparison = tensor
+        if reference_state_dict is not None:
+            if name not in reference_state_dict:
+                raise KeyError(f"reference checkpoint is missing text tensor {name!r}")
+            comparison = reference_state_dict[name]
+            if tuple(comparison.shape) != tuple(tensor.shape):
+                raise ValueError(
+                    f"reference shape for {name} is {tuple(comparison.shape)}, "
+                    f"expected {tuple(tensor.shape)}"
+                )
+        delta = comparison.detach().float().cpu() - reconstructed.float().cpu()
+        abs_error = float(delta.abs().max()) if delta.numel() else 0.0
+        rmse = float(delta.square().mean().sqrt()) if delta.numel() else 0.0
+        if abs_error > max_abs_error:
+            max_abs_error = abs_error
+            worst_tensor = name
+        max_rmse = max(max_rmse, rmse)
+
+    unused = set(selected_bits) - reported_names
+    if unused:
+        preview = ", ".join(sorted(unused)[:8])
+        raise KeyError(f"allocation contains {len(unused)} absent tensors: {preview}")
+
+    actual_bytes = tensor_payload_bytes(tensors)
+    charged_bytes = int(allocation_report.get("total_bytes", -1))
+    if actual_bytes != charged_bytes:
+        raise ValueError(
+            "physical tensor bytes do not match allocation: "
+            f"packed={actual_bytes:,}, charged={charged_bytes:,}"
+        )
+    selector_bytes = sum(
+        tensor.numel() * tensor.element_size()
+        for key, tensor in tensors.items()
+        if key.endswith(".selector")
+    )
+    expected_selectors = int(allocation_report.get("selector_bytes", -1))
+    if selector_bytes != expected_selectors:
+        raise ValueError(
+            f"selector bytes {selector_bytes:,} do not match report {expected_selectors:,}"
+        )
+
+    report = {
+        "format": FORMAT,
+        "canonical_tensors": len(entries),
+        "aliases": len(aliases),
+        "flat_payload_tensors": len(tensors),
+        "tensor_bytes": actual_bytes,
+        "selector_bytes": selector_bytes,
+        "max_abs_repack_error": max_abs_error,
+        "max_repack_rmse": max_rmse,
+        "max_abs_scale_storage_error": max_abs_scale_error,
+        "max_scale_storage_rmse": max_scale_rmse,
+        "worst_tensor": worst_tensor,
+        "error_reference": (
+            "expanded_fake_quant" if reference_state_dict is not None else "packing_source"
+        ),
+        "row_block": block,
+        "group_size": group_size,
+        "scale_normalization": scale_normalization,
+        "source_budget": int(allocation_report.get("budget", charged_bytes)),
+        "source_margin": int(allocation_report.get("margin_bytes", 0)),
+    }
+    payload = {
+        "format": FORMAT,
+        "model": allocation_report.get("model"),
+        "state_dict": tensors,
+        "entries": entries,
+        "aliases": aliases,
+        "report": report,
+    }
+    return payload, report
+
+
+def unpack_block_adaptive_state(payload: Mapping[str, Any]) -> dict[str, torch.Tensor]:
+    """Restore every packed text tensor without loading the base architecture."""
+    if payload.get("format") != FORMAT:
+        raise ValueError(f"unsupported block-adaptive artifact: {payload.get('format')!r}")
+    tensors = payload.get("state_dict")
+    entries = payload.get("entries")
+    if not isinstance(tensors, dict) or not isinstance(entries, dict):
+        raise ValueError("artifact is missing flat state_dict or entries metadata")
+    restored: dict[str, torch.Tensor] = {}
+    for name, metadata in entries.items():
+        kind = metadata.get("kind")
+        if kind == "tensor":
+            value = unpack_tensor(_load_entry(metadata["entry"], tensors))
+        elif kind == "adaptive_rows":
+            shape = tuple(int(value) for value in metadata["shape"])
+            widths = tensors[metadata["selector_key"]].to(torch.uint8)
+            value = torch.empty(shape, dtype=torch.bfloat16)
+            for bit_text, group_metadata in metadata["groups"].items():
+                bit_width = int(bit_text)
+                ids = _row_ids(widths, bit_width, shape[0], int(metadata["block"]))
+                if ids.numel() != int(group_metadata["rows"]):
+                    raise ValueError(f"row count mismatch while restoring {name} at {bit_width} bits")
+                rows = unpack_tensor(_load_entry(group_metadata["entry"], tensors))
+                value.index_copy_(0, ids, rows.to(value.dtype))
+        else:
+            raise ValueError(f"unknown entry kind for {name}: {kind!r}")
+        restored[name] = value
+    for alias, canonical in payload.get("aliases", {}).items():
+        if canonical not in restored:
+            raise KeyError(f"alias target is absent: {alias} -> {canonical}")
+        restored[alias] = restored[canonical]
+    return restored
+
+
+def save_block_adaptive_artifact(payload: Mapping[str, Any], path: str | Path) -> dict[str, Any]:
+    output = Path(path)
+    if output.exists():
+        raise FileExistsError(f"refusing to overwrite artifact: {output}")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    started = time.time()
+    torch.save(dict(payload), output)
+    report = dict(payload.get("report", {}))
+    report.update(
+        {
+            "disk_bytes": output.stat().st_size,
+            "sha256": sha256_file(output),
+            "save_seconds": round(time.time() - started, 3),
+        }
+    )
+    output.with_suffix(output.suffix + ".json").write_text(json.dumps(report, indent=2))
+    return report
+
+
+def load_block_adaptive_artifact(path: str | Path) -> dict[str, Any]:
+    payload = torch.load(Path(path), map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict) or payload.get("format") != FORMAT:
+        raise ValueError(f"unsupported block-adaptive artifact in {path}")
+    actual = tensor_payload_bytes(payload.get("state_dict", {}))
+    expected = int(payload.get("report", {}).get("tensor_bytes", -1))
+    if actual != expected:
+        raise ValueError(f"artifact tensor bytes changed: {actual:,} != {expected:,}")
+    return payload
+
+
+def restore_block_adaptive_artifact(
+    model_name: str,
+    artifact_path: str | Path,
+    output_path: str | Path,
+    *,
+    cache_dir: str | None = None,
+) -> dict[str, Any]:
+    """Expand a physical artifact into the BF16 HF checkpoint used by evaluation."""
+    out = Path(output_path)
+    if out.exists() and (not out.is_dir() or any(out.iterdir())):
+        raise FileExistsError(f"refusing to overwrite restore target: {out}")
+    started = time.time()
+    payload = load_block_adaptive_artifact(artifact_path)
+    restored = unpack_block_adaptive_state(payload)
+    model = load_model(
+        model_name,
+        dtype="bfloat16",
+        device_map=None,
+        cache_dir=cache_dir,
+        multimodal=True,
+    )
+    targets = model.state_dict()
+    with torch.no_grad():
+        for name, target in targets.items():
+            if EXCLUDED_RE.search(name):
+                target.zero_()
+                continue
+            if name not in restored:
+                raise KeyError(f"artifact did not restore text tensor: {name}")
+            value = restored[name]
+            if tuple(value.shape) != tuple(target.shape):
+                raise ValueError(
+                    f"shape mismatch for {name}: {tuple(value.shape)} != {tuple(target.shape)}"
+                )
+            target.copy_(value.to(dtype=target.dtype))
+    copied = save_checkpoint(model, out, source_model=model_name, cache_dir=cache_dir)
+    report = {
+        "format": FORMAT,
+        "artifact": str(artifact_path),
+        "artifact_sha256": sha256_file(artifact_path),
+        "tensor_bytes": int(payload["report"]["tensor_bytes"]),
+        "restored_text_tensors": len(restored),
+        "zeroed_excluded_tensors": sum(1 for name in targets if EXCLUDED_RE.search(name)),
+        "auxiliary_files": copied,
+        "restore_seconds": round(time.time() - started, 1),
+    }
+    (out / "block_adaptive_restore.json").write_text(json.dumps(report, indent=2))
+    return report

src/eaimath/answers.py

@@ -0,0 +1,173 @@
+"""Extraction and equivalence checking for free-form math answers.
+
+The hidden leaderboard eval is math, so answers are LaTeX expressions inside
+``\\boxed{...}`` rather than multiple-choice letters. Extraction therefore needs
+real brace matching (``\\boxed{\\frac{1}{2}}`` breaks any naive regex), and
+comparison needs symbolic equivalence (``0.5`` == ``\\frac{1}{2}`` == ``1/2``).
+"""
+
+from __future__ import annotations
+
+import re
+from fractions import Fraction
+
+BOXED_TOKENS = ("\\boxed", "\\fbox")
+
+# Applied in order when falling back to string comparison.
+_STRIP_WRAPPERS = (
+    (r"\\left", ""),
+    (r"\\right", ""),
+    (r"\\!", ""),
+    (r"\\,", ""),
+    (r"\\;", ""),
+    (r"\\ ", " "),
+    (r"\\dfrac", r"\\frac"),
+    (r"\\tfrac", r"\\frac"),
+    (r"\\cdot", "*"),
+    (r"\\times", "*"),
+    (r"\^\{\\circ\}", ""),
+    (r"\^\\circ", ""),
+    (r"\\%", ""),
+    (r"\\\$", ""),
+)
+
+
+def extract_boxed(text: str) -> str | None:
+    """Return the contents of the LAST ``\\boxed{...}`` in ``text``.
+
+    Uses brace-depth matching so nested LaTeX survives. Returns ``None`` when no
+    boxed expression is present or the braces never close (a truncated
+    generation), which the caller should treat as an unparsed answer.
+    """
+    if not text:
+        return None
+
+    start_idx = -1
+    token_len = 0
+    for token in BOXED_TOKENS:
+        idx = text.rfind(token)
+        if idx > start_idx:
+            start_idx, token_len = idx, len(token)
+    if start_idx == -1:
+        return None
+
+    i = start_idx + token_len
+    while i < len(text) and text[i].isspace():
+        i += 1
+    if i >= len(text):
+        return None
+
+    # `\boxed 42` (no braces) — take the next non-space run.
+    if text[i] != "{":
+        match = re.match(r"[^\s$\\]+", text[i:])
+        return match.group(0) if match else None
+
+    depth = 0
+    content_start = i + 1
+    while i < len(text):
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+            if depth == 0:
+                return text[content_start:i]
+        i += 1
+    return None  # unclosed brace -> truncated mid-answer
+
+
+def extract_answer(text: str) -> str | None:
+    """Extract a final answer, preferring ``\\boxed{}`` with light fallbacks."""
+    boxed = extract_boxed(text)
+    if boxed is not None:
+        return boxed.strip()
+
+    # Fallback: an explicit "answer is X" on one of the last few lines.
+    tail = [ln.strip() for ln in text.strip().splitlines() if ln.strip()][-5:]
+    pattern = re.compile(
+        r"(?:final\s+answer|answer)\s*(?:is|:)\s*\$?([^\s$.,]+)", re.IGNORECASE
+    )
+    for line in reversed(tail):
+        match = pattern.search(line)
+        if match:
+            return match.group(1).strip()
+    return None
+
+
+def normalize(expr: str) -> str:
+    """Aggressively normalize a LaTeX answer for string comparison."""
+    if expr is None:
+        return ""
+    out = expr.strip()
+    out = re.sub(r"^\$+|\$+$", "", out).strip()
+    out = re.sub(r"^\\\[|\\\]$", "", out).strip()
+    for pattern, repl in _STRIP_WRAPPERS:
+        out = re.sub(pattern, repl, out)
+    out = re.sub(r"\\text\{([^}]*)\}", r"\1", out)
+    out = re.sub(r"\\mathrm\{([^}]*)\}", r"\1", out)
+    # 1,234,567 -> 1234567 (thousands separators only, not tuples)
+    if re.fullmatch(r"-?\d{1,3}(,\d{3})+(\.\d+)?", out):
+        out = out.replace(",", "")
+    out = out.replace(" ", "").rstrip(".")
+    # Trailing ".0" and redundant leading "+"
+    out = re.sub(r"^\+", "", out)
+    if re.fullmatch(r"-?\d+\.0+", out):
+        out = out.split(".")[0]
+    return out
+
+
+def _as_number(expr: str) -> float | None:
+    """Best-effort numeric value for simple ints, decimals and \\frac{a}{b}."""
+    text = normalize(expr)
+    try:
+        return float(Fraction(text))
+    except (ValueError, ZeroDivisionError):
+        pass
+    frac = re.fullmatch(r"\\frac\{(-?[\d.]+)\}\{(-?[\d.]+)\}", text)
+    if frac:
+        try:
+            denom = float(frac.group(2))
+            return float(frac.group(1)) / denom if denom else None
+        except (ValueError, ZeroDivisionError):
+            return None
+    return None
+
+
+_MATH_VERIFY_STATE: dict[str, object] = {}
+
+
+def _math_verify():
+    """Lazily import math_verify; cache the failure so we only warn once."""
+    if "fn" not in _MATH_VERIFY_STATE:
+        try:
+            from math_verify import parse, verify
+
+            _MATH_VERIFY_STATE["fn"] = (parse, verify)
+        except Exception:  # pragma: no cover - depends on optional install
+            _MATH_VERIFY_STATE["fn"] = None
+    return _MATH_VERIFY_STATE["fn"]
+
+
+def answers_match(predicted: str | None, gold: str) -> bool:
+    """True when ``predicted`` is mathematically equivalent to ``gold``.
+
+    Tries symbolic equivalence via ``math_verify`` first, then numeric
+    comparison, then normalized string equality. Each stage is a superset of
+    what the next can catch, so a False here means all three disagreed.
+    """
+    if predicted is None:
+        return False
+
+    checker = _math_verify()
+    if checker is not None:
+        parse, verify = checker
+        try:
+            if verify(parse(f"${gold}$"), parse(f"${predicted}$")):
+                return True
+        except Exception:
+            pass  # malformed LaTeX -> fall through to the cheaper checks
+
+    pred_num, gold_num = _as_number(predicted), _as_number(gold)
+    if pred_num is not None and gold_num is not None:
+        return abs(pred_num - gold_num) < 1e-6
+
+    return normalize(predicted) == normalize(gold)

src/eaimath/artifact.py

@@ -0,0 +1,686 @@
+"""Submission-shaped packing and restoration for the Round 9 Bucket-B artifact.
+
+The calibrated GPTQ checkpoint is an inference diagnostic, not the object the
+course grades: modules skipped by GPTQ are still BF16 in that directory.  This
+module dequantizes the diagnostic values once, repacks every text parameter at
+the target mixed precision, stores only selected embedding rows, and restores a
+normal BF16 Hugging Face checkpoint.
+
+The artifact deliberately preserves original tokenizer ids.  The supplied
+evaluator copies the original tokenizer into the restored directory, so a dense
+id remap would be overwritten and would make prompts invalid.
+"""
+
+from __future__ import annotations
+
+import json
+import lzma
+import re
+import tempfile
+import time
+import zlib
+from pathlib import Path
+from typing import Any, Iterable
+
+import numpy as np
+import torch
+
+from .embedding_predictor import predict_token_rows
+from .model import DEFAULT_MODEL, load_model, load_tokenizer, save_checkpoint
+from .pack import (
+    pack_bits,
+    pack_rows,
+    pack_tensor,
+    state_dict_bytes,
+    unpack_bits,
+    unpack_adaptive_codebook,
+    unpack_rows,
+    unpack_tensor,
+)
+
+FORMAT = "eaimath-r09-v1"
+ZLIB_MAGIC = b"EAIMATH_ZLIB_V1\n"
+LZMA_MAGIC = b"EAIMATH_LZMA_V1\n"
+EXCLUDED = re.compile(r"visual|vision_tower|(^|\.)mtp\.")
+PROTECTED = re.compile(r"A_log|dt_bias|conv1d|in_proj_a|in_proj_b|norm|bias")
+EMBED = re.compile(r"embed_tokens|lm_head")
+LINEAR_ATTN = re.compile(r"linear_attn")
+FULL_ATTN = re.compile(r"self_attn")
+MLP = re.compile(r"\.mlp\.")
+
+
+def component_and_bits(
+    name: str,
+    recipe_bits: dict[str, int] | None = None,
+) -> tuple[str, int]:
+    """Return component and stored width; protected parameters always stay BF16."""
+    bits = recipe_bits or {
+        "mlp": 3,
+        "linear_attn": 3,
+        "full_attn": 8,
+        "embed": 8,
+    }
+    if PROTECTED.search(name):
+        return "protected", 16
+    if EMBED.search(name):
+        return "embed", bits["embed"]
+    if LINEAR_ATTN.search(name):
+        return "linear_attn", bits["linear_attn"]
+    if FULL_ATTN.search(name):
+        return "full_attn", bits["full_attn"]
+    if MLP.search(name):
+        return "mlp", bits["mlp"]
+    return "other", 16
+
+
+def _tensor_error(original: torch.Tensor, restored: torch.Tensor) -> tuple[float, float]:
+    if original.numel() == 0:
+        return 0.0, 0.0
+    delta = original.detach().float().cpu() - restored.detach().float().cpu()
+    return float(delta.abs().max()), float(delta.square().mean().sqrt())
+
+
+def read_trace_corpus(
+    paths: Iterable[Path],
+) -> tuple[list[str], list[str], dict[str, int]]:
+    """Read verified train-only problems and completions without mixing their roles."""
+    problems: list[str] = []
+    completions: list[str] = []
+    stats = {"rows": 0, "usable": 0, "bad_json": 0, "characters": 0}
+    for path in paths:
+        if not path.is_file():
+            raise FileNotFoundError(f"vocabulary corpus does not exist: {path}")
+        with path.open(encoding="utf-8") as handle:
+            for line in handle:
+                stats["rows"] += 1
+                try:
+                    row = json.loads(line)
+                except json.JSONDecodeError:
+                    stats["bad_json"] += 1
+                    continue
+                problem = str(row.get("problem") or "").strip()
+                completion = str(row.get("completion") or row.get("response") or "").strip()
+                if not problem or not completion:
+                    continue
+                if row.get("correct") is False or row.get("finished") is False:
+                    continue
+                problems.append(problem)
+                completions.append(completion)
+                stats["usable"] += 1
+                stats["characters"] += len(problem) + len(completion)
+    if stats["usable"] < 128 or stats["characters"] < 1_000_000:
+        raise ValueError(
+            "vocabulary ranking needs at least 128 verified traces and 1M characters; "
+            f"found {stats['usable']} traces / {stats['characters']:,} characters"
+        )
+    return problems, completions, stats
+
+
+def read_trace_texts(paths: Iterable[Path]) -> tuple[list[str], dict[str, int]]:
+    """Read disjoint training problems and verified completions for vocab ranking."""
+    problems, completions, stats = read_trace_corpus(paths)
+    return problems + completions, stats
+
+
+def _decode_gptq_qzeros(module) -> torch.Tensor:
+    """Decode GPTQ v2 zero points without relying on a kernel-private helper.
+
+    Three-bit GPTQ does not have an integer ``pack_factor``: 32 values span
+    three 32-bit words, with values 10 and 21 crossing word boundaries.
+    GPTQModel's TorchLinear streaming helper uses the ordinary pack-factor path
+    and therefore cannot decode this layout.  Its reference dequantizer has the
+    special case reproduced here.
+    """
+    qzeros = module.qzeros.detach().cpu()
+    bits = int(module.bits)
+    word_bits = int(module.pack_dtype_bits)
+    maxq = int(module.maxq)
+
+    if bits in {2, 4, 8}:
+        shifts = torch.arange(0, word_bits, bits, dtype=torch.int32).reshape(1, 1, -1)
+        zeros = torch.bitwise_right_shift(qzeros.unsqueeze(2), shifts)
+        zeros = torch.bitwise_and(zeros, maxq)
+    elif bits == 3:
+        if word_bits != 32 or qzeros.shape[1] % 3:
+            raise RuntimeError(
+                "unsupported 3-bit GPTQ zero layout: "
+                f"word_bits={word_bits}, qzeros={tuple(qzeros.shape)}"
+            )
+        shifts = torch.tensor(
+            [
+                [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 0],
+                [0, 1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31],
+                [0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 0],
+            ],
+            dtype=torch.int32,
+        ).reshape(1, 1, 3, 12)
+        zeros = qzeros.reshape(qzeros.shape[0], qzeros.shape[1] // 3, 3, 1)
+        zeros = torch.bitwise_right_shift(zeros.expand(-1, -1, -1, 12), shifts)
+        zeros[:, :, 0, 10] = (zeros[:, :, 0, 10] & 0x3) | (
+            (zeros[:, :, 1, 0] << 2) & 0x4
+        )
+        zeros[:, :, 1, 11] = (zeros[:, :, 1, 11] & 0x1) | (
+            (zeros[:, :, 2, 0] << 1) & 0x6
+        )
+        zeros = torch.bitwise_and(zeros, maxq)
+        zeros = torch.cat(
+            [zeros[:, :, 0, :11], zeros[:, :, 1, 1:12], zeros[:, :, 2, 1:11]],
+            dim=2,
+        )
+    else:
+        raise RuntimeError(f"unsupported GPTQ zero-point width: {bits}")
+
+    expected = int(module.scales.numel())
+    if zeros.numel() != expected:
+        raise RuntimeError(
+            "decoded GPTQ zero-point count does not match scales: "
+            f"zeros={zeros.numel()}, scales={expected}"
+        )
+    return zeros.reshape(module.scales.shape).to(torch.int64)
+
+
+def _capture_gptq_entry(module) -> tuple[dict[str, Any], float, torch.Tensor]:
+    """Store GPTQ's exact codes/scales instead of re-quantizing its values."""
+    dequantized = module.dequantize_weight().detach().cpu()
+    scales = module.scales.detach().cpu().contiguous()
+    g_idx = module.g_idx.detach().cpu().long()
+    expected_shape = (module.in_features, module.out_features)
+    if tuple(dequantized.shape) != expected_shape or g_idx.numel() != module.in_features:
+        raise RuntimeError(
+            "GPTQ auto-padding was not removed by dequantization: "
+            f"weight={tuple(dequantized.shape)}, g_idx={g_idx.numel()}, "
+            f"expected={expected_shape}"
+        )
+    zeros = _decode_gptq_qzeros(module)
+    scale_rows = scales.float().index_select(0, g_idx)
+    zero_rows = zeros.index_select(0, g_idx)
+    codes = torch.round(dequantized.float() / scale_rows + zero_rows.float()).to(torch.int64)
+    codes.clamp_(0, 2**module.bits - 1)
+    reconstructed = scales.index_select(0, g_idx) * (
+        codes - zero_rows
+    ).to(scales.dtype)
+    target = dequantized.to(torch.bfloat16)
+    max_error = float((reconstructed.to(torch.bfloat16) - target).float().abs().max())
+
+    sequential = torch.div(
+        torch.arange(module.in_features, dtype=torch.int64),
+        module.group_size,
+        rounding_mode="floor",
+    ).clamp_max(scales.shape[0] - 1)
+    zero_flat = zeros.reshape(-1)
+    uniform_zero = bool(torch.all(zero_flat == zero_flat[0]))
+    entry: dict[str, Any] = {
+        "kind": "gptq_codes",
+        "packed": torch.from_numpy(pack_bits(codes.numpy().astype(np.uint32), module.bits)),
+        "scale": scales,
+        "shape": (module.out_features, module.in_features),
+        "code_shape": tuple(codes.shape),
+        "bits": int(module.bits),
+        "group": int(module.group_size),
+        "sym": bool(module.sym),
+        "dtype": "torch.bfloat16",
+        "zero_value": int(zero_flat[0]) if uniform_zero else None,
+    }
+    if not uniform_zero:
+        entry["zero_packed"] = torch.from_numpy(
+            pack_bits(zeros.numpy().astype(np.uint32), module.bits)
+        )
+        entry["zero_shape"] = tuple(zeros.shape)
+    if not torch.equal(g_idx, sequential):
+        entry["g_idx"] = g_idx.to(torch.int32)
+    return entry, max_error, target.T.contiguous()
+
+
+def unpack_gptq_entry(entry: dict[str, Any]) -> torch.Tensor:
+    """Expand an exact code/scale GPTQ entry to its BF16 linear weight."""
+    in_features, out_features = tuple(entry["code_shape"])
+    count = in_features * out_features
+    codes = unpack_bits(entry["packed"].numpy(), entry["bits"], count)
+    codes_tensor = torch.from_numpy(codes.astype(np.int16)).reshape(in_features, out_features)
+    scales = entry["scale"]
+    if entry.get("g_idx") is None:
+        g_idx = torch.div(
+            torch.arange(in_features, dtype=torch.int64),
+            entry["group"],
+            rounding_mode="floor",
+        ).clamp_max(scales.shape[0] - 1)
+    else:
+        g_idx = entry["g_idx"].long()
+    if entry.get("zero_value") is not None:
+        zero_rows = int(entry["zero_value"])
+    else:
+        zero_shape = tuple(entry["zero_shape"])
+        zero_count = int(np.prod(zero_shape))
+        zeros = unpack_bits(entry["zero_packed"].numpy(), entry["bits"], zero_count)
+        zero_tensor = torch.from_numpy(zeros.astype(np.int16)).reshape(zero_shape)
+        zero_rows = zero_tensor.index_select(0, g_idx)
+    values = scales.index_select(0, g_idx) * (
+        codes_tensor - zero_rows
+    ).to(scales.dtype)
+    return values.T.contiguous().to(torch.bfloat16)
+
+
+def dequantize_gptq_model(checkpoint: str | Path):
+    """Capture exact GPTQ entries, then replace linears for a normal state dict."""
+    import torch.nn as nn
+    from gptqmodel import GPTQModel
+    from gptqmodel.nn_modules.qlinear import BaseQuantLinear
+    from gptqmodel.nn_modules.qlinear.torch import TorchLinear
+    from gptqmodel.utils.backend import BACKEND
+
+    loaded = GPTQModel.load(
+        str(checkpoint),
+        backend=BACKEND.GPTQ_TORCH,
+        device="cpu",
+        trust_remote_code=True,
+    )
+    model = loaded.model
+    named = dict(model.named_modules())
+    replaced = 0
+    exact_weights: dict[str, dict[str, Any]] = {}
+    capture_error = 0.0
+    for name, module in list(model.named_modules()):
+        if isinstance(module, BaseQuantLinear) and not isinstance(module, TorchLinear):
+            raise TypeError(
+                f"{name} loaded as {type(module).__name__}; the submission packer "
+                "requires BACKEND.GPTQ_TORCH"
+            )
+        if not isinstance(module, TorchLinear):
+            continue
+        exact_entry, error, weight = _capture_gptq_entry(module)
+        exact_weights[f"{name}.weight"] = exact_entry
+        capture_error = max(capture_error, error)
+        replacement = nn.Linear(
+            module.in_features,
+            module.out_features,
+            bias=module.bias is not None,
+            device="cpu",
+            dtype=torch.bfloat16,
+        )
+        replacement.weight = nn.Parameter(weight)
+        if module.bias is not None:
+            replacement.bias = nn.Parameter(module.bias.detach().cpu().to(torch.bfloat16))
+        parent_name, child_name = name.rsplit(".", 1)
+        setattr(named[parent_name], child_name, replacement)
+        named[name] = replacement
+        replaced += 1
+    if replaced == 0:
+        raise RuntimeError("no GPTQ TorchLinear modules were found in the source checkpoint")
+    return loaded, model, replaced, exact_weights, capture_error
+
+
+def pack_model_state(
+    state_dict: dict[str, torch.Tensor],
+    keep_ids: list[int],
+    *,
+    group_size: int = 128,
+    exact_entries: dict[str, dict[str, Any]] | None = None,
+    recipe_bits: dict[str, int] | None = None,
+    recipe_name: str = "m3l3a8e8",
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Pack unique text tensors and validate every local round trip."""
+    packed: dict[str, dict[str, Any]] = {}
+    aliases: dict[str, str] = {}
+    seen: dict[int, str] = {}
+    groups: dict[str, dict[str, float | int]] = {}
+    max_abs_error = 0.0
+    max_rmse = 0.0
+    worst_tensor = ""
+    keep_tensor = torch.tensor(keep_ids, dtype=torch.int64)
+    exact_entries = exact_entries or {}
+
+    for index, (name, tensor) in enumerate(state_dict.items(), 1):
+        if EXCLUDED.search(name) or not isinstance(tensor, torch.Tensor):
+            continue
+        pointer = tensor.data_ptr()
+        if pointer and pointer in seen:
+            aliases[name] = seen[pointer]
+            continue
+        if pointer:
+            seen[pointer] = name
+
+        component, bits = component_and_bits(name, recipe_bits)
+        if name in exact_entries:
+            entry = exact_entries[name]
+            restored = unpack_gptq_entry(entry)
+            original = tensor.detach().cpu()
+        elif component == "embed" and tensor.ndim == 2 and tensor.shape[0] > len(keep_ids):
+            entry = pack_rows(tensor, keep_tensor, bits, group_size=group_size)
+            restored = unpack_rows(entry, fill="zero").index_select(0, keep_tensor)
+            original = tensor.detach().cpu().index_select(0, keep_tensor)
+        else:
+            entry = pack_tensor(tensor, bits, group_size=group_size)
+            restored = unpack_tensor(entry)
+            original = tensor.detach().cpu()
+        abs_error, rmse = _tensor_error(original, restored)
+        if abs_error > max_abs_error:
+            max_abs_error = abs_error
+            worst_tensor = name
+        max_rmse = max(max_rmse, rmse)
+        packed[name] = entry
+
+        group = groups.setdefault(
+            component,
+            {"tensors": 0, "source_parameters": 0, "stored_bytes": 0, "bits": bits},
+        )
+        group["tensors"] = int(group["tensors"]) + 1
+        group["source_parameters"] = int(group["source_parameters"]) + tensor.numel()
+        group["stored_bytes"] = int(group["stored_bytes"]) + state_dict_bytes({name: entry})
+        if index % 100 == 0:
+            print(f"  packed through state tensor {index}/{len(state_dict)}: {name}")
+
+    if not packed:
+        raise RuntimeError("source state dict yielded no text tensors")
+    if not any(entry["kind"].endswith("_rows") for entry in packed.values()):
+        raise RuntimeError("embedding row packing was not applied")
+    missing_exact = sorted(set(exact_entries) - set(packed))
+    if missing_exact:
+        raise RuntimeError(f"{len(missing_exact)} captured GPTQ weights were not packed: {missing_exact[:3]}")
+
+    report = {
+        "format": FORMAT,
+        "recipe": recipe_name,
+        "recipe_bits": recipe_bits,
+        "group_size": group_size,
+        "vocab_size": len(keep_ids),
+        "packed_tensors": len(packed),
+        "aliases": len(aliases),
+        "tensor_bytes": state_dict_bytes(packed),
+        "groups": groups,
+        "max_abs_repack_error": max_abs_error,
+        "max_repack_rmse": max_rmse,
+        "worst_tensor": worst_tensor,
+    }
+    payload = {
+        "format": FORMAT,
+        "recipe": recipe_name,
+        "group_size": group_size,
+        "keep_ids": keep_tensor.to(torch.int32),
+        "state_dict": packed,
+        "aliases": aliases,
+        "report": report,
+    }
+    return payload, report
+
+
+def add_input_embedding_extension(
+    payload: dict[str, Any],
+    source_weight: torch.Tensor,
+    input_ids: Iterable[int],
+    *,
+    bits: int,
+    group_size: int = 128,
+) -> dict[str, Any]:
+    """Store extra input-only rows while retaining the smaller output allow-set."""
+    output_ids = set(int(value) for value in payload["keep_ids"].tolist())
+    input_set = set(int(value) for value in input_ids)
+    if not output_ids.issubset(input_set):
+        raise ValueError("input vocabulary must contain every output token id")
+    extra = sorted(input_set - output_ids)
+    if not extra:
+        raise ValueError("input extension has no rows beyond the output vocabulary")
+    row_names = [
+        name
+        for name, entry in payload["state_dict"].items()
+        if entry["kind"] in {"quant_rows", "raw_rows"}
+    ]
+    if len(row_names) != 1:
+        raise ValueError(f"expected one canonical row-packed embedding, found {row_names}")
+    entry = pack_rows(
+        source_weight,
+        torch.tensor(extra, dtype=torch.int64),
+        bits,
+        group_size=group_size,
+    )
+    extension = {
+        "target": row_names[0],
+        "entry": entry,
+        "bits": bits,
+        "output_vocab_size": len(output_ids),
+        "input_vocab_size": len(input_set),
+        "extra_rows": len(extra),
+    }
+    payload["input_embedding_extension"] = extension
+    extension_bytes = state_dict_bytes({"input_embedding_extension": entry})
+    payload["report"]["tensor_bytes"] += extension_bytes
+    payload["report"]["output_vocab_size"] = len(output_ids)
+    payload["report"]["input_vocab_size"] = len(input_set)
+    payload["report"]["input_extension_bits"] = bits
+    payload["report"]["input_extension_rows"] = len(extra)
+    payload["report"]["input_extension_bytes"] = extension_bytes
+    return extension
+
+
+def save_artifact(
+    payload: dict[str, Any],
+    output_path: Path,
+    *,
+    lossless_codec: str = "none",
+    zlib_level: int = 6,
+    lzma_preset: int = 6,
+) -> dict[str, Any]:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    if output_path.exists():
+        raise FileExistsError(f"refusing to overwrite existing artifact: {output_path}")
+    if lossless_codec not in {"none", "zlib", "lzma"}:
+        raise ValueError(f"unsupported lossless artifact codec: {lossless_codec}")
+    raw_disk_bytes = None
+    if lossless_codec == "none":
+        torch.save(payload, output_path)
+    else:
+        temporary_path: Path | None = None
+        try:
+            with tempfile.NamedTemporaryFile(
+                dir=output_path.parent,
+                prefix=f".{output_path.name}.",
+                suffix=".raw",
+                delete=False,
+            ) as handle:
+                temporary_path = Path(handle.name)
+            torch.save(payload, temporary_path)
+            raw_disk_bytes = temporary_path.stat().st_size
+            compressor = (
+                zlib.compressobj(level=zlib_level)
+                if lossless_codec == "zlib"
+                else lzma.LZMACompressor(preset=lzma_preset)
+            )
+            magic = ZLIB_MAGIC if lossless_codec == "zlib" else LZMA_MAGIC
+            with temporary_path.open("rb") as source, output_path.open("xb") as target:
+                target.write(magic)
+                while block := source.read(8 * 1024 * 1024):
+                    target.write(compressor.compress(block))
+                target.write(compressor.flush())
+        except Exception:
+            if output_path.exists():
+                output_path.unlink()
+            raise
+        finally:
+            if temporary_path is not None and temporary_path.exists():
+                temporary_path.unlink()
+    disk_bytes = output_path.stat().st_size
+    report = {
+        **payload["report"],
+        "artifact_path": str(output_path),
+        "disk_bytes": disk_bytes,
+        "lossless_codec": lossless_codec,
+    }
+    if raw_disk_bytes is not None:
+        report.update(
+            {
+                "raw_artifact_bytes": raw_disk_bytes,
+                "lossless_savings_bytes": raw_disk_bytes - disk_bytes,
+                "codec_level": zlib_level if lossless_codec == "zlib" else lzma_preset,
+            }
+        )
+        if lossless_codec == "zlib":
+            report["zlib_level"] = zlib_level
+        else:
+            report["lzma_preset"] = lzma_preset
+    output_path.with_suffix(output_path.suffix + ".json").write_text(json.dumps(report, indent=2))
+    return report
+
+
+def load_artifact(path: str | Path) -> dict[str, Any]:
+    artifact = Path(path)
+    with artifact.open("rb") as source:
+        magic = source.read(len(ZLIB_MAGIC))
+        if magic in {ZLIB_MAGIC, LZMA_MAGIC}:
+            decompressor = zlib.decompressobj() if magic == ZLIB_MAGIC else lzma.LZMADecompressor()
+            # Keep the restored stream beside the artifact so a multi-GB load
+            # cannot unexpectedly exhaust a small system /tmp partition.
+            with tempfile.TemporaryFile(dir=artifact.parent) as restored:
+                while block := source.read(8 * 1024 * 1024):
+                    restored.write(decompressor.decompress(block))
+                if magic == ZLIB_MAGIC:
+                    restored.write(decompressor.flush())
+                if not decompressor.eof:
+                    codec = "zlib" if magic == ZLIB_MAGIC else "lzma"
+                    raise ValueError(f"truncated {codec} artifact: {path}")
+                restored.seek(0)
+                payload = torch.load(restored, map_location="cpu", weights_only=False)
+        else:
+            source.seek(0)
+            payload = torch.load(source, map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict) or payload.get("format") != FORMAT:
+        raise ValueError(f"unsupported artifact format in {path}")
+    return payload
+
+
+def restore_artifact(
+    model_name: str,
+    artifact_path: str | Path,
+    output_path: str | Path,
+    *,
+    embedding_fill: str = "zero",
+    cache_dir: str | None = None,
+) -> dict[str, Any]:
+    """Expand the artifact into the evaluator's required full BF16 checkpoint."""
+    if embedding_fill not in {"zero", "mean", "token", "base"}:
+        raise ValueError("embedding_fill must be zero, mean, token, or base")
+    out = Path(output_path)
+    if out.exists():
+        if not out.is_dir() or any(out.iterdir()):
+            raise FileExistsError(f"refusing to overwrite restore target: {out}")
+
+    payload = load_artifact(artifact_path)
+    started = time.time()
+    model = load_model(
+        model_name,
+        dtype="bfloat16",
+        device_map=None,
+        cache_dir=cache_dir,
+        multimodal=True,
+    )
+    targets = model.state_dict()
+    restored_names: set[str] = set()
+    token_prediction = None
+    if embedding_fill == "token":
+        predictor = payload.get("embedding_predictor")
+        if predictor is None:
+            raise ValueError("artifact has no token-string embedding predictor")
+        tokenizer = load_tokenizer(model_name, cache_dir=cache_dir)
+        token_prediction = predict_token_rows(predictor, tokenizer)
+
+    with torch.no_grad():
+        for name, entry in payload["state_dict"].items():
+            if name not in targets:
+                raise KeyError(f"artifact tensor is absent from base architecture: {name}")
+            if entry["kind"] == "gptq_codes":
+                value = unpack_gptq_entry(entry)
+            elif entry["kind"] == "adaptive_codebook":
+                value = unpack_adaptive_codebook(entry)
+            elif entry["kind"] in {"quant_rows", "raw_rows"}:
+                if embedding_fill == "base":
+                    selected_entry = {
+                        **entry,
+                        "kind": "quant" if entry["kind"] == "quant_rows" else "raw",
+                    }
+                    selected = unpack_tensor(selected_entry)
+                    value = targets[name].detach().cpu().clone()
+                    value.index_copy_(0, entry["row_ids"].long(), selected)
+                elif embedding_fill == "token":
+                    if token_prediction is None or tuple(token_prediction.shape) != tuple(
+                        entry["full_shape"]
+                    ):
+                        raise RuntimeError(
+                            "token prediction shape does not match row-packed embedding: "
+                            f"prediction={None if token_prediction is None else tuple(token_prediction.shape)}, "
+                            f"expected={tuple(entry['full_shape'])}"
+                        )
+                    selected_entry = {
+                        **entry,
+                        "kind": "quant" if entry["kind"] == "quant_rows" else "raw",
+                    }
+                    selected = unpack_tensor(selected_entry)
+                    value = token_prediction.clone()
+                    value.index_copy_(0, entry["row_ids"].long(), selected)
+                else:
+                    value = unpack_rows(entry, fill=embedding_fill)
+            else:
+                value = unpack_tensor(entry)
+            extension = payload.get("input_embedding_extension")
+            if extension is not None and extension.get("target") == name:
+                extension_entry = extension["entry"]
+                dense_entry = {
+                    **extension_entry,
+                    "kind": "quant"
+                    if extension_entry["kind"] == "quant_rows"
+                    else "raw",
+                }
+                extra_values = unpack_tensor(dense_entry)
+                value.index_copy_(0, extension_entry["row_ids"].long(), extra_values)
+            targets[name].copy_(value.to(dtype=targets[name].dtype))
+            restored_names.add(name)
+
+        for alias, canonical in payload["aliases"].items():
+            if alias not in targets or canonical not in targets:
+                raise KeyError(f"artifact alias cannot be resolved: {alias} -> {canonical}")
+            if targets[alias].data_ptr() != targets[canonical].data_ptr():
+                targets[alias].copy_(targets[canonical])
+            restored_names.add(alias)
+
+    missing_text = [
+        name
+        for name in targets
+        if not EXCLUDED.search(name) and name not in restored_names
+    ]
+    if missing_text:
+        preview = ", ".join(missing_text[:8])
+        raise RuntimeError(f"artifact did not restore {len(missing_text)} text tensors: {preview}")
+
+    copied = save_checkpoint(model, out, source_model=model_name, cache_dir=cache_dir)
+    keep = set(int(token_id) for token_id in payload["keep_ids"].tolist())
+    original_vocab = max(
+        int(entry["full_shape"][0])
+        for entry in payload["state_dict"].values()
+        if entry["kind"] in {"quant_rows", "raw_rows"}
+    )
+    suppressed = [token_id for token_id in range(original_vocab) if token_id not in keep]
+    try:
+        generation_config = model.generation_config
+        generation_config.suppress_tokens = suppressed
+        generation_config.save_pretrained(out)
+    except Exception as exc:  # noqa: BLE001 - explicit manifest still records the failure
+        raise RuntimeError(f"could not save suppress_tokens generation config: {exc}") from exc
+
+    report = {
+        "format": FORMAT,
+        "artifact": str(artifact_path),
+        "model_name": model_name,
+        "embedding_fill": embedding_fill,
+        "submission_valid": embedding_fill in {"zero", "mean", "token"},
+        "restored_text_tensors": len(restored_names),
+        "suppressed_tokens": len(suppressed),
+        "kept_tokens": len(keep),
+        "input_vocab_size": payload.get("report", {}).get("input_vocab_size", len(keep)),
+        "auxiliary_files": copied,
+        "restore_seconds": round(time.time() - started, 1),
+    }
+    (out / "artifact_restore.json").write_text(json.dumps(report, indent=2))
+    (out / "keep_ids.json").write_text(json.dumps(sorted(keep)))
+    return report
+
+
+def default_model() -> str:
+    """Small indirection retained for the standalone submission wrapper."""
+    return DEFAULT_MODEL

src/eaimath/artifact_refine.py

@@ -0,0 +1,248 @@
+"""Physical code-grid refinement helpers for the Bucket-B artifact."""
+
+from __future__ import annotations
+
+import copy
+import math
+import re
+from typing import Any
+
+import numpy as np
+import torch
+
+from .artifact import unpack_gptq_entry
+from .pack import pack_bits, unpack_bits, unpack_tensor
+
+MLP = re.compile(r"\.mlp\.")
+GDN = re.compile(r"linear_attn")
+FULL_ATTN = re.compile(r"self_attn")
+
+
+def resolved_recipe_from_report(report: dict[str, Any]) -> str:
+    """Recover the exact payload recipe from a physical artifact sidecar.
+
+    Early Round-21 sidecars inherited the base report's recipe even though the
+    payload recipe correctly included its refinement suffix.  The refinement
+    fields are complete, so existing artifacts can be validated without being
+    rebuilt or decompressed an extra time.
+    """
+    recipe = str(report.get("recipe") or "")
+    refinement = report.get("refinement")
+    if refinement == "frozen_grid_adapter_projection_v1":
+        if "+grid-" in recipe:
+            return recipe
+        profile = report.get("refinement_profile")
+        alpha = report.get("refinement_alpha")
+        if not recipe or profile not in {"all", "attention", "gdn"} or alpha is None:
+            raise ValueError("incomplete frozen-grid recipe metadata")
+        return f"{recipe}+grid-{profile}-a{float(alpha):g}"
+    if refinement == "dual_adaptive_codebook_v1":
+        if "+aaac-" in recipe:
+            return recipe
+        mode = report.get("refinement_mode")
+        if not recipe or mode not in {"activation", "weight"}:
+            raise ValueError("incomplete adaptive-codebook recipe metadata")
+        return f"{recipe}+aaac-gdn-{mode}"
+    if refinement == "stochastic_grid_adapter_projection_v1":
+        if "+stochastic-" in recipe:
+            return recipe
+        profile = report.get("refinement_profile")
+        gain = report.get("refinement_gain")
+        seed = report.get("refinement_seed")
+        if not recipe or profile not in {"all", "attention", "gdn"} or gain is None or seed is None:
+            raise ValueError("incomplete stochastic-grid recipe metadata")
+        return f"{recipe}+stochastic-{profile}-g{float(gain):g}-s{int(seed)}"
+    if not recipe:
+        raise ValueError("artifact sidecar has no recipe")
+    return recipe
+
+
+def component(name: str) -> str | None:
+    if MLP.search(name):
+        return "mlp"
+    if GDN.search(name):
+        return "gdn"
+    if FULL_ATTN.search(name):
+        return "full_attn"
+    return None
+
+
+def profile_matches(name: str, profile: str) -> bool:
+    value = component(name)
+    if profile == "all":
+        return value is not None
+    if profile == "attention":
+        return value in {"gdn", "full_attn"}
+    if profile == "gdn":
+        return value == "gdn"
+    raise ValueError(f"unknown projection profile: {profile}")
+
+
+def _gptq_codes_and_zeros(entry: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    in_features, out_features = tuple(entry["code_shape"])
+    count = in_features * out_features
+    codes = unpack_bits(entry["packed"].numpy(), entry["bits"], count)
+    codes_tensor = torch.from_numpy(codes.astype(np.int64)).reshape(in_features, out_features)
+    scales = entry["scale"].float()
+    if entry.get("g_idx") is None:
+        g_idx = torch.div(
+            torch.arange(in_features, dtype=torch.int64),
+            int(entry["group"]),
+            rounding_mode="floor",
+        ).clamp_max(scales.shape[0] - 1)
+    else:
+        g_idx = entry["g_idx"].long()
+    if entry.get("zero_value") is not None:
+        zeros = torch.full_like(scales, int(entry["zero_value"]), dtype=torch.int64)
+    else:
+        zero_shape = tuple(entry["zero_shape"])
+        zero_count = int(np.prod(zero_shape))
+        values = unpack_bits(entry["zero_packed"].numpy(), entry["bits"], zero_count)
+        zeros = torch.from_numpy(values.astype(np.int64)).reshape(zero_shape)
+    return codes_tensor, scales.index_select(0, g_idx), zeros.index_select(0, g_idx)
+
+
+def project_gptq_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    base = unpack_gptq_entry(entry).float()
+    target = base + alpha * (adapted.detach().float().cpu() - base)
+    old_codes, scale_rows, zero_rows = _gptq_codes_and_zeros(entry)
+    maxq = 2 ** int(entry["bits"]) - 1
+    new_codes = torch.round(target.T / scale_rows.clamp_min(1e-12) + zero_rows).clamp(0, maxq)
+    new_codes = new_codes.to(torch.int64)
+    projected = (scale_rows * (new_codes - zero_rows).float()).T
+    changed = new_codes.ne(old_codes)
+    updated = copy.copy(entry)
+    updated["packed"] = torch.from_numpy(
+        pack_bits(new_codes.numpy().astype(np.uint32), int(entry["bits"]))
+    )
+    return updated, {
+        "weights": int(new_codes.numel()),
+        "changed_codes": int(changed.sum()),
+        "target_delta_rmse": float((target - base).square().mean().sqrt()),
+        "projection_rmse": float((target - projected).square().mean().sqrt()),
+    }
+
+
+def project_quant_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    base = unpack_tensor(entry).float()
+    target = base + alpha * (adapted.detach().float().cpu() - base)
+    bits, group = int(entry["bits"]), int(entry["group"])
+    count = int(np.prod(entry["shape"]))
+    old = unpack_bits(entry["packed"].numpy(), bits, count)
+    old_codes = torch.from_numpy(old.astype(np.int64)).reshape(-1, group)
+    scale = entry["scale"].float().reshape(-1, 1)
+    scale = scale * math.ldexp(1.0, -int(entry.get("scale_exponent", 0)))
+    flat = target.reshape(-1, group)
+    if entry["symmetric"]:
+        qmax = 2 ** (bits - 1) - 1
+        new_codes = torch.round(flat / scale.clamp_min(1e-12)) + qmax + 1
+        new_codes = new_codes.clamp(0, 2**bits - 1).to(torch.int64)
+        projected = (new_codes - qmax - 1).float() * scale
+    else:
+        zero = entry["zero"].float().reshape(-1, 1)
+        new_codes = (torch.round(flat / scale.clamp_min(1e-12)) + zero).clamp(0, 2**bits - 1)
+        new_codes = new_codes.to(torch.int64)
+        projected = (new_codes.float() - zero) * scale
+    updated = copy.copy(entry)
+    updated["packed"] = torch.from_numpy(pack_bits(new_codes.numpy().astype(np.uint32), bits))
+    return updated, {
+        "weights": int(new_codes.numel()),
+        "changed_codes": int(new_codes.ne(old_codes).sum()),
+        "target_delta_rmse": float((target - base).square().mean().sqrt()),
+        "projection_rmse": float((target - projected.reshape_as(target)).square().mean().sqrt()),
+    }
+
+
+def project_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    if entry.get("kind") == "gptq_codes":
+        return project_gptq_entry(entry, adapted, alpha)
+    if entry.get("kind") == "quant":
+        return project_quant_entry(entry, adapted, alpha)
+    raise ValueError(f"entry kind cannot be grid-projected: {entry.get('kind')!r}")
+
+
+def _stochastic_integer_round(desired: torch.Tensor, *, generator: torch.Generator) -> torch.Tensor:
+    """Unbiased stochastic rounding to integers with a deterministic generator."""
+    lower = torch.floor(desired)
+    probability = desired - lower
+    draw = torch.rand(desired.shape, generator=generator, dtype=torch.float32)
+    return lower + draw.lt(probability)
+
+
+def stochastic_project_gptq_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float, seed: int
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    base = unpack_gptq_entry(entry).float()
+    delta = adapted.detach().float().cpu() - base
+    target = base + alpha * delta
+    old_codes, scale_rows, zero_rows = _gptq_codes_and_zeros(entry)
+    desired = old_codes.float() + alpha * delta.T / scale_rows.clamp_min(1e-12)
+    generator = torch.Generator().manual_seed(seed)
+    maxq = 2 ** int(entry["bits"]) - 1
+    new_codes = _stochastic_integer_round(desired, generator=generator)
+    new_codes = new_codes.clamp(0, maxq).to(torch.int64)
+    projected = (scale_rows * (new_codes - zero_rows).float()).T
+    changed = new_codes.ne(old_codes)
+    updated = copy.copy(entry)
+    updated["packed"] = torch.from_numpy(
+        pack_bits(new_codes.numpy().astype(np.uint32), int(entry["bits"]))
+    )
+    return updated, {
+        "weights": int(new_codes.numel()),
+        "changed_codes": int(changed.sum()),
+        "absolute_code_steps": int((new_codes - old_codes).abs().sum()),
+        "target_delta_rmse": float((target - base).square().mean().sqrt()),
+        "projection_rmse": float((target - projected).square().mean().sqrt()),
+    }
+
+
+def stochastic_project_quant_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float, seed: int
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    base = unpack_tensor(entry).float()
+    delta = adapted.detach().float().cpu() - base
+    target = base + alpha * delta
+    bits, group = int(entry["bits"]), int(entry["group"])
+    count = int(np.prod(entry["shape"]))
+    old = unpack_bits(entry["packed"].numpy(), bits, count)
+    old_codes = torch.from_numpy(old.astype(np.int64)).reshape(-1, group)
+    scale = entry["scale"].float().reshape(-1, 1)
+    scale = scale * math.ldexp(1.0, -int(entry.get("scale_exponent", 0)))
+    flat_delta = delta.reshape(-1, group)
+    if entry["symmetric"]:
+        qmax = 2 ** (bits - 1) - 1
+        zero = qmax + 1
+    else:
+        zero = entry["zero"].float().reshape(-1, 1)
+    desired = old_codes.float() + alpha * flat_delta / scale.clamp_min(1e-12)
+    generator = torch.Generator().manual_seed(seed)
+    new_codes = _stochastic_integer_round(desired, generator=generator)
+    new_codes = new_codes.clamp(0, 2**bits - 1).to(torch.int64)
+    projected = (new_codes.float() - zero) * scale
+    changed = new_codes.ne(old_codes)
+    updated = copy.copy(entry)
+    updated["packed"] = torch.from_numpy(pack_bits(new_codes.numpy().astype(np.uint32), bits))
+    return updated, {
+        "weights": int(new_codes.numel()),
+        "changed_codes": int(changed.sum()),
+        "absolute_code_steps": int((new_codes - old_codes).abs().sum()),
+        "target_delta_rmse": float((target - base).square().mean().sqrt()),
+        "projection_rmse": float((target - projected.reshape_as(target)).square().mean().sqrt()),
+    }
+
+
+def stochastic_project_entry(
+    entry: dict[str, Any], adapted: torch.Tensor, alpha: float, seed: int
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    """Stochastically preserve a sub-grid learned update in physical codes."""
+    if entry.get("kind") == "gptq_codes":
+        return stochastic_project_gptq_entry(entry, adapted, alpha, seed)
+    if entry.get("kind") == "quant":
+        return stochastic_project_quant_entry(entry, adapted, alpha, seed)
+    raise ValueError(f"entry kind cannot be grid-projected: {entry.get('kind')!r}")

src/eaimath/buckets.py

@@ -0,0 +1,91 @@
+"""Submission buckets — three per week, scored on accuracy *within* a band.
+
+The leaderboard ranks by accuracy inside each compression band, not on the
+size/accuracy frontier. That inverts the optimization: there is **no reward for
+being smaller than your bucket allows**, so every submission should sit as close
+to its upper size limit as it can. A 2.0 GB model and a 3.3 GB model compete on
+equal terms in bucket C; the 3.3 GB one simply gets more bits to spend.
+
+Three bands × two tracks = six submissions selected at the end of each week.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Conservative denominator: the text-only weight budget we actually compress.
+# If the graders measure against the full HF repo (9.31 GB, vision tower
+# included) every limit below is ~11% more generous -- worth confirming, but
+# aiming at the tighter number can only be safe.
+ORIGINAL_GB = 8.412
+ORIGINAL_GB_FULL_REPO = 9.313
+
+
+@dataclass(frozen=True)
+class Bucket:
+    name: str
+    label: str
+    low_pct: float
+    high_pct: float
+
+    @property
+    def low_gb(self) -> float:
+        return self.low_pct / 100 * ORIGINAL_GB
+
+    @property
+    def high_gb(self) -> float:
+        return self.high_pct / 100 * ORIGINAL_GB
+
+    @property
+    def target_gb(self) -> float:
+        """Where a submission should aim: just under the ceiling."""
+        return self.high_gb * 0.98
+
+    @property
+    def bits_per_weight(self) -> tuple[float, float]:
+        return (16 * self.low_pct / 100, 16 * self.high_pct / 100)
+
+
+BUCKETS: tuple[Bucket, ...] = (
+    Bucket("A", "≤10%", 0.0, 10.0),
+    Bucket("B", "10–20%", 10.0, 20.0),
+    Bucket("C", "20–40%", 20.0, 40.0),
+)
+
+BY_NAME = {b.name: b for b in BUCKETS}
+
+
+def classify(size_gb: float | None) -> Bucket | None:
+    """Which bucket a checkpoint qualifies for, or None if it is too large."""
+    if size_gb is None:
+        return None
+    pct = 100 * size_gb / ORIGINAL_GB
+    for bucket in BUCKETS:
+        if bucket.low_pct <= pct < bucket.high_pct:
+            return bucket
+    return None
+
+
+def pct_of_original(size_gb: float) -> float:
+    return 100 * size_gb / ORIGINAL_GB
+
+
+def describe() -> str:
+    lines = [
+        f"{'bucket':<8}{'band':<10}{'size range':>18}{'bits/weight':>14}{'aim for':>10}",
+        "-" * 60,
+    ]
+    for b in BUCKETS:
+        lo, hi = b.bits_per_weight
+        lines.append(
+            f"{b.name:<8}{b.label:<10}{f'{b.low_gb:.2f}–{b.high_gb:.2f} GB':>18}"
+            f"{f'{lo:.2f}–{hi:.2f}':>14}{b.target_gb:>9.2f} GB"
+        )
+    lines.append("")
+    lines.append(
+        f"Scored on accuracy within a band, so aim at the ceiling — being "
+        f"smaller earns nothing.\nDenominator: {ORIGINAL_GB} GB (text-only). "
+        f"If graders use the full repo ({ORIGINAL_GB_FULL_REPO} GB) every limit "
+        f"is ~11% looser."
+    )
+    return "\n".join(lines)

src/eaimath/data.py

@@ -0,0 +1,304 @@
+"""Math evaluation sets and calibration data.
+
+The leaderboard set is hidden and probably postdates the model, so we keep two
+tiers deliberately separate:
+
+* **gate**  — small, fast, run on every recipe. Cheap signal for iteration.
+* **holdout** — recent competitions we never tune against. The honest estimate.
+
+AIME 2024 is deliberately excluded from the holdout: it is measurably
+contaminated (inflating scores 10-20 points over clean contests), so it flatters
+every recipe equally and discriminates between none of them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Iterable, Sequence
+
+
+@dataclass
+class MathExample:
+    example_id: str
+    problem: str
+    answer: str
+    source: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class DatasetSpec:
+    """How to pull one benchmark off the Hub.
+
+    Field names differ between mirrors of the same benchmark, so each role lists
+    candidate column names tried in order.
+    """
+
+    name: str
+    hf_id: str
+    split: str = "test"
+    config: str | None = None
+    problem_fields: Sequence[str] = ("problem", "Problem", "question", "Question")
+    answer_fields: Sequence[str] = ("answer", "Answer", "solution", "expected_answer")
+    tier: str = "gate"
+    filters: tuple[tuple[str, str, Any], ...] = ()
+    max_examples: int | None = None
+    note: str = ""
+
+
+# Three tiers, by how often we run them and how much tuning pressure they can
+# absorb before their numbers stop meaning anything.
+#
+#   gate       every experiment. Cheap, tuned against freely.
+#   checkpoint before a weekly leaderboard submission. Moderate tuning risk.
+#   holdout    the two graded checkpoints only. NEVER tuned against — these are
+#              post-release contests and the closest proxy we have for a hidden
+#              eval that "is not available in public domain today".
+REGISTRY: dict[str, DatasetSpec] = {
+    "math500_hard": DatasetSpec(
+        name="math500_hard",
+        hf_id="HuggingFaceH4/MATH-500",
+        split="test",
+        tier="gate",
+        filters=(("level", "gte", 4),),
+        max_examples=100,
+        note="MATH-500 levels 4-5, deterministic 100-problem subsample. The "
+        "fast regression signal: sensitive enough to catch damage, cheap "
+        "enough to run on every recipe.",
+    ),
+    "math500": DatasetSpec(
+        name="math500",
+        hf_id="HuggingFaceH4/MATH-500",
+        split="test",
+        tier="checkpoint",
+        note="Full 500. Largely saturated for this model (~84.5 bf16), so it "
+        "detects collapse but not subtle reasoning damage.",
+    ),
+    "aime25": DatasetSpec(
+        name="aime25",
+        hf_id="MathArena/aime_2025",
+        split="train",
+        tier="checkpoint",
+        note="30 problems. Hard tail — where quantization damage actually shows.",
+    ),
+    "hmmt_feb25": DatasetSpec(
+        name="hmmt_feb25",
+        hf_id="MathArena/hmmt_feb_2025",
+        split="train",
+        tier="checkpoint",
+        note="30 problems. Reported on the model card (74.0), so we have a "
+        "published bf16 reference to validate our harness against.",
+    ),
+    "aime26": DatasetSpec(
+        name="aime26",
+        hf_id="MathArena/aime_2026",
+        split="train",
+        tier="holdout",
+        note="30 problems, Feb 2026 contest. Post-dates most training data.",
+    ),
+    "hmmt_feb26": DatasetSpec(
+        name="hmmt_feb26",
+        hf_id="MathArena/hmmt_feb_2026",
+        split="train",
+        tier="holdout",
+        note="33 problems, Feb 2026 contest. Cleanest proxy for the hidden eval.",
+    ),
+    "aime24": DatasetSpec(
+        name="aime24",
+        hf_id="Maxwell-Jia/AIME_2024",
+        split="train",
+        tier="diagnostic",
+        note="Measurably contaminated — inflates scores 10-20 points over clean "
+        "contests. Diagnostic only, never for recipe selection.",
+    ),
+}
+
+# Training pools — problems with known answers, used to generate our own
+# reasoning traces. Disjoint from every eval set above: MATH-500 is drawn from
+# the MATH *test* split, so the MATH train split cannot leak into it.
+TRAIN_REGISTRY: dict[str, DatasetSpec] = {
+    "math_train": DatasetSpec(
+        name="math_train",
+        hf_id="EleutherAI/hendrycks_math",
+        split="train",
+        config="algebra",
+        answer_fields=("solution",),  # gold answer is the \boxed{} in the solution
+        tier="train",
+        note="MATH train split. Pass --config to pick a subject.",
+    ),
+    "openr1": DatasetSpec(
+        name="openr1",
+        hf_id="open-r1/OpenR1-Math-220k",
+        split="train",
+        answer_fields=("answer", "solution"),
+        tier="train",
+        note="220k competition problems with verified answers.",
+    ),
+}
+
+MATH_SUBJECTS = (
+    "algebra", "counting_and_probability", "geometry", "intermediate_algebra",
+    "number_theory", "prealgebra", "precalculus",
+)
+
+SUITES: dict[str, list[str]] = {
+    "gate": ["math500_hard"],
+    "checkpoint": ["math500", "aime25", "hmmt_feb25"],
+    "holdout": ["aime26", "hmmt_feb26"],
+}
+
+
+def _passes_filters(row: dict[str, Any], filters: Sequence[tuple[str, str, Any]]) -> bool:
+    for field_name, op, value in filters:
+        actual = row.get(field_name)
+        if actual is None:
+            return False
+        if op == "gte" and not actual >= value:
+            return False
+        if op == "lte" and not actual <= value:
+            return False
+        if op == "eq" and actual != value:
+            return False
+        if op == "in" and actual not in value:
+            return False
+    return True
+
+
+def _subsample(examples: list[MathExample], n: int, seed: int = 0) -> list[MathExample]:
+    """Deterministic subsample, stable across runs and machines.
+
+    Shuffles with a fixed seed rather than taking a prefix, because these sets
+    are ordered by subject/difficulty and a prefix would be badly skewed. The
+    gate set must be identical across every recipe or the comparison is
+    meaningless.
+    """
+    import random
+
+    if len(examples) <= n:
+        return examples
+    indices = sorted(range(len(examples)))
+    random.Random(seed).shuffle(indices)
+    return [examples[i] for i in sorted(indices[:n])]
+
+
+def _resolve_field(row: dict[str, Any], candidates: Iterable[str]) -> str | None:
+    lowered = {k.lower(): k for k in row}
+    for candidate in candidates:
+        key = lowered.get(candidate.lower())
+        if key is not None and row[key] is not None:
+            return str(row[key])
+    return None
+
+
+def load_dataset_examples(
+    spec: DatasetSpec | str,
+    limit: int | None = None,
+    cache_dir: str | None = None,
+) -> list[MathExample]:
+    """Load one benchmark into ``MathExample`` records.
+
+    Raises with the observed column names when a field cannot be resolved, so a
+    schema change on the Hub produces an actionable error instead of silently
+    empty problems.
+    """
+    from datasets import load_dataset
+
+    if isinstance(spec, str):
+        table = {**REGISTRY, **TRAIN_REGISTRY}
+        if spec not in table:
+            raise KeyError(f"Unknown dataset {spec!r}. Known: {sorted(table)}")
+        spec = table[spec]
+
+    kwargs: dict[str, Any] = {"split": spec.split}
+    if spec.config:
+        kwargs["name"] = spec.config
+    if cache_dir:
+        kwargs["cache_dir"] = cache_dir
+
+    dataset = load_dataset(spec.hf_id, **kwargs)
+
+    examples: list[MathExample] = []
+    for i, row in enumerate(dataset):
+        if not _passes_filters(row, spec.filters):
+            continue
+        problem = _resolve_field(row, spec.problem_fields)
+        answer = _resolve_field(row, spec.answer_fields)
+        if problem is None or answer is None:
+            raise ValueError(
+                f"{spec.name}: could not resolve problem/answer fields. "
+                f"Available columns: {sorted(row)}. "
+                f"Tried problem={list(spec.problem_fields)}, answer={list(spec.answer_fields)}."
+            )
+        if "\\boxed" in answer:
+            from .answers import extract_boxed
+
+            boxed = extract_boxed(answer)
+            if boxed is None:
+                continue  # unparseable gold: drop rather than train on it
+            answer = boxed
+
+        examples.append(
+            MathExample(
+                example_id=f"{spec.name}:{i}",
+                problem=problem,
+                answer=answer,
+                source=spec.name,
+                metadata={
+                    k: row[k]
+                    for k in ("level", "subject", "type", "url", "id", "problem_idx")
+                    if k in row
+                },
+            )
+        )
+
+    # Spec cap first (defines the canonical set), then the ad-hoc --limit.
+    if spec.max_examples is not None:
+        examples = _subsample(examples, spec.max_examples)
+    if limit is not None:
+        examples = examples[:limit]
+    return examples
+
+
+def load_suite(
+    names: Sequence[str],
+    limit: int | None = None,
+    cache_dir: str | None = None,
+) -> list[MathExample]:
+    """Load and concatenate several benchmarks. ``limit`` applies per dataset.
+
+    Accepts tier names (``gate``/``checkpoint``/``holdout``) as shorthand for
+    the datasets in that tier.
+    """
+    resolved: list[str] = []
+    for name in names:
+        resolved.extend(SUITES[name] if name in SUITES else [name])
+
+    out: list[MathExample] = []
+    for name in resolved:
+        out.extend(load_dataset_examples(name, limit=limit, cache_dir=cache_dir))
+    return out
+
+
+def describe_registry() -> str:
+    lines = []
+    for tier in ("gate", "checkpoint", "holdout", "diagnostic"):
+        members = [s for s in REGISTRY.values() if s.tier == tier]
+        if not members:
+            continue
+        lines.append(f"[{tier}]")
+        for spec in members:
+            cap = f" (capped at {spec.max_examples})" if spec.max_examples else ""
+            lines.append(f"  {spec.name:<14} {spec.hf_id}{cap}")
+            lines.append(f"  {'':<14} {spec.note}")
+    return "\n".join(lines)
+
+
+MATH_PROMPT = (
+    "Solve the following math problem. Put your final answer inside "
+    "\\boxed{{}} on the last line.\n\n"
+    "Problem:\n{problem}"
+)
+
+
+def build_prompt(example: MathExample) -> str:
+    return MATH_PROMPT.format(problem=example.problem.strip())

src/eaimath/delta_sharing.py

@@ -0,0 +1,144 @@
+"""Cross-layer MLP sharing with low-rank, layer-specific weight deltas.
+
+Hard tying erased every layer's identity and scored zero in Round 8.  This
+module keeps one full MLP per contiguous Middle-Cycle group and approximates
+each follower as ``shared_base + U @ V.T``.  It is the directly testable
+initialization used by DeltaLLM-style recovery and the same broad idea as Basis
+Sharing: shared storage plus small layer-specific coefficients.
+
+The evaluation checkpoint expands the deltas back to ordinary dense weights so
+vLLM needs no custom kernel.  The report separately counts the compressed
+representation; a submission decompressor would perform the same expansion.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from .sharing import MLP_PROJECTIONS, find_decoder_layers, sharing_groups
+
+
+def low_rank_delta(
+    anchor: torch.Tensor,
+    target: torch.Tensor,
+    rank: int,
+    *,
+    device: str,
+    niter: int = 1,
+    seed: int = 0,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Return ``anchor + rank-r(target-anchor)`` and reconstruction metrics."""
+    if rank < 1 or rank > min(target.shape):
+        raise ValueError(f"rank {rank} is invalid for matrix {tuple(target.shape)}")
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    base = anchor.detach().to(device=device, dtype=torch.float32)
+    wanted = target.detach().to(device=device, dtype=torch.float32)
+    delta = wanted - base
+    delta_norm = torch.linalg.vector_norm(delta)
+    target_norm = torch.linalg.vector_norm(wanted)
+
+    # Oversampling materially improves randomized SVD at almost no storage cost;
+    # only the requested rank is retained in the representation.
+    q = min(rank + 8, min(delta.shape))
+    u, s, v = torch.svd_lowrank(delta, q=q, niter=niter)
+    approximation = (u[:, :rank] * s[:rank]) @ v[:, :rank].T
+    residual = delta - approximation
+    residual_norm = torch.linalg.vector_norm(residual)
+    restored = (base + approximation).to(dtype=target.dtype, device="cpu")
+
+    metrics = {
+        "relative_delta_error": float(residual_norm / delta_norm.clamp_min(1e-12)),
+        "relative_weight_error": float(residual_norm / target_norm.clamp_min(1e-12)),
+        "delta_norm": float(delta_norm),
+        "residual_norm": float(residual_norm),
+    }
+    del base, wanted, delta, approximation, residual, u, s, v
+    if str(device).startswith("cuda"):
+        torch.cuda.empty_cache()
+    return restored, metrics
+
+
+def apply_delta_sharing(
+    model: torch.nn.Module,
+    factor: int,
+    rank: int,
+    *,
+    device: str,
+    niter: int = 1,
+    seed: int = 0,
+) -> dict[str, Any]:
+    """Approximate follower MLPs in place and count the compressed form."""
+    layers = find_decoder_layers(model)
+    groups = sharing_groups(len(layers), factor)
+    dense_mlp_params = sum(
+        getattr(layer.mlp, projection).weight.numel()
+        for layer in layers
+        for projection in MLP_PROJECTIONS
+    )
+    basis_layers: set[int] = {0, 1, len(layers) - 1}
+    follower_layers: set[int] = set()
+    rows: list[dict[str, Any]] = []
+    delta_params = 0
+
+    with torch.no_grad():
+        for group_index, group in enumerate(groups):
+            if len(group) < 2:
+                basis_layers.update(group)
+                continue
+            anchor_index = group[0]
+            basis_layers.add(anchor_index)
+            for follower_index in group[1:]:
+                follower_layers.add(follower_index)
+                for projection_index, projection in enumerate(MLP_PROJECTIONS):
+                    anchor = getattr(layers[anchor_index].mlp, projection).weight
+                    follower = getattr(layers[follower_index].mlp, projection).weight
+                    restored, metrics = low_rank_delta(
+                        anchor,
+                        follower,
+                        rank,
+                        device=device,
+                        niter=niter,
+                        seed=seed + group_index * 10_000 + follower_index * 10 + projection_index,
+                    )
+                    follower.copy_(restored)
+                    delta_params += rank * (follower.shape[0] + follower.shape[1])
+                    rows.append(
+                        {
+                            "group": group_index,
+                            "anchor_layer": anchor_index,
+                            "follower_layer": follower_index,
+                            "projection": projection,
+                            "shape": list(follower.shape),
+                            **metrics,
+                        }
+                    )
+
+    basis_params = sum(
+        getattr(layers[index].mlp, projection).weight.numel()
+        for index in sorted(basis_layers)
+        for projection in MLP_PROJECTIONS
+    )
+    represented = basis_params + delta_params
+    return {
+        "factor": factor,
+        "rank": rank,
+        "niter": niter,
+        "groups": [group for group in groups if len(group) > 1],
+        "basis_layers": sorted(basis_layers),
+        "follower_layers": sorted(follower_layers),
+        "dense_mlp_params": dense_mlp_params,
+        "basis_params": basis_params,
+        "delta_params": delta_params,
+        "represented_mlp_params": represented,
+        "mlp_parameter_fraction": represented / dense_mlp_params,
+        "bf16_representation_bytes": represented * 2,
+        "max_relative_weight_error": max(row["relative_weight_error"] for row in rows),
+        "mean_relative_weight_error": sum(row["relative_weight_error"] for row in rows) / len(rows),
+        "max_relative_delta_error": max(row["relative_delta_error"] for row in rows),
+        "rows": rows,
+    }

src/eaimath/embedding_predictor.py

@@ -0,0 +1,129 @@
+"""Tiny deterministic predictor for vocabulary rows omitted from an artifact.
+
+The predictor borrows the idea of intra-frame prediction from codecs: a token's
+string is side information already present in the tokenizer, so only a small
+linear dictionary has to be stored. Exact retained rows are scattered over the
+prediction during restoration.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+
+from .pack import pack_tensor, state_dict_bytes, unpack_tensor
+
+FEATURE_DIM = 272
+
+
+def _token_strings(tokenizer: Any, ids: list[int]) -> list[str]:
+    tokens = tokenizer.convert_ids_to_tokens(ids)
+    if isinstance(tokens, str):
+        tokens = [tokens]
+    return ["" if token is None else str(token) for token in tokens]
+
+
+def token_string_features(tokenizer: Any, ids: torch.Tensor) -> torch.Tensor:
+    """Return tokenizer-derived features without storing anything per token.
+
+    Layout: 256 normalized UTF-8 byte counts, eight log-length buckets, four
+    first-byte classes and four last-byte classes. The last two blocks retain a
+    little order information while keeping the dictionary under one megabyte.
+    """
+    flat_ids = ids.detach().cpu().long().reshape(-1)
+    strings = _token_strings(tokenizer, [int(value) for value in flat_ids.tolist()])
+    features = torch.zeros((len(strings), FEATURE_DIM), dtype=torch.float32)
+    for row, token in enumerate(strings):
+        encoded = token.encode("utf-8", errors="replace") or b"\x00"
+        byte_ids = torch.tensor(list(encoded), dtype=torch.int64)
+        counts = torch.bincount(byte_ids, minlength=256).float()
+        features[row, :256] = counts / counts.square().sum().sqrt().clamp_min(1.0)
+        length_bucket = min(7, int(math.log2(max(1, len(encoded)))))
+        features[row, 256 + length_bucket] = 1.0
+        features[row, 264 + encoded[0] // 64] = 1.0
+        features[row, 268 + encoded[-1] // 64] = 1.0
+    return features
+
+
+def _sample_ids(vocab_size: int, sample_size: int, offset: float = 0.0) -> torch.Tensor:
+    count = min(vocab_size, sample_size)
+    if count == vocab_size:
+        return torch.arange(vocab_size, dtype=torch.int64)
+    positions = (torch.arange(count, dtype=torch.float64) + offset) * vocab_size / count
+    return positions.floor().clamp_max(vocab_size - 1).long().unique()
+
+
+def fit_token_predictor(
+    weight: torch.Tensor,
+    tokenizer: Any,
+    *,
+    sample_size: int = 65_536,
+    heldout_size: int = 4_096,
+    ridge: float = 1e-2,
+    device: str = "cpu",
+) -> tuple[dict[str, Any], dict[str, float | int | str]]:
+    """Fit and INT8-pack a ridge dictionary for a full embedding matrix."""
+    if weight.ndim != 2 or weight.shape[0] < 2:
+        raise ValueError(f"expected a vocabulary matrix, found {tuple(weight.shape)}")
+    if ridge <= 0:
+        raise ValueError("ridge must be positive")
+    train_ids = _sample_ids(weight.shape[0], sample_size)
+    x = token_string_features(tokenizer, train_ids).to(device)
+    y = weight.detach().cpu().index_select(0, train_ids).float().to(device)
+    gram = x.T @ x
+    gram.diagonal().add_(ridge)
+    basis = torch.linalg.solve(gram, x.T @ y).cpu()
+
+    heldout_ids = _sample_ids(weight.shape[0], heldout_size, offset=0.5)
+    # Exclude any collision with the evenly spaced training grid.
+    train_set = set(int(value) for value in train_ids.tolist())
+    heldout_list = [int(value) for value in heldout_ids.tolist() if int(value) not in train_set]
+    if not heldout_list:
+        heldout_list = [int(train_ids[-1])]
+    heldout_ids = torch.tensor(heldout_list, dtype=torch.int64)
+    hx = token_string_features(tokenizer, heldout_ids)
+    target = weight.detach().cpu().index_select(0, heldout_ids).float()
+    prediction = hx @ basis
+    cosine = torch.nn.functional.cosine_similarity(prediction, target, dim=1).mean()
+    denominator = target.square().mean().sqrt().clamp_min(1e-12)
+    relative_rmse = (prediction - target).square().mean().sqrt() / denominator
+
+    packed_basis = pack_tensor(basis, bits=8, group_size=256)
+    entry: dict[str, Any] = {
+        "kind": "token_string_ridge_v1",
+        "feature_dim": FEATURE_DIM,
+        "basis": packed_basis,
+        "vocab_size": int(weight.shape[0]),
+        "hidden_size": int(weight.shape[1]),
+        "ridge": float(ridge),
+        "sample_size": int(train_ids.numel()),
+    }
+    report: dict[str, float | int | str] = {
+        "kind": entry["kind"],
+        "stored_bytes": state_dict_bytes({"basis": packed_basis}),
+        "sample_size": int(train_ids.numel()),
+        "heldout_size": int(heldout_ids.numel()),
+        "heldout_mean_cosine": float(cosine),
+        "heldout_relative_rmse": float(relative_rmse),
+    }
+    return entry, report
+
+
+def predict_token_rows(
+    entry: dict[str, Any],
+    tokenizer: Any,
+    *,
+    batch_size: int = 8_192,
+) -> torch.Tensor:
+    """Decode a predictor into the original dense vocabulary matrix."""
+    if entry.get("kind") != "token_string_ridge_v1":
+        raise ValueError(f"unsupported token predictor: {entry.get('kind')!r}")
+    basis = unpack_tensor(entry["basis"]).float()
+    rows = []
+    for start in range(0, int(entry["vocab_size"]), batch_size):
+        stop = min(int(entry["vocab_size"]), start + batch_size)
+        ids = torch.arange(start, stop, dtype=torch.int64)
+        rows.append((token_string_features(tokenizer, ids) @ basis).to(torch.bfloat16))
+    return torch.cat(rows, dim=0)

src/eaimath/evaluate.py

@@ -0,0 +1,222 @@
+"""The math eval loop, and the metrics that actually discriminate recipes.
+
+Accuracy alone is not enough. Published measurements on this exact model family
+show W4A16 retaining ~99% of MATH-500 while losing ~11-20 points of AIME, and
+show quantization roughly doubling the truncation rate on AIME under a fixed
+token cap. The mechanism is that low-bit weights perturb high-entropy
+"branching" tokens, the model rambles, and it never emits its closing tag.
+
+So every run reports three things:
+
+* ``accuracy``          — did it get the answer right
+* ``truncation_rate``   — did it run out of budget instead of stopping
+* ``mean_generated_tokens`` / ``think_close_rate`` — is CoT inflating
+
+A recipe that holds accuracy while inflating tokens is not safe; it is a recipe
+that will collapse the moment the grader's token cap is tighter than ours.
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+from .answers import answers_match, extract_answer
+from .data import MathExample, build_prompt
+from .generate import GenerationOutput, generate
+
+THINK_CLOSE_TAG = "</think>"
+
+
+@dataclass
+class MathPrediction:
+    example_id: str
+    source: str
+    gold: str
+    predicted: str | None
+    correct: bool
+    finished: bool
+    think_closed: bool
+    num_generated_tokens: int
+    num_prompt_tokens: int
+    response: str
+
+
+def score_generation(
+    example: MathExample, output: GenerationOutput
+) -> MathPrediction:
+    predicted = extract_answer(output.text)
+    return MathPrediction(
+        example_id=example.example_id,
+        source=example.source,
+        gold=example.answer,
+        predicted=predicted,
+        correct=answers_match(predicted, example.answer),
+        finished=output.finished,
+        # A thinking model that never closes its tag has looped, even if it
+        # somehow stopped afterwards.
+        think_closed=THINK_CLOSE_TAG in output.text,
+        num_generated_tokens=output.num_generated_tokens,
+        num_prompt_tokens=output.num_prompt_tokens,
+        response=output.text,
+    )
+
+
+def summarize(predictions: Sequence[MathPrediction]) -> dict[str, Any]:
+    n = len(predictions)
+    if n == 0:
+        return {"num_examples": 0}
+
+    lengths = [p.num_generated_tokens for p in predictions]
+    finished = [p for p in predictions if p.finished]
+    truncated = [p for p in predictions if not p.finished]
+
+    return {
+        "num_examples": n,
+        "accuracy": sum(p.correct for p in predictions) / n,
+        "parse_rate": sum(p.predicted is not None for p in predictions) / n,
+        # The headline risk metric: budget exhaustion, not wrong answers.
+        "truncation_rate": len(truncated) / n,
+        "think_close_rate": sum(p.think_closed for p in predictions) / n,
+        "mean_generated_tokens": statistics.mean(lengths),
+        "median_generated_tokens": statistics.median(lengths),
+        "max_generated_tokens": max(lengths),
+        # Splitting accuracy by termination separates "reasoned badly" from
+        # "never got to answer" — they need different fixes.
+        "accuracy_when_finished": (
+            sum(p.correct for p in finished) / len(finished) if finished else None
+        ),
+        "accuracy_when_truncated": (
+            sum(p.correct for p in truncated) / len(truncated) if truncated else None
+        ),
+    }
+
+
+def summarize_by_source(predictions: Sequence[MathPrediction]) -> dict[str, Any]:
+    sources = sorted({p.source for p in predictions})
+    return {s: summarize([p for p in predictions if p.source == s]) for s in sources}
+
+
+def run_eval(
+    model,
+    tokenizer,
+    examples: Sequence[MathExample],
+    *,
+    max_new_tokens: int = 65536,
+    temperature: float = 0.0,
+    top_p: float = 0.95,
+    top_k: int = 20,
+    presence_penalty: float = 0.0,
+    repetition_penalty: float = 1.0,
+    batch_size: int = 8,
+    enable_thinking: bool | None = None,
+) -> list[MathPrediction]:
+    outputs = generate(
+        model,
+        tokenizer,
+        [build_prompt(ex) for ex in examples],
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        batch_size=batch_size,
+        enable_thinking=enable_thinking,
+        desc="math eval",
+    )
+    return [score_generation(ex, out) for ex, out in zip(examples, outputs)]
+
+
+def run_eval_vllm(
+    model_path: str,
+    tokenizer,
+    examples: Sequence[MathExample],
+    *,
+    max_new_tokens: int = 65536,
+    temperature: float = 0.0,
+    top_p: float = 0.95,
+    top_k: int = 20,
+    presence_penalty: float = 0.0,
+    repetition_penalty: float = 1.0,
+    enable_thinking: bool | None = None,
+    gpu_memory_utilization: float = 0.90,
+    allowed_token_ids: Sequence[int] | None = None,
+    llm=None,
+) -> tuple[list[MathPrediction], object]:
+    """Same scoring, vLLM engine. Returns the engine so it can be reused."""
+    from .vllm_backend import generate_vllm
+
+    outputs, llm = generate_vllm(
+        model_path,
+        tokenizer,
+        [build_prompt(ex) for ex in examples],
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        presence_penalty=presence_penalty,
+        repetition_penalty=repetition_penalty,
+        enable_thinking=enable_thinking,
+        gpu_memory_utilization=gpu_memory_utilization,
+        allowed_token_ids=allowed_token_ids,
+        llm=llm,
+    )
+    return [score_generation(ex, out) for ex, out in zip(examples, outputs)], llm
+
+
+def save_results(
+    output_dir: str | Path,
+    *,
+    run_name: str,
+    config: dict[str, Any],
+    predictions: Sequence[MathPrediction],
+) -> Path:
+    """Write ``summary.json`` (tracked) and ``generations.jsonl`` (gitignored)."""
+    out = Path(output_dir) / run_name
+    out.mkdir(parents=True, exist_ok=True)
+
+    summary = {
+        "run_name": run_name,
+        "config": config,
+        "overall": summarize(predictions),
+        "by_source": summarize_by_source(predictions),
+    }
+    (out / "summary.json").write_text(json.dumps(summary, indent=2))
+
+    with (out / "generations.jsonl").open("w") as fh:
+        for prediction in predictions:
+            fh.write(json.dumps(asdict(prediction)) + "\n")
+
+    return out
+
+
+def format_summary(summary: dict[str, Any]) -> str:
+    overall = summary["overall"]
+    if not overall.get("num_examples"):
+        return "no examples evaluated"
+
+    lines = [
+        f"  examples            {overall['num_examples']}",
+        f"  accuracy            {overall['accuracy']:.3f}",
+        f"  parse rate          {overall['parse_rate']:.3f}",
+        f"  TRUNCATION RATE     {overall['truncation_rate']:.3f}   <- budget exhaustion",
+        f"  think-close rate    {overall['think_close_rate']:.3f}",
+        f"  mean gen tokens     {overall['mean_generated_tokens']:.0f}",
+        f"  median gen tokens   {overall['median_generated_tokens']:.0f}",
+        f"  max gen tokens      {overall['max_generated_tokens']}",
+    ]
+    if overall.get("accuracy_when_finished") is not None:
+        lines.append(f"  acc | finished      {overall['accuracy_when_finished']:.3f}")
+    if overall.get("accuracy_when_truncated") is not None:
+        lines.append(f"  acc | truncated     {overall['accuracy_when_truncated']:.3f}")
+
+    lines.append("")
+    for source, stats in summary["by_source"].items():
+        lines.append(
+            f"  [{source}] n={stats['num_examples']} acc={stats['accuracy']:.3f} "
+            f"trunc={stats['truncation_rate']:.3f} tok={stats['mean_generated_tokens']:.0f}"
+        )
+    return "\n".join(lines)

src/eaimath/extreme_quant.py

@@ -0,0 +1,582 @@
+"""Training-free extreme weight compression primitives for Bucket A.
+
+The functions in this module produce dense BF16 values for ordinary Hugging
+Face/vLLM evaluation while accounting for the compact representation that
+would be stored in a submission artifact.  They deliberately avoid custom
+inference kernels: Round 11 asks which representation still reasons, not which
+kernel is fastest.
+
+Three representation families are implemented:
+
+* binary and ternary scalar grids with one BF16 scale per group;
+* binary grids with a fixed number of exact sparse residuals per group;
+* residual vector codebooks (an intentionally small AQLM-style screen).
+
+The vector-codebook implementation is not claimed to reproduce AQLM.  AQLM
+jointly optimizes block outputs and codebooks; this is a deterministic,
+weight-only cold-start test that tells us whether paying for that optimization
+is justified on Qwen3.5's hybrid architecture.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import re
+from dataclasses import asdict, dataclass, replace
+from typing import Any, Sequence
+
+import torch
+
+from .model import BUDGET_GROUPS
+from .quantize import EXCLUDED, PROTECTED, QuantSpec, quantize_dequantize
+
+
+@dataclass(frozen=True)
+class ExtremeSpec:
+    kind: str
+    group_size: int = 256
+    outliers: int = 0
+    vector_dim: int = 8
+    codebooks: int = 0
+    codebook_size: int = 16
+    kmeans_iterations: int = 4
+    kmeans_samples: int = 32_768
+    assign_chunk: int = 65_536
+    rank: int = 0
+
+
+def _component(name: str) -> str:
+    for group, pattern in BUDGET_GROUPS:
+        if re.search(pattern, name):
+            return group
+    return "other"
+
+
+def _seed_for(name: str, seed: int) -> int:
+    digest = hashlib.sha256(f"{seed}:{name}".encode()).digest()
+    return int.from_bytes(digest[:8], "little") % (2**31)
+
+
+def _grouped(weight: torch.Tensor, group_size: int) -> tuple[torch.Tensor, tuple[int, ...]]:
+    shape = tuple(weight.shape)
+    width = shape[-1]
+    if group_size <= 0 or width % group_size:
+        group_size = width
+    return weight.reshape(-1, group_size).float(), shape
+
+
+def _normalized_fwht(values: torch.Tensor) -> torch.Tensor:
+    """Orthonormal Walsh-Hadamard transform over the last dimension.
+
+    The transform is its own inverse and fixed by shape, so a submission stores
+    no transform matrix. This is the transform-coding analogue of using a DCT
+    before scalar quantization, but every operation is an add/subtract and all
+    power-of-two Qwen group widths are supported.
+    """
+    width = values.shape[-1]
+    if width < 1 or width & (width - 1):
+        raise ValueError(f"Hadamard width must be a power of two, found {width}")
+    # ``float().contiguous()`` can alias an already-contiguous float input.
+    # The butterfly is in-place, so clone to keep this a pure transform.
+    output = values.float().contiguous().clone()
+    stride = 1
+    while stride < width:
+        blocks = output.reshape(-1, width // (2 * stride), 2, stride)
+        left = blocks[:, :, 0, :].clone()
+        right = blocks[:, :, 1, :].clone()
+        blocks[:, :, 0, :] = left + right
+        blocks[:, :, 1, :] = left - right
+        output = blocks.reshape(values.shape)
+        stride *= 2
+    return output / math.sqrt(width)
+
+
+def _hadamard_quantize(weight: torch.Tensor, spec: ExtremeSpec) -> tuple[torch.Tensor, int]:
+    """Transform-code independent groups and invert to dense evaluation values."""
+    base_kind = spec.kind.removeprefix("hadamard_")
+    if base_kind == spec.kind:
+        raise ValueError(f"not a Hadamard spec: {spec.kind}")
+    groups, shape = _grouped(weight, spec.group_size)
+    if groups.shape[-1] & (groups.shape[-1] - 1):
+        raise ValueError("Hadamard transform requires a power-of-two effective group")
+    coefficients = _normalized_fwht(groups)
+    base_spec = replace(spec, kind=base_kind, group_size=groups.shape[-1])
+    if base_kind == "binary":
+        coded = _binary(coefficients, base_spec)
+        codebook_values = 0
+    elif base_kind == "ternary":
+        coded = _ternary(coefficients, base_spec)
+        codebook_values = 0
+    elif base_kind == "sparse_binary":
+        coded = _sparse_binary(coefficients, base_spec)
+        codebook_values = 0
+    elif base_kind.startswith("int"):
+        bits = int(base_kind.removeprefix("int"))
+        coded = quantize_dequantize(
+            coefficients,
+            QuantSpec(pattern=".*", bits=bits, group_size=groups.shape[-1]),
+        )
+        codebook_values = 0
+    else:
+        raise ValueError(f"unsupported Hadamard coefficient codec: {base_kind}")
+    restored = _normalized_fwht(coded).reshape(shape).to(weight.dtype)
+    return restored, codebook_values
+
+
+def _binary(weight: torch.Tensor, spec: ExtremeSpec) -> torch.Tensor:
+    groups, shape = _grouped(weight, spec.group_size)
+    scale = groups.abs().mean(dim=1, keepdim=True).clamp_min(1e-8)
+    signs = torch.where(groups >= 0, 1.0, -1.0)
+    return (signs * scale).reshape(shape).to(weight.dtype)
+
+
+def _ternary(weight: torch.Tensor, spec: ExtremeSpec) -> torch.Tensor:
+    """MSE-select a symmetric ternary threshold independently per group."""
+    groups, shape = _grouped(weight, spec.group_size)
+    magnitude = groups.abs()
+    mean_abs = magnitude.mean(dim=1, keepdim=True).clamp_min(1e-8)
+    best = torch.zeros_like(groups)
+    best_error = torch.full(
+        (groups.shape[0], 1), float("inf"), device=groups.device, dtype=groups.dtype
+    )
+    # A compact grid is enough for a diagnostic and avoids training on eval data.
+    for multiplier in (0.35, 0.45, 0.55, 0.65, 0.75, 0.90):
+        mask = magnitude >= mean_abs * multiplier
+        count = mask.sum(dim=1, keepdim=True).clamp_min(1)
+        scale = (magnitude * mask).sum(dim=1, keepdim=True) / count
+        candidate = torch.sign(groups) * mask * scale
+        error = (groups - candidate).square().mean(dim=1, keepdim=True)
+        improve = error < best_error
+        best = torch.where(improve, candidate, best)
+        best_error = torch.where(improve, error, best_error)
+    return best.reshape(shape).to(weight.dtype)
+
+
+def _activation_weighted_ternary(
+    weight: torch.Tensor,
+    spec: ExtremeSpec,
+    importance: torch.Tensor | None,
+) -> torch.Tensor:
+    """Ternarize against diagonal activation covariance instead of weight MSE."""
+    if importance is None or weight.ndim != 2 or importance.numel() != weight.shape[-1]:
+        return _ternary(weight, replace(spec, kind="ternary"))
+    groups, shape = _grouped(weight, spec.group_size)
+    group = groups.shape[-1]
+    coordinate_weight = importance.detach().float().clamp_min(0)
+    coordinate_weight = coordinate_weight / coordinate_weight.mean().clamp_min(1e-12)
+    weights = coordinate_weight.reshape(1, -1).expand(shape[0], -1).reshape(-1, group)
+    magnitude = groups.abs()
+    mean_abs = (magnitude * weights).sum(dim=1, keepdim=True) / weights.sum(
+        dim=1, keepdim=True
+    ).clamp_min(1e-12)
+    best = torch.zeros_like(groups)
+    best_error = torch.full(
+        (groups.shape[0], 1), float("inf"), device=groups.device, dtype=groups.dtype
+    )
+    for multiplier in (0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.90, 1.10):
+        mask = magnitude >= mean_abs * multiplier
+        selected_weight = weights * mask
+        scale = (magnitude * selected_weight).sum(dim=1, keepdim=True) / selected_weight.sum(
+            dim=1, keepdim=True
+        ).clamp_min(1e-12)
+        candidate = torch.sign(groups) * mask * scale
+        error = ((groups - candidate).square() * weights).sum(dim=1, keepdim=True)
+        improve = error < best_error
+        best = torch.where(improve, candidate, best)
+        best_error = torch.where(improve, error, best_error)
+    return best.reshape(shape).to(weight.dtype)
+
+
+def _sigma_delta_ternary(
+    weight: torch.Tensor,
+    spec: ExtremeSpec,
+    importance: torch.Tensor | None,
+) -> torch.Tensor:
+    """First-order error-feedback ternary in activation-energy order.
+
+    Coordinates with lower observed energy are visited first, allowing their
+    quantization residue to influence later, more important coordinates. The
+    traversal is reconstructed from the same train-only sensitivity vector, so
+    it requires no stored permutation or side information.
+    """
+    if importance is None or weight.ndim != 2 or importance.numel() != weight.shape[-1]:
+        return _ternary(weight, replace(spec, kind="ternary"))
+    groups, shape = _grouped(weight, spec.group_size)
+    group = groups.shape[-1]
+    column_groups = shape[-1] // group
+    energy = importance.detach().float().clamp_min(0).reshape(column_groups, group)
+    order = energy.argsort(dim=1).repeat(shape[0], 1).to(groups.device)
+    ordinary = _ternary(weight, replace(spec, kind="ternary")).reshape_as(groups).float()
+    scale = ordinary.abs().amax(dim=1).clamp_min(1e-8)
+    output = torch.zeros_like(groups)
+    feedback = torch.zeros(groups.shape[0], device=groups.device, dtype=groups.dtype)
+    rows = torch.arange(groups.shape[0], device=groups.device)
+    for step in range(group):
+        index = order[:, step]
+        value = groups[rows, index] + feedback
+        quantized = torch.where(
+            value > scale * 0.5,
+            scale,
+            torch.where(value < -scale * 0.5, -scale, torch.zeros_like(value)),
+        )
+        output[rows, index] = quantized
+        feedback = value - quantized
+    return output.reshape(shape).to(weight.dtype)
+
+
+def _structured_residual_ternary(
+    weight: torch.Tensor,
+    spec: ExtremeSpec,
+    *,
+    seed: int,
+) -> tuple[torch.Tensor, int]:
+    """Preserve a low-rank subspace and ternarize only the residual."""
+    if weight.ndim != 2 or spec.rank < 1:
+        raise ValueError("srr_ternary requires a matrix and rank >= 1")
+    rank = min(spec.rank, min(weight.shape))
+    target = weight.float()
+    devices = [target.device.index or 0] if target.is_cuda else []
+    with torch.random.fork_rng(devices=devices):
+        torch.manual_seed(seed)
+        if target.is_cuda:
+            torch.cuda.manual_seed_all(seed)
+        u, singular, v = torch.svd_lowrank(
+            target,
+            q=min(rank + 4, min(target.shape)),
+            niter=2,
+        )
+    left = u[:, :rank] * singular[:rank]
+    right = v[:, :rank].T
+    residual = target - left @ right
+    coded = _ternary(residual, replace(spec, kind="ternary"))
+    # BF16 left/right factors are the complete side information.
+    factor_values = rank * (weight.shape[0] + weight.shape[1])
+    return (coded.float() + left @ right).to(weight.dtype), factor_values
+
+
+def _sparse_binary(weight: torch.Tensor, spec: ExtremeSpec) -> torch.Tensor:
+    """Binary base plus exact, fixed-rate residual exceptions in each group."""
+    groups, shape = _grouped(weight, spec.group_size)
+    if not 0 < spec.outliers < groups.shape[1]:
+        raise ValueError("sparse_binary requires 0 < outliers < group_size")
+
+    first_scale = groups.abs().mean(dim=1, keepdim=True).clamp_min(1e-8)
+    first = torch.where(groups >= 0, first_scale, -first_scale)
+    indices = (groups - first).abs().topk(spec.outliers, dim=1).indices
+    residual_mask = torch.zeros_like(groups, dtype=torch.bool)
+    residual_mask.scatter_(1, indices, True)
+
+    base_mask = ~residual_mask
+    scale = (
+        (groups.abs() * base_mask).sum(dim=1, keepdim=True)
+        / base_mask.sum(dim=1, keepdim=True).clamp_min(1)
+    ).clamp_min(1e-8)
+    output = torch.where(groups >= 0, scale, -scale)
+    output[residual_mask] = groups[residual_mask]
+    return output.reshape(shape).to(weight.dtype)
+
+
+def _assign(points: torch.Tensor, centers: torch.Tensor, chunk: int) -> torch.Tensor:
+    assignments = []
+    center_norm = centers.square().sum(dim=1).unsqueeze(0)
+    for start in range(0, points.shape[0], chunk):
+        block = points[start : start + chunk]
+        distances = block.square().sum(dim=1, keepdim=True) + center_norm
+        distances.addmm_(block, centers.T, beta=1.0, alpha=-2.0)
+        assignments.append(distances.argmin(dim=1))
+    return torch.cat(assignments)
+
+
+def _fit_codebook(
+    points: torch.Tensor,
+    *,
+    size: int,
+    iterations: int,
+    max_samples: int,
+    chunk: int,
+    seed: int,
+) -> torch.Tensor:
+    if points.shape[0] < size:
+        raise ValueError(f"need at least {size} vectors, found {points.shape[0]}")
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+    if points.shape[0] > max_samples:
+        # randperm(points.shape[0]) would allocate gigabytes for the largest
+        # MLP matrices. Sampling with replacement keeps this bounded at O(sample).
+        sample_idx = torch.randint(
+            points.shape[0], (max_samples,), generator=generator
+        )
+        sample = points.index_select(0, sample_idx.to(points.device))
+    else:
+        sample = points
+    init_idx = torch.randperm(sample.shape[0], generator=generator)[:size]
+    centers = sample.index_select(0, init_idx.to(sample.device)).clone()
+
+    for _ in range(iterations):
+        assignment = _assign(sample, centers, chunk)
+        sums = torch.zeros_like(centers)
+        sums.index_add_(0, assignment, sample)
+        counts = torch.bincount(assignment, minlength=size).to(sample.dtype).unsqueeze(1)
+        nonempty = counts.squeeze(1) > 0
+        centers[nonempty] = sums[nonempty] / counts[nonempty]
+    return centers
+
+
+def _vector_additive(
+    weight: torch.Tensor, spec: ExtremeSpec, *, seed: int
+) -> tuple[torch.Tensor, int]:
+    if spec.codebooks < 1:
+        raise ValueError("vector quantization requires at least one codebook")
+    if spec.codebook_size < 2 or spec.codebook_size & (spec.codebook_size - 1):
+        raise ValueError("codebook_size must be a power of two")
+
+    original_shape = tuple(weight.shape)
+    flat = weight.float().reshape(-1)
+    padding = (-flat.numel()) % spec.vector_dim
+    if padding:
+        flat = torch.nn.functional.pad(flat, (0, padding))
+    points = flat.reshape(-1, spec.vector_dim)
+    residual = points.clone()
+    reconstructed = torch.zeros_like(points)
+    codebook_values = 0
+
+    for book in range(spec.codebooks):
+        centers = _fit_codebook(
+            residual,
+            size=spec.codebook_size,
+            iterations=spec.kmeans_iterations,
+            max_samples=spec.kmeans_samples,
+            chunk=spec.assign_chunk,
+            seed=seed + book * 104_729,
+        )
+        assignment = _assign(residual, centers, spec.assign_chunk)
+        contribution = centers.index_select(0, assignment)
+        reconstructed.add_(contribution)
+        residual.sub_(contribution)
+        codebook_values += centers.numel()
+
+    restored = reconstructed.reshape(-1)[: weight.numel()].reshape(original_shape)
+    return restored.to(weight.dtype), codebook_values
+
+
+def _representation_bits(
+    numel: int, spec: ExtremeSpec, *, codebook_values: int = 0
+) -> int:
+    if spec.kind.startswith("hadamard_"):
+        # The fixed normalized transform is reconstructed from group width and
+        # costs no payload bytes. Only its coefficient codec is charged.
+        return _representation_bits(
+            numel,
+            replace(spec, kind=spec.kind.removeprefix("hadamard_")),
+            codebook_values=codebook_values,
+        )
+    if spec.kind == "aw_ternary":
+        return _representation_bits(
+            numel, replace(spec, kind="ternary"), codebook_values=codebook_values
+        )
+    if spec.kind == "sigma_ternary":
+        return _representation_bits(
+            numel, replace(spec, kind="ternary"), codebook_values=codebook_values
+        )
+    if spec.kind == "binary":
+        groups = math.ceil(numel / spec.group_size)
+        return numel + groups * 16
+    if spec.kind == "ternary":
+        # Five base-3 digits fit in one byte. Account in independently packed groups.
+        groups = math.ceil(numel / spec.group_size)
+        trit_bytes = groups * math.ceil(spec.group_size / 5)
+        return trit_bytes * 8 + groups * 16
+    if spec.kind == "sparse_binary":
+        groups = math.ceil(numel / spec.group_size)
+        position_bits = math.ceil(math.log2(spec.group_size))
+        return numel + groups * (16 + spec.outliers * (16 + position_bits))
+    if spec.kind == "vector":
+        vectors = math.ceil(numel / spec.vector_dim)
+        index_bits = int(math.log2(spec.codebook_size))
+        return vectors * index_bits * spec.codebooks + codebook_values * 16
+    if spec.kind.startswith("int"):
+        bits = int(spec.kind.removeprefix("int"))
+        groups = math.ceil(numel / spec.group_size)
+        return numel * bits + groups * 16
+    if spec.kind == "bf16":
+        return numel * 16
+    raise ValueError(f"unknown extreme quantizer kind: {spec.kind}")
+
+
+def quantize_tensor(
+    weight: torch.Tensor,
+    spec: ExtremeSpec,
+    *,
+    seed: int,
+    importance: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, int, dict[str, Any]]:
+    factor_values = 0
+    if spec.kind.startswith("hadamard_"):
+        restored, codebook_values = _hadamard_quantize(weight, spec)
+    elif spec.kind == "binary":
+        restored = _binary(weight, spec)
+        codebook_values = 0
+    elif spec.kind == "ternary":
+        restored = _ternary(weight, spec)
+        codebook_values = 0
+    elif spec.kind == "aw_ternary":
+        restored = _activation_weighted_ternary(weight, spec, importance)
+        codebook_values = 0
+    elif spec.kind == "sigma_ternary":
+        restored = _sigma_delta_ternary(weight, spec, importance)
+        codebook_values = 0
+    elif spec.kind == "srr_ternary":
+        restored, factor_values = _structured_residual_ternary(weight, spec, seed=seed)
+        codebook_values = 0
+    elif spec.kind == "sparse_binary":
+        restored = _sparse_binary(weight, spec)
+        codebook_values = 0
+    elif spec.kind == "vector":
+        restored, codebook_values = _vector_additive(weight, spec, seed=seed)
+    elif spec.kind.startswith("int"):
+        bits = int(spec.kind.removeprefix("int"))
+        restored = quantize_dequantize(
+            weight, QuantSpec(pattern=".*", bits=bits, group_size=spec.group_size)
+        )
+        codebook_values = 0
+    elif spec.kind == "bf16":
+        restored = weight
+        codebook_values = 0
+    else:
+        raise ValueError(f"unknown extreme quantizer kind: {spec.kind}")
+
+    delta = weight.float() - restored.float()
+    denominator = weight.float().square().sum().sqrt().clamp_min(1e-12)
+    info = {
+        "relative_weight_error": float(delta.square().sum().sqrt() / denominator),
+        "max_abs_error": float(delta.abs().max()),
+        "codebook_values": codebook_values,
+        "factor_values": factor_values,
+        "activation_weighted": spec.kind == "aw_ternary" and importance is not None,
+        "error_feedback": spec.kind == "sigma_ternary" and importance is not None,
+    }
+    if spec.kind == "srr_ternary":
+        bits = _representation_bits(
+            weight.numel(), replace(spec, kind="ternary"), codebook_values=0
+        ) + factor_values * 16
+    else:
+        bits = _representation_bits(weight.numel(), spec, codebook_values=codebook_values)
+    return restored, bits, info
+
+
+def apply_extreme_recipe(
+    model: torch.nn.Module,
+    specs: dict[str, ExtremeSpec],
+    *,
+    device: str,
+    stored_vocab_rows: int,
+    stored_vocab_ids: Sequence[int] | None = None,
+    input_second_moments: dict[str, torch.Tensor] | None = None,
+    seed: int = 6013,
+) -> dict[str, Any]:
+    """Quantize the text model in place and return physical representation accounting."""
+    excluded = re.compile(EXCLUDED)
+    protected = re.compile(PROTECTED.pattern)
+    seen: set[int] = set()
+    tensors: list[dict[str, Any]] = []
+    component_bits: dict[str, int] = {}
+    component_params: dict[str, int] = {}
+    total_bits = 0
+    total_counted_params = 0
+    weighted_squared_error = 0.0
+    weighted_squared_norm = 0.0
+    vocab_ids = None
+    sensitivity_hits = 0
+    if stored_vocab_ids is not None:
+        vocab_ids = torch.tensor(sorted(set(int(value) for value in stored_vocab_ids)))
+        if vocab_ids.numel() != stored_vocab_rows:
+            raise ValueError(
+                "stored_vocab_ids must contain exactly stored_vocab_rows unique ids: "
+                f"ids={vocab_ids.numel()}, rows={stored_vocab_rows}"
+            )
+
+    with torch.no_grad():
+        for name, param in model.named_parameters():
+            if excluded.search(name):
+                continue
+            pointer = param.data_ptr()
+            if pointer in seen:
+                continue
+            seen.add(pointer)
+
+            component = _component(name)
+            spec = ExtremeSpec("bf16") if protected.search(name) else specs.get(
+                component, ExtremeSpec("bf16")
+            )
+            target = param.data.to(device) if spec.kind != "bf16" else param.data
+            original = target.float()
+            module_name = name.removesuffix(".weight")
+            importance = None
+            if input_second_moments is not None and module_name in input_second_moments:
+                importance = input_second_moments[module_name].to(target.device)
+                sensitivity_hits += 1
+            restored, bits, info = quantize_tensor(
+                target,
+                spec,
+                seed=_seed_for(name, seed),
+                importance=importance,
+            )
+            if component == "embed" and param.ndim == 2 and stored_vocab_rows < param.shape[0]:
+                if vocab_ids is None:
+                    raise ValueError(
+                        "a compact embedding claim requires exact stored_vocab_ids; "
+                        "output masking alone does not reconstruct omitted input rows"
+                    )
+                if int(vocab_ids[0]) < 0 or int(vocab_ids[-1]) >= param.shape[0]:
+                    raise ValueError("stored vocabulary ids are outside the embedding table")
+                keep_mask = torch.zeros(param.shape[0], dtype=torch.bool, device=restored.device)
+                keep_mask[vocab_ids.to(restored.device)] = True
+                restored[~keep_mask] = 0
+            if spec.kind != "bf16" or (
+                component == "embed" and param.ndim == 2 and stored_vocab_rows < param.shape[0]
+            ):
+                param.data.copy_(restored.to(param.device))
+
+            counted_params = param.numel()
+            # The physical A artifact will store selected original-id rows only.
+            # Quantize the full table for an input-compatible diagnostic, but do
+            # not charge rows that the compact representation omits.
+            if component == "embed" and param.ndim == 2:
+                counted_params = min(stored_vocab_rows, param.shape[0]) * param.shape[1]
+                ratio = counted_params / param.numel()
+                bits = math.ceil(bits * ratio)
+
+            component_bits[component] = component_bits.get(component, 0) + bits
+            component_params[component] = component_params.get(component, 0) + counted_params
+            total_bits += bits
+            total_counted_params += counted_params
+            weighted_squared_error += float((original - restored.float()).square().sum())
+            weighted_squared_norm += float(original.square().sum())
+            tensors.append(
+                {
+                    "name": name,
+                    "component": component,
+                    "shape": list(param.shape),
+                    "spec": asdict(spec),
+                    "counted_params": counted_params,
+                    "representation_bytes": math.ceil(bits / 8),
+                    **info,
+                }
+            )
+
+    return {
+        "format": "eaimath-extreme-screen-v1",
+        "stored_vocab_rows": stored_vocab_rows,
+        "omitted_embedding_rows_zeroed": vocab_ids is not None,
+        "total_params_in_representation": total_counted_params,
+        "total_bytes": math.ceil(total_bits / 8),
+        "total_gb": total_bits / 8e9,
+        "effective_bits_per_weight": total_bits / total_counted_params,
+        "relative_weight_error": math.sqrt(weighted_squared_error / weighted_squared_norm),
+        "component_bytes": {name: math.ceil(bits / 8) for name, bits in component_bits.items()},
+        "component_params": component_params,
+        "specs": {name: asdict(spec) for name, spec in specs.items()},
+        "sensitivity_tensors_used": sensitivity_hits,
+        "tensors": tensors,
+    }

src/eaimath/gauge.py

@@ -0,0 +1,136 @@
+"""Exact MLP gauge transformations used as quantization preconditioners.
+
+For a gated MLP ``down(silu(gate(x)) * up(x))``, every intermediate channel has
+an exact positive scaling symmetry: multiply one row of ``up`` by ``c`` and the
+matching column of ``down`` by ``1/c``. Gauge fixing chooses ``c`` so the two
+channel norms are balanced before low-bit compression. The factors are fused
+into the weights, so neither the compact artifact nor inference stores them.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+
+def balance_up_down(
+    up: torch.Tensor,
+    down: torch.Tensor,
+    *,
+    min_scale: float = 1 / 16,
+    max_scale: float = 16,
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+    """Return a function-equivalent, norm-balanced pair of MLP weights."""
+    if up.ndim != 2 or down.ndim != 2 or up.shape[0] != down.shape[1]:
+        raise ValueError(f"incompatible up/down shapes: {tuple(up.shape)}, {tuple(down.shape)}")
+    up32 = up.float()
+    down32 = down.float()
+    up_rms = up32.square().mean(dim=1).sqrt().clamp_min(1e-12)
+    down_rms = down32.square().mean(dim=0).sqrt().clamp_min(1e-12)
+    scale = torch.sqrt(down_rms / up_rms).clamp(min_scale, max_scale)
+    balanced_up = up32 * scale[:, None]
+    balanced_down = down32 / scale[None, :]
+    before_ratio = torch.maximum(up_rms / down_rms, down_rms / up_rms)
+    after_up = balanced_up.square().mean(dim=1).sqrt().clamp_min(1e-12)
+    after_down = balanced_down.square().mean(dim=0).sqrt().clamp_min(1e-12)
+    after_ratio = torch.maximum(after_up / after_down, after_down / after_up)
+    quantiles = torch.tensor([0.01, 0.5, 0.99], device=scale.device)
+    scale_quantiles = torch.quantile(scale, quantiles).cpu().tolist()
+    report = {
+        "channels": scale.numel(),
+        "scale_min": float(scale.min()),
+        "scale_p01": scale_quantiles[0],
+        "scale_median": scale_quantiles[1],
+        "scale_p99": scale_quantiles[2],
+        "scale_max": float(scale.max()),
+        "mean_norm_ratio_before": float(before_ratio.mean()),
+        "mean_norm_ratio_after": float(after_ratio.mean()),
+        "clamped_low": int((scale == min_scale).sum()),
+        "clamped_high": int((scale == max_scale).sum()),
+        "payload_bytes": 0,
+    }
+    return balanced_up.to(up.dtype), balanced_down.to(down.dtype), report
+
+
+def gated_mlp_output(
+    inputs: torch.Tensor,
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    down: torch.Tensor,
+) -> torch.Tensor:
+    return (F.silu(inputs @ gate.T) * (inputs @ up.T)) @ down.T
+
+
+def apply_mlp_gauge(
+    model: torch.nn.Module,
+    *,
+    device: str,
+    verification_vectors: int = 2,
+    seed: int = 6013,
+) -> dict[str, Any]:
+    """Fuse the exact up/down gauge into every dense text MLP in ``model``."""
+    rows = []
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+    with torch.no_grad():
+        for name, module in model.named_modules():
+            if "visual" in name or "mtp" in name:
+                continue
+            if not all(hasattr(module, attr) for attr in ("gate_proj", "up_proj", "down_proj")):
+                continue
+            gate_param = module.gate_proj.weight
+            up_param = module.up_proj.weight
+            down_param = module.down_proj.weight
+            gate = gate_param.detach().to(device)
+            up = up_param.detach().to(device)
+            down = down_param.detach().to(device)
+            inputs = torch.randn(
+                verification_vectors,
+                up.shape[1],
+                generator=generator,
+                dtype=torch.float32,
+            ).to(device=device, dtype=up.dtype)
+            before = gated_mlp_output(inputs, gate, up, down).float()
+            balanced_up, balanced_down, report = balance_up_down(up, down)
+            after = gated_mlp_output(inputs, gate, balanced_up, balanced_down).float()
+            denominator = before.square().sum().sqrt().clamp_min(1e-12)
+            report.update(
+                {
+                    "module": name,
+                    "up_shape": list(up.shape),
+                    "down_shape": list(down.shape),
+                    "verification_relative_l2": float(
+                        (before - after).square().sum().sqrt() / denominator
+                    ),
+                    "verification_max_abs": float((before - after).abs().max()),
+                }
+            )
+            up_param.copy_(balanced_up.to(up_param.device, dtype=up_param.dtype))
+            down_param.copy_(balanced_down.to(down_param.device, dtype=down_param.dtype))
+            if getattr(module.up_proj, "bias", None) is not None:
+                bias = module.up_proj.bias
+                # Qwen uses bias-free MLPs, but preserve the symmetry for any
+                # compatible checkpoint that does carry an up-projection bias.
+                up_rms = up.float().square().mean(dim=1).sqrt().clamp_min(1e-12)
+                down_rms = down.float().square().mean(dim=0).sqrt().clamp_min(1e-12)
+                scale = torch.sqrt(down_rms / up_rms).clamp(1 / 16, 16)
+                bias.mul_(scale.to(bias.device, dtype=bias.dtype))
+            rows.append(report)
+
+    if not rows:
+        raise RuntimeError("no dense gate/up/down MLP modules were found")
+    squared = sum(row["verification_relative_l2"] ** 2 for row in rows)
+    return {
+        "format": "eaimath-mlp-gauge-v1",
+        "method": "exact positive up/down channel gauge balancing",
+        "modules": len(rows),
+        "payload_bytes": 0,
+        "rms_verification_relative_l2": math.sqrt(squared / len(rows)),
+        "max_verification_relative_l2": max(
+            row["verification_relative_l2"] for row in rows
+        ),
+        "rows": rows,
+    }

src/eaimath/generate.py

@@ -0,0 +1,162 @@
+"""Batched generation with the bookkeeping our metrics need.
+
+``transformers.generate`` gives back padded token ids and nothing about *why*
+each sequence stopped. For a thinking model that distinction is the whole game:
+a sequence that ran out of budget is a different failure from one that finished
+and got the answer wrong, and only the first is fixable by shortening CoT.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+import torch
+from tqdm.auto import tqdm
+
+
+@dataclass
+class GenerationOutput:
+    text: str
+    num_prompt_tokens: int
+    num_generated_tokens: int
+    finished: bool  # True = emitted EOS; False = hit max_new_tokens
+
+
+def build_chat_prompts(tokenizer, prompts: Sequence[str], enable_thinking: bool | None = None):
+    """Apply the chat template, passing ``enable_thinking`` only if supported."""
+    kwargs = {"tokenize": False, "add_generation_prompt": True}
+    if enable_thinking is not None:
+        kwargs["enable_thinking"] = enable_thinking
+
+    messages = [[{"role": "user", "content": p}] for p in prompts]
+    try:
+        return [tokenizer.apply_chat_template(m, **kwargs) for m in messages]
+    except TypeError:
+        kwargs.pop("enable_thinking", None)
+        return [tokenizer.apply_chat_template(m, **kwargs) for m in messages]
+
+
+@torch.inference_mode()
+def generate(
+    model,
+    tokenizer,
+    prompts: Sequence[str],
+    *,
+    max_new_tokens: int = 65536,
+    temperature: float = 0.0,
+    top_p: float = 0.95,
+    top_k: int = 20,
+    repetition_penalty: float = 1.0,
+    batch_size: int = 8,
+    enable_thinking: bool | None = None,
+    desc: str = "generating",
+) -> list[GenerationOutput]:
+    chat_texts = build_chat_prompts(tokenizer, prompts, enable_thinking)
+    do_sample = temperature > 0
+    eos_ids = _eos_token_ids(model, tokenizer)
+
+    gen_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "pad_token_id": tokenizer.pad_token_id,
+        "do_sample": do_sample,
+    }
+    # HF has no presence_penalty; repetition_penalty is the available analogue.
+    if repetition_penalty != 1.0:
+        gen_kwargs["repetition_penalty"] = repetition_penalty
+    if do_sample:
+        gen_kwargs.update(temperature=temperature, top_p=top_p, top_k=top_k)
+
+    results: list[GenerationOutput] = []
+    progress = tqdm(total=len(chat_texts), desc=desc, unit="ex")
+    index = 0
+    current_batch = batch_size
+
+    try:
+        while index < len(chat_texts):
+            batch = chat_texts[index : index + current_batch]
+            try:
+                results.extend(
+                    _run_batch(model, tokenizer, batch, gen_kwargs, eos_ids)
+                )
+            except Exception as exc:
+                # A shared box means free VRAM changes under us, and long-CoT KV
+                # cache grows with max_new_tokens. Back off rather than losing
+                # an hour of completed work -- but only for genuine OOM; anything
+                # else is a real bug and must surface.
+                if not _is_oom(exc) or current_batch == 1:
+                    raise
+                torch.cuda.empty_cache()
+                current_batch = max(1, current_batch // 2)
+                progress.write(f"  OOM -> retrying at batch size {current_batch}")
+                continue
+            index += len(batch)
+            progress.update(len(batch))
+    finally:
+        progress.close()
+
+    return results
+
+
+_OOM_TYPES = tuple(
+    err
+    for err in (
+        getattr(torch, "OutOfMemoryError", None),
+        getattr(torch.cuda, "OutOfMemoryError", None),
+    )
+    if err is not None
+)
+
+
+def _is_oom(exc: BaseException) -> bool:
+    """True for out-of-memory failures only, across torch versions."""
+    if _OOM_TYPES and isinstance(exc, _OOM_TYPES):
+        return True
+    return isinstance(exc, RuntimeError) and "out of memory" in str(exc).lower()
+
+
+@torch.inference_mode()
+def _run_batch(model, tokenizer, batch, gen_kwargs, eos_ids) -> list[GenerationOutput]:
+    inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=False)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    prompt_len = inputs["input_ids"].shape[1]
+
+    outputs = model.generate(**inputs, **gen_kwargs)
+
+    out: list[GenerationOutput] = []
+    for row, prompt_tokens in zip(outputs, inputs["attention_mask"].sum(dim=1).tolist()):
+        generated = row[prompt_len:]
+        num_generated, finished = _measure_completion(generated, eos_ids)
+        out.append(
+            GenerationOutput(
+                text=tokenizer.decode(generated[:num_generated], skip_special_tokens=True),
+                num_prompt_tokens=int(prompt_tokens),
+                num_generated_tokens=num_generated,
+                finished=finished,
+            )
+        )
+    return out
+
+
+def _eos_token_ids(model, tokenizer) -> set[int]:
+    ids: set[int] = set()
+    config_eos = getattr(getattr(model, "generation_config", None), "eos_token_id", None)
+    for source in (config_eos, tokenizer.eos_token_id):
+        if isinstance(source, int):
+            ids.add(source)
+        elif isinstance(source, (list, tuple)):
+            ids.update(int(i) for i in source)
+    return ids
+
+
+def _measure_completion(generated: torch.Tensor, eos_ids: set[int]) -> tuple[int, bool]:
+    """Length up to and including the first EOS, and whether one was found.
+
+    Sequences that finish early are right-padded by ``generate``, so scanning for
+    the first EOS is what separates a real stop from budget exhaustion.
+    """
+    if eos_ids:
+        for position, token in enumerate(generated.tolist()):
+            if token in eos_ids:
+                return position + 1, True
+    return int(generated.shape[0]), False

src/eaimath/gptq_backend.py

@@ -0,0 +1,40 @@
+"""GPTQModel loading adapter for the existing Hugging Face evaluation loop.
+
+GPTQModel's public object deliberately wraps the underlying Transformers model.
+It exposes ``generate`` but does not promise every attribute that
+``eaimath.generate`` uses (notably ``device`` and ``generation_config``).  This
+small adapter keeps that compatibility boundary in one place.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class GPTQGenerationAdapter:
+    """Expose GPTQModel through the subset of ``PreTrainedModel`` we evaluate."""
+
+    def __init__(self, wrapped: Any, device: str):
+        self.wrapped = wrapped
+        self.device = torch.device(device)
+        inner = getattr(wrapped, "model", None)
+        self.generation_config = getattr(inner, "generation_config", None)
+
+    def generate(self, *args, **kwargs):
+        return self.wrapped.generate(*args, **kwargs)
+
+
+def load_gptq_for_generation(model_path: str, device: str) -> GPTQGenerationAdapter:
+    """Load a packed GPTQModel checkpoint on one device for deterministic eval."""
+    try:
+        from gptqmodel import GPTQModel
+    except ImportError as exc:
+        raise RuntimeError(
+            "The gptqmodel backend is not installed. Run "
+            "`./scripts/setup_round8.sh` in the project environment."
+        ) from exc
+
+    wrapped = GPTQModel.load(model_path, device=device)
+    return GPTQGenerationAdapter(wrapped, device)

src/eaimath/ledger.py

@@ -0,0 +1,224 @@
+"""The results ledger: one row per checkpoint, not per experiment.
+
+Storage is **one small file per checkpoint**, under `experiments/records/`.
+
+This is the third design. A single JSON array conflicted whenever two machines
+recorded anything. Append-only JSONL was better but still conflicted, because
+both sides append at the same end-of-file position — verified with
+`git merge-file`, which reported a conflict even though the two added rows were
+different. One file per checkpoint is the version that actually works: two
+machines recording different checkpoints create different files, which git
+merges without even noticing. Recording the *same* checkpoint twice still
+conflicts, and should — that is a real collision.
+
+`RESULTS.md` is a **build artifact and is gitignored**, fully derived from these
+records. Rebuild or print it with `python scripts/results.py`.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+RECORDS = Path("experiments/records")  # one JSON file per checkpoint
+RESULTS_MD = Path("RESULTS.md")        # derived, gitignored
+
+# Commentary lives in code so nobody has to hand-edit a generated file.
+# The reference every result is quoted against. It must be a row that a run
+# directory can produce: the previous reference claimed 0.950 at rp=1.10, matched
+# no run on any machine, and disagreed with both measurements that did exist.
+BASELINE = "bf16_rp120"
+
+CURATED_NOTES: dict[str, str] = {
+    "bf16_rp120": "**REFERENCE** bf16 at its actual optimum, rp=1.20. acc|finished 1.000",
+    "bf16_rp110": "bf16 at rp=1.10 — identical to w16 on all three metrics, so the "
+                  "quantize/save/load path is exactly lossless",
+    "bf16-greedy-nopenalty": "same weights, no penalty — decoding is worth ~+30 pts",
+    "w16": "identity check: quantize/save/eval path is lossless",
+    "mlp4": "Red Hat's recipe (GDN untouched). No loss vs w16.",
+    "mlp4_attn4": "full_attn@4 costs 21 pts for 0.44 GB — worst trade found",
+    "mlp4_lin8": "linear_attn@8 free on top of attn@4 — GDN survives INT8",
+    "all4_lin8": "embed@4 looked free here (−3 pts); round 4 revised this",
+    "mlp4_lin4": "linear_attn@4 cliff — later shown to be a full_attn@4 artifact",
+    "all4": "acc|finished 1.000 — the loss is pure truncation",
+    "mlp3_lin8": "mlp@3 cliff (RTN; untested with a calibrated quantizer)",
+    "mlp3_lin4": "98% truncated, yet acc|finished still 1.000",
+    "all3": "collapsed — this is the RTN row of the published recovery ladder",
+    "mlp2_lin8": "distinct failure: gibberish, stops at 1k tokens, 0% parse",
+    "m4l16a16e4": "embed@4 free here; GDN and attention untouched",
+    "m4l8a16e4": "linear_attn@8 free AGAIN (+0.010) for −1.00 GB",
+    "m4l8a8e8": "3.14 GB / 0.890 at rp=1.10",
+    "m4l8a8e4": "acc|finished 0.988 — the 8 pt gap was all truncation",
+    "m4l4a8e4": "GDN@4 at rp=1.10",
+    "m3l8a8e4": "MLP cliff between 4 and 3 bits (RTN)",
+    "m4l8a8e4_rp12": "2.83 GB at w16 accuracy — rp=1.20 not 1.10, zero bytes",
+    "m4l8a8e4_g64": "g=64 does not help (−0.030) and costs +0.06 GB",
+    "m4l8a16e4_g64": "g=64 does not help (−0.080)",
+    "m4l8a8e8_g64": "g=64 flat (+0.010) for +0.07 GB — finer groups rejected",
+    "m4l6a8e4_g64": "linear_attn@6 at g=64",
+    "m5l6a8e4_rp12": "**BEST OVERALL** 2.86 GB / 0.950 = bf16 exactly, 2.94x, fewer tokens than bf16",
+    "m5l6a8e8_rp12": "3.18 GB / 0.950 — same score as embed@4 above: embed is free once mlp is 5",
+    "m4l8a8e8_rp12": "0.890 → 0.920 from rp alone; mlp@4 caps out here regardless of the other bits",
+    "m4l6a8e8_rp12": "2.89 GB / 0.920 — superseded by m5l6a8e4: fewer bytes AND 3 pts more",
+    "m4l5a8e4_rp12": "2.45 GB / 0.860 — GDN survives 5 bits",
+    "m4l6a8e4_rp12": "2.57 GB / 0.840 — embed@4 costs ~8 pts vs embed@8 here",
+    "m4l4a8e4_rp12": "0.670 → 0.750 purely from retuning the penalty",
+    # Round 7 — the overnight sweep on buckets A and B.
+    "all3_rp13": "rp=1.30 made truncation WORSE (1.000) — past the useful range",
+    "all3_rp14": "trunc 0.900, still 0.000 — finishing more does not mean answering",
+    "all3_rp15": "trunc 0.660 and tokens −23%, accuracy STILL 0.000. Law 1 breaks here",
+    "m3l6a8e8_rp12": "mlp@3, everything else healthy: **0.300** — damaged, NOT dead",
+    "m3l6a8e8_rp135": "same weights at rp=1.35: 0.300 → 0.030. The penalty optimum has a ceiling",
+    "m3l3a8e8_rp12": "0.000 — adding GDN@3 on top of mlp@3 is what actually kills it",
+    "m3l3a8e8_g64_rp12": "g=64 does not rescue 3-bit either — 0.000",
+    "m4l6a8e4_rp12_sh2": "2x MLP sharing, no recovery: 0.000. Even the gentle factor collapses",
+    "m4l6a8e4_rp12_sh4": "4x MLP sharing, no recovery: 0.000 at trunc 0.46 — finishes, wrong",
+    "m5l6a8e4_rp12_v64k": "INVALID — probe read the wrong JSON key and kept only ~4k tokens",
+    "m5l6a8e4_rp12_v32k": "INVALID — same bug as the 64k row",
+    # Round 8 — calibrated GPTQ and clean BF16 sharing controls.
+    "gptq_mlp3_gar_rp12": "GPTQ MLP@3 with GAR: 0.870; 3-bit weights are healthy under calibration",
+    "gptq_mlp3_descact_rp12": "GPTQ MLP@3 with desc_act: ties GAR at 0.870",
+    "gptq_m3l4a8e8_gar_rp12": "**BEST GPTQ PROBE** 2.36 GB / 0.910; only 3 pts below re-measured bf16",
+    "gptq_m3l4a6e4_gar_rp12": "1.97 GB / 0.860; +30k vocab projects to 1.678 GB but leaves only 4 MB slack",
+    "gptq_m3l3a8e8_gar_rp12": "**B PRECURSOR** exact RTN recipe rose 0.000→0.860; +30k vocab projects to 1.664 GB",
+    "w16_sh2_mean_rp12": "BF16 2x MLP sharing, mean init: 0.000. Quantization was not the cause",
+    "w16_sh2_first_rp12": "BF16 2x MLP sharing, first init: finishes 77%, still 0.000",
+    "w16_sh4_mean_rp12": "BF16 4x MLP sharing, mean init: 0.000; untrained sharing closed",
+    # Round 9 — physical Bucket-B artifact and embedding-row controls.
+    "b0_zero_rp12": "**PHYSICAL B CANDIDATE** 1.664 GB / 0.780; exact GPTQ codes + 30k stored rows",
+    "b0_zero_checkpoint_rp12": "same physical B0 on n=560: 0.786, acc|finished 0.971",
+    "b0_zero_nomask_n30_rp12": "n=30 deployment control without an explicit output keep-set: ties zero fill",
+    "b0_mean_n30_rp12": "submission-valid mean omitted-row fill: no gain over zero on n=30",
+    "b0_base_n30_rp12": "INVALID diagnostic: original omitted rows improve 0.767→0.833 on n=30",
+    # Round 10 — B vocabulary/decoding plateau and A delta-sharing capacity.
+    "r10_b1_pf30_rp115": "30k problem-first physical B, rp=1.15: pure termination loss",
+    "r10_b1_pf30_rp120": "30k problem-first physical B, rp=1.20; no material gain over Round 9",
+    "r10_b1_pf30_rp125": "30k physical B over-penalized: 0.800→0.670",
+    "r10_b2_pf36_rp115": "36k problem-first physical B, rp=1.15",
+    "r10_b2_pf36_rp120": "**ROUND 10 GATE WINNER** 1.680 GB / 0.810; must defer to n=560",
+    "r10_b2_pf36_rp125": "lower truncation, but acc|finished collapses 0.988→0.860",
+    "r10_b2_pf36_checkpoint_rp120": "**PHYSICAL B FRONTIER** n=560 0.787: one problem above Round 9, effectively tied",
+    "r10_a_delta_f2_r64_n30": "factor 2/rank 64 delta sharing: finishes 43%, still 0/30",
+    "r10_a_delta_f8_r64_n30": "factor 8/rank 64 delta sharing: 100% truncation, 0/30",
+    "r10_a_delta_f8_r32_n30": "factor 8/rank 32 delta sharing: 100% truncation, 0/30",
+    "r10_a_delta_f8_r16_n30": "factor 8/rank 16 size endpoint: 100% truncation, 0/30",
+    # Round 12 — learned cross-band transfer and creative controls.
+    "r12_short2500_c_checkpoint": "**NEW C FRONTIER** short2500 + m5l6a8e4: n=560 0.916",
+    "r12_balanced6000_c_checkpoint": "balanced learned source: healthy 0.914, below short2500",
+    "r12_balanced6000_gauge_c_checkpoint": "gauge cuts truncation but trades conditional correctness; accuracy stays 0.914",
+    "r12_short2500_gauge_c_checkpoint": "exact SwiGLU gauge control: 0.909 vs ordinary short C 0.916",
+    "r12_hadamard_c_checkpoint": "fixed-transform codec: 0.887; rejected",
+    "r12_paro_int4_checkpoint": "public Paro method control: 0.886 and repository outside C",
+    "r12_short2500_b_zero_checkpoint": "short-source physical B zero fill: n=560 0.814",
+    "r12_short2500_b_token_checkpoint": "**NEW B FRONTIER** token predictor: n=560 0.834 vs zero 0.814",
+}
+
+HEADER = """# Results ledger
+
+One row per **checkpoint**. Sweeps and ablations live in `experiments/`.
+
+**Generated file — do not edit.** Rebuild with `python scripts/results.py`.
+Source of truth is one file per checkpoint under `experiments/records/`, so two
+machines recording different results never touch the same file.
+
+Eval: `gate` (MATH-500 levels 4-5, n=100), 32,768-token cap, greedy. Size is the
+compressed size from the quantization recipe; `pack.py` verifies those are real
+bytes to 0.0%.
+
+Noise floor ~±5 points unpaired — use `scripts/05_compare_runs.py` for paired
+McNemar before believing a small delta.
+
+"""
+
+COLUMNS = [
+    ("checkpoint", "checkpoint", "s"),
+    ("size GB", "size_gb", ".2f"),
+    ("bpw", "bits_per_weight", ".2f"),
+    ("acc", "accuracy", ".3f"),
+    ("trunc", "truncation_rate", ".3f"),
+    ("mean tok", "mean_generated_tokens", ".0f"),
+    ("notes", "notes", "s"),
+]
+
+
+def load() -> list[dict[str, Any]]:
+    """Every recorded checkpoint, largest first."""
+    if not RECORDS.exists():
+        return []
+    rows = [json.loads(p.read_text()) for p in sorted(RECORDS.glob("*.json"))]
+    return sorted(rows, key=lambda r: -(r.get("size_gb") or 0))
+
+
+def _safe_name(checkpoint: str) -> str:
+    return "".join(c if c.isalnum() or c in "-_" else "_" for c in checkpoint)
+
+
+def record(
+    checkpoint: str,
+    *,
+    size_gb: float,
+    bits_per_weight: float,
+    summary: dict[str, Any],
+    notes: str = "",
+    run_name: str = "",
+    suite: str = "gate",
+) -> None:
+    """Write one checkpoint's result to its own file.
+
+    ``suite`` is part of the identity, not decoration. The same recipe evaluated
+    on `gate` (n=100) and on `checkpoint` (n=560) produces two different numbers,
+    and without this they collide on one filename -- the second silently
+    replacing the first in a table whose header promises every row is `gate`.
+    Callers must fold the suite into ``checkpoint`` for anything but `gate`.
+    """
+    RECORDS.mkdir(parents=True, exist_ok=True)
+    row = {
+        "checkpoint": checkpoint,
+        "suite": suite,
+        "size_gb": size_gb,
+        "bits_per_weight": bits_per_weight,
+        "accuracy": summary.get("accuracy"),
+        "truncation_rate": summary.get("truncation_rate"),
+        "mean_generated_tokens": summary.get("mean_generated_tokens"),
+        "notes": notes,
+        "run_name": run_name,
+    }
+    (RECORDS / f"{_safe_name(checkpoint)}.json").write_text(json.dumps(row, indent=2))
+
+
+def render(rows: list[dict[str, Any]]) -> str:
+    head = "| " + " | ".join(name for name, *_ in COLUMNS) + " |"
+    align = "|" + "|".join(":--" if f == "s" else "--:" for _, _, f in COLUMNS) + "|"
+    lines = [HEADER, head, align]
+
+    for row in rows:
+        cells = []
+        for _, key, fmt in COLUMNS:
+            value = row.get(key)
+            if key == "notes":
+                cells.append(CURATED_NOTES.get(row["checkpoint"], str(value or "")) or "—")
+            elif value is None:
+                cells.append("—")
+            elif fmt == "s":
+                cells.append(str(value))
+            else:
+                cells.append(f"{value:{fmt}}")
+        lines.append("| " + " | ".join(cells) + " |")
+
+    baseline = next((r for r in rows if r["checkpoint"] == BASELINE), None)
+    if baseline and baseline.get("accuracy") is not None:
+        lines += [
+            "",
+            f"Baseline: **bf16 = {baseline['accuracy']:.3f} acc, "
+            f"{baseline['truncation_rate']:.3f} trunc, "
+            f"{baseline['mean_generated_tokens']:.0f} tokens, "
+            f"{baseline['size_gb']:.2f} GB** "
+            f"(`{BASELINE}`, greedy + repetition_penalty=1.20 — bf16's actual optimum).",
+        ]
+    return "\n".join(lines) + "\n"
+
+
+def write_results_md() -> Path:
+    RESULTS_MD.write_text(render(load()))
+    return RESULTS_MD

src/eaimath/model.py

@@ -0,0 +1,440 @@
+"""Loading, inspecting and slimming the Qwen3.5-4B checkpoint.
+
+Qwen3.5-4B is a hybrid multimodal model — 24 Gated DeltaNet (linear attention)
+layers interleaved with 8 full-attention layers, plus a 24-block vision tower
+and a multi-token-prediction head. For a text-only math benchmark the vision
+tower and MTP head are never executed, so they are ~455M parameters (9.7% of the
+checkpoint) of pure dead weight.
+
+Exact class names and module paths for this architecture vary across
+transformers releases, so everything here introspects the loaded module tree
+rather than hardcoding paths. ``describe_model`` exists to dump ground truth
+before we commit to any of it.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+
+DEFAULT_MODEL = "Qwen/Qwen3.5-4B"
+
+# First matching pattern wins, so order matters: `visual`/`mtp` must precede the
+# generic `mlp`/`norm` patterns since those submodules contain MLPs too.
+BUDGET_GROUPS: tuple[tuple[str, str], ...] = (
+    ("vision", r"visual|vision"),
+    ("mtp", r"(^|\.)mtp\."),
+    ("embed", r"embed_tokens|lm_head"),
+    ("linear_attn", r"linear_attn"),
+    ("full_attn", r"self_attn"),
+    ("mlp", r"\.mlp\."),
+    ("norm", r"norm"),
+)
+
+# Tiny but precision-critical: the SSM recurrence dynamics. Keeping every one of
+# these in fp32 costs ~3 MB, and quantizing them is what makes low-bit SSM
+# quantization collapse.
+SSM_SENSITIVE = r"A_log|dt_bias|conv1d|in_proj_a|in_proj_b"
+
+
+@dataclass
+class GroupStats:
+    name: str
+    num_tensors: int = 0
+    num_params: int = 0  # unique storage only
+    num_bytes: int = 0  # unique storage only
+    raw_params: int = 0  # counting tied aliases twice
+    num_aliases: int = 0
+    examples: list[str] = field(default_factory=list)
+
+    def add(self, name: str, tensor: torch.Tensor, *, counted: bool) -> None:
+        self.num_tensors += 1
+        self.raw_params += tensor.numel()
+        if counted:
+            self.num_params += tensor.numel()
+            self.num_bytes += tensor.numel() * tensor.element_size()
+        else:
+            self.num_aliases += 1
+        if len(self.examples) < 3:
+            self.examples.append(name + ("  (tied alias)" if not counted else ""))
+
+
+def _group_for(name: str) -> str:
+    for group, pattern in BUDGET_GROUPS:
+        if re.search(pattern, name):
+            return group
+    return "other"
+
+
+def parameter_budget(state_dict: dict[str, torch.Tensor]) -> dict[str, Any]:
+    """Group every tensor by component and report params / bytes per group.
+
+    Tied weights (Qwen3.5 ties ``lm_head`` to ``embed_tokens``) appear twice in a
+    ``state_dict`` while sharing one allocation. Counting both would overstate
+    the checkpoint by 636M params / 1.27 GB here — and checkpoint bytes is the
+    metric we are graded on — so aliases are detected by storage pointer and
+    counted once.
+    """
+    groups: dict[str, GroupStats] = {}
+    seen_storage: dict[int, str] = {}
+    aliases: list[dict[str, str]] = []
+
+    for name, tensor in state_dict.items():
+        if not isinstance(tensor, torch.Tensor):
+            continue
+        pointer = tensor.data_ptr()
+        is_alias = pointer != 0 and pointer in seen_storage
+        if is_alias:
+            aliases.append({"name": name, "aliases": seen_storage[pointer]})
+        else:
+            seen_storage[pointer] = name
+        groups.setdefault(_group_for(name), GroupStats(_group_for(name))).add(
+            name, tensor, counted=not is_alias
+        )
+
+    total_params = sum(g.num_params for g in groups.values())
+    total_bytes = sum(g.num_bytes for g in groups.values())
+    return {
+        "total_params": total_params,
+        "total_bytes": total_bytes,
+        "total_gb": total_bytes / 1e9,
+        "raw_params_with_aliases": sum(g.raw_params for g in groups.values()),
+        "tied_aliases": aliases,
+        "groups": {
+            name: {
+                "num_tensors": g.num_tensors,
+                "num_params": g.num_params,
+                "num_bytes": g.num_bytes,
+                "gb": g.num_bytes / 1e9,
+                "num_aliases": g.num_aliases,
+                "pct_params": 100 * g.num_params / total_params if total_params else 0.0,
+                "examples": g.examples,
+            }
+            for name, g in sorted(groups.items(), key=lambda kv: -kv[1].num_params)
+        },
+    }
+
+
+def checkpoint_size_bytes(state_dict: dict[str, torch.Tensor]) -> int:
+    """Unique bytes in a state_dict — the number the leaderboard scores."""
+    return parameter_budget(state_dict)["total_bytes"]
+
+
+def format_budget(budget: dict[str, Any]) -> str:
+    lines = [
+        f"{'component':<16}{'params':>16}{'% params':>11}{'GB':>9}{'tensors':>9}",
+        "-" * 61,
+    ]
+    for name, stats in budget["groups"].items():
+        suffix = f"  ({stats['num_aliases']} tied)" if stats["num_aliases"] else ""
+        lines.append(
+            f"{name:<16}{stats['num_params']:>16,}{stats['pct_params']:>10.1f}%"
+            f"{stats['gb']:>9.3f}{stats['num_tensors']:>9}{suffix}"
+        )
+    lines.append("-" * 61)
+    lines.append(
+        f"{'TOTAL':<16}{budget['total_params']:>16,}{100.0:>10.1f}%"
+        f"{budget['total_gb']:>9.3f}"
+    )
+    if budget.get("tied_aliases"):
+        excess = budget["raw_params_with_aliases"] - budget["total_params"]
+        lines.append(
+            f"  ({len(budget['tied_aliases'])} tied alias tensor(s) counted once; "
+            f"naive state_dict sum would overstate by {excess:,} params)"
+        )
+    return "\n".join(lines)
+
+
+def load_model(
+    model_id: str = DEFAULT_MODEL,
+    dtype: str = "bfloat16",
+    device_map: str | None = "auto",
+    cache_dir: str | None = None,
+    device: str | None = None,
+    multimodal: bool = False,
+):
+    """Load the model, trying each plausible auto-class for this architecture.
+
+    On transformers >= 5, ``AutoModelForCausalLM`` resolves Qwen3.5 to
+    ``Qwen3_5ForCausalLM`` and materializes only the language model — the vision
+    tower and MTP head are never allocated, so the 455M-param text-only saving
+    happens for free at load time.
+
+    Pass ``device`` to pin the whole model to one GPU. At 8.4 GB it fits on any
+    of ours, and ``device_map="auto"`` across several GPUs would pipeline-shard
+    it instead, adding cross-device hops on every layer for no benefit.
+    """
+    import transformers
+
+    torch_dtype = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }[dtype]
+
+    kwargs: dict[str, Any] = {
+        "dtype": torch_dtype,
+        "trust_remote_code": True,
+    }
+    if device is not None:
+        kwargs["device_map"] = {"": device}
+    elif device_map is not None:
+        kwargs["device_map"] = device_map
+    if cache_dir:
+        kwargs["cache_dir"] = cache_dir
+
+    candidates = [
+        "AutoModelForCausalLM",
+        "AutoModelForImageTextToText",
+        "AutoModelForVision2Seq",
+        "AutoModel",
+    ]
+    if multimodal:
+        # vLLM only registers Qwen3_5ForConditionalGeneration and rejects a
+        # text-only Qwen3_5TextConfig (vllm#39231). Any checkpoint we intend to
+        # evaluate with vLLM must therefore keep the multimodal wrapper, even
+        # though the vision tower is never executed for a math prompt.
+        candidates = [
+            "Qwen3_5ForConditionalGeneration",
+            "AutoModelForImageTextToText",
+            "AutoModelForVision2Seq",
+            "AutoModelForCausalLM",
+        ]
+    errors: list[str] = []
+    for class_name in candidates:
+        auto_class = getattr(transformers, class_name, None)
+        if auto_class is None:
+            continue
+        try:
+            model = auto_class.from_pretrained(model_id, **kwargs)
+            model.eval()
+            print(f"[model] loaded {model_id} via {class_name}")
+            return model
+        except Exception as exc:  # noqa: BLE001 - we want the full error list
+            errors.append(f"  {class_name}: {type(exc).__name__}: {exc}")
+
+    raise RuntimeError(
+        f"Could not load {model_id} with any auto-class.\n" + "\n".join(errors)
+    )
+
+
+def load_tokenizer(model_id: str = DEFAULT_MODEL, cache_dir: str | None = None):
+    from transformers import AutoTokenizer
+
+    kwargs: dict[str, Any] = {"trust_remote_code": True}
+    if cache_dir:
+        kwargs["cache_dir"] = cache_dir
+    tokenizer = AutoTokenizer.from_pretrained(model_id, **kwargs)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"  # required for correct batched generation
+    return tokenizer
+
+
+def estimate_memory(
+    model: torch.nn.Module, batch_size: int, seq_len: int
+) -> dict[str, float]:
+    """Predict peak VRAM for a generation run, so OOM is caught before it costs an hour.
+
+    Only the 8 full-attention layers hold a growing KV cache; the 24 Gated
+    DeltaNet layers keep a fixed-size recurrent state regardless of length,
+    which is why this model's cache is far cheaper than a dense transformer's.
+    """
+    config = getattr(model, "config", None)
+    config = getattr(config, "text_config", config)
+
+    layer_types = getattr(config, "layer_types", None)
+    num_layers = getattr(config, "num_hidden_layers", 32)
+    if layer_types:
+        num_full = sum(1 for t in layer_types if "full" in str(t))
+    else:
+        interval = getattr(config, "full_attention_interval", 4)
+        num_full = num_layers // interval
+
+    kv_heads = getattr(config, "num_key_value_heads", 4)
+    head_dim = getattr(config, "head_dim", 256)
+    dtype_size = 2
+
+    kv_bytes_per_token = num_full * kv_heads * head_dim * 2 * dtype_size
+    kv_bytes = kv_bytes_per_token * seq_len * batch_size
+
+    # DynamicCache grows by torch.cat on every decode step: allocate (n+1),
+    # copy, free n. Peak is therefore ~2x the steady-state cache, which is
+    # exactly the "ladder" you see climbing in nvtop until it OOMs.
+    # expandable_segments lets the allocator grow a segment in place, cutting
+    # the transient sharply -- but we still budget for it.
+    growth_factor = 1.4 if "expandable_segments" in os.environ.get(
+        "PYTORCH_CUDA_ALLOC_CONF", ""
+    ) else 2.0
+    kv_peak_bytes = kv_bytes * growth_factor
+
+    # Gated DeltaNet recurrent state: [v_heads, k_dim, v_dim] per layer, length-independent.
+    v_heads = getattr(config, "linear_num_value_heads", 32)
+    k_dim = getattr(config, "linear_key_head_dim", 128)
+    v_dim = getattr(config, "linear_value_head_dim", 128)
+    state_bytes = (num_layers - num_full) * v_heads * k_dim * v_dim * 4 * batch_size
+
+    weight_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
+
+    return {
+        "weights_gb": weight_bytes / 1e9,
+        "kv_cache_gb": kv_bytes / 1e9,
+        "kv_peak_gb": kv_peak_bytes / 1e9,
+        "ssm_state_gb": state_bytes / 1e9,
+        "steady_gb": (weight_bytes + kv_bytes + state_bytes) / 1e9,
+        # What the preflight must check against: peak, not steady state.
+        "total_gb": (weight_bytes + kv_peak_bytes + state_bytes) / 1e9,
+        "kv_bytes_per_token": kv_bytes_per_token,
+        "growth_factor": growth_factor,
+    }
+
+
+def free_vram_gb(device: str | torch.device) -> tuple[float, float]:
+    """(free, total) GB on ``device`` — reflects other users' jobs on a shared box."""
+    free, total = torch.cuda.mem_get_info(torch.device(device))
+    return free / 1e9, total / 1e9
+
+
+def describe_model(model: torch.nn.Module, max_depth: int = 3) -> dict[str, Any]:
+    """Summarize the module tree — the ground truth we build stripping on."""
+    tree: list[dict[str, Any]] = []
+    for name, module in model.named_modules():
+        depth = name.count(".")
+        if name and depth <= max_depth:
+            n_params = sum(p.numel() for p in module.parameters(recurse=True))
+            tree.append(
+                {
+                    "path": name,
+                    "class": type(module).__name__,
+                    "depth": depth,
+                    "num_params": n_params,
+                }
+            )
+
+    state_dict = model.state_dict()
+    sensitive = [n for n in state_dict if re.search(SSM_SENSITIVE, n)]
+    return {
+        "model_class": type(model).__name__,
+        "num_parameters": sum(p.numel() for p in model.parameters()),
+        "module_tree": tree,
+        "budget": parameter_budget(state_dict),
+        "ssm_sensitive_tensors": {
+            "count": len(sensitive),
+            "num_params": sum(state_dict[n].numel() for n in sensitive),
+            "names": sensitive[:20],
+        },
+        "state_dict_prefixes": sorted(
+            {".".join(n.split(".")[:2]) for n in state_dict}
+        ),
+    }
+
+
+def find_submodule(model: torch.nn.Module, pattern: str) -> list[str]:
+    """Paths of modules whose name matches ``pattern`` and whose parent does not."""
+    regex = re.compile(pattern)
+    hits = [name for name, _ in model.named_modules() if name and regex.search(name)]
+    # Keep only the shallowest match on each branch.
+    return [h for h in hits if not any(h.startswith(o + ".") for o in hits)]
+
+
+def strip_unused_modules(
+    model: torch.nn.Module,
+    drop_vision: bool = True,
+    drop_mtp: bool = True,
+) -> dict[str, Any]:
+    """Delete text-irrelevant submodules in place; return what was removed.
+
+    The vision tower and MTP head together are ~455M params. A text-only math
+    eval never executes either: the ViT has no image inputs, and the MTP head is
+    a speculative-decoding draft head that HF ``generate()`` does not call.
+    """
+    removed: list[dict[str, Any]] = []
+    targets: list[str] = []
+    if drop_vision:
+        targets.append(r"(^|\.)(visual|vision_tower)$")
+    if drop_mtp:
+        targets.append(r"(^|\.)mtp$")
+
+    for pattern in targets:
+        for path in find_submodule(model, pattern):
+            parent = model
+            parts = path.split(".")
+            for part in parts[:-1]:
+                parent = getattr(parent, part)
+            child = getattr(parent, parts[-1], None)
+            if child is None:
+                continue
+            n_params = sum(p.numel() for p in child.parameters(recurse=True))
+            setattr(parent, parts[-1], None)
+            removed.append({"path": path, "num_params": n_params})
+
+    # transformers >= 5 loads Qwen3.5 through Qwen3_5ForCausalLM, which never
+    # allocates the vision tower or MTP head. Finding nothing to remove is the
+    # expected outcome there, not a failure.
+    already_text_only = not removed and not find_submodule(model, r"visual|vision_tower|mtp")
+
+    return {
+        "removed": removed,
+        "num_params_removed": sum(r["num_params"] for r in removed),
+        "bytes_removed_bf16": 2 * sum(r["num_params"] for r in removed),
+        "already_text_only": already_text_only,
+    }
+
+
+# Files the model repo carries that `save_pretrained` does not reproduce.
+# vLLM's multimodal path builds an image processor even for a text-only prompt,
+# so a checkpoint missing `preprocessor_config.json` fails to load outright.
+AUXILIARY_FILES = (
+    "preprocessor_config.json",
+    "video_preprocessor_config.json",
+    "chat_template.jinja",
+)
+
+
+def save_checkpoint(
+    model: torch.nn.Module,
+    out_dir: str | Path,
+    source_model: str = DEFAULT_MODEL,
+    cache_dir: str | None = None,
+) -> list[str]:
+    """Write a checkpoint that vLLM and transformers can both load.
+
+    ``model.save_pretrained`` emits weights + config, and the tokenizer covers
+    vocab/merges, but the processor and chat-template files come from the source
+    repo and must be copied across explicitly.
+    """
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out)
+    load_tokenizer(source_model, cache_dir=cache_dir).save_pretrained(out)
+
+    copied: list[str] = []
+    source_dir = Path(source_model)
+    for filename in AUXILIARY_FILES:
+        target = out / filename
+        if target.exists():
+            copied.append(filename)
+            continue
+        try:
+            if source_dir.is_dir():
+                candidate = source_dir / filename
+                if not candidate.is_file():
+                    continue
+                shutil.copy(candidate, target)
+            else:
+                from huggingface_hub import hf_hub_download
+
+                downloaded = hf_hub_download(
+                    source_model, filename, cache_dir=cache_dir
+                )
+                shutil.copy(downloaded, target)
+            copied.append(filename)
+        except Exception as exc:  # noqa: BLE001 - optional files; report and continue
+            print(f"[save_checkpoint] could not copy {filename}: {exc}")
+    return copied

src/eaimath/pack.py

@@ -0,0 +1,397 @@
+"""Real bit-packing — turning our analytic sizes into actual bytes on disk.
+
+Everything so far has used *simulated* quantization: weights rounded onto a
+coarse grid but stored back as bf16, with the compressed size computed
+analytically from the recipe. That is the right tool for research — it isolates
+accuracy effects without committing to a format — but a submission is scored on
+bytes that actually exist.
+
+This module closes that gap. `pack_tensor` stores an n-bit quantized tensor as
+a flat `uint8` buffer plus fp16 scales, so `numel * element_size` over the saved
+state dict equals the number we have been reporting. `unpack_tensor` inverts it.
+
+Sub-byte widths are packed densely (two 4-bit values per byte, eight 3-bit
+values per three bytes, and so on) rather than padded to a byte, because padding
+would silently inflate 3-bit to 4-bit and make our size claims wrong.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+
+# Packing 2.26B values at once would materialise a ~9 GB bit array, so the
+# general path works in chunks.
+CHUNK = 1 << 24
+
+
+def _pack_uint8_chunked(values: np.ndarray, bits: int) -> np.ndarray:
+    out = []
+    for start in range(0, values.size, CHUNK):
+        chunk = values[start : start + CHUNK]
+        bit_planes = ((chunk[:, None] >> np.arange(bits - 1, -1, -1)) & 1).astype(np.uint8)
+        out.append(np.packbits(bit_planes.reshape(-1)))
+    return np.concatenate(out) if out else np.zeros(0, dtype=np.uint8)
+
+
+def _unpack_uint8_chunked(packed: np.ndarray, bits: int, count: int) -> np.ndarray:
+    bit_stream = np.unpackbits(packed)[: count * bits].reshape(count, bits)
+    weights = (1 << np.arange(bits - 1, -1, -1)).astype(np.uint32)
+    return (bit_stream * weights).sum(axis=1).astype(np.uint32)
+
+
+def pack_bits(values: np.ndarray, bits: int) -> np.ndarray:
+    """Pack unsigned integers of width ``bits`` into a dense uint8 buffer."""
+    flat = np.ascontiguousarray(values.reshape(-1)).astype(np.uint32)
+    if bits == 8:
+        return flat.astype(np.uint8)
+    if bits == 4:
+        # Fast path: two nibbles per byte. Pad to even length.
+        padded = np.pad(flat, (0, flat.size % 2)).astype(np.uint8)
+        return (padded[0::2] << 4 | padded[1::2]).astype(np.uint8)
+    if bits == 16:
+        return flat.astype(np.uint16).view(np.uint8)
+    # Preserve widths above eight bits. The general bit-plane packer accepts
+    # uint32 values; narrowing here silently wrapped 9--15 bit codes.
+    return _pack_uint8_chunked(flat, bits)
+
+
+def unpack_bits(packed: np.ndarray, bits: int, count: int) -> np.ndarray:
+    if bits == 8:
+        return packed[:count].astype(np.uint32)
+    if bits == 4:
+        high, low = packed >> 4, packed & 0x0F
+        return np.stack([high, low], axis=1).reshape(-1)[:count].astype(np.uint32)
+    if bits == 16:
+        return packed.view(np.uint16)[:count].astype(np.uint32)
+    return _unpack_uint8_chunked(packed, bits, count)
+
+
+def _store_scale(
+    scale: torch.Tensor,
+    normalization: str,
+) -> tuple[torch.Tensor, int]:
+    """Encode a positive scale field in two bytes per group.
+
+    ``pow2`` is a block-floating representation: one exact power-of-two
+    exponent is shared by the tensor entry, while every group retains an FP16
+    mantissa.  It costs the same tensor bytes as raw FP16 scales but prevents
+    the tiny scales used by 9--12 bit weights from falling into FP16's
+    subnormal range.
+    """
+    if normalization == "none":
+        return scale.reshape(-1).to(torch.float16), 0
+    if normalization != "pow2":
+        raise ValueError(
+            f"unsupported scale normalization {normalization!r}; expected none or pow2"
+        )
+    maximum = float(scale.detach().abs().max()) if scale.numel() else 0.0
+    exponent = 0 if maximum <= 0.0 else -math.floor(math.log2(maximum))
+    # A normalized maximum in [1, 2) is comfortably inside FP16's normal
+    # range.  The exponent is metadata (one small integer per entry), while
+    # the charged scale stream remains one FP16 value per group.
+    normalized = scale * math.ldexp(1.0, exponent)
+    return normalized.reshape(-1).to(torch.float16), int(exponent)
+
+
+def pack_tensor(
+    weight: torch.Tensor,
+    bits: int,
+    group_size: int = 128,
+    symmetric: bool = True,
+    scale_normalization: str = "none",
+) -> dict[str, Any]:
+    """Quantize and pack one tensor into buffers whose bytes are the real cost."""
+    if bits >= 16:
+        return {"kind": "raw", "data": weight.to(torch.bfloat16).cpu()}
+
+    shape = tuple(weight.shape)
+    group = group_size if group_size > 0 and shape[-1] % group_size == 0 else shape[-1]
+    flat = weight.detach().float().cpu().reshape(-1, group)
+
+    if symmetric:
+        qmax = 2 ** (bits - 1) - 1
+        scale = (flat.abs().amax(dim=1, keepdim=True) / qmax).clamp(min=1e-8)
+        q = torch.round(flat / scale).clamp(-qmax - 1, qmax)
+        codes = (q + (qmax + 1)).to(torch.int64)  # shift to unsigned for packing
+        zero = None
+    else:
+        qmax = 2**bits - 1
+        wmin, wmax = flat.amin(dim=1, keepdim=True), flat.amax(dim=1, keepdim=True)
+        scale = ((wmax - wmin) / qmax).clamp(min=1e-8)
+        zero = torch.round(-wmin / scale)
+        codes = torch.clamp(torch.round(flat / scale) + zero, 0, qmax).to(torch.int64)
+
+    packed = pack_bits(codes.numpy().astype(np.uint32), bits)
+    stored_scale, scale_exponent = _store_scale(scale, scale_normalization)
+    entry = {
+        "kind": "quant",
+        "packed": torch.from_numpy(packed),  # uint8 -> 1 byte each, honestly counted
+        "scale": stored_scale,
+        "shape": shape,
+        "bits": bits,
+        "group": group,
+        "symmetric": symmetric,
+        "dtype": str(weight.dtype),
+        "scale_normalization": scale_normalization,
+    }
+    entry["scale_exponent"] = scale_exponent
+    if zero is not None:
+        entry["zero"] = zero.reshape(-1).to(torch.float16)
+    return entry
+
+
+def unpack_tensor(entry: dict[str, Any]) -> torch.Tensor:
+    if entry["kind"] == "raw":
+        return entry["data"]
+
+    bits, group, shape = entry["bits"], entry["group"], tuple(entry["shape"])
+    count = int(np.prod(shape))
+    codes = unpack_bits(entry["packed"].numpy(), bits, count)
+    codes = torch.from_numpy(codes.astype(np.int64)).reshape(-1, group)
+    scale = entry["scale"].float().reshape(-1, 1)
+    scale = scale * math.ldexp(1.0, -int(entry.get("scale_exponent", 0)))
+
+    if entry["symmetric"]:
+        qmax = 2 ** (bits - 1) - 1
+        values = (codes - (qmax + 1)).float() * scale
+    else:
+        values = (codes.float() - entry["zero"].float().reshape(-1, 1)) * scale
+
+    dtype = getattr(torch, entry["dtype"].replace("torch.", ""))
+    return values.reshape(shape).to(dtype)
+
+
+def _nearest_code_indices(values: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor:
+    """Nearest scalar-codebook entry without materialising a huge 3-D tensor."""
+    best_distance = torch.full_like(values, float("inf"), dtype=torch.float32)
+    best_index = torch.zeros_like(values, dtype=torch.int64)
+    for index, center in enumerate(codebook):
+        distance = (values - center).square()
+        replace = distance < best_distance
+        best_distance = torch.where(replace, distance, best_distance)
+        best_index = torch.where(replace, index, best_index)
+    return best_index
+
+
+def pack_adaptive_codebook(
+    weight: torch.Tensor,
+    *,
+    bits: int = 4,
+    group_size: int = 128,
+    importance: torch.Tensor | None = None,
+    codebook_count: int = 2,
+    iterations: int = 4,
+) -> tuple[dict[str, Any], dict[str, float | int]]:
+    """Pack groups with tiny learned scalar codebooks and sign-bit selectors.
+
+    The sign of each otherwise-positive FP16 group scale encodes which of two
+    codebooks is used.  The absolute scale and packed indices are ordinary
+    physical tensors, so byte accounting includes every codebook value while
+    the selector itself costs no additional storage.
+    """
+    if bits != 4:
+        raise ValueError("adaptive codebooks currently require four-bit indices")
+    if codebook_count != 2:
+        raise ValueError("scale-sign encoding requires exactly two codebooks")
+    if iterations < 1:
+        raise ValueError("iterations must be positive")
+    shape = tuple(weight.shape)
+    if weight.ndim != 2 or shape[-1] % group_size:
+        raise ValueError(
+            f"adaptive codebooks need a 2-D matrix divisible by group size; got {shape}"
+        )
+    device = weight.device
+    flat = weight.detach().float().reshape(-1, group_size)
+    groups = flat.shape[0]
+    if importance is None:
+        feature_weight = torch.ones(group_size, device=device, dtype=torch.float32)
+    else:
+        feature_weight = importance.detach().to(device=device, dtype=torch.float32)
+        if feature_weight.numel() != shape[-1]:
+            raise ValueError(
+                f"importance has {feature_weight.numel()} values, expected {shape[-1]}"
+            )
+        feature_weight = feature_weight.reshape(-1, group_size).repeat(shape[0], 1)
+        feature_weight = feature_weight / feature_weight.mean().clamp_min(1e-12)
+    if feature_weight.ndim == 1:
+        feature_weight = feature_weight.expand(groups, -1)
+
+    scale = flat.abs().amax(dim=1).clamp_min(1e-8)
+    normalized = flat / scale[:, None]
+    shape_score = normalized.abs().mean(dim=1)
+    split = shape_score.median()
+    partitions = (shape_score > split).long()
+    levels = 2**bits
+    quantiles = torch.linspace(0, 1, levels + 2, device=device)[1:-1]
+    books = []
+    for book in range(codebook_count):
+        values = normalized[partitions == book].reshape(-1)
+        if values.numel() < levels:
+            values = normalized.reshape(-1)
+        # Bound initialization cost for unusually large matrices while keeping
+        # deterministic coverage across their complete flattened range.
+        if values.numel() > 2_000_000:
+            stride = max(1, values.numel() // 2_000_000)
+            values = values[::stride][:2_000_000]
+        books.append(torch.quantile(values, quantiles).sort().values)
+    codebooks = torch.stack(books)
+
+    selector = partitions
+    indices = torch.zeros_like(flat, dtype=torch.int64)
+    for _ in range(iterations):
+        errors = []
+        candidate_indices = []
+        for book in range(codebook_count):
+            codes = _nearest_code_indices(flat / scale[:, None], codebooks[book])
+            reconstructed = codebooks[book][codes] * scale[:, None]
+            errors.append(((flat - reconstructed).square() * feature_weight).sum(dim=1))
+            candidate_indices.append(codes)
+        selector = torch.stack(errors, dim=1).argmin(dim=1)
+        indices = torch.where(
+            selector[:, None].eq(0), candidate_indices[0], candidate_indices[1]
+        )
+        chosen = codebooks[selector[:, None], indices]
+        numerator = (feature_weight * flat * chosen).sum(dim=1)
+        denominator = (feature_weight * chosen.square()).sum(dim=1).clamp_min(1e-12)
+        scale = (numerator / denominator).abs().clamp_min(1e-8)
+        normalized = flat / scale[:, None]
+        for book in range(codebook_count):
+            selected_groups = selector.eq(book)
+            if not bool(selected_groups.any()):
+                continue
+            selected_codes = indices[selected_groups].reshape(-1)
+            selected_values = normalized[selected_groups].reshape(-1)
+            selected_weights = feature_weight[selected_groups].reshape(-1)
+            sums = torch.zeros(levels, device=device).scatter_add_(
+                0, selected_codes, selected_values * selected_weights
+            )
+            counts = torch.zeros(levels, device=device).scatter_add_(
+                0, selected_codes, selected_weights
+            )
+            updated = torch.where(
+                counts > 0, sums / counts.clamp_min(1e-12), codebooks[book]
+            )
+            codebooks[book] = updated.sort().values.clamp(-1.5, 1.5)
+
+    # Recompute once after the final codebook update and then charge the exact
+    # FP16 scale/codebook representation used by the decoder.
+    errors = []
+    candidate_indices = []
+    for book in range(codebook_count):
+        codes = _nearest_code_indices(flat / scale[:, None], codebooks[book])
+        reconstructed = codebooks[book][codes] * scale[:, None]
+        errors.append(((flat - reconstructed).square() * feature_weight).sum(dim=1))
+        candidate_indices.append(codes)
+    selector = torch.stack(errors, dim=1).argmin(dim=1)
+    indices = torch.where(selector[:, None].eq(0), candidate_indices[0], candidate_indices[1])
+    signed_scale = scale.to(torch.float16)
+    signed_scale = torch.where(selector.bool(), -signed_scale.abs(), signed_scale.abs())
+    stored_books = codebooks.to(torch.float16)
+    packed = pack_bits(indices.detach().cpu().numpy().astype(np.uint32), bits)
+    entry: dict[str, Any] = {
+        "kind": "adaptive_codebook",
+        "packed": torch.from_numpy(packed),
+        "scale": signed_scale.detach().cpu(),
+        "codebooks": stored_books.detach().cpu(),
+        "shape": shape,
+        "bits": bits,
+        "group": group_size,
+        "dtype": str(weight.dtype),
+        "selector_encoding": "scale_sign",
+    }
+    restored = unpack_adaptive_codebook(entry).to(device=device, dtype=torch.float32)
+    delta = weight.detach().float() - restored
+    weighted = feature_weight * delta.reshape(-1, group_size).square()
+    report = {
+        "groups": int(groups),
+        "selector_one_groups": int(selector.sum()),
+        "codebook_bytes": int(stored_books.numel() * stored_books.element_size()),
+        "weighted_sse": float(weighted.sum()),
+        "weighted_mse": float(weighted.mean()),
+        "rmse": float(delta.square().mean().sqrt()),
+        "max_abs_error": float(delta.abs().max()),
+    }
+    return entry, report
+
+
+def unpack_adaptive_codebook(entry: dict[str, Any]) -> torch.Tensor:
+    if entry.get("kind") != "adaptive_codebook":
+        raise ValueError(f"not an adaptive-codebook entry: {entry.get('kind')!r}")
+    bits, group, shape = int(entry["bits"]), int(entry["group"]), tuple(entry["shape"])
+    count = int(np.prod(shape))
+    codes = unpack_bits(entry["packed"].numpy(), bits, count)
+    codes_tensor = torch.from_numpy(codes.astype(np.int64)).reshape(-1, group)
+    signed_scale = entry["scale"].float().reshape(-1)
+    selector = torch.signbit(signed_scale).long()
+    scale = signed_scale.abs().reshape(-1, 1)
+    books = entry["codebooks"].float()
+    chosen = books[selector[:, None], codes_tensor]
+    values = chosen * scale
+    dtype = getattr(torch, entry["dtype"].replace("torch.", ""))
+    return values.reshape(shape).to(dtype)
+
+
+def pack_rows(
+    weight: torch.Tensor,
+    row_ids: torch.Tensor,
+    bits: int,
+    group_size: int = 128,
+    symmetric: bool = True,
+) -> dict[str, Any]:
+    """Pack selected rows while retaining the full tensor shape in metadata.
+
+    The original token ids are preserved.  This matters for the course evaluator,
+    which restores the original tokenizer after decompression.  At restore time the
+    omitted rows can be filled deterministically and the selected rows scattered
+    back to their original ids.
+    """
+    if weight.ndim != 2:
+        raise ValueError(f"row packing requires a matrix, got shape {tuple(weight.shape)}")
+    ids = row_ids.detach().to(device="cpu", dtype=torch.int64).sort().values
+    if ids.numel() == 0:
+        raise ValueError("row_ids must not be empty")
+    if int(ids[0]) < 0 or int(ids[-1]) >= weight.shape[0]:
+        raise ValueError(
+            f"row ids [{int(ids[0])}, {int(ids[-1])}] outside 0..{weight.shape[0] - 1}"
+        )
+    if torch.unique_consecutive(ids).numel() != ids.numel():
+        raise ValueError("row_ids contains duplicates")
+
+    selected = weight.detach().cpu().index_select(0, ids)
+    entry = pack_tensor(selected, bits, group_size=group_size, symmetric=symmetric)
+    entry["kind"] = "quant_rows" if entry["kind"] == "quant" else "raw_rows"
+    entry["full_shape"] = tuple(weight.shape)
+    entry["row_ids"] = ids.to(torch.int32)
+    return entry
+
+
+def unpack_rows(entry: dict[str, Any], fill: str = "zero") -> torch.Tensor:
+    """Expand a row-packed matrix with a deterministic omitted-row fill."""
+    if entry["kind"] not in {"quant_rows", "raw_rows"}:
+        raise ValueError(f"not a row-packed entry: {entry.get('kind')!r}")
+    dense_entry = {**entry, "kind": "quant" if entry["kind"] == "quant_rows" else "raw"}
+    selected = unpack_tensor(dense_entry)
+    full_shape = tuple(entry["full_shape"])
+    if fill == "zero":
+        restored = torch.zeros(full_shape, dtype=selected.dtype)
+    elif fill == "mean":
+        mean_row = selected.float().mean(dim=0).to(selected.dtype)
+        restored = mean_row.expand(full_shape[0], -1).clone()
+    else:
+        raise ValueError(f"unsupported row fill {fill!r}; expected zero or mean")
+    restored.index_copy_(0, entry["row_ids"].long(), selected)
+    return restored
+
+
+def state_dict_bytes(packed: dict[str, dict[str, Any]]) -> int:
+    """Real bytes of a packed state dict — the number a submission is scored on."""
+    total = 0
+    for entry in packed.values():
+        for value in entry.values():
+            if isinstance(value, torch.Tensor):
+                total += value.numel() * value.element_size()
+    return total

src/eaimath/peft_compat.py

@@ -0,0 +1,54 @@
+"""Narrow PEFT compatibility helpers for a shared quantization environment."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def is_runtime_quantized(model: Any) -> bool:
+    """Return whether PEFT should treat ``model`` as a quantized runtime model."""
+    if bool(getattr(model, "is_quantized", False)):
+        return True
+    if getattr(model, "quantization_method", None):
+        return True
+    config = getattr(model, "config", None)
+    return bool(getattr(config, "quantization_config", None))
+
+
+QUANT_DISPATCHERS = (
+    "dispatch_eetq",
+    "dispatch_aqlm",
+    "dispatch_awq",
+    "dispatch_gptq",
+    "dispatch_hqq",
+    "dispatch_inc",
+    "dispatch_torchao",
+)
+
+
+def prepare_dense_lora_dispatch(model: Any, *, dispatch_module: Any | None = None) -> bool:
+    """Bypass PEFT's irrelevant quantized dispatchers for a dense BF16 model.
+
+    PEFT 0.18 calls every available quantization backend before its ordinary
+    ``torch.nn.Linear`` fallback. Merely having GPTQModel or TorchAO installed
+    can therefore raise an optional-backend compatibility error even though the
+    target model is not quantized. Round 12 needs those packages elsewhere, so
+    disabling their dispatch functions in this one dense-training process is
+    safer than mutating the shared environment.
+    """
+    if is_runtime_quantized(model):
+        return False
+    if dispatch_module is None:
+        from peft.tuners.lora import model as dispatch_module
+
+    def not_applicable(*_args: Any, **_kwargs: Any) -> None:
+        return None
+
+    for name in QUANT_DISPATCHERS:
+        if not hasattr(dispatch_module, name):
+            raise RuntimeError(f"PEFT dense dispatch compatibility is missing {name}")
+        setattr(dispatch_module, name, not_applicable)
+    # Avoid eager BitsAndBytes imports at the start of PEFT's dispatcher list.
+    dispatch_module.is_bnb_available = lambda: False
+    dispatch_module.is_bnb_4bit_available = lambda: False
+    return True

src/eaimath/quantize.py

@@ -0,0 +1,233 @@
+"""Group-wise weight quantization with exact size accounting.
+
+Simulated ("fake") quantization: weights are quantized then dequantized back to
+bf16, so the saved checkpoint is still a normal HF model that vLLM and
+`transformers` load unchanged. That is not a hack — it mirrors the submission
+pipeline exactly, where ``convert_to_hf_checkpoint`` must emit a full bf16 model.
+The *scored* size is computed analytically from the recipe, not from the bytes on
+disk.
+
+This lets us answer the question that decides our compression ceiling — how far
+can each layer group be pushed before long chain-of-thought breaks — without
+first committing to a packed format or a quantization toolchain.
+
+Round-to-nearest is deliberate as a starting point: it is the floor, needs no
+calibration data, and is enough to *rank* layer-group sensitivity.
+
+But it is a worse floor than we assumed. arXiv 2606.25519 measures CoT *token
+inflation* by quantizer on Qwen3-4B at group size 128, and RTN is the worst of
+the lot: +42.5% at INT4 against GPTQ's +12.0% and rotation-based ParoQuant's
++4.7%. At INT3 the spread across methods reaches 10x. Since our accuracy is
+gated by truncation, token inflation is the quantity that actually costs us
+points -- so moving off RTN is worth considerably more here than the "1-2
+points" that reconstruction-error framing would suggest.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any, Iterable
+
+import torch
+
+KEEP_BITS = 16  # bf16 passthrough
+
+
+@dataclass(frozen=True)
+class QuantSpec:
+    """Quantization for every parameter whose name matches ``pattern``.
+
+    ``group_size`` groups along the input dimension (the last axis); 0 means one
+    group per output channel. Smaller groups cost more scale bytes but track the
+    weight distribution better.
+    """
+
+    pattern: str
+    bits: int
+    group_size: int = 128
+    symmetric: bool = True
+
+    def bits_per_weight(self) -> float:
+        """Effective stored bits per weight, including scale/zero-point overhead."""
+        if self.bits >= KEEP_BITS:
+            return float(KEEP_BITS)
+        if self.group_size <= 0:
+            return float(self.bits)
+        # fp16 scale per group, plus an int zero-point per group when asymmetric.
+        overhead = 16 + (0 if self.symmetric else self.bits)
+        return self.bits + overhead / self.group_size
+
+
+def quantize_dequantize(
+    weight: torch.Tensor, spec: QuantSpec
+) -> torch.Tensor:
+    """Round-trip ``weight`` through ``spec``'s grid, returning the same dtype."""
+    if spec.bits >= KEEP_BITS:
+        return weight
+
+    original_dtype, original_shape = weight.dtype, weight.shape
+    group_size = spec.group_size if spec.group_size > 0 else original_shape[-1]
+    if original_shape[-1] % group_size != 0:
+        # Fall back to per-channel rather than silently mis-grouping.
+        group_size = original_shape[-1]
+
+    flat = weight.reshape(-1, group_size).float()
+
+    if spec.symmetric:
+        qmax = 2 ** (spec.bits - 1) - 1
+        scale = (flat.abs().amax(dim=1, keepdim=True) / qmax).clamp(min=1e-8)
+        q = torch.round(flat / scale).clamp(-qmax - 1, qmax)
+        out = q * scale
+    else:
+        qmax = 2**spec.bits - 1
+        wmin = flat.amin(dim=1, keepdim=True)
+        wmax = flat.amax(dim=1, keepdim=True)
+        scale = ((wmax - wmin) / qmax).clamp(min=1e-8)
+        zero = torch.round(-wmin / scale)
+        q = torch.clamp(torch.round(flat / scale) + zero, 0, qmax)
+        out = (q - zero) * scale
+
+    return out.reshape(original_shape).to(original_dtype)
+
+
+class Recipe:
+    """Ordered specs; first match wins, anything unmatched stays bf16."""
+
+    def __init__(self, specs: Iterable[QuantSpec]):
+        self.specs = list(specs)
+        self._compiled = [(re.compile(s.pattern), s) for s in self.specs]
+
+    def spec_for(self, name: str) -> QuantSpec | None:
+        for regex, spec in self._compiled:
+            if regex.search(name):
+                return spec
+        return None
+
+
+# The vision tower and MTP head are never executed by a text-only math eval, and
+# a real submission simply drops them (llama.cpp's text-only conversion of this
+# model does exactly that). They are excluded from both quantization and the size
+# accounting. We nonetheless keep them *present* in the saved checkpoint, because
+# vLLM refuses to load a text-only Qwen3_5 config (vllm#39231).
+EXCLUDED = r"visual|vision_tower|(^|\.)mtp\."
+
+
+def apply_recipe(
+    model: torch.nn.Module, recipe: Recipe, device: str | None = None
+) -> dict[str, Any]:
+    """Fake-quantize ``model`` in place; return per-group size accounting.
+
+    Tied tensors are counted once — Qwen3.5 ties ``lm_head`` to ``embed_tokens``,
+    and double-counting would overstate the checkpoint by 1.27 GB.
+    """
+    excluded_regex = re.compile(EXCLUDED)
+    seen_storage: set[int] = set()
+    groups: dict[str, dict[str, float]] = {}
+    total_bits = 0.0
+    total_params = 0
+    excluded_params = 0
+
+    with torch.no_grad():
+        for name, param in model.named_parameters():
+            if excluded_regex.search(name):
+                excluded_params += param.numel()
+                continue
+            spec = recipe.spec_for(name)
+            bits = spec.bits_per_weight() if spec else float(KEEP_BITS)
+
+            if spec is not None and spec.bits < KEEP_BITS:
+                target = param.data.to(device) if device else param.data
+                quantized = quantize_dequantize(target, spec)
+                param.data.copy_(quantized.to(param.data.device))
+
+            pointer = param.data_ptr()
+            if pointer in seen_storage:
+                continue  # tied alias: already counted
+            seen_storage.add(pointer)
+
+            label = _group_label(name, spec)
+            entry = groups.setdefault(
+                label, {"params": 0, "bits_per_weight": bits, "bytes": 0.0}
+            )
+            entry["params"] += param.numel()
+            entry["bytes"] += param.numel() * bits / 8
+            total_params += param.numel()
+            total_bits += param.numel() * bits
+
+    total_bytes = total_bits / 8
+    return {
+        "total_params": total_params,
+        "total_bytes": int(total_bytes),
+        "total_gb": total_bytes / 1e9,
+        "effective_bits_per_weight": total_bits / total_params if total_params else 0.0,
+        "compression_vs_bf16": (total_params * 2) / total_bytes if total_bytes else 0.0,
+        "groups": groups,
+        # Present in the saved checkpoint for vLLM compatibility, excluded from
+        # the score because a real submission drops them.
+        "excluded_params": excluded_params,
+        "excluded_gb_bf16": excluded_params * 2 / 1e9,
+    }
+
+
+def _group_label(name: str, spec: QuantSpec | None) -> str:
+    from .model import BUDGET_GROUPS
+
+    for group, pattern in BUDGET_GROUPS:
+        if re.search(pattern, name):
+            return f"{group}@{spec.bits if spec else KEEP_BITS}b"
+    return f"other@{spec.bits if spec else KEEP_BITS}b"
+
+
+def format_size_report(report: dict[str, Any]) -> str:
+    lines = [
+        f"{'group':<22}{'params':>15}{'bpw':>7}{'GB':>9}",
+        "-" * 53,
+    ]
+    for label, stats in sorted(report["groups"].items(), key=lambda kv: -kv[1]["params"]):
+        lines.append(
+            f"{label:<22}{stats['params']:>15,}"
+            f"{stats['bits_per_weight']:>7.2f}{stats['bytes'] / 1e9:>9.3f}"
+        )
+    lines.append("-" * 53)
+    lines.append(
+        f"{'TOTAL':<22}{report['total_params']:>15,}"
+        f"{report['effective_bits_per_weight']:>7.2f}{report['total_gb']:>9.3f}"
+    )
+    lines.append(f"  compression vs bf16: {report['compression_vs_bf16']:.2f}x")
+    if report.get("excluded_params"):
+        lines.append(
+            f"  excluded (vision+MTP, present on disk but not scored): "
+            f"{report['excluded_params']:,} params / {report['excluded_gb_bf16']:.2f} GB"
+        )
+    return "\n".join(lines)
+
+
+# The SSM recurrence dynamics must stay high precision: error compounds through
+# the linear recurrence instead of being renormalized each step, and naive
+# low-bit PTQ on these collapses SSM models entirely. They are ~0.02% of
+# parameters, so protecting all of them is nearly free.
+PROTECTED = QuantSpec(
+    pattern=r"A_log|dt_bias|conv1d|in_proj_a|in_proj_b|norm|bias",
+    bits=KEEP_BITS,
+)
+
+
+def build_recipe(
+    mlp_bits: int = 16,
+    linear_attn_bits: int = 16,
+    full_attn_bits: int = 16,
+    embed_bits: int = 16,
+    group_size: int = 128,
+    symmetric: bool = True,
+) -> Recipe:
+    """Per-component recipe. PROTECTED comes first so it always wins."""
+    return Recipe(
+        [
+            PROTECTED,
+            QuantSpec(r"embed_tokens|lm_head", embed_bits, group_size, symmetric),
+            QuantSpec(r"linear_attn", linear_attn_bits, group_size, symmetric),
+            QuantSpec(r"self_attn", full_attn_bits, group_size, symmetric),
+            QuantSpec(r"\.mlp\.", mlp_bits, group_size, symmetric),
+        ]
+    )

src/eaimath/rate_distortion.py

@@ -0,0 +1,84 @@
+"""Exact-byte utilities for monotone rate--distortion allocation."""
+
+from __future__ import annotations
+
+import math
+import heapq
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Choice:
+    bits: int
+    bytes: int
+    distortion: float
+
+
+def quantized_bytes(shape: tuple[int, ...], bits: int, group_size: int = 128) -> int:
+    numel = math.prod(shape)
+    if bits >= 16 or len(shape) < 2:
+        return numel * 2
+    group = group_size if group_size > 0 and shape[-1] % group_size == 0 else shape[-1]
+    groups = math.ceil(numel / group)
+    return math.ceil(numel * bits / 8) + groups * 2
+
+
+def allocate(
+    options: dict[str, list[Choice]],
+    budget: int,
+) -> tuple[dict[str, Choice], int]:
+    """Greedily buy monotone next upgrades by distortion reduction per byte."""
+    if not options:
+        raise ValueError("no rate--distortion options")
+    ordered = {
+        name: sorted(rows, key=lambda row: (row.bytes, row.distortion))
+        for name, rows in options.items()
+    }
+    if any(not rows for rows in ordered.values()):
+        raise ValueError("every tensor requires at least one choice")
+    index = {name: 0 for name in ordered}
+    total = sum(rows[0].bytes for rows in ordered.values())
+    if total > budget:
+        raise ValueError(f"minimum representation is {total:,} bytes, over budget {budget:,}")
+
+    # Only one next upgrade per item is eligible at a time. The original
+    # implementation rebuilt and scanned that full set after buying every
+    # individual upgrade: O(items * upgrades), which becomes effectively
+    # quadratic for Round 14's 280k row blocks. A heap preserves the exact
+    # greedy ordering in O((items + upgrades) log items).
+    name_rank = {name: rank for rank, name in enumerate(sorted(ordered))}
+    heap: list[tuple[float, float, int, int, str, int]] = []
+
+    def push_next(name: str) -> None:
+        current_index = index[name]
+        rows = ordered[name]
+        if current_index + 1 >= len(rows):
+            return
+        current = rows[current_index]
+        target = rows[current_index + 1]
+        added = target.bytes - current.bytes
+        gain = current.distortion - target.distortion
+        if added <= 0 or gain <= 0:
+            return
+        # This is the min-heap inverse of the old max tuple:
+        # (gain / added, gain, -added, name).
+        heapq.heappush(
+            heap,
+            (-gain / added, -gain, added, -name_rank[name], name, current_index + 1),
+        )
+
+    for name in ordered:
+        push_next(name)
+
+    while heap:
+        _ratio, _gain, added, _rank, winner, target_index = heapq.heappop(heap)
+        if total + added > budget:
+            # Total never decreases, so this upgrade (and its dependent chain)
+            # can never fit later. Smaller independent upgrades remain eligible.
+            continue
+        total += added
+        index[winner] = target_index
+        push_next(winner)
+
+    selected = {name: ordered[name][position] for name, position in index.items()}
+    return selected, total

src/eaimath/sharing.py

@@ -0,0 +1,196 @@
+"""MLP layer sharing — buying size structurally instead of with bits.
+
+Bucket A (<=0.84 GB) is out of reach by quantization alone. ParetoQ measures 3-
+and 4-bit QAT saturating at 10 B tokens while sub-2-bit needs 30 B, because
+sub-2-bit "breaks the grid" and has to form new representations rather than nudge
+within one. Our token budget is 1-3 B. So the bits have to stay at 3-4 and the
+size has to come from somewhere else.
+
+The MLPs are 53.9% of the model (2.26 B of 4.21 B). Tying one MLP block across
+every k-th layer divides that by k while keeping every forward pass intact --
+unlike depth pruning, which deletes the pass as well and collapses this workload
+(GSM8K 79.3 -> 0.9 at 25% removal; AIME24 gone after a single layer).
+
+**Why this is measurable tonight with no new inference code.** We tie the
+*Parameters*, so `apply_recipe`'s existing data_ptr dedup counts the block once
+and the size report is the real shared size. Then `untie` clones them back into
+separate storage before saving, so what lands on disk is an ordinary checkpoint
+that vLLM loads unchanged. Accuracy measured on that checkpoint is exactly the
+shared model's accuracy -- the tying is mathematically identical, only the
+storage differs. A submission would ship the deduped form plus an expansion
+step, which is precisely the "compression + decompression" pair the course asks
+for.
+
+**What this does NOT do.** There is no recovery training here. Averaging k
+learned blocks into one and running it cold measures the *raw damage*, which is
+the number we need to size the LoRA/distillation budget against -- not a result
+to submit. Expect it to look bad.
+
+Middle-Cycle only: layers 0, 1 and the last are left alone. They are
+independently found to be functionally distinct in Qwen3-4B, and tying them is a
+known way to destroy a model for no size win worth having.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+# The projections inside one MLP block. Named rather than discovered so that a
+# renamed submodule fails loudly instead of silently sharing nothing.
+MLP_PROJECTIONS = ("gate_proj", "up_proj", "down_proj")
+
+PROTECTED_HEAD = 2   # never tie layers 0-1
+PROTECTED_TAIL = 1   # never tie the final layer
+
+
+def find_decoder_layers(model: torch.nn.Module) -> torch.nn.ModuleList:
+    """The decoder layer list, wherever this transformers release put it.
+
+    Qwen3.5 is multimodal, so the text stack sits under a `language_model` in
+    some releases and directly under `model` in others. Both appear in the
+    versions we have run against, so this introspects instead of guessing.
+    """
+    candidates: list[torch.nn.ModuleList] = []
+    for name, module in model.named_modules():
+        if not isinstance(module, torch.nn.ModuleList):
+            continue
+        if "visual" in name or "vision" in name or ".mtp" in name:
+            continue
+        if not name.endswith("layers"):
+            continue
+        first = module[0] if len(module) else None
+        if first is not None and hasattr(first, "mlp"):
+            candidates.append(module)
+    if not candidates:
+        raise RuntimeError(
+            "Could not locate the decoder layer list. Run "
+            "`python scripts/01_inspect_model.py` and check the module tree."
+        )
+    # The text stack is the longest one that is not the vision tower.
+    return max(candidates, key=len)
+
+
+def sharing_groups(num_layers: int, factor: int) -> list[list[int]]:
+    """Contiguous groups of layer indices to tie, honouring Middle-Cycle."""
+    if factor < 2:
+        return []
+    body = list(range(PROTECTED_HEAD, num_layers - PROTECTED_TAIL))
+    return [body[i : i + factor] for i in range(0, len(body), factor)]
+
+
+def tie_mlps(
+    model: torch.nn.Module, factor: int, *, init: str = "mean"
+) -> dict[str, Any]:
+    """Tie each group of ``factor`` MLP blocks to one shared block, in place.
+
+    ``init='mean'`` averages the group's weights into the shared block;
+    ``init='first'`` keeps the first block and discards the rest. Mean is the
+    default because it is the standard cheap initializer and discarding 3 of 4
+    learned blocks is strictly worse. Neither is permutation-aligned merging,
+    which is what a real attempt would use -- see PLAN.md W6.
+
+    Returns accounting: how many MLP params exist after tying, and how many the
+    dense model had.
+    """
+    if init not in ("mean", "first"):
+        raise ValueError(f"init must be 'mean' or 'first', got {init!r}")
+
+    layers = find_decoder_layers(model)
+    groups = sharing_groups(len(layers), factor)
+    if not groups:
+        return {"factor": factor, "groups": 0, "tied_layers": 0}
+
+    dense_params = sum(
+        p.numel()
+        for layer in layers
+        for name in MLP_PROJECTIONS
+        for p in getattr(layer.mlp, name).parameters()
+    )
+
+    tied_layers = 0
+    with torch.no_grad():
+        for group in groups:
+            if len(group) < 2:
+                continue  # a trailing singleton shares with nobody
+            anchor = layers[group[0]]
+            for proj_name in MLP_PROJECTIONS:
+                anchor_proj = getattr(anchor.mlp, proj_name)
+                if init == "mean":
+                    stacked = torch.stack([
+                        getattr(layers[i].mlp, proj_name).weight.data.float()
+                        for i in group
+                    ])
+                    anchor_proj.weight.data.copy_(
+                        stacked.mean(dim=0).to(anchor_proj.weight.dtype)
+                    )
+                # Point every follower at the anchor's Parameter object. Same
+                # storage => apply_recipe's data_ptr dedup counts it once.
+                for i in group[1:]:
+                    getattr(layers[i].mlp, proj_name).weight = anchor_proj.weight
+            tied_layers += len(group) - 1
+
+    shared_params = _unique_mlp_params(layers)
+    return {
+        "factor": factor,
+        "init": init,
+        "groups": len([g for g in groups if len(g) > 1]),
+        "tied_layers": tied_layers,
+        "mlp_params_dense": dense_params,
+        "mlp_params_shared": shared_params,
+        "mlp_params_saved": dense_params - shared_params,
+    }
+
+
+def _unique_mlp_params(layers: torch.nn.ModuleList) -> int:
+    seen: set[int] = set()
+    total = 0
+    for layer in layers:
+        for proj_name in MLP_PROJECTIONS:
+            weight = getattr(layer.mlp, proj_name).weight
+            if weight.data_ptr() in seen:
+                continue
+            seen.add(weight.data_ptr())
+            total += weight.numel()
+    return total
+
+
+def untie(model: torch.nn.Module) -> int:
+    """Give every tied MLP weight its own storage again.
+
+    Call this after quantizing and before saving. `safetensors` refuses to write
+    tensors that share memory, and transformers' workaround silently drops the
+    duplicates -- which would produce a checkpoint with missing keys. Cloning
+    costs disk, not accuracy: the values are identical, so the model evaluated
+    is exactly the shared one.
+    """
+    layers = find_decoder_layers(model)
+    seen: set[int] = set()
+    cloned = 0
+    with torch.no_grad():
+        for layer in layers:
+            for proj_name in MLP_PROJECTIONS:
+                proj = getattr(layer.mlp, proj_name)
+                pointer = proj.weight.data_ptr()
+                if pointer not in seen:
+                    seen.add(pointer)
+                    continue
+                proj.weight = torch.nn.Parameter(
+                    proj.weight.data.clone(), requires_grad=proj.weight.requires_grad
+                )
+                cloned += 1
+    return cloned
+
+
+def format_sharing_report(report: dict[str, Any]) -> str:
+    if not report.get("groups"):
+        return "[sharing] none"
+    saved_gb = report["mlp_params_saved"] * 2 / 1e9
+    return (
+        f"[sharing] {report['factor']}x MLP, init={report['init']}: "
+        f"{report['groups']} groups, {report['tied_layers']} layers tied\n"
+        f"          MLP params {report['mlp_params_dense']:,} -> "
+        f"{report['mlp_params_shared']:,} "
+        f"(-{report['mlp_params_saved']:,}, -{saved_gb:.2f} GB at bf16)"
+    )

src/eaimath/vllm_backend.py

@@ -0,0 +1,144 @@
+"""vLLM generation backend — the fix for our eval throughput problem.
+
+The HF path runs a fixed batch until *every* sequence in it stops. With a 43%
+truncation rate and a 32k cap, one long sequence pins the whole batch for 32,768
+steps while the other 15 sit finished and padded. That is why a 50-problem shard
+took ~4 hours.
+
+vLLM's continuous batching retires finished sequences immediately and admits new
+ones, and paged attention removes the `torch.cat` cache-regrowth spike entirely.
+For this workload — high variance in output length, long tails — that is worth
+roughly an order of magnitude.
+
+Optional dependency: ``pip install -e '.[vllm]'``. The HF backend in
+``generate.py`` stays the reference implementation and always works.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Sequence
+
+from .generate import GenerationOutput, build_chat_prompts
+
+
+def _checkpoint_dtype(model_path: str) -> str:
+    """Use FP16 activations for GPTQ; retain BF16 for the established RTN path."""
+    # Official ParoQuant checkpoints use an FP16 activation kernel and can be
+    # addressed by a remote Hugging Face id, before a local config exists.
+    if "paro" in model_path.lower():
+        return "float16"
+    root = Path(model_path)
+    quantize_config = root / "quantize_config.json"
+    if quantize_config.is_file():
+        try:
+            config = json.loads(quantize_config.read_text())
+            if config.get("bits") in (2, 3, 4, 8):
+                return "float16"
+        except (OSError, json.JSONDecodeError):
+            pass
+
+    model_config = root / "config.json"
+    if model_config.is_file():
+        try:
+            config = json.loads(model_config.read_text())
+            quantization = config.get("quantization_config") or {}
+            if str(quantization.get("quant_method", "")).lower() in {"gptq", "paroquant"}:
+                return "float16"
+        except (OSError, json.JSONDecodeError):
+            pass
+    return "bfloat16"
+
+
+def is_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("vllm") is not None
+
+
+def generate_vllm(
+    model_path: str,
+    tokenizer,
+    prompts: Sequence[str],
+    *,
+    max_new_tokens: int = 65536,
+    temperature: float = 0.0,
+    top_p: float = 0.95,
+    top_k: int = 20,
+    presence_penalty: float = 0.0,
+    repetition_penalty: float = 1.0,
+    enable_thinking: bool | None = None,
+    gpu_memory_utilization: float = 0.90,
+    max_model_len: int | None = None,
+    seed: int = 0,
+    allowed_token_ids: Sequence[int] | None = None,
+    llm=None,
+) -> tuple[list[GenerationOutput], object]:
+    """Generate with vLLM, returning the same records as the HF backend.
+
+    Returns ``(outputs, llm)`` so the engine can be reused across calls — engine
+    startup costs ~1-2 min, which matters when sweeping recipes.
+    """
+    from vllm import LLM, SamplingParams
+
+    chat_texts = build_chat_prompts(tokenizer, prompts, enable_thinking)
+
+    if llm is None:
+        longest_prompt = max(len(tokenizer(t)["input_ids"]) for t in chat_texts)
+        if max_model_len is None:
+            # Headroom over the longest prompt so no request is rejected outright.
+            max_model_len = max_new_tokens + longest_prompt + 256
+        llm = LLM(
+            model=model_path,
+            # vLLM 0.19 GPTQ kernels support FP16 activations. The dense base
+            # model and our RTN-value checkpoints remain on their proven BF16 path.
+            dtype=_checkpoint_dtype(model_path),
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len,
+            trust_remote_code=True,
+            seed=seed,
+        )
+
+    sampling_kwargs = dict(
+        max_tokens=max_new_tokens,
+        temperature=temperature,
+        # vLLM ignores top_p/top_k at temperature 0, but passing them when
+        # sampling keeps parity with Qwen's recommended thinking-mode settings.
+        top_p=top_p if temperature > 0 else 1.0,
+        top_k=top_k if temperature > 0 else -1,
+        presence_penalty=presence_penalty,
+        repetition_penalty=repetition_penalty,
+    )
+    if allowed_token_ids:
+        # Restricting generation to a token subset is how we measure the cost of
+        # a trimmed vocabulary without rebuilding the tokenizer. Not every vLLM
+        # release accepts it, and an overnight sweep must not die at 3 a.m. over
+        # a kwarg -- so fail loudly here, before the engine is built.
+        try:
+            SamplingParams(**sampling_kwargs, allowed_token_ids=[0])
+        except TypeError as exc:
+            raise RuntimeError(
+                "This vLLM build does not accept SamplingParams(allowed_token_ids=...), "
+                "so the vocabulary-restriction probe cannot run here. Drop "
+                "--vocab-size to run everything else.\n"
+                f"  underlying error: {exc}"
+            ) from exc
+        sampling_kwargs["allowed_token_ids"] = list(allowed_token_ids)
+
+    sampling = SamplingParams(**sampling_kwargs)
+
+    results: list[GenerationOutput] = []
+    for request, chat_text in zip(llm.generate(list(chat_texts), sampling), chat_texts):
+        completion = request.outputs[0]
+        results.append(
+            GenerationOutput(
+                text=completion.text,
+                num_prompt_tokens=len(request.prompt_token_ids),
+                num_generated_tokens=len(completion.token_ids),
+                # vLLM reports this directly, so unlike the HF path we do not
+                # have to infer termination by scanning for EOS.
+                finished=completion.finish_reason != "length",
+            )
+        )
+    return results, llm

src/eaimath/vocab.py

@@ -0,0 +1,161 @@
+"""Vocabulary trimming — 285 MB, the single largest untouched lever.
+
+The tied embedding is 635.7 M params, 15.1% of the model. A math-only model does
+not need 248,320 tokens: a 13 k vocabulary retains 98.6% coverage on MATH-500 and
+AIME, and math is the easiest domain to trim (far easier than code or
+multilingual).
+
+    vocab      params     @4.125 bpw
+    248,320    635.7 M    327.8 MB
+    32,768      83.9 M     43.3 MB
+
+That 285 MB is 34% of the entire bucket-A ceiling. It is also what makes bucket B
+reachable *with a healthy attention stack* -- see `scripts/plan_configs.py`:
+without trimming, fitting under 1.68 GB forces full_attn down to ~3 bits, and
+full_attn is the most fragile component we have measured (-21 points at 4 bits).
+
+**Two separate jobs, deliberately separated.**
+
+1. `allowed_token_ids` (here) measures the *accuracy cost* of a smaller
+   vocabulary with no tokenizer surgery at all: generation is restricted to the
+   kept set, so the model produces exactly what a trimmed model could produce.
+   This validates the idea before anyone builds the real thing.
+2. Actually rebuilding the tokenizer -- filtering merges so the vocabulary stays
+   closed, remapping ids, slicing the embedding -- is the shipping step, and it
+   only makes sense once (1) says the cost is acceptable.
+
+Doing (1) first also dissolves the zero-logit trap. Zeroing a dropped embedding
+row does not remove the token: its logit becomes exactly 0, which wins the argmax
+whenever every real logit is negative. Masking at sampling time cannot have that
+failure mode.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from typing import Any, Iterable, Sequence
+
+# Every byte-level fallback token, so any string stays encodable, plus the
+# specials the chat template depends on. Termination is our failure mode and
+# </think> / EOS are the mechanism, so they are never droppable.
+ALWAYS_KEEP_TOKENS = (
+    "<|endoftext|>", "<|im_start|>", "<|im_end|>",
+    "<think>", "</think>",
+)
+
+
+def token_frequencies(texts: Iterable[str], tokenizer) -> Counter:
+    """Token id -> count over a corpus, using the model's own tokenizer."""
+    counts: Counter = Counter()
+    for text in texts:
+        counts.update(tokenizer(text, add_special_tokens=False)["input_ids"])
+    return counts
+
+
+def build_keep_set(
+    counts: Counter,
+    tokenizer,
+    target_size: int,
+    *,
+    priority_counts: Counter | None = None,
+    coverage_report: bool = True,
+) -> dict[str, Any]:
+    """The ``target_size`` most useful token ids, plus everything unconditional.
+
+    Byte-fallback tokens are kept whole rather than by frequency: dropping one
+    makes some strings unencodable, and the storage is 256 rows.
+    """
+    keep: set[int] = set()
+
+    for token in ALWAYS_KEEP_TOKENS:
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        if token_id is not None and token_id >= 0:
+            keep.add(token_id)
+    for token_id in getattr(tokenizer, "all_special_ids", []) or []:
+        keep.add(token_id)
+
+    # Byte-level fallbacks: single-byte pieces the BPE falls back to.
+    vocab = tokenizer.get_vocab()
+    for token, token_id in vocab.items():
+        if len(token) == 1 or (token.startswith("<0x") and token.endswith(">")):
+            keep.add(token_id)
+
+    total = sum(counts.values())
+    covered = sum(counts[i] for i in keep if i in counts)
+    priority_added = 0
+    if priority_counts:
+        # Round 9 showed that omitted *input* rows are the clearest recoverable
+        # loss. Preserve every token seen in disjoint training problems before
+        # the much longer completions consume the frequency budget.
+        for token_id, _ in priority_counts.most_common():
+            if len(keep) >= target_size:
+                break
+            if token_id not in keep:
+                keep.add(token_id)
+                covered += counts[token_id]
+                priority_added += 1
+    for token_id, count in counts.most_common():
+        if len(keep) >= target_size:
+            break
+        if token_id not in keep:
+            keep.add(token_id)
+            covered += count
+
+    if len(keep) < target_size:
+        # The corpus did not contain enough distinct tokens to fill the budget.
+        # Silently shipping the short set is what turned the first vocabulary
+        # probe into a measurement of "restrict the model to 4k tokens drawn
+        # from problem statements" -- which unsurprisingly scored 0.010, and
+        # said nothing at all about vocabulary trimming.
+        #
+        # BPE vocabularies are ordered by merge priority, so low ids are the
+        # most generally useful pieces. Filling from the bottom is a defensible
+        # completion, and the caller is told it happened.
+        filled = 0
+        for token_id in range(len(vocab)):
+            if len(keep) >= target_size:
+                break
+            if token_id not in keep:
+                keep.add(token_id)
+                filled += 1
+    else:
+        filled = 0
+
+    result: dict[str, Any] = {
+        "keep_ids": sorted(keep),
+        "kept": len(keep),
+        "original_vocab": len(vocab),
+        "target_size": target_size,
+        "observed_distinct": len(counts),
+        "priority_distinct": len(priority_counts or {}),
+        "priority_added": priority_added,
+        "filled_by_id_order": filled,
+    }
+    if coverage_report and total:
+        result["token_coverage"] = covered / total
+        result["dropped_token_mass"] = 1.0 - covered / total
+    return result
+
+
+def embedding_savings(
+    original_vocab: int, kept: int, hidden_size: int, bits_per_weight: float
+) -> dict[str, float]:
+    """Bytes reclaimed by a trim, for the tied embedding + lm_head (counted once)."""
+    dropped_params = (original_vocab - kept) * hidden_size
+    return {
+        "dropped_rows": original_vocab - kept,
+        "dropped_params": dropped_params,
+        "saved_bytes": dropped_params * bits_per_weight / 8,
+        "saved_mb": dropped_params * bits_per_weight / 8 / 1e6,
+    }
+
+
+def sampling_kwargs(keep_ids: Sequence[int] | None) -> dict[str, Any]:
+    """SamplingParams kwargs restricting generation to ``keep_ids``.
+
+    Returns empty when there is nothing to restrict, so callers can splat this
+    unconditionally.
+    """
+    if not keep_ids:
+        return {}
+    return {"allowed_token_ids": list(keep_ids)}

src/eaimath/workspace.py

@@ -0,0 +1,40 @@
+"""Workspace layout helpers shared by runners, collectors, and hygiene tools."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from pathlib import Path
+
+
+def round_log_files(
+    round_name: str,
+    patterns: Iterable[str] = (),
+    experiments: Path = Path("experiments"),
+) -> list[Path]:
+    """Return round-local logs plus legacy flat/archive locations.
+
+    Round 15 onward writes logs beside its round. Rounds 8--14 historically
+    wrote into ``experiments/logs``; the hygiene tool archives those files
+    without making old evidence collectors blind to them.
+    """
+    local = experiments / "rounds" / round_name / "logs"
+    legacy = experiments / "logs"
+    candidates = [local / "runner.log", legacy / f"{round_name}_runner.log"]
+    for pattern in patterns:
+        candidates.extend(sorted(local.glob(pattern)))
+        candidates.extend(sorted(legacy.glob(pattern)))
+
+    archive = legacy / "archive"
+    if archive.is_dir():
+        candidates.extend(sorted(archive.rglob(f"{round_name}_runner.log")))
+        for pattern in patterns:
+            candidates.extend(sorted(archive.rglob(pattern)))
+
+    unique: list[Path] = []
+    seen: set[Path] = set()
+    for path in candidates:
+        resolved = path.resolve()
+        if path.is_file() and resolved not in seen:
+            seen.add(resolved)
+            unique.append(path)
+    return unique