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