From 97c756a06472b2f632453f35a541b2f0553bec00 Mon Sep 17 00:00:00 2001 From: Dennis Paul Date: Mon, 13 Jul 2026 14:54:55 +0200 Subject: [PATCH] =?UTF-8?q?tools:=20quant=5Fablation.py=20=E2=80=94=20engi?= =?UTF-8?q?ne-free=20A/B=20of=20any=20quantization=20scheme=20vs=20fp16=20?= =?UTF-8?q?(fake-quant,=20isolates=20weight=20error)=20(#108,=20#115)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measuring "what does int4 cost?" by comparing colibri's score to a published model-card number does not work: this harness scores 0-shot log-likelihood while published numbers are few-shot/CoT, and that protocol gap can swamp the quantization effect entirely (#108). This removes the confound by construction: take an fp16 model, push its weights through colibri's own quantizer (quantize -> dequantize, in place), and score both with the SAME harness on the SAME questions. The only variable is the quantizer, so the delta IS the quantization cost. Runs on OLMoE in minutes, so a scheme can be ranked BEFORE committing to a multi-hour GLM conversion. Quantizer math is replicated from tools/convert_fp8_to_int4.py (symmetric absmax, per-row scales) and generalised with an optional group size, so grouped/finer schemes can be compared directly against what ships today. Measured on OLMoE-1B-7B, n=200/task (#108): scheme hellaswag arc_c mmlu mean delta fp16 77.0% 47.0% 47.0% 57.0% -- int4 (shipped) 74.0% 41.0% 31.5% 48.8% -8.2pp int4-nohead 73.5% 40.5% 37.5% 50.5% -6.5pp int4-g128 78.5% 45.5% 38.0% 54.0% -3.0pp int4-g128-nohead 78.5% 46.5% 38.0% 54.3% -2.7pp The per-row int4 container costs ~8pp, concentrated on the HARD task: MMLU falls to 31.5% against a 25% random baseline while easy HellaSwag barely moves -- per-row scales eat the small logit margins that hard questions depend on (the same margin erosion that flips near-tie tokens in #100). group=128 recovers ~63% of the loss. Keeping lm_head/embed in fp16 is NOT the fix (+1.7pp alone, +0.3pp atop grouping). Includes a coverage assert: transformers fuses MoE experts into 3D tensors, so a ndim==2 filter silently skips every expert and leaves ~85% of the model in fp16 while appearing to work. The tool fails loudly instead of reporting fiction. Dev-only tool (torch + transformers); the engine's dependency-free path is untouched. --- c/tools/quant_ablation.py | 235 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 c/tools/quant_ablation.py diff --git a/c/tools/quant_ablation.py b/c/tools/quant_ablation.py new file mode 100644 index 0000000..c3c5bf9 --- /dev/null +++ b/c/tools/quant_ablation.py @@ -0,0 +1,235 @@ +""" +A/B any quantization scheme against fp16 โ€” WITHOUT converting a 370 GB model first. + +Why this exists +--------------- +Measuring "what does int4 cost us?" by comparing colibri's score to a published model-card +number does not work: the harness scores 0-shot log-likelihood, published numbers are +few-shot/CoT, and that protocol gap can swamp the quantization effect (see #108). + +This tool removes the confound by construction. It takes an fp16 model, pushes its weights +through colibri's OWN quantizer (quantize -> dequantize, in place), and scores both with the +SAME harness, on the SAME questions, on the SAME machine. The only variable is the quantizer, +so the delta IS the quantization cost. + +It runs on a small model (OLMoE) in minutes, so a scheme can be ranked BEFORE committing to +a multi-hour GLM conversion. The quantizer math is replicated from tools/convert_fp8_to_int4.py +(symmetric absmax, per-row scales) and generalised with an optional group size. + +Measured with this tool on OLMoE-1B-7B, n=200/task (issue #108): + + scheme hellaswag arc_c mmlu mean delta + fp16 77.0% 47.0% 47.0% 57.0% -- + int4 (shipped) 74.0% 41.0% 31.5% 48.8% -8.2pp + int4-nohead 73.5% 40.5% 37.5% 50.5% -6.5pp + int4-g128 78.5% 45.5% 38.0% 54.0% -3.0pp + int4-g128-nohead 78.5% 46.5% 38.0% 54.3% -2.7pp + + -> the per-row int4 container costs ~8pp, concentrated on the HARD task (MMLU falls to + 31.5% against a 25% random baseline, while easy HellaSwag barely moves): per-row scales + eat the small logit margins hard questions depend on. + -> group=128 recovers ~63% of that loss for ~+0.25 bits/weight. + -> leaving lm_head/embed in fp16 is NOT the fix (+1.7pp alone, +0.3pp on top of grouping). + +Usage +----- + pip install torch transformers # dev-only; the engine stays dependency-free + python tools/fetch_benchmarks.py --out ./bench --tasks hellaswag,arc_challenge,mmlu --limit 200 + python tools/quant_ablation.py --model allenai/OLMoE-1B-7B-0924 --data ./bench \ + --tasks hellaswag,arc_challenge,mmlu --limit 200 \ + --schemes fp16,int4,int4-g64,int4-g128 + +Scheme grammar: fp16 | int{2,4,8}[-g][-nohead] + int4 per-row absmax int4 -- what the converter ships today + int4-g64 one scale per 64 input weights instead of per row + int4-nohead as int4, but lm_head/embed kept in fp16 +""" +import argparse +import json +import random +import re +import sys + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + + +# -------------------------------------------------------------------------------------- +# colibri's quantizer (tools/convert_fp8_to_int4.py:32-52), generalised with a group size. +# +# LAYOUT NOTE, and it is a trap: transformers fuses MoE experts into 3D tensors +# (mlp.experts.gate_up_proj = [n_experts, in, out]) even when the checkpoint stores one 2D +# matrix per expert. A `p.ndim == 2` filter therefore skips EVERY expert and silently leaves +# ~85% of an MoE in fp16 while appearing to work. Both layouts must be handled, and the +# coverage assert below exists to make that failure loud instead of plausible. +# -------------------------------------------------------------------------------------- +def _quant_last_dim(x, bits, group): + """Symmetric absmax quantize->dequantize along the last (input) dim.""" + qmax = (1 << (bits - 1)) - 1 # int4 -> 7, int8 -> 127, int2 -> 1 + qmin = -(qmax + 1) # int4 -> -8 (nibble [-8,7], as the converter does) + if group: + if x.shape[-1] % group: + raise SystemExit(f"group {group} does not divide input dim {x.shape[-1]}") + x = x.reshape(*x.shape[:-1], x.shape[-1] // group, group) + amax = x.abs().amax(dim=-1, keepdim=True) + s = torch.clamp(amax / qmax, min=1e-8) + q = torch.clamp(torch.round(x / s), qmin, qmax) + out = q * s + return out.reshape(*out.shape[:-2], -1) if group else out + + +def quantize_param(w, bits, group): + if w.ndim == 3: # fused experts [E, in, out] -> move input last + x = w.transpose(1, 2).contiguous() + return _quant_last_dim(x, bits, group).transpose(1, 2).contiguous() + return _quant_last_dim(w, bits, group) # nn.Linear [out, in] -- input already last + + +SCHEME_RE = re.compile(r"^int(2|4|8)(?:-g(\d+))?(-nohead)?$") + + +def parse_scheme(name): + """'int4-g128-nohead' -> (bits=4, group=128, skip_head=True). 'fp16' -> None.""" + if name == "fp16": + return None + m = SCHEME_RE.match(name) + if not m: + raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,4,8}}[-g][-nohead])") + return int(m.group(1)), int(m.group(2) or 0), bool(m.group(3)) + + +def is_router(name): + # The router (mlp.gate.weight) stays f32 in the converter -- convert_fp8_to_int4.py:14. + # Careful: expert weights are gate_proj/up_proj/down_proj and DO get quantized. + return name.endswith("mlp.gate.weight") + + +def is_head_or_embed(name): + return "embed_tokens" in name or "lm_head" in name + + +def apply_scheme(model, scheme): + """Quantize the tensor classes the converter hits (attn/mlp/expert/embed/lm_head); + norms, router and biases stay float. Returns (n_tensors, quantized_params, total).""" + total = sum(p.numel() for p in model.parameters()) + spec = parse_scheme(scheme) + if spec is None: + return 0, 0, total + bits, group, skip_head = spec + n = qp = 0 + with torch.no_grad(): + for name, p in model.named_parameters(): + if p.ndim < 2 or is_router(name): + continue + if skip_head and is_head_or_embed(name): + continue + p.data.copy_(quantize_param(p.data.float(), bits, group).to(p.dtype)) + n += 1 + qp += p.numel() + return n, qp, total + + +# -------------------------------------------------------------------------------------- +# Scoring โ€” mirrors tools/eval_glm.py exactly: +# acc = argmax over options of sum(logprob of continuation tokens) +# acc_norm = argmax over options of sum(logprob) / len(continuation string in CHARACTERS) +# -------------------------------------------------------------------------------------- +def load_docs(task, data_dir, limit, seed): + path = f"{data_dir}/{task}.jsonl" + try: + docs = [json.loads(l) for l in open(path) if l.strip()] + except FileNotFoundError: + raise SystemExit(f"missing {path} โ€” run: python tools/fetch_benchmarks.py --out {data_dir} --tasks {task}") + random.Random(seed).shuffle(docs) # same seed/shuffle convention as eval_glm.py + return docs[:limit] if limit else docs + + +@torch.no_grad() +def score(model, tk, docs, device): + acc = accn = 0 + for d in docs: + ctx, choices, gold = d["ctx"], d["choices"], int(d["gold"]) + ctx_ids = tk(ctx, add_special_tokens=False).input_ids + lps, norms = [], [] + for cont in choices: + full = tk(ctx + cont, add_special_tokens=False).input_ids + cl = len(ctx_ids) + while cl > 0 and (cl > len(full) or full[:cl] != ctx_ids[:cl]): + cl -= 1 + if not full[cl:]: + full = ctx_ids + tk(cont, add_special_tokens=False).input_ids + cl = len(ctx_ids) + cl = max(1, cl) + ids = torch.tensor([full], device=device) + logprobs = torch.log_softmax(model(ids).logits.float()[0, :-1], dim=-1) + tgt = ids[0, 1:] + lps.append(logprobs[torch.arange(cl - 1, len(full) - 1), tgt[cl - 1:]].sum().item()) + norms.append(max(1, len(cont))) # CHARACTER length, like eval_glm.py + acc += max(range(len(lps)), key=lambda i: lps[i]) == gold + accn += max(range(len(lps)), key=lambda i: lps[i] / norms[i]) == gold + n = len(docs) + return 100 * acc / n, 100 * accn / n + + +def main(): + ap = argparse.ArgumentParser(description="A/B a quantization scheme against fp16, engine-free") + ap.add_argument("--model", default="allenai/OLMoE-1B-7B-0924", help="HF repo id or local dir") + ap.add_argument("--data", default="./bench") + ap.add_argument("--tasks", default="hellaswag,arc_challenge,mmlu") + ap.add_argument("--limit", type=int, default=200) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--schemes", default="fp16,int4,int4-g128", + help="comma list: fp16 | int{2,4,8}[-g][-nohead]") + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--min-coverage", type=float, default=95.0, + help="fail if a scheme quantized less than this %% of params (catches the " + "3D-fused-expert trap, where a ndim==2 filter skips every expert)") + a = ap.parse_args() + + tasks = a.tasks.split(",") + schemes = a.schemes.split(",") + for s in schemes: + parse_scheme(s) # fail fast on a typo + + tk = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True) + docs = {t: load_docs(t, a.data, a.limit, a.seed) for t in tasks} + + means, rows = {}, {} + for scheme in schemes: + model = AutoModelForCausalLM.from_pretrained( + a.model, dtype=torch.float16, low_cpu_mem_usage=True, + device_map={"": 0} if a.device == "cuda" else None, trust_remote_code=True) + model.eval() + if a.device != "cuda": + model.to(a.device) + + n, qp, tp = apply_scheme(model, scheme) + cov = 100 * qp / tp if tp else 0.0 + print(f"[{scheme}] {n} tensors ยท {qp/1e9:.2f}B/{tp/1e9:.2f}B params ({cov:.1f}% coverage)", + flush=True) + if scheme != "fp16" and cov < a.min_coverage: + raise SystemExit( + f"ERROR: {scheme} quantized only {cov:.1f}% of parameters (< {a.min_coverage}%).\n" + f" The experts are probably being skipped: transformers fuses MoE experts\n" + f" into 3D tensors, so a ndim==2 filter silently leaves them in fp16.") + + rows[scheme] = {t: score(model, tk, docs[t], a.device) for t in tasks} + means[scheme] = sum(v[1] for v in rows[scheme].values()) / len(tasks) + for t in tasks: + print(f" {t:<16} n={len(docs[t]):<4} acc {rows[scheme][t][0]:5.1f}%" + f" acc_norm {rows[scheme][t][1]:5.1f}%", flush=True) + print(f" {'MEAN acc_norm':<16} {means[scheme]:5.1f}%\n", flush=True) + del model + torch.cuda.empty_cache() + + base = means.get("fp16") + print(f"{'scheme':<20}{'mean acc_norm':>14}{'delta vs fp16':>16}") + for scheme in schemes: + d = f"{means[scheme]-base:+.1f}pp" if base is not None and scheme != "fp16" else "--" + print(f"{scheme:<20}{means[scheme]:>13.1f}%{d:>16}") + if base is None: + print("\n(no fp16 baseline in --schemes, so no deltas)", file=sys.stderr) + + +if __name__ == "__main__": + main()