url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tools: repair_mtp_int8.py — fix int4-converted MTP heads in place; warn on --mtp --ebits <8

Browse Source 
Per-row int4 quantization of the MTP head's eh_proj [D,2D] zeroes its ENTIRE
embedding half: the tensor's two column halves differ in scale by ~20-30x per
row (embedding-half absmax ~0.05, hidden-half ~1.5 on GLM-5.2), so a single
per-row scale (absmax/7) puts every embedding-half weight below half a
quantization step and np.rint rounds it to exact zero (packed bytes 0x88).
The draft head then cannot see the input token and MTP acceptance collapses
to ~0% — the mechanism behind issue #8's measurement (int4: 0-4%; int8:
39-59%), and the reason --mtp already defaults to --ebits 8.

This showed up in the wild: a published pre-converted container had its MTP
shards made with an explicit --ebits 4 and drafts were pure garbage on every
backend (deterministic, data-side; verified byte-for-byte by reproducing the
published bytes with quant_int4 on the official BF16 rows).

- tools/repair_mtp_int8.py: repairs such a container IN PLACE — finds the MTP
  layer's per-row-int4 dense tensors (eh_proj, q/kv/o projections, shared
  experts; routed experts untouched), re-downloads only those from the FP8
  source repo (~355 MB of HTTP range reads, no torch, no token), requantizes
  at int8 with quant_int8's exact math, and rewrites the shards atomically
  with *.bak-int4 backups. --dry-run to inspect. The engine picks up int8
  automatically (qt_from_disk detects format by blob size).
  Validated on GLM-5.2 744B: MTP acceptance 0% -> 100% (greedy), 1.11 -> 3.20
  tokens/forward, decode 0.18 -> 0.30 tok/s (Metal, 4-bit KV).

- convert_fp8_to_int4.py: print a warning when --mtp is combined with
  --ebits <8 and per-row scales (the exact footgun above); --group-size 128
  remains a valid int4 alternative since group scales give the embedding
  half its own scale.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
