src/eaimath/quantize.py

"""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),
        ]
    )