tools: quant_ablation.py — engine-free A/B of any quantization scheme vs fp16 (fake-quant, isolates weight error) (#108, #115)
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.
This commit is contained in:
Dennis Paul
@@ -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<N>][-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<N>][-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<N>][-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()
|
||||||
Reference in New Issue
Block a user