recviking
2026-07-18 16:27:29 -04:00
parent 72d3d37231
commit d6676c10e9
2 changed files with 240 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/tools/convert_fp8_to_int4.py
+11
View File
@@ -286,6 +286,17 @@ def main():
# testa MTP a int4 = acceptance ~0-4% (misurato, issue #8): il draft sbaglia sempre
# e la speculazione non parte mai. A int8: 39-59%, 2.2-2.8 token/forward.
a.ebits = 8 if (a.mtp or a.indexer) else 4
if a.mtp and a.ebits < 8 and a.group_size <= 0:
# Non solo lossy: eh_proj ha ~20-30x di asimmetria di scala fra le due meta' di
# colonna, quindi l'int4 per-riga (UNA scala per riga) arrotonda a ZERO l'intera
# meta' embedding -> il draft non vede il token -> acceptance ~0% (issue #8).
# EN: not merely lossy: eh_proj has ~20-30x column-scale asymmetry, so per-row
# EN: int4 rounds its ENTIRE embedding half to exact zeros -> the draft cannot
# EN: see the input token -> ~0% acceptance (issue #8). A container converted
# EN: this way is repairable in place with tools/repair_mtp_int8.py.
print(f"WARNING: --mtp with --ebits {a.ebits} and per-row scales ZEROES eh_proj's "
"embedding half -> MTP acceptance ~0% (issue #8). Use the default --ebits 8, "
"or add --group-size 128 for group-scaled int4.")
if a.xbits is None: a.xbits = a.ebits
# Build per-type bits map. If a type-specific arg is set, use it; otherwise the
c/tools/repair_mtp_int8.py
+229
View File
@@ -0,0 +1,229 @@
"""Repair an existing colibri int4 container whose MTP head was quantized at int4.
WHY THIS EXISTS
`model.layers.<N>.eh_proj.weight` [D, 2D] multiplies the MTP concat
[embedding_norm ; hidden_norm], and its two column halves differ in scale by
~20-30x per row (embedding-half absmax ~0.05, hidden-half ~1.5 on GLM-5.2).
Per-row int4 uses ONE scale (= absmax/7) per row, so every embedding-half
weight lands below half a quantization step and np.rint rounds the ENTIRE
embedding half to exact zeros (packed bytes 0x88). The MTP head then drafts
garbage: acceptance ~0% (issue #8 measured 0-4% at int4; 39-59% at int8 —
which is why `convert_fp8_to_int4.py --mtp` defaults to --ebits 8).
A container converted (or downloaded) with an int4 MTP head does not need a
full re-conversion: this script re-downloads ONLY the affected dense tensors
(~355 MB of HTTP range reads against the FP8 source repo), requantizes them
at int8 with the converter's exact math, and patches the local shards in
place. Originals are kept beside as *.bak-int4.
WHAT IT TOUCHES
The MTP layer's dense tensors only (eh_proj, q_a/q_b/kv_a/kv_b/o_proj,
shared_experts.*): the ones that stream into RAM once and stay resident.
Routed experts (model.layers.<N>.mlp.experts.*) are NOT touched — they are
statistically like the main layers' experts and int4 is acceptable there.
The engine auto-detects int8 vs int4 by blob size (qt_from_disk), so no
engine or config change is needed. Cost: ~+133 MB on disk / resident RAM.
USAGE
python3 tools/repair_mtp_int8.py --snap /path/to/glm52_i4 # repair
python3 tools/repair_mtp_int8.py --snap /path/to/glm52_i4 --dry-run # inspect only
--source-repo defaults to zai-org/GLM-5.2-FP8 (the checkpoint the public
int4 containers were converted from). Requires numpy and network access;
no torch, no HF token (public repo, anonymous range reads).
"""
import argparse, glob, json, os, ssl, struct, sys, urllib.request
import numpy as np
# macOS python.org builds ship no CA bundle: use certifi when available (Linux
# system Pythons generally have working system certs and skip this).
try:
import certifi
_SSL_CTX = ssl.create_default_context(cafile=certifi.where())
except ImportError:
_SSL_CTX = ssl.create_default_context()
# ---------- HTTP range reads against the source repo ----------
def http_range(url, start, length, tries=5):
req = urllib.request.Request(url, headers={"User-Agent": "colibri-mtp-repair",
"Range": f"bytes={start}-{start+length-1}"})
for attempt in range(tries):
try:
with urllib.request.urlopen(req, timeout=30, context=_SSL_CTX) as r:
data = r.read()
if len(data) == length:
return data
except KeyboardInterrupt:
raise
except Exception as ex:
if attempt == tries - 1:
raise RuntimeError(f"range read failed for {url}: {ex}")
raise RuntimeError(f"short range read for {url}")
class SourceRepo:
def __init__(self, repo, revision="main"):
self.base = f"https://huggingface.co/{repo}/resolve/{revision}/"
with urllib.request.urlopen(self.base + "model.safetensors.index.json", timeout=30, context=_SSL_CTX) as r:
self.wmap = json.loads(r.read())["weight_map"]
self._hdr = {}
def _shard_header(self, shard):
if shard not in self._hdr:
n = struct.unpack("<Q", http_range(self.base + shard, 0, 8))[0]
self._hdr[shard] = (json.loads(http_range(self.base + shard, 8, n)), 8 + n)
return self._hdr[shard]
def meta(self, name):
shard = self.wmap.get(name)
if not shard:
return None
hdr, _ = self._shard_header(shard)
return hdr[name]
def fetch_f32(self, name):
"""Download one tensor and dequantize to f32 [O, I] (BF16 or FP8+block scales)."""
shard = self.wmap[name]
hdr, base = self._shard_header(shard)
m = hdr[name]
o0, o1 = m["data_offsets"]
raw = http_range(self.base + shard, base + o0, o1 - o0)
if m["dtype"] == "BF16":
u = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 16
return u.view(np.float32).reshape(m["shape"]).astype(np.float32)
if m["dtype"] == "F8_E4M3":
b = np.frombuffer(raw, dtype=np.uint8).astype(np.uint16)
sign = np.where(b & 0x80, -1.0, 1.0)
e = (b >> 3) & 0xF
mant = (b & 7).astype(np.float64)
v = (sign * np.where(e > 0, (1 + mant / 8) * np.exp2(e.astype(np.float64) - 7),
mant / 8 * np.exp2(-6.0))).reshape(m["shape"])
sn = name + "_scale_inv"
sshard = self.wmap[sn]
shdr, sbase = self._shard_header(sshard)
sm = shdr[sn]
so0, so1 = sm["data_offsets"]
sc = np.frombuffer(http_range(self.base + sshard, sbase + so0, so1 - so0),
dtype=np.float32).reshape(sm["shape"])
O, I = m["shape"]
scf = np.repeat(np.repeat(sc, 128, axis=0)[:O], 128, axis=1)[:, :I]
return (v * scf).astype(np.float32)
raise ValueError(f"{name}: unsupported source dtype {m['dtype']}")
# ---------- quantization: identical to convert_fp8_to_int4.quant_int8 ----------
def quant_int8(w):
amax = np.abs(w).max(axis=1, keepdims=True)
s = np.maximum(amax / 127, 1e-8)
q = np.clip(np.rint(w / s), -128, 127).astype(np.int8)
return q.reshape(-1).view(np.uint8).copy(), s[:, 0].astype(np.float32)
# ---------- local safetensors IO (no deps; preserves byte-identity of untouched tensors) ----------
def read_shard(path):
with open(path, "rb") as fh:
n = struct.unpack("<Q", fh.read(8))[0]
hdr = json.loads(fh.read(n))
base = 8 + n
order = sorted(((k, v) for k, v in hdr.items() if k != "__metadata__"),
key=lambda kv: kv[1]["data_offsets"][0])
out = {}
for k, v in order:
fh.seek(base + v["data_offsets"][0])
out[k] = (v["dtype"], v["shape"], fh.read(v["data_offsets"][1] - v["data_offsets"][0]))
return out, hdr.get("__metadata__")
def write_shard(path, tensors, meta):
hdr = {}
off = 0
for k, (dt, shape, raw) in tensors.items():
hdr[k] = {"dtype": dt, "shape": shape, "data_offsets": [off, off + len(raw)]}
off += len(raw)
if meta:
hdr["__metadata__"] = meta
hj = json.dumps(hdr).encode()
hj += b" " * ((8 - len(hj) % 8) % 8)
with open(path, "wb") as fh:
fh.write(struct.pack("<Q", len(hj)))
fh.write(hj)
for _, (_, _, raw) in tensors.items():
fh.write(raw)
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--snap", required=True, help="local colibri int4 container directory")
ap.add_argument("--source-repo", default="zai-org/GLM-5.2-FP8")
ap.add_argument("--revision", default="main")
ap.add_argument("--dry-run", action="store_true", help="report what would change, touch nothing")
a = ap.parse_args()
cfg = json.load(open(os.path.join(a.snap, "config.json")))
L = cfg["num_hidden_layers"]
pref = f"model.layers.{L}."
src = SourceRepo(a.source_repo, a.revision)
# find the MTP layer's dense int4 tensors across the local shards
plan = {} # shard path -> [tensor names to repair]
already_ok, skipped_experts = [], 0
for f in sorted(glob.glob(os.path.join(a.snap, "*.safetensors"))):
with open(f, "rb") as fh:
n = struct.unpack("<Q", fh.read(8))[0]
hdr = json.loads(fh.read(n))
for name, v in hdr.items():
if not name.startswith(pref) or name.endswith(".qs") or v.get("dtype") != "U8":
continue
if ".mlp.experts." in name:
skipped_experts += 1
continue
m = src.meta(name)
if m is None:
print(f" ?? {name}: not in source repo, skipping")
continue
O, I = m["shape"]
nb = v["data_offsets"][1] - v["data_offsets"][0]
if nb == O * I:
already_ok.append(name)
elif nb == O * ((I + 1) // 2):
plan.setdefault(f, []).append((name, O, I))
else:
print(f" ?? {name}: unexpected blob size {nb} for shape {O}x{I}, skipping")
n_fix = sum(len(v) for v in plan.values())
print(f"MTP layer {L}: {n_fix} dense tensor(s) at per-row int4 to repair, "
f"{len(already_ok)} already int8, {skipped_experts} routed-expert tensors left as-is")
if not n_fix:
print("nothing to do."); return
if a.dry_run:
for f, names in plan.items():
for nm, O, I in names:
print(f" would repair {nm} [{O},{I}] in {os.path.basename(f)}")
return
for f, names in plan.items():
print(f"patching {os.path.basename(f)} ({len(names)} tensors)")
tensors, meta = read_shard(f)
for nm, O, I in names:
print(f" {nm}: fetching source + requantizing at int8...", flush=True)
w = src.fetch_f32(nm)
assert w.shape == (O, I), f"{nm}: source shape {w.shape} != container {O}x{I}"
q, s = quant_int8(w)
tensors[nm] = ("U8", [len(q)], q.tobytes())
tensors[nm + ".qs"] = ("F32", [O], s.tobytes())
# verify first row against the source
loc = q[:I].view(np.int8).astype(np.float64) * s[0]
ref = w[0].astype(np.float64)
cos = float(ref @ loc / (np.linalg.norm(ref) * np.linalg.norm(loc) + 1e-30))
print(f" row-0 cosine vs source: {cos:.5f}")
bak = f + ".bak-int4"
write_shard(f + ".new", tensors, meta)
os.replace(f, bak)
os.replace(f + ".new", f)
print(f" saved; original kept as {os.path.basename(bak)}")
print("done. Re-run with --dry-run to confirm (should report 'already int8').")
if __name__ == "__main__":
main